[Backport branch/3.2.x] [libcu++] Correctly handle extended lambda in cuda::launch

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<transformed_device_argument_t<_Args>>...>)
+                                            ::cuda::std::decay_t<transformed_device_argument_t<_Args>>...>
+#    if _CCCL_CUDA_COMPILER(NVCC)
+                && !__nv_is_extended_device_lambda_closure_type(_Kernel)
+#    endif
+  )
   {
     auto __launcher =
       __kernel_launcher<decltype(__combined), _Kernel, ::cuda::std::decay_t<transformed_device_argument_t<_Args>>...>;

libcudacxx/test/libcudacxx/cuda/ccclrt/common/host_device.cuh

@@ -13,7 +13,7 @@
 
 #include <cuda/hierarchy>
 
-#include "utility.cuh"
+#include "testing.cuh"
 
 template <typename Dims, typename Lambda>
 void __global__ lambda_launcher(const Dims dims, const Lambda lambda)

libcudacxx/test/libcudacxx/cuda/ccclrt/common/testing.cuh

@@ -20,31 +20,11 @@
 #include <exception> // IWYU pragma: keep
 #include <sstream>
 
+#include "utility.cuh"
 #include <c2h/catch2_test_helper.h>
 
 #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)

libcudacxx/test/libcudacxx/cuda/ccclrt/common/utility.cuh

@@ -22,7 +22,26 @@
 
 #include <new> // 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

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 <cuda/devices>
+#include <cuda/launch>
+#include <cuda/stream>
+
+#include "../common/utility.cuh"
+
+__host__ void test_extended_lambda()
+{
+  cuda::stream stream{cuda::devices[0]};
+  test::pinned<int> 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;
+}

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 <typename DefaultConfig>
 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 <typename Config>
+  __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<decltype(kernel.default_config())>);
     static_assert(cuda::__kernel_has_default_config<decltype(kernel)>);
 
-    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)

libcudacxx/test/libcudacxx/cuda/ccclrt/stream/stream_smoke.cu

@@ -13,7 +13,7 @@
 #include <cuda/std/utility>
 #include <cuda/stream>
 
-#include <utility.cuh>
+#include <testing.cuh>
 
 C2H_CCCLRT_TEST("Can create a stream and launch work into it", "[stream]")
 {

libcudacxx/test/libcudacxx/cuda/memory_resource/resources/common_tests.cuh

@@ -10,6 +10,7 @@
 
 #pragma once
 
+#include <testing.cuh>
 #include <utility.cuh>
 
 template <typename ResourceType>