Fused L2 1-NN based on cutlass 3xTF32 / DMMA #1118

mdoijade · 2022-12-30T15:40:34Z

-- 3xTF32 & DMMA cutlass based persistent FusedL2NN kernel version loosely based on grouped gemm but customized for single problem size.
-- as the value of k increases the performance benefit of this implementation gets better.
for k==64 upto 1.3x, for k ==128 upto 1.53x, k == 256, up to 1.67x.
-- for all the sizes of k this kernel out performs previous implementation.
-- attaching the results of FusedL2NN Benchmark of previous implementation with this cutlass version.

…ck/warp shape size, this now touches the perf of fusedL2NN simt kernel

… from per row to multi-rows. now this kernel is 1.3x to 1.8x faster as k value increases perf gets better than fusedL2NN simt kernel

rapids-bot · 2022-12-30T15:40:41Z

Pull requests from external contributors require approval from a rapidsai organization member with write or admin permissions before CI can begin.

… to predicated tile iterator

… than per warp multi row lock, cleanup and doc update

mdoijade · 2023-01-04T10:51:41Z

FusedL2NN cutlass vs fp32 simt kernel.xlsx

mdoijade · 2023-01-04T10:53:25Z

@teju85 @tfeher for visibility and review.

cjnolet · 2023-01-19T19:07:30Z

/okay to test

cjnolet · 2023-01-19T19:10:47Z

/ok to test

…reviously needed

…election, add comments on register spills tile shape, add test case for veclen=2

tfeher

Thanks @mdoijade for the updates! The PR has changed significantly since the previous review round to enable an impressive speedup. Here are the first batch of comments of my second review.

This PR adds customized cutlass headers to implement the fused L2 NN operation.

While the changeset seem to be large, it is actually much smaller, if we compare the new cutlass headers to their original version (based on the reference added to files, and after applying raft formatting to the originals). Comparing that way makes it easier to follow how the cutlass code was adapted to our needs, and reveals a clean implementation. Great work, I have only smaller comments (so far)!

cpp/include/raft/util/cutlass_utils.cuh

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

cpp/include/raft/distance/detail/fused_distance_nn/custom_epilogue_with_broadcast.h

cpp/include/raft/distance/detail/fused_distance_nn/epilogue.cuh

cpp/include/raft/distance/detail/fused_distance_nn/epilogue_elementwise.cuh

cpp/include/raft/distance/detail/fused_distance_nn/epilogue.cuh

cpp/include/raft/distance/detail/fused_distance_nn/epilogue_elementwise.cuh

…usedL2NN_cutlass

…fix cutlass_utils.h comments

…ing 2 copies of the same code

tfeher

Thanks Mahesh for the updates so far. Here is my second batch of comments.

cpp/include/raft/distance/detail/fused_distance_nn/cutlass_base.cuh

cpp/include/raft/distance/detail/fused_distance_nn/gemm.h

cpp/include/raft/distance/detail/fused_distance_nn/persistent_gemm.h

cpp/include/raft/distance/detail/fused_distance_nn/predicated_tile_iterator_normvec_smem.h

…ntime

…usedL2NN_cutlass

cjnolet · 2023-05-15T21:22:50Z

I've lost track of where we are with this PR. Do you guys think this will make it into 23.06? (Burndown is in 3 days).

tfeher

Thanks Mahesh for the updates!

@cjnolet, the PR is in a good shape, most of the issues have been addressed, it shall make it to 23.06.

cpp/include/raft/distance/detail/fused_distance_nn/epilogue_elementwise.cuh

cpp/include/raft/distance/detail/fused_distance_nn/cutlass_base.cuh

cpp/include/raft/distance/detail/fused_distance_nn/persistent_gemm.h

…usedL2NN_cutlass

tfeher

Thanks Mahesh for resolving the issues! The PR looks good to me.

cjnolet · 2023-05-16T22:04:46Z

/merge

mdoijade added 5 commits November 18, 2022 20:40

add cutlass source files for fusedL2NN with initial set of changes

fb3e03c

temp commit to store current progress

7039a50

working code with atomicCAS based reduction from each thread

356c2e1

improve perf by warp reduce and reduce register by adjusting gemm blo…

6663d4c

…ck/warp shape size, this now touches the perf of fusedL2NN simt kernel

improve the perf of fusedL2NN cutlass kernel by reducing atomic locks…

ca88ad1

… from per row to multi-rows. now this kernel is 1.3x to 1.8x faster as k value increases perf gets better than fusedL2NN simt kernel

github-actions bot added the cpp label Dec 30, 2022

mdoijade added 3 commits December 30, 2022 21:13

add the custom gemm fused epilogue header required for passing params…

4e4e6ff

… to predicated tile iterator

rename final_op as cg_reduce_op, cleanup

8af3ae7

add whole block multi-row single lock impl which performs 5-7% faster…

d2c3833

… than per warp multi row lock, cleanup and doc update

cjnolet assigned mdoijade Jan 3, 2023

mdoijade added 5 commits January 4, 2023 13:35

merge branch-23.02

ea9e62a

fix connected components reduction functor for working with cutlass

b38414f

