Use JIT-LTO kernel handles through an API the toolkit actually provides

[drzraf]

Use JIT-LTO kernel handles through an API the toolkit actually provides

Part 1 of 2. On its own this commit changes nothing observable: it is a prerequisite for building cuVS with a CUDA toolkit older than 12.8, and that becomes possible only in combination with rapidsai/librtcx#17 - "Support CUDA toolkits older than 12.8 via the driver API"

Neither is useful without the other, and either may merge first. The split follows the code: since #2311 the JIT-LTO plumbing lives in librtcx, which loads the linked image (cudaLibraryLoadData / cudaLibraryGetKernel / cudaLibraryUnload) and launches it (cudaLaunchKernelExC); cuVS keeps only what it does with the resulting handle. Both halves hit the same wall -- the runtime library-management API arrived in CUDA 12.8 -- but they are in different repositories, so they are fixed separately. Without librtcx#17 every cuVS translation unit that includes the rtcx launcher still fails with

rtcx/algorithm_launcher.hpp:24: error: identifier "cudaLibrary_t" is undefined

(this commit does not change that)

The cuVS half. cuVS queries and sets attributes on the cudaKernel_t handles that rtcx hands out for run-time-linked kernels. Both things it does with those handles are CUDA 12.8 features:

• Passing a cudaKernel_t where a "const void* func" entry point is expected.
  The runtime documents this from 12.8 onwards -- "If the specified function
  does not exist, then it is assumed to be a cudaKernel_t and used as is",
  which appears in the 12.8 documentation of cudaFuncGetAttributes,
  cudaFuncSetAttribute and cudaOccupancyMaxActiveBlocksPerMultiprocessor and in
  no earlier version. On an older toolkit the handle is not a valid entry
  address and the call fails with cudaErrorInvalidDeviceFunction.

• cudaKernelSetAttributeForDevice, which does not exist before 12.8 in either the headers or libcudart.

The driver API has always taken the equivalent CUkernel / CUfunction, and cudaKernel_t is a typedef for struct CUkern_st*, i.e. exactly CUkernel, so the handles carry over unchanged. Add src/util/jit_kernel_compat.hpp, which forwards to the runtime API on 12.8+ and to cuKernelSetAttribute, cuFuncGetAttribute and cuOccupancyMaxActiveBlocksPerMultiprocessor otherwise, and route the six call sites through it. On 12.8+ every wrapper compiles down to the original runtime call and CUDA::cuda_driver is not linked, so nothing changes for the toolkits cuVS supports today.

Two details worth recording:

• The runtime's primary context is forced to exist before any driver call,
  because cuKernelGetFunction resolves against the current context. Skipping
  this yields CUDA_ERROR_INVALID_CONTEXT from an otherwise correct sequence.

• Driver statuses are translated to the nearest runtime status so that existing
  RAFT_CUDA_TRY diagnostics stay meaningful.

In launchConfigGenerator the choice between the two occupancy APIs is an if constexpr on the argument type, because that helper is also called with genuine __global__ function pointers, for which the runtime API is correct.

The IVF-PQ kernel selector needed one further change. It asserted that a failed cudaFuncSetAttribute was also observable through cudaGetLastError(), which only holds for the runtime API; the driver path never sets the runtime's sticky error. The intent -- skip a kernel candidate that cannot get the shared memory it wants -- is preserved, and the sticky error is cleared explicitly.

Verified with CUDA 12.6.85: with rapidsai/librtcx#17 applied (-DCPM_rtcx_SOURCE=...),
libcuvs.so builds and links for sm_50, and CAGRA build and search run correctly on a Maxwell device. Without it, the same tree fails only inside the rtcx headers, which is what "part 1 of 2" means in practice.

Note that dependencies.yaml still lists CUDA 12.2 and 12.5 in its cuda_version matrix, so the toolkit range the project declares and the one it can actually build have disagreed since #1405.

(Opus 5)

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.