Replace dots_along_rows with rowNorm and improve coalescedReduction performance

[Nyrio]

dots_along_rows in ann_utils.cuh was in some cases more performant than the corresponding raft primitive rowNorm, so I have improved that primitive in order to replace dots_along_rows without performance regressions. rowNorm for a row-major matrix calls coalescedReduction, which I have modified to conditionally select one of the following code paths based on the input dimensions:

• Thin: for matrices with many small rows, one block processes multiple rows, with 2 to 32 threads collaborating on each row using a shuffle-based reduction.
• Medium: the existing cub-based implementation with one block per row (I have only changed the reduction algorithm to raking which is more performant provided that the workload is big enough)
• Thick: two-step implementation. In the first step, multiple blocks per row reducing to an intermediate buffer (main_op is applied but not final_op). In the second step, reduces the intermediate buffer using the thin kernel (this time final_op is applied but not main_op).

Other changes included in this PR:

• In order to properly support shuffle-based reductions, I have added generic shuffle helpers that support arbitrary types by cutting them into chunks (based on size/alignment). This was adapted from similar helpers in CUB.
• I have added a helper for "logical" warp reduction, i.e sub-warps of 2, 4, 8, 16 or 32 threads, and added support for arbitrary reduction operations in the warp reduction.
• I have consolidated tests with support for arbitrary types and operations and tested some operations that in particular use the index argument of main_op such as an argmax, and only for the coalesced reduction I have added test cases with raft::KeyValuePair

[Nyrio]

Note to reviewers: I am aware that the reduction currently doesn't compile with non-trivial types such as cub pairs due to the shuffle-based reductions. Working on a fix.

[Nyrio]

I have fixed support for non-trivial types, please have a detailed look at the last commit and in particular changes to cuda_utils.cuh.

[tfeher]

After these changes, is the following comment still valid?

raft/cpp/include/raft/linalg/norm.cuh

Line 41 in 355f693

* current implementation is optimized only for bigger values of 'D'.

[Nyrio]

After these changes, is the following comment still valid?

Removed.

[tfeher]

Hi Louis, it is nice to see further improvements in our prims. I see that the bulk of the changes are the updates in the tests cases, thanks for the thorough work!

I have just few smaller comments for the code.

Please update the PR description:

• mention adding general shuffle and reduction op
• move detailed description about performance of different kernels into a separate comment.

If you have any measurements/notes on why is this approach better than cub segmented reduction, then please add a comment.

[Nyrio]

@tfeher cub::DeviceSegmentedReduce is a more generic primitive and my expectation was that it would not perform better, but I should run some benchmarks against it to make sure of that. Segmented reduce can work with segments of arbitrary lengths, and reads the start and end offsets of segments from arrays. I haven't read the implementation but my guess is that it does a BlockReduce per segment, in which case we have no reason to pay the price of creating and reading these offsets.

[Nyrio]

Some notes on the performance of the thick vs medium kernel:

• For the thick implementation, I considered using atomics but for a generic reduction, it requires a pre-step to initialize the output, allocate and initialize mutexes, and a post-step for the final op. That is altogether much costlier than the two-step approach I ended up using.
• The prim is heavily memory-bound, so only one block per SM is enough to reach near-SOL global memory bandwidth, meaning the medium kernel will perform better if the number of rows is near or greater than the number of SMs, i.e anything more than a few dozens of rows.
• If the number of rows is small and the number of columns is up to a few thousand, we should also prefer the one-kernel approach because the bottleneck is the launch latencies.

Visual demonstration of the performance of the medium vs thick implementations (y-axis is time in ms, lower is better):

[tfeher]

Thanks Louis for the update LGTM!

[cjnolet]

rerun tests

[cjnolet]

@Nyrio i suspect maybe the CI checks aren’t being executed because of the conflicts in your branch.

[Nyrio]

@cjnolet I was waiting for my local compilation and test run to succeed before pushing, but as you can expect, compiling the neighbors test took a few hours.

[cjnolet]

Wait, a few hours?!?! What type of environment / configuration are you using? How many cores are you using to compile?

[cjnolet]

@gpucibot merge

[cjnolet]

@gpucibot merge

[cjnolet]

rerun tests

[Nyrio]

@cjnolet It looks like the CI errors are unrelated to the contents of this PR.

[cjnolet]

rerun tests

[cjnolet]

@Nyrio yep you are right about that. @ajschmidt8 has fixed the issue so we should be able to get this in today, assuming it passes.

[cjnolet]

rerun tests

[cjnolet]

rerun tests

[cjnolet]

rerun tests