merge branch-23.02

68b75c8

fix clang format and copyright year

c89f1c1

add comments to cutlass headers which are customized

207a964

mdoijade marked this pull request as ready for review January 4, 2023 10:43

mdoijade requested a review from a team as a code owner January 4, 2023 10:43

mdoijade changed the title ~~WIP: Fused L2 1-NN based on cutlass 3xTF32 / DMMA~~ Fused L2 1-NN based on cutlass 3xTF32 / DMMA Jan 4, 2023

fix clang format issues

ff22154

cjnolet added improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels Jan 19, 2023

merge branch-23.04

cc3e669

mdoijade force-pushed the fusedL2NN_cutlass branch from 8169088 to cc3e669 Compare February 3, 2023 09:27

mdoijade requested review from a team as code owners February 3, 2023 09:27

mdoijade and others added 5 commits May 5, 2023 13:05

Merge branch 'branch-23.06' into fusedL2NN_cutlass

65d9c8e

remove the redundant header inclusion in fusedl2knn tests which was p…

1fce36a

…reviously needed

remove commented code and fix formating

b072d80

add larger input test cases to test fusedL2nn all code paths

a7c7303

fix launch config for small input sizes, fix atomicCAS optimal path s…

d741c4a

…election, add comments on register spills tile shape, add test case for veclen=2

tfeher requested changes May 8, 2023

View reviewed changes

mdoijade and others added 6 commits May 10, 2023 02:15

fix formatting issues

7e4b298

Merge remote-tracking branch 'mdoijade_fork/fusedL2NN_cutlass' into f…

c3ab218

…usedL2NN_cutlass

Merge branch 'branch-23.06' into fusedL2NN_cutlass

bc1bfad

move raft copyright below cutlass's and fix start year to be 2023

7f1d30d

add specialization for veclen=1 with 32x128x16 having no reg spills, …

514bd1e

…fix cutlass_utils.h comments

move smem init code to their respective tile iterators instead of hav…

cba6bbc

…ing 2 copies of the same code

tfeher requested changes May 11, 2023

View reviewed changes

mdoijade and others added 7 commits May 12, 2023 03:04

combine the optimal path and non-optimal path gmem writes

71b91e9

use the new dispatch mechanism to select the appropriate kernel at ru…

ad2ce75

…ntime

fix comments and small cleanup

6b87cc9

add comment about persistent_gemm.h mapping to its cutlass version

383f5bd

fix all formatting issues

3f4ceff

Merge remote-tracking branch 'mdoijade_fork/fusedL2NN_cutlass' into f…

5dbf438

…usedL2NN_cutlass

Merge branch 'branch-23.06' into fusedL2NN_cutlass

268c218

tfeher reviewed May 15, 2023

View reviewed changes

cpp/include/raft/distance/detail/fused_distance_nn/epilogue_elementwise.cuh Show resolved Hide resolved

cpp/include/raft/distance/detail/fused_distance_nn/cutlass_base.cuh Show resolved Hide resolved

cpp/include/raft/distance/detail/fused_distance_nn/persistent_gemm.h Show resolved Hide resolved

cjnolet approved these changes May 16, 2023

View reviewed changes

cjnolet and others added 3 commits May 15, 2023 22:18

Merge branch 'branch-23.06' into fusedL2NN_cutlass

e21a8da

remove dead/commented code from epilogue broadcast header

5f8f33b

Merge remote-tracking branch 'mdoijade_fork/fusedL2NN_cutlass' into f…

d603f81

…usedL2NN_cutlass

tfeher approved these changes May 16, 2023

View reviewed changes

fix formatting

4140734

rapids-bot bot merged commit a1d1fd6 into rapidsai:branch-23.06 May 16, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fused L2 1-NN based on cutlass 3xTF32 / DMMA #1118

Fused L2 1-NN based on cutlass 3xTF32 / DMMA #1118

mdoijade commented Dec 30, 2022 •

edited

Loading

rapids-bot bot commented Dec 30, 2022

mdoijade commented Jan 4, 2023

mdoijade commented Jan 4, 2023

cjnolet commented Jan 19, 2023

cjnolet commented Jan 19, 2023

tfeher left a comment

tfeher left a comment

cjnolet commented May 15, 2023

tfeher left a comment

tfeher left a comment

cjnolet commented May 16, 2023

Fused L2 1-NN based on cutlass 3xTF32 / DMMA #1118

Fused L2 1-NN based on cutlass 3xTF32 / DMMA #1118

Conversation

mdoijade commented Dec 30, 2022 • edited Loading

rapids-bot bot commented Dec 30, 2022

mdoijade commented Jan 4, 2023

mdoijade commented Jan 4, 2023

cjnolet commented Jan 19, 2023

cjnolet commented Jan 19, 2023

tfeher left a comment

Choose a reason for hiding this comment

tfeher left a comment

Choose a reason for hiding this comment

cjnolet commented May 15, 2023

tfeher left a comment

Choose a reason for hiding this comment

tfeher left a comment

Choose a reason for hiding this comment

cjnolet commented May 16, 2023

mdoijade commented Dec 30, 2022 •

edited

Loading