CAGRA - separable compilation for distance computation

[achirkin]

Factor the compute_distance function and related template parameters out of the CAGRA search kernels.
This reduces the total number of kernel instances, thus reducing the binary size and the compile time.

The change, however, has a few drawbacks:

• CUDA separable compilation needs to be enabled to allow compute_distance functions being compiled in separate object files. I introduced a static library component for the affected sources to minimize the impact of the change.
• The separable compilation and dynamic dispatch of compute_distance function means the compiler cannot optimize across the kernel-compute_distance boundary, which results in higher register usage and occasional register spilling. Most of the cases are optimized in this PR, but some compromises seem unavoidable.
• Dynamic dispatch (constructing a dataset descriptor) requires an extra kernel call (xxx_init_kernel) to get the function pointer, which adds extra latency. This is mitigated to some extent by caching the constructed descriptor using raft custom resource.

[achirkin]

Current WIP status, with full functionality restored after the refactoring:

• CI build size 880 -> 663MB.
• Slowdown is up to 2x
• Bonus: multi-kernel version now naturally supports CAGRA-Q compression.

[achirkin]

Current WIP status:

• CI build size 788 -> 600MB.
• Slowdown is up to 20% for standard distance and up to 30% for VPQ distance. The worst case is the single-cta kernel on a big batch; the multi-kernel version actually sees some speedup in a few cases.
• Multi-kernel version now supports CAGRA-Q compression.
• A few tests failing
• Disabled target_link_options(cuvs PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/fatbin.ld") in CMakeLists.txt due to a linker error. Not sure how to resolve this at the moment.

[achirkin]

Update on performance: the worst-case slowdown now is around 6-7% on the deep-100M and wiki-all datasets.
Comparison against #324 (comment)
https://docs.google.com/spreadsheets/d/191f1sYsAUwPidncV4xDzgNtOhKx9sp6jJFvOfpB-nxs

[tfeher]

Thank you Artem for this PR! It is great what this PR achieves: finally the distance computation functions are decoupled from the search kernels and that reduces the number of times the search kernels are compiled. This significantly decreases of the binary size.

This comes at a price: for some parameter combinations cagra::search will become up to 6% slower. Thank you for investigating another solution in #324. Since that has similar impact on runtime, but results in less reduction in binary size, I am in favor of the current PR.

Fixing CAGRA's binary size will enable us to add more features that will improve the performance (persistent kernel, fp16). Therefore would recommend that we go ahead and merge the current PR.

There is still one aspect that I find unfortunate: the complexity of selecting and dispatching the cagra kernels (and distance functions) is significantly increased. The persistent kernel PR will add another set of complications on top of this. To compensate, we shall improve the developer documentation: I have left a few comments along this line.

Still, I suspect that we could do better, therefore please open an issue (as a follow-up of this an #215) to re-evaluate and simplify CAGRA kernel/distance selection and dispatch logic.

[tfeher]

Thanks Artem for fixing the issues, the added comments are really useful! The PR looks good to me. Please create a follow up issue to improve naming of descriptors and simplify call hierarchy.

[cjnolet]

Was building a new artifact really necessary to improve the perf / binary size? Or was this just done to make the build more modular?

[achirkin]

Generally, separable compilation affects performance negatively, so I reduced the AOE by setting cmake CUDA_SEPARABLE_COMPILATION on the relevant files only - via this component.

[cjnolet]

/merge