Improve the performance of radix top-k

[yong-wang]

The main changes are:

• Add a one-block version. It uses single thread block for one row of a batch and is used when len is relatively small (<= 16384)
• Avoid writing candidates to buffers when the number of candidates is larger than buffer length.
• Add a parameter to control whether to use a fused filter in the last pass or use a standalone filter kernel. The later case is preferable when the leading bits of inputs are almost same.
• Early stopping: when the target bucket contains k values, we can stop the computation earlier
• Many implementation details are polished, like the initialization of counter, calculation of kernel launch parameters, and the scan step
• Tests and benchmarks are updated to include the new implementations. New benchmarks are added to demonstrate the advantage of adaptive version.

[rapids-bot]

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

[yong-wang]

More details about the benchmark changes and results.

The adaptive version is most useful when the most significant bits of input data are almost the same. That is, when the value range of input data is narrow. So, some new benchmarks are added to demonstrate its advantage.

• For float input, the value range of input is set to [1.0, 1.00003], and the corresponding binaries are 0x3F800000 and 0x3F8000FF, respectively. So the leading 3 bytes are all the same.
• For double input, the value range of input is set to [1.0, 1.000015], and the corresponding binaries are 0x3FF0000000000000 and 0x3FF0000FFFFFFFFF, respectively.

Benchmark results (using A100 GPU and CUDA 12.0):

• For existing benchmarks, on average, kRadix11bitsUpdated is 1.82X, 2.19X and 1.34X faster than kRadix11bits for float-int (float value & int index), double-int and double-size_t inputs, respectively.
• For the new added benchmarks, which have narrow range of inputs, kRadix11bitsAdaptive is 9.76X, 1.66X and 1X faster than kRadix11bitsUpdated for float-int, double-int and double-size_t inputs, respectively.

[cjnolet]

/ok to test

[codecov-commenter]

Codecov Report

❗ No coverage uploaded for pull request base (branch-23.04@88cb31d). Click here to learn what that means.
Patch has no changes to coverable lines.

Additional details and impacted files

@@               Coverage Diff               @@
##             branch-23.04    #1175   +/-   ##
===============================================
  Coverage                ?   87.99%           
===============================================
  Files                   ?       21           
  Lines                   ?      483           
  Branches                ?        0           
===============================================
  Hits                    ?      425           
  Misses                  ?       58           
  Partials                ?        0           

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[cjnolet]

/okay to test

[cjnolet]

/add to allowlist

[cjnolet]

/okay to test

[cjnolet]

/ok to test

[cjnolet]

/ok to test

[cjnolet]

@yong-wang just a heads up that we are reaching the halfway point for release 23.04 and are trying to avoid having too many last-minute changes merged during burndown. It looks like this has a conflict to resolve and it’s otherwise approved to go in. Do you think we can get this in for 23.04?

[cjnolet]

/ok to test

[cjnolet]

/merge

[tfeher]

@cjnolet there seems to be an issue with the input types passed to select_k while calling from here

raft/cpp/include/raft/spatial/knn/knn.cuh

Lines 156 to 167 in e4aec7b

	matrix::detail::select::radix::select_k<value_t, idx_t, 8, 512>(
	in_keys, in_values, n_inputs, input_len, k, out_keys, out_values, select_min, stream);
	break;
	case SelectKAlgo::RADIX_11_BITS:
	matrix::detail::select::radix::select_k<value_t, idx_t, 11, 512>(
	in_keys, in_values, n_inputs, input_len, k, out_keys, out_values, select_min, stream);
	break;
	case SelectKAlgo::WARP_SORT:
	matrix::detail::select::warpsort::select_k<value_t, idx_t>(
	in_keys, in_values, n_inputs, input_len, k, out_keys, out_values, select_min, stream);

I am investigating and I will send a commit to fix this. The question is why does it pass other tests? I could reproduce the CI error locally.

[yong-wang]

Thanks Tamas for locating the bug. I had misread the error message and missed the bug.

[cjnolet]

/ok to test

[cjnolet]

/ok to test

[tfeher]

There is an error with some of the tests:

./MATRIX_TEST --gtest_filter=*ReferencedRandomDoubleSizeT*

[  FAILED  ] 4 tests, listed below:
[  FAILED  ] SelectK/ReferencedRandomDoubleSizeT.Run/77, where GetParam() = (params{batch_size: 20, len: 700, k: 10, asc, no-input-index}, kRadix8bits)
[  FAILED  ] SelectK/ReferencedRandomDoubleSizeT.Run/112, where GetParam() = (params{batch_size: 100, len: 1700, k: 31, asc, no-input-index}, kRadix8bits)
[  FAILED  ] SelectK/ReferencedRandomDoubleSizeT.Run/140, where GetParam() = (params{batch_size: 100, len: 1700, k: 64, dsc, no-input-index}, kRadix8bits)
[  FAILED  ] SelectK/ReferencedRandomDoubleSizeT.Run/182, where GetParam() = (params{batch_size: 100, len: 1700, k: 1023, dsc, no-input-index}, kRadix8bits)

[yong-wang]

Got it. I'll take a look.

[cjnolet]

/okay to test

[cjnolet]

/okay to test