From d5a855c7eba7f50502c6b282c09a1f0be4b80129 Mon Sep 17 00:00:00 2001 From: achirkin Date: Fri, 28 Aug 2026 16:44:45 +0200 Subject: [PATCH 1/2] Extend kernel_launch with kernel_ref and attributes --- cpp/include/raft/util/kernel_launch.hpp | 145 ++++++++++++-- cpp/tests/util/kernel_launch.cu | 248 ++++++++++++++++++++++++ 2 files changed, 379 insertions(+), 14 deletions(-) diff --git a/cpp/include/raft/util/kernel_launch.hpp b/cpp/include/raft/util/kernel_launch.hpp index 3acc3a22d4..8bfe5cdd65 100644 --- a/cpp/include/raft/util/kernel_launch.hpp +++ b/cpp/include/raft/util/kernel_launch.hpp @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -71,14 +72,44 @@ void dispatch(cudaLaunchConfig_t const& config, */ /** - * @brief Where a kernel is launched: the stream, the dynamic shared memory size, and the call site - * to blame for launch errors. + * @brief Launch attribute: the kernel synchronizes across the whole grid. + * + * Such a launch fails unless the whole grid is resident on the device at once, so its grid size has + * to come from an occupancy query rather than from the problem size. + */ +inline auto cooperative() -> cudaLaunchAttribute +{ + cudaLaunchAttribute attr{}; + attr.id = cudaLaunchAttributeCooperative; + attr.val.cooperative = 1; + return attr; +} + +/** + * @brief Launch attribute: preferred share of the combined L1/shared memory to use as shared + * memory, in percent. + * + * Only a hint; the driver may pick a different split. Unrelated to the cap on dynamic shared + * memory, which is a property of the kernel rather than of the launch. + */ +inline auto shmem_carveout(unsigned percent) -> cudaLaunchAttribute +{ + cudaLaunchAttribute attr{}; + attr.id = cudaLaunchAttributePreferredSharedMemoryCarveout; + attr.val.sharedMemCarveout = percent; + return attr; +} + +/** + * @brief How and where a kernel is launched: the stream, the dynamic shared memory size, the launch + * attributes, and the call site to blame for launch errors. * * Converts implicitly from raft resources or from a stream, so that a launch reads as a single * call and the diagnostics of a failed launch point at the launch expression: * @code * raft::launch_kernel(res, grid, block, my_kernel, arg0, arg1); * raft::launch_kernel({stream, smem}, grid, block, my_kernel, arg0, arg1); + * raft::launch_kernel({res, smem, {raft::cooperative()}}, grid, block, my_kernel, arg0, arg1); * @endcode * * Launching on raft resources is dry run compliant: the kernel does not run when the handle has the @@ -87,7 +118,9 @@ void dispatch(cudaLaunchConfig_t const& config, * * Copy and move are deleted and @c launch_kernel takes this by value, so the parameter can only be * initialized from a prvalue: an instance stored in a variable can never be launched, and the - * captured location is therefore always the one of the launch expression. + * captured location is therefore always the one of the launch expression. That is also what makes + * it safe for @c config to point at @p attrs, whose backing array lives until the end of that + * expression. */ struct launch_on { public: @@ -99,15 +132,18 @@ struct launch_on { * * @param[in] res raft resources providing the stream to launch on * @param[in] smem dynamic shared memory size in bytes + * @param[in] attrs launch attributes, e.g. @c raft::cooperative() * @param[in] loc call site to blame for launch errors; leave at its default */ launch_on( // NOLINT(google-explicit-constructor) resources const& res, - std::size_t smem = 0, - std::source_location loc = std::source_location::current()) + std::size_t smem = 0, + std::initializer_list attrs = {}, + std::source_location loc = std::source_location::current()) : launch_on{resource::get_cuda_stream(res).value(), smem, resource::get_dry_run_flag(res) ? detail::kSkipExecution : detail::launch_flags{}, + attrs, loc} { } @@ -122,14 +158,16 @@ struct launch_on { * @param[in] stream stream to launch on * @param[in] smem dynamic shared memory size in bytes * @param[in] kSkipExecution whether to skip the launch, e.g. a dry-run flag plumbed by the caller + * @param[in] attrs launch attributes, e.g. @c raft::cooperative() * @param[in] loc call site to blame for launch errors; leave at its default */ launch_on( // NOLINT(google-explicit-constructor) rmm::cuda_stream_view stream, - std::size_t smem = 0, - bool kSkipExecution = false, - std::source_location loc = std::source_location::current()) - : launch_on{stream.value(), smem, kSkipExecution, loc} + std::size_t smem = 0, + bool kSkipExecution = false, + std::initializer_list attrs = {}, + std::source_location loc = std::source_location::current()) + : launch_on{stream.value(), smem, kSkipExecution, attrs, loc} { } @@ -143,14 +181,17 @@ struct launch_on { * @param[in] stream stream to launch on * @param[in] smem dynamic shared memory size in bytes * @param[in] kSkipExecution whether to skip the launch, e.g. a dry-run flag plumbed by the caller + * @param[in] attrs launch attributes, e.g. @c raft::cooperative() * @param[in] loc call site to blame for launch errors; leave at its default */ launch_on( // NOLINT(google-explicit-constructor) cudaStream_t stream, - std::size_t smem = 0, - bool kSkipExecution = false, - std::source_location loc = std::source_location::current()) - : launch_on{stream, smem, kSkipExecution ? detail::kSkipExecution : detail::launch_flags{}, loc} + std::size_t smem = 0, + bool kSkipExecution = false, + std::initializer_list attrs = {}, + std::source_location loc = std::source_location::current()) + : launch_on{ + stream, smem, kSkipExecution ? detail::kSkipExecution : detail::launch_flags{}, attrs, loc} { } @@ -170,19 +211,58 @@ struct launch_on { private: /** * The flags are private so that they stay a property of the resources: a call site names a - * stream, a shared memory size and at most a dry-run flag, never a launch mode. + * stream, a shared memory size, the launch attributes and at most a dry-run flag, never a launch + * mode. Attributes differ from flags in exactly that respect: they describe the launch itself, so + * they are given at the call site and go straight into @c config. */ launch_on(cudaStream_t stream, std::size_t smem, detail::launch_flags launch_with, + std::initializer_list attrs, std::source_location loc) : location{loc}, flags{launch_with} { config.stream = stream; config.dynamicSmemBytes = smem; + config.numAttrs = static_cast(attrs.size()); + // cudaLaunchConfig_t::attrs is not const, although cudaLaunchKernelExC takes the configuration + // by const pointer and never writes to the list. + config.attrs = const_cast(attrs.begin()); } }; +/** + * @brief A kernel that exists only at run time, together with the signature it was compiled with. + * + * A kernel loaded from a runtime-linked library (@c cudaLibraryGetKernel, e.g. after a JIT LTO + * link) has no @c __global__ function pointer for @c launch_kernel to read the parameter types + * from, so the signature is named explicitly: + * @code + * using scan_kernel_t = void(float const*, std::uint32_t); + * raft::launch_kernel({res, smem}, grid, block, + * raft::kernel_ref{handle}, queries, n_queries); + * @endcode + * + * Whether the handle really has that signature is on whoever loaded it. Given the signature, the + * launch converts each argument to its parameter type, so a call site does not need casts to make + * the argument types match the kernel exactly. + * + * @tparam Signature the kernel's function type, e.g. @c void(float const*, std::uint32_t) + */ +template +struct kernel_ref { + static_assert(sizeof(Signature) == 0, "kernel_ref needs a function type, e.g. void(float*, int)"); +}; + +template +struct kernel_ref { + /** @param[in] kernel handle to a loaded kernel whose signature is @c void(Params...) */ + explicit kernel_ref(cudaKernel_t kernel) : handle{kernel} {} + + /** Handle to the loaded kernel. */ + cudaKernel_t handle; +}; + /** * @brief Launch @p kernel with @p args, which already have the kernel parameter types. * @@ -258,6 +338,43 @@ requires(sizeof...(Params) == sizeof...(Args) && static_cast(std::forward(args))...); } +/** + * @brief Launch a kernel named by a runtime handle, converting @p args to its parameter types. + * + * Behaves like the converting overload above, except that the kernel and its parameter types come + * from @p kernel rather than from a @c __global__ function pointer: + * @code + * raft::launch_kernel({res, smem}, grid, block, + * raft::kernel_ref{launcher->get_kernel()}, queries, n); + * @endcode + * + * Unlike the two function-pointer overloads, this one accepts arguments that already have the + * parameter types too, because there is no exactly-matching overload for them to prefer. + * + * @param[in] where stream to launch on, dynamic shared memory size, attributes, and the call site + * @param[in] grid grid dimensions + * @param[in] block block dimensions + * @param[in] kernel handle to the loaded kernel, with the signature to launch it by + * @param[in] args arguments to convert and pass to @p kernel + */ +template +requires(sizeof...(Params) == sizeof...(Args)) void launch_kernel( + launch_on where, dim3 grid, dim3 block, kernel_ref kernel, Args&&... args) +{ + static_assert((std::is_convertible_v && ...), + "Each launch argument must be convertible to the corresponding kernel parameter"); + + where.config.gridDim = grid; + where.config.blockDim = block; + // A cudaKernel_t is an object pointer, so it needs no cast to reach dispatch, which launches it + // with the same cudaLaunchKernelExC that a __global__ function pointer goes through. + detail::dispatch(where.config, + kernel.handle, + where.flags, + where.location, + static_cast(std::forward(args))...); +} + /** @} */ // end group kernel_launch } // namespace raft diff --git a/cpp/tests/util/kernel_launch.cu b/cpp/tests/util/kernel_launch.cu index c34efb506b..ac01ecf8e9 100644 --- a/cpp/tests/util/kernel_launch.cu +++ b/cpp/tests/util/kernel_launch.cu @@ -14,6 +14,7 @@ #include +#include #include #include #include @@ -29,6 +30,25 @@ RAFT_KERNEL write_one_kernel(int* out) if (threadIdx.x == 0 && blockIdx.x == 0) { *out = 1; } } +RAFT_KERNEL write_count_kernel(int* out, std::uint32_t n) +{ + if (threadIdx.x == 0 && blockIdx.x == 0) { *out = static_cast(n); } +} + +RAFT_KERNEL copy_one_kernel(int const* in, int* out) +{ + if (threadIdx.x == 0 && blockIdx.x == 0) { *out = *in; } +} + +/** The handle a runtime-linked library would hand out, for a kernel this test compiled itself. */ +template +auto handle_of(Kernel* kernel) -> cudaKernel_t +{ + cudaKernel_t handle{}; + RAFT_CUDA_TRY(cudaGetKernel(&handle, reinterpret_cast(kernel))); + return handle; +} + RAFT_KERNEL copy_restricted_kernel(int const* __restrict__ in, int* out) { if (threadIdx.x == 0 && blockIdx.x == 0) { *out = *in; } @@ -76,6 +96,22 @@ static_assert(!launchable_as_named, "an lvalue launch_on must static_assert(!launchable_when_moved, "a moved-from launch_on must not be launchable"); +/** Whether a kernel named at run time by `Signature` can be launched with `Args`. */ +template +concept launchable_at_runtime = requires(raft::resources & res, cudaKernel_t handle, Args... args) +{ + raft::launch_kernel(res, dim3{}, dim3{}, raft::kernel_ref{handle}, args...); +}; + +// Unlike the function-pointer overloads, the runtime-kernel one has no exactly-matching sibling to +// defer to, so it must accept arguments that already have the parameter types. +static_assert(launchable_at_runtime, + "an argument that already has the parameter type must be accepted"); +static_assert(launchable_at_runtime, + "an argument that converts to its parameter must be accepted"); +static_assert(!launchable_at_runtime, "too few arguments must not compile"); +static_assert(!launchable_at_runtime, "too many arguments must not compile"); + } // namespace TEST(KernelLaunch, SuccessfulLaunch) @@ -226,4 +262,216 @@ TEST(KernelLaunch, SkipExecutionOnStream) EXPECT_EQ(host_out, 0) << "skip_execution must suppress the launch"; } +TEST(KernelLaunch, CooperativeLaunch) +{ + raft::resources res; + rmm::device_uvector out(1, resource::get_cuda_stream(res)); + RAFT_CUDA_TRY(cudaMemsetAsync(out.data(), 0, sizeof(int), resource::get_cuda_stream(res))); + + raft::launch_kernel({res, 0, {raft::cooperative()}}, 1, 32, write_one_kernel, out.data()); + resource::sync_stream(res); + + int host_out = 0; + RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); + EXPECT_EQ(host_out, 1); +} + +TEST(KernelLaunch, CooperativeLaunchRejectsNonResidentGrid) +{ + raft::resources res; + + // A cooperative launch requires the whole grid to be resident at once, so a grid that launches + // fine on its own must fail once the attribute reaches the driver. Without this asymmetry the + // test could not tell a plumbed attribute from a dropped one. + constexpr int k_huge_grid = 1 << 20; + constexpr int k_block = 1024; + EXPECT_NO_THROW(raft::launch_kernel(res, k_huge_grid, k_block, noop_kernel)); + EXPECT_THROW( + raft::launch_kernel({res, 0, {raft::cooperative()}}, k_huge_grid, k_block, noop_kernel), + raft::cuda_error); + resource::sync_stream(res); +} + +TEST(KernelLaunch, SharedMemoryCarveout) +{ + raft::resources res; + auto stream = resource::get_cuda_stream(res); + rmm::device_uvector out(1, stream); + RAFT_CUDA_TRY(cudaMemsetAsync(out.data(), 0, sizeof(int), stream)); + + raft::launch_kernel( + {res, sizeof(int), {raft::shmem_carveout(100)}}, 1, 32, smem_kernel, out.data()); + resource::sync_stream(res); + + int host_out = 0; + RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); + EXPECT_EQ(host_out, 1); +} + +TEST(KernelLaunch, MultipleAttributes) +{ + raft::resources res; + auto stream = resource::get_cuda_stream(res); + rmm::device_uvector out(1, stream); + RAFT_CUDA_TRY(cudaMemsetAsync(out.data(), 0, sizeof(int), stream)); + + raft::launch_kernel({res, sizeof(int), {raft::cooperative(), raft::shmem_carveout(50)}}, + 1, + 32, + smem_kernel, + out.data()); + resource::sync_stream(res); + + int host_out = 0; + RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); + EXPECT_EQ(host_out, 1); +} + +TEST(KernelLaunch, AttributesInDryRunAreSkipped) +{ + raft::resources res; + auto stream = resource::get_cuda_stream(res); + rmm::device_uvector out(1, stream); + RAFT_CUDA_TRY(cudaMemsetAsync(out.data(), 0, sizeof(int), stream)); + resource::sync_stream(res); + + { + raft::dry_run_resources dry_res(res); + raft::launch_kernel({dry_res, 0, {raft::cooperative()}}, 1, 32, write_one_kernel, out.data()); + resource::sync_stream(dry_res); + } + + int host_out = -1; + RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); + EXPECT_EQ(host_out, 0) << "attributes must not make a dry-run launch execute"; +} + +TEST(KernelLaunch, RuntimeKernelLaunch) +{ + raft::resources res; + rmm::device_uvector out(1, resource::get_cuda_stream(res)); + RAFT_CUDA_TRY(cudaMemsetAsync(out.data(), 0, sizeof(int), resource::get_cuda_stream(res))); + + raft::launch_kernel( + res, 1, 32, raft::kernel_ref{handle_of(write_one_kernel)}, out.data()); + resource::sync_stream(res); + + int host_out = 0; + RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); + EXPECT_EQ(host_out, 1); +} + +TEST(KernelLaunch, RuntimeKernelConvertsArguments) +{ + raft::resources res; + auto stream = resource::get_cuda_stream(res); + rmm::device_uvector out(1, stream); + RAFT_CUDA_TRY(cudaMemsetAsync(out.data(), 0, sizeof(int), stream)); + + // A std::size_t into a std::uint32_t parameter: the conversion is what lets a call site drop the + // casts that a launch taking the addresses of its arguments would need for the sizes to match. + std::size_t const n = 7; + raft::launch_kernel(res, + 1, + 32, + raft::kernel_ref{handle_of(write_count_kernel)}, + out.data(), + n); + resource::sync_stream(res); + + int host_out = 0; + RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); + EXPECT_EQ(host_out, 7); +} + +TEST(KernelLaunch, RuntimeKernelConvertsPointerArgument) +{ + raft::resources res; + auto stream = resource::get_cuda_stream(res); + rmm::device_uvector in(1, stream); + rmm::device_uvector out(1, stream); + int host_in = 1; + RAFT_CUDA_TRY(cudaMemcpyAsync(in.data(), &host_in, sizeof(int), cudaMemcpyHostToDevice, stream)); + + // `int*` into a `int const*` parameter. + raft::launch_kernel(res, + 1, + 32, + raft::kernel_ref{handle_of(copy_one_kernel)}, + in.data(), + out.data()); + resource::sync_stream(res); + + int host_out = 0; + RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); + EXPECT_EQ(host_out, 1); +} + +TEST(KernelLaunch, CooperativeRuntimeKernel) +{ + raft::resources res; + rmm::device_uvector out(1, resource::get_cuda_stream(res)); + RAFT_CUDA_TRY(cudaMemsetAsync(out.data(), 0, sizeof(int), resource::get_cuda_stream(res))); + + // The two features are independent: a runtime kernel takes its attributes from `launch_on` just + // like a statically compiled one. + raft::launch_kernel({res, 0, {raft::cooperative()}}, + 1, + 32, + raft::kernel_ref{handle_of(write_one_kernel)}, + out.data()); + resource::sync_stream(res); + + int host_out = 0; + RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); + EXPECT_EQ(host_out, 1); +} + +TEST(KernelLaunch, RuntimeKernelDryRunSkipsLaunch) +{ + raft::resources res; + auto stream = resource::get_cuda_stream(res); + rmm::device_uvector out(1, stream); + RAFT_CUDA_TRY(cudaMemsetAsync(out.data(), 0, sizeof(int), stream)); + resource::sync_stream(res); + + auto handle = handle_of(write_one_kernel); + { + raft::dry_run_resources dry_res(res); + raft::launch_kernel(dry_res, 1, 32, raft::kernel_ref{handle}, out.data()); + resource::sync_stream(dry_res); + } + + int host_out = -1; + RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); + EXPECT_EQ(host_out, 0) << "a runtime kernel must not run in dry-run mode"; +} + +TEST(KernelLaunch, RuntimeKernelErrorReportsCallSite) +{ + raft::resources res; + auto handle = handle_of(noop_kernel); + + constexpr int k_bad_block = 2048; + std::string caught; + int launch_line = 0; + try { + launch_line = __LINE__ + 1; + raft::launch_kernel(res, 1, k_bad_block, raft::kernel_ref{handle}); + FAIL() << "Expected cuda_error from invalid launch configuration"; + } catch (raft::cuda_error const& e) { + caught = e.what(); + } + + EXPECT_EQ(caught.find("kernel_launch.hpp"), std::string::npos) << caught; + EXPECT_NE(caught.find("kernel_launch.cu"), std::string::npos) << caught; + + std::string re_exp{R"(CUDA error encountered at: file=.*kernel_launch\.cu line=)"}; + re_exp += std::to_string(launch_line); + re_exp += + R"( function=.*RuntimeKernelErrorReportsCallSite.*: call='cudaLaunchKernelExC', Reason=.*)"; + EXPECT_TRUE(std::regex_search(caught, std::regex(re_exp))) + << "message:'" << caught << "'\nexpected regex:'" << re_exp << "'"; +} + } // namespace raft From 09d3ecb00d4da5745c4342d9bd5bda5253551ccc Mon Sep 17 00:00:00 2001 From: achirkin Date: Fri, 28 Aug 2026 17:20:05 +0200 Subject: [PATCH 2/2] Add dry run checks --- cpp/tests/util/kernel_launch.cu | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/cpp/tests/util/kernel_launch.cu b/cpp/tests/util/kernel_launch.cu index ac01ecf8e9..3bca911184 100644 --- a/cpp/tests/util/kernel_launch.cu +++ b/cpp/tests/util/kernel_launch.cu @@ -3,6 +3,8 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "../test_utils.cuh" + #include #include #include @@ -335,15 +337,24 @@ TEST(KernelLaunch, AttributesInDryRunAreSkipped) RAFT_CUDA_TRY(cudaMemsetAsync(out.data(), 0, sizeof(int), stream)); resource::sync_stream(res); + auto launch = [&](raft::resources const& h) { + raft::launch_kernel({h, 0, {raft::cooperative()}}, 1, 32, write_one_kernel, out.data()); + }; + { raft::dry_run_resources dry_res(res); - raft::launch_kernel({dry_res, 0, {raft::cooperative()}}, 1, 32, write_one_kernel, out.data()); + launch(dry_res); resource::sync_stream(dry_res); } int host_out = -1; RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); EXPECT_EQ(host_out, 0) << "attributes must not make a dry-run launch execute"; + + raft::execute_with_dry_run_check(res, launch, raft::alloc_behavior::NO_ALLOCATIONS); + + RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); + EXPECT_EQ(host_out, 1) << "the real pass must execute the kernel"; } TEST(KernelLaunch, RuntimeKernelLaunch) @@ -436,15 +447,24 @@ TEST(KernelLaunch, RuntimeKernelDryRunSkipsLaunch) resource::sync_stream(res); auto handle = handle_of(write_one_kernel); + auto launch = [&](raft::resources const& h) { + raft::launch_kernel(h, 1, 32, raft::kernel_ref{handle}, out.data()); + }; + { raft::dry_run_resources dry_res(res); - raft::launch_kernel(dry_res, 1, 32, raft::kernel_ref{handle}, out.data()); + launch(dry_res); resource::sync_stream(dry_res); } int host_out = -1; RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); EXPECT_EQ(host_out, 0) << "a runtime kernel must not run in dry-run mode"; + + raft::execute_with_dry_run_check(res, launch, raft::alloc_behavior::NO_ALLOCATIONS); + + RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); + EXPECT_EQ(host_out, 1) << "the real pass must execute the kernel"; } TEST(KernelLaunch, RuntimeKernelErrorReportsCallSite)