Add cutlass 3xTF32,DMMA based L2/cosine distance kernels for SM 8.0 or higher #939

mdoijade · 2022-10-21T16:04:31Z

-- 3xTF32 cutlass based L2 exp/cosine kernel provides 3.5x speedup for fp32 inputs compared to existing pairwise distance kernel for ampere or higher.
-- DMMA cutlass based implementation for L2 exp/cosine provides 2.6x speedup instead of existing double precision FMA pipeline based kernel.
-- add cutlass as header only dependency to RAFT.

…ompilation, fix clang formating and correct line endings

…ng issues

GPUtester · 2022-11-03T14:34:49Z

Can one of the admins verify this patch?

Admins can comment ok to test to allow this one PR to run or add to allowlist to allow all future PRs from the same author to run.

cpp/cmake/thirdparty/get_cutlass.cmake

robertmaynard

By these lines it looks like cutlass is now a hard requirement to use RAFT no matter the configuration, since the dependency is on the core raft library and not raft::distance.

I expect that is wrong, and this all needs to be moved to raft::distance

cpp/CMakeLists.txt

cpp/cmake/thirdparty/get_cutlass.cmake

…ectly

…required now

tfeher

Thanks @mdoijade for the PR! It is exciting to see these changes! The PR looks great overall, I have mostly requested documenting from where (and how) the CUTLASS code was adapted from. The goal with this to make it easier to adapt to changes in future CUTLASS releases.

Please remove the [WIP] tag and add a PR title and please add a description: it would be great to have a few lines on how the normalization factors are handled using iterators and epilogue functions.

cpp/include/raft/util/cudart_utils.hpp

cpp/test/distance/dist_adj.cu

cpp/include/raft/distance/detail/predicated_tile_iterator_normvec.h

cpp/include/raft/distance/detail/euclidean.cuh

cpp/include/raft/distance/detail/cosine.cuh

cpp/include/raft/distance/detail/pairwise_distance_epilogue.h

cpp/include/raft/distance/detail/pairwise_distance_epilogue_elementwise.h

cpp/include/raft/distance/detail/pairwise_distance_gemm.h

cpp/cmake/thirdparty/get_cutlass.cmake

…nstead of cutlass for its cmake config file

cpp/cmake/thirdparty/get_cutlass.cmake

tfeher

Thanks Mahesh for addressing the issues. The C++ changes look good to me.

robertmaynard · 2022-11-14T20:44:59Z

cpp/CMakeLists.txt

@@ -217,6 +218,7 @@ target_link_libraries(
            CUDA::cusolver${_ctk_static_suffix}
            CUDA::cusparse${_ctk_static_suffix}
            $<$<BOOL:${RAFT_ENABLE_thrust_DEPENDENCY}>:raft::Thrust>
+            nvidia::cutlass::cutlass


This dependency needs to be on the raft-distance target and not the raft target

@robertmaynard I see that several tests and core algorithms depend on cosine/euclidean distance headers where we are using cutlass, hence I think it needs to be dependency on raft target. without it I am seeing several build failures when those sources are built.
I've modified get_cutlass.cmake from raft-distance-exports to raft-exports. Can this resolve the build issue in pylibraft?
I've submitted this change and waiting to see if CI passes.

@cjnolet Are you okay with cutlass being a hard requirement for raft?

I'd prefer if we could make CUTLASS a dependency only of raft::distance (which pylibraft uses).

@mdoijade any tests/benchmarks and downstream projects which depend on distances also specify and use the raft::distance target.

I see, I've tried to add dependency on raft_distance wherever required and got the build working locally, please review if it looks good now.
also I've enabled the cutlass path only till cuda 11.x.

…nents also use distance headers

…pendency wherever needed in downstream tests/lib

… and enable non-cutlass version for ampere+ temporarily

robertmaynard · 2022-11-16T14:04:24Z

cpp/CMakeLists.txt

@@ -422,7 +424,7 @@ if(RAFT_COMPILE_NN_LIBRARY)
               INTERFACE_POSITION_INDEPENDENT_CODE ON
  )

-  target_link_libraries(raft_nn_lib PUBLIC faiss::faiss raft::raft)
+  target_link_libraries(raft_nn_lib PUBLIC faiss::faiss raft::raft raft_distance)


Why does raft_nn_lib need the raft_distance.

raft_nn_lib sources make use of distance headers such as cosine.cuh/euclidean.cuh where it encounters build failure with missing cutlass headers.

If it is raft_nn_lib srcs that have this dependency, the target_link_libraries usage should be PRIVATE. Secondly it should be a dependency on the cutlass target and not raft_distance.

This currently creates a hard dependency for downstream consumers of raft, which breaks requests like find_package(raft COMPONENTS nn) since the raft_distance target won't exist for those users.

so can I make it target_link_libraries(raft_nn_lib PUBLIC faiss::faiss raft::raft nvidia::cutlass::cutlass) ?
and then will I need to modify get_cutlass.cmake to add rapids_export_package to raft-nn-lib-exports ?

you would want target_link_libraries(raft_nn_lib PUBLIC faiss::faiss raft::raft PRIVATE nvidia::cutlass::cutlass) and no other change would be required.

I've made these changes and pushed them.

Thanks @cjnolet and @robertmaynard for quick help and fixes!

robertmaynard

CMake code looks good!

cjnolet

LGTM!

cjnolet · 2022-11-16T20:02:39Z

rerun tests

cjnolet · 2022-11-16T22:55:26Z

@gpucibot merge

cjnolet · 2022-11-16T22:57:23Z

@gpucibot merge

PR #939 introduced CUTLASS dependency. When compiled in debug mode, this leads to the following error: ``` ptxas error : Stack size for entry function '_ZN12raft_cutlass6KernelINS_...' cannot be statically determined ``` This would be normally just a warning, but we treat warnings as errors. This PR disables the warning in Debug mode. Authors: - Tamas Bela Feher (https://github.com/tfeher) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: #1033

mdoijade added 4 commits October 20, 2022 23:02

cutlass based euclidean expanded, cosine kernels

20648c5

add prior ampere pairwisedistmat kernel to prevent redundant kernel c…

a9dabc8

…ompilation, fix clang formating and correct line endings

add noexcept to the functor methods

1a45bfa

merge branch 22.12 and resolve conflicts

c6f091b

mdoijade requested review from a team as code owners October 21, 2022 16:04

github-actions bot added CMake cpp labels Oct 21, 2022

fix comments, remove redundant code and fix formatting issues

7786fcb

cjnolet added improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels Oct 27, 2022

mdoijade added 2 commits October 28, 2022 16:41

add cutlass cmake support for raft with custom namespace, fix formati…

181fc40

…ng issues

fix formatting issues

3d34545

cjnolet assigned mdoijade Nov 1, 2022

fix the cutlass_include_dir path in cmake

02c23ed

robertmaynard requested changes Nov 3, 2022

View reviewed changes

cpp/CMakeLists.txt Outdated Show resolved Hide resolved

cpp/CMakeLists.txt Outdated Show resolved Hide resolved

robertmaynard requested changes Nov 3, 2022

View reviewed changes

cpp/cmake/thirdparty/get_cutlass.cmake Outdated Show resolved Hide resolved

fix bugs in get_cutlass cmake to use cutlass provided properties corr…

7933436

…ectly

mdoijade requested a review from robertmaynard November 4, 2022 09:41

mdoijade added 3 commits November 4, 2022 15:44

remove the cutlass namespace setting in test cmakefiles as it is not …

d4bdec5

…required now

temp remove dist dependency from cutlass to check if it works in ci/cd

d26bcef

merge branch-22.12 latest changes

4df4185

tfeher requested changes Nov 9, 2022

View reviewed changes

robertmaynard requested changes Nov 9, 2022

View reviewed changes

cpp/cmake/thirdparty/get_cutlass.cmake Outdated Show resolved Hide resolved

fix get_cutlass.cmake to work with pylibraft by using NvidiaCutlass i…

451c3c0

…nstead of cutlass for its cmake config file

robertmaynard requested changes Nov 10, 2022

View reviewed changes

cpp/cmake/thirdparty/get_cutlass.cmake Outdated Show resolved Hide resolved

fix get_cutlass install path, make changes as per review comments

7b512f9

mdoijade added 3 commits November 11, 2022 01:20

merge branch-22.12

a05e1e2

fix clang format issues

d32b4c0

temp fix to check if python build works

f7c440a

mdoijade changed the title ~~[WIP] Add cutlass 3xTF32,DMMA based L2/cosine distance kernels for SM 8.0 or higher~~ Add cutlass 3xTF32,DMMA based L2/cosine distance kernels for SM 8.0 or higher Nov 11, 2022

tfeher approved these changes Nov 14, 2022

View reviewed changes

robertmaynard requested changes Nov 14, 2022

View reviewed changes

mdoijade added 4 commits November 15, 2022 19:12

add raft-exports instead of raft-distance-exports as other raft compo…

b1a1fd7

…nents also use distance headers

make cutlass to depend only on raft_distance and add raft_distance de…

4ef44e7

…pendency wherever needed in downstream tests/lib

fix cmake formatting issues

186fcc7

prevent cutlass based pairwise dist kernels to be disabled on cuda 12…

8aa8909

… and enable non-cutlass version for ampere+ temporarily

robertmaynard requested changes Nov 16, 2022

View reviewed changes

cjnolet added 7 commits November 16, 2022 12:18

Moving cutlass dependency to distance and nn to keep them separate.

abfd493

Adding CUTLASS to build docs as dependency

f1b1239

Updating to export to both distance and nn

32e6052

Adding cutlass as private dependency

f6de9ee

Making cutlass INTERFACE in raft::nn and raft::distance

9bf0647

Using proper exports per Robert Maynard's suggestion.

8f0119a

Adding cutlass as private dependency of lib targets

6ad4fd1

robertmaynard approved these changes Nov 16, 2022

View reviewed changes

cjnolet approved these changes Nov 16, 2022

View reviewed changes

rapids-bot bot merged commit 611abc7 into rapidsai:branch-22.12 Nov 16, 2022

tfeher mentioned this pull request Nov 18, 2022

Suppres ptxas stack size warning in Debug mode #1033

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add cutlass 3xTF32,DMMA based L2/cosine distance kernels for SM 8.0 or higher #939

Add cutlass 3xTF32,DMMA based L2/cosine distance kernels for SM 8.0 or higher #939

mdoijade commented Oct 21, 2022 •

edited

Loading

GPUtester commented Nov 3, 2022

robertmaynard left a comment

tfeher left a comment

tfeher left a comment

robertmaynard Nov 14, 2022

mdoijade Nov 15, 2022

robertmaynard Nov 15, 2022

cjnolet Nov 15, 2022

cjnolet Nov 15, 2022

mdoijade Nov 16, 2022

robertmaynard Nov 16, 2022

mdoijade Nov 16, 2022

robertmaynard Nov 16, 2022

mdoijade Nov 16, 2022 •

edited

Loading

robertmaynard Nov 16, 2022

cjnolet Nov 16, 2022

mdoijade Nov 16, 2022

robertmaynard left a comment

cjnolet left a comment

cjnolet commented Nov 16, 2022

cjnolet commented Nov 16, 2022

cjnolet commented Nov 16, 2022

Add cutlass 3xTF32,DMMA based L2/cosine distance kernels for SM 8.0 or higher #939

Add cutlass 3xTF32,DMMA based L2/cosine distance kernels for SM 8.0 or higher #939

Conversation

mdoijade commented Oct 21, 2022 • edited Loading

GPUtester commented Nov 3, 2022

robertmaynard left a comment

Choose a reason for hiding this comment

tfeher left a comment

Choose a reason for hiding this comment

tfeher left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mdoijade Nov 16, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

robertmaynard left a comment

Choose a reason for hiding this comment

cjnolet left a comment

Choose a reason for hiding this comment

cjnolet commented Nov 16, 2022

cjnolet commented Nov 16, 2022

cjnolet commented Nov 16, 2022

mdoijade commented Oct 21, 2022 •

edited

Loading

mdoijade Nov 16, 2022 •

edited

Loading