From 06e7bf1dd8246062a2b30f39fc739e1e0e5e000b Mon Sep 17 00:00:00 2001 From: pciolkosz Date: Fri, 19 Dec 2025 01:34:23 -0800 Subject: [PATCH] [libcu++] Correctly handle extended lambda in cuda::launch (#6987) * Don't set current device in CUDA 13 and handle extended lambda * Add extended lambda test * Compiler workarounds * Waive extended lambda test on NVRTC * Apply suggestion from @davebayer --------- Co-authored-by: David Bayer <48736217+davebayer@users.noreply.github.com> (cherry picked from commit 8928e6fbe917d55662fa058a8d70b26f9a0c0ec6) --- libcudacxx/include/cuda/__launch/launch.h | 6 +- .../cuda/ccclrt/common/host_device.cuh | 2 +- .../libcudacxx/cuda/ccclrt/common/testing.cuh | 22 +----- .../libcudacxx/cuda/ccclrt/common/utility.cuh | 23 +++++- .../ccclrt/launch/extended_lambda.pass.cpp | 46 +++++++++++ .../cuda/ccclrt/launch/launch_smoke.cu | 77 +++++++++---------- .../cuda/ccclrt/stream/stream_smoke.cu | 2 +- .../resources/common_tests.cuh | 1 + 8 files changed, 111 insertions(+), 68 deletions(-) create mode 100644 libcudacxx/test/libcudacxx/cuda/ccclrt/launch/extended_lambda.pass.cpp diff --git a/libcudacxx/include/cuda/__launch/launch.h b/libcudacxx/include/cuda/__launch/launch.h index edf3e3e1b8a1..9c7303b3d8e5 100644 --- a/libcudacxx/include/cuda/__launch/launch.h +++ b/libcudacxx/include/cuda/__launch/launch.h @@ -191,7 +191,11 @@ _CCCL_HOST_API auto launch(_Submitter&& __submitter, auto __combined = __conf.combine_with_default(__kernel); if constexpr (::cuda::std::is_invocable_v<_Kernel, kernel_config<_Dimensions, _Config...>, - ::cuda::std::decay_t>...>) + ::cuda::std::decay_t>...> +# if _CCCL_CUDA_COMPILER(NVCC) + && !__nv_is_extended_device_lambda_closure_type(_Kernel) +# endif + ) { auto __launcher = __kernel_launcher>...>; diff --git a/libcudacxx/test/libcudacxx/cuda/ccclrt/common/host_device.cuh b/libcudacxx/test/libcudacxx/cuda/ccclrt/common/host_device.cuh index 424cd7fe0248..22f8d599f4da 100644 --- a/libcudacxx/test/libcudacxx/cuda/ccclrt/common/host_device.cuh +++ b/libcudacxx/test/libcudacxx/cuda/ccclrt/common/host_device.cuh @@ -13,7 +13,7 @@ #include -#include "utility.cuh" +#include "testing.cuh" template void __global__ lambda_launcher(const Dims dims, const Lambda lambda) diff --git a/libcudacxx/test/libcudacxx/cuda/ccclrt/common/testing.cuh b/libcudacxx/test/libcudacxx/cuda/ccclrt/common/testing.cuh index fdec338ad36e..8cf7bc9b7f91 100644 --- a/libcudacxx/test/libcudacxx/cuda/ccclrt/common/testing.cuh +++ b/libcudacxx/test/libcudacxx/cuda/ccclrt/common/testing.cuh @@ -20,31 +20,11 @@ #include // IWYU pragma: keep #include +#include "utility.cuh" #include #define CUDART(call) REQUIRE((call) == cudaSuccess) -__device__ inline void ccclrt_require_impl( - bool condition, const char* condition_text, const char* filename, unsigned int linenum, const char* funcname) -{ - if (!condition) - { - // TODO do warp aggregate prints for easier readability? - printf("%s:%u: %s: block: [%d,%d,%d], thread: [%d,%d,%d] Condition `%s` failed.\n", - filename, - linenum, - funcname, - blockIdx.x, - blockIdx.y, - blockIdx.z, - threadIdx.x, - threadIdx.y, - threadIdx.z, - condition_text); - __trap(); - } -} - // There is a problem with clang-cuda and nv/target, but we don't need the device side macros yet, // disable them for now #if _CCCL_CUDA_COMPILER(CLANG) diff --git a/libcudacxx/test/libcudacxx/cuda/ccclrt/common/utility.cuh b/libcudacxx/test/libcudacxx/cuda/ccclrt/common/utility.cuh index d7c7e477042f..14bb07ad6e9b 100644 --- a/libcudacxx/test/libcudacxx/cuda/ccclrt/common/utility.cuh +++ b/libcudacxx/test/libcudacxx/cuda/ccclrt/common/utility.cuh @@ -22,7 +22,26 @@ #include // IWYU pragma: keep (needed for placement new) -#include "testing.cuh" +__device__ inline void ccclrt_require_impl( + bool condition, const char* condition_text, const char* filename, unsigned int linenum, const char* funcname) +{ + if (!condition) + { + // TODO do warp aggregate prints for easier readability? + printf("%s:%u: %s: block: [%d,%d,%d], thread: [%d,%d,%d] Condition `%s` failed.\n", + filename, + linenum, + funcname, + blockIdx.x, + blockIdx.y, + blockIdx.z, + threadIdx.x, + threadIdx.y, + threadIdx.z, + condition_text); + __trap(); + } +} namespace { @@ -157,7 +176,7 @@ void launch_kernel_single_thread(cuda::stream_ref stream, Fn fn, Args... args) { cuda::__ensure_current_context guard(stream); kernel_launcher<<<1, 1, 0, stream.get()>>>(fn, args...); - CUDART(cudaGetLastError()); + assert(cudaGetLastError() == cudaSuccess); } } // namespace test } // namespace diff --git a/libcudacxx/test/libcudacxx/cuda/ccclrt/launch/extended_lambda.pass.cpp b/libcudacxx/test/libcudacxx/cuda/ccclrt/launch/extended_lambda.pass.cpp new file mode 100644 index 000000000000..755cd6d81bf7 --- /dev/null +++ b/libcudacxx/test/libcudacxx/cuda/ccclrt/launch/extended_lambda.pass.cpp @@ -0,0 +1,46 @@ +//===----------------------------------------------------------------------===// +// +// Part of libcu++, the C++ Standard Library for your entire system, +// under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// + +// ADDITIONAL_COMPILE_FLAGS: --extended-lambda +// UNSUPPORTED: nvrtc + +#include +#include +#include + +#include "../common/utility.cuh" + +__host__ void test_extended_lambda() +{ + cuda::stream stream{cuda::devices[0]}; + test::pinned i(0); + auto config = cuda::block_dims<32>() & cuda::grid_dims<1>(); + auto assign_42_lambda = [] __device__(int* pi) { + *pi = 42; + }; + cuda::launch(stream, config, assign_42_lambda, i.get()); + stream.sync(); + assert(*i == 42); + + auto assign_1337_lambda = [] __device__(auto config, int* pi) { + static_assert(config.dims.count(cuda::gpu_thread, cuda::block) == 32); + static_assert(config.dims.count(cuda::block) == 1); + *pi = 1337; + }; + cuda::launch(stream, config, assign_1337_lambda, config, i.get()); + stream.sync(); + assert(*i == 1337); +} + +int main(int, char**) +{ + NV_IF_TARGET(NV_IS_HOST, test_extended_lambda();) + return 0; +} diff --git a/libcudacxx/test/libcudacxx/cuda/ccclrt/launch/launch_smoke.cu b/libcudacxx/test/libcudacxx/cuda/ccclrt/launch/launch_smoke.cu index 232df2f0a96a..fc2e4a9726fa 100644 --- a/libcudacxx/test/libcudacxx/cuda/ccclrt/launch/launch_smoke.cu +++ b/libcudacxx/test/libcudacxx/cuda/ccclrt/launch/launch_smoke.cu @@ -215,20 +215,6 @@ void launch_smoke_test(cudaStream_t dst) } } - /* Comment out for now until I figure how to enable extended lambda for only this file - // Lambda - { - cuda::launch(dst, cuda::block_dims<256>() & cuda::grid_dims(1), - [] __device__(auto config) { - if (config.dims.rank(cuda::gpu_thread, cuda::block) == 0) { - printf("Hello from the GPU\n"); - kernel_run_proof = true; - } - }); - check_kernel_run(dst); - } - */ - // Dynamic shared memory option { auto config = cuda::block_dims<32>() & cuda::grid_dims<1>(); @@ -264,20 +250,25 @@ void launch_smoke_test(cudaStream_t dst) } } -C2H_TEST("Launch smoke stream", "[launch]") +C2H_CCCLRT_TEST("Launch smoke stream", "[launch]") { // Use raw stream to make sure it can be implicitly converted on call to // launch cudaStream_t stream; - CUDART(cudaStreamCreate(&stream)); + { + ::cuda::__ensure_current_context guard(cuda::device_ref{0}); + CUDART(cudaStreamCreate(&stream)); + } launch_smoke_test(stream); - CUDART(cudaStreamSynchronize(stream)); - CUDART(cudaStreamDestroy(stream)); + { + ::cuda::__ensure_current_context guard(cuda::device_ref{0}); + CUDART(cudaStreamSynchronize(stream)); + CUDART(cudaStreamDestroy(stream)); + } } -#endif // !_CCCL_CUDA_COMPILER(CLANG) template struct kernel_with_default_config @@ -300,42 +291,44 @@ struct kernel_with_default_config } }; -/* Comment out for now until I figure how to enable extended lambda for only this file -void test_default_config() { - cuda::stream stream{cuda::device_ref{0}}; - auto grid = cuda::grid_dims(4); - auto block = cuda::block_dims<256>; - - auto verify_lambda = [] __device__(auto config) { +struct verify_callable +{ + template + __device__ void operator()(Config config) + { static_assert(config.dims.count(cuda::gpu_thread, cuda::block) == 256); CCCLRT_REQUIRE(config.dims.count(cuda::block) == 4); cooperative_groups::this_grid().sync(); - }; + } +}; + +C2H_CCCLRT_TEST("Launch with default config", "") +{ + cuda::stream stream{cuda::device_ref{0}}; + auto grid = cuda::grid_dims(4); + auto block = cuda::block_dims<256>; - SECTION("Combine with empty") { - kernel_with_default_config kernel{ - cuda::make_config(block, grid, cuda::cooperative_launch())}; + SECTION("Combine with empty") + { + kernel_with_default_config kernel{cuda::make_config(block, grid, cuda::cooperative_launch())}; static_assert(cuda::__is_kernel_config); static_assert(cuda::__kernel_has_default_config); - cuda::launch(stream, cuda::make_config(), kernel, verify_lambda); + cuda::launch(stream, cuda::make_config(), kernel, verify_callable{}); stream.sync(); } - SECTION("Combine with no overlap") { + SECTION("Combine with no overlap") + { kernel_with_default_config kernel{cuda::make_config(block)}; - cuda::launch(stream, cuda::make_config(grid, cuda::cooperative_launch()), - kernel, verify_lambda); + cuda::launch(stream, cuda::make_config(grid, cuda::cooperative_launch()), kernel, verify_callable{}); stream.sync(); } - SECTION("Combine with overlap") { - kernel_with_default_config kernel{ - cuda::make_config(cuda::block_dims<1>, cuda::cooperative_launch())}; - cuda::launch(stream, - cuda::make_config(block, grid, cuda::cooperative_launch()), - kernel, verify_lambda); + SECTION("Combine with overlap") + { + kernel_with_default_config kernel{cuda::make_config(cuda::block_dims<1>(), cuda::cooperative_launch())}; + cuda::launch(stream, cuda::make_config(block, grid, cuda::cooperative_launch()), kernel, verify_callable{}); stream.sync(); } } -C2H_TEST("Launch with default config", "") { test_default_config(); } -*/ +#endif // !_CCCL_CUDA_COMPILER(CLANG) diff --git a/libcudacxx/test/libcudacxx/cuda/ccclrt/stream/stream_smoke.cu b/libcudacxx/test/libcudacxx/cuda/ccclrt/stream/stream_smoke.cu index c42bf45b6439..04a361749dc0 100644 --- a/libcudacxx/test/libcudacxx/cuda/ccclrt/stream/stream_smoke.cu +++ b/libcudacxx/test/libcudacxx/cuda/ccclrt/stream/stream_smoke.cu @@ -13,7 +13,7 @@ #include #include -#include +#include C2H_CCCLRT_TEST("Can create a stream and launch work into it", "[stream]") { diff --git a/libcudacxx/test/libcudacxx/cuda/memory_resource/resources/common_tests.cuh b/libcudacxx/test/libcudacxx/cuda/memory_resource/resources/common_tests.cuh index b6617bf30881..ed2eaf022f9e 100644 --- a/libcudacxx/test/libcudacxx/cuda/memory_resource/resources/common_tests.cuh +++ b/libcudacxx/test/libcudacxx/cuda/memory_resource/resources/common_tests.cuh @@ -10,6 +10,7 @@ #pragma once +#include #include template