Test proclaim_copyable_arguments for lambdas

NVIDIA · Nov 23, 2024 · 60ea954 · 60ea954
1 parent 83d180f
commit 60ea954
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 2 deletions.
diff --git a/libcudacxx/include/cuda/__functional/address_stability.h b/libcudacxx/include/cuda/__functional/address_stability.h
@@ -48,6 +48,13 @@ _CCCL_INLINE_VAR constexpr bool proclaims_copyable_arguments_v = proclaims_copya
 template <typename F>
 struct __callable_permitting_copied_arguments : F
 {
+#if _CCCL_STD_VER < 2017
+  template <typename G>
+  _LIBCUDACXX_HIDE_FROM_ABI constexpr __callable_permitting_copied_arguments(G&& g)
+      : F(::cuda::std::forward<G>(g))
+  {}
+#endif
+
   using F::operator();
 };
 
@@ -61,9 +68,9 @@ struct proclaims_copyable_arguments<__callable_permitting_copied_arguments<F>> :
 //! @see proclaims_copyable_arguments
 template <typename F>
 _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr auto
-proclaim_copyable_arguments(F f) -> __callable_permitting_copied_arguments<F>
+proclaim_copyable_arguments(F&& f) -> __callable_permitting_copied_arguments<::cuda::std::decay_t<F>>
 {
-  return __callable_permitting_copied_arguments<F>{_CUDA_VSTD::move(f)};
+  return {::cuda::std::forward<F>(f)};
 }
 
 // Specializations for libcu++ function objects are provided here to not pull this include into `<cuda/std/...>` headers

diff --git a/thrust/testing/address_stability.cmake b/thrust/testing/address_stability.cmake
@@ -0,0 +1 @@
+target_compile_options(${test_target} PRIVATE $<$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>:--extended-lambda>)
diff --git a/thrust/testing/address_stability.cu b/thrust/testing/address_stability.cu
@@ -83,3 +83,40 @@ void TestAddressStabilityUserDefinedFunctionObject()
   static_assert(proclaims_copyable_arguments<decltype(proclaim_copyable_arguments(my_plus<const int&&>{}))>::value, "");
 }
 DECLARE_UNITTEST(TestAddressStabilityUserDefinedFunctionObject);
+
+void TestAddressStabilityLambda()
+{
+  using ::cuda::proclaim_copyable_arguments;
+  using ::cuda::proclaims_copyable_arguments;
+
+  {
+    auto l = [](const int& i) {
+      return i + 2;
+    };
+    static_assert(!proclaims_copyable_arguments<decltype(l)>::value, "");
+    auto pr_l = proclaim_copyable_arguments(l);
+    ASSERT_EQUAL(pr_l(3), 5);
+    static_assert(proclaims_copyable_arguments<decltype(pr_l)>::value, "");
+  }
+
+  {
+    auto l = [] _CCCL_DEVICE(const int& i) {
+      return i + 2;
+    };
+    static_assert(!proclaims_copyable_arguments<decltype(l)>::value, "");
+    auto pr_device_l = proclaim_copyable_arguments(l);
+    (void) &pr_device_l;
+    static_assert(proclaims_copyable_arguments<decltype(pr_device_l)>::value, "");
+  }
+
+  {
+    auto l = [] _CCCL_HOST_DEVICE(const int& i) {
+      return i + 2;
+    };
+    static_assert(!proclaims_copyable_arguments<decltype(l)>::value, "");
+    auto pr_l = proclaim_copyable_arguments(l);
+    ASSERT_EQUAL(pr_l(3), 5);
+    static_assert(proclaims_copyable_arguments<decltype(pr_l)>::value, "");
+  }
+}
+DECLARE_UNITTEST(TestAddressStabilityLambda);