From 457f0c9475613d42a3c2333cf6be88d021943935 Mon Sep 17 00:00:00 2001
From: Michael Schellenberger Costa <miscco@nvidia.com>
Date: Tue, 25 Nov 2025 12:56:32 +0100
Subject: [PATCH] Ensure that we properly warn about device lambdas that need
 to query the return type (#6765)

* Ensure that we properly check if a device lambda is wrapped

We were checking `__nv_is_extended_device_lambda_closure_type(_Fp)` and so on, but that does not work if the device lambda is cvref-qualified

Fix this by checking `remove_cvref_t`

Fixes #6762

* Ensure that `thrust::transform_iterator` properly checks for device lambdas

Fixes [BUG]: `thrust::transform_iterator` no longer emits diagnostic when used with __device__ lambda
Fixes #6763

(cherry picked from commit 8588acb2b3fb252bcc8be959469021dd710df006)
---
 libcudacxx/include/cuda/std/__functional/invoke.h | 8 +++++---
 thrust/thrust/iterator/transform_iterator.h       | 8 ++++----
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/libcudacxx/include/cuda/std/__functional/invoke.h b/libcudacxx/include/cuda/std/__functional/invoke.h
index 08d66bcf6b27..0fc242a2babd 100644
--- a/libcudacxx/include/cuda/std/__functional/invoke.h
+++ b/libcudacxx/include/cuda/std/__functional/invoke.h
@@ -33,6 +33,7 @@
 #include <cuda/std/__type_traits/is_same.h>
 #include <cuda/std/__type_traits/is_void.h>
 #include <cuda/std/__type_traits/nat.h>
+#include <cuda/std/__type_traits/remove_cvref.h>
 #include <cuda/std/__utility/declval.h>
 #include <cuda/std/__utility/forward.h>
 
@@ -192,15 +193,16 @@ struct _CCCL_TYPE_VISIBILITY_DEFAULT invoke_result //
 {
 #if _CCCL_CUDA_COMPILER(NVCC) && defined(__CUDACC_EXTENDED_LAMBDA__) && !_CCCL_DEVICE_COMPILATION()
 #  if _CCCL_CUDACC_BELOW(12, 3)
-  static_assert(!__nv_is_extended_device_lambda_closure_type(_Fp),
+  static_assert(!__nv_is_extended_device_lambda_closure_type(remove_cvref_t<_Fp>),
                 "Attempt to use an extended __device__ lambda in a context "
                 "that requires querying its return type in host code. Use a "
                 "named function object, an extended __host__ __device__ lambda, or "
                 "cuda::proclaim_return_type instead.");
 #  else // ^^^ _CCCL_CUDACC_BELOW(12, 3) ^^^ / vvv _CCCL_CUDACC_AT_LEAST(12, 3) vvv
   static_assert(
-    !__nv_is_extended_device_lambda_closure_type(_Fp) || __nv_is_extended_host_device_lambda_closure_type(_Fp)
-      || __nv_is_extended_device_lambda_with_preserved_return_type(_Fp),
+    !__nv_is_extended_device_lambda_closure_type(remove_cvref_t<_Fp>)
+      || __nv_is_extended_host_device_lambda_closure_type(remove_cvref_t<_Fp>)
+      || __nv_is_extended_device_lambda_with_preserved_return_type(remove_cvref_t<_Fp>),
     "Attempt to use an extended __device__ lambda in a context "
     "that requires querying its return type in host code. Use a "
     "named function object, an extended __host__ __device__ lambda, "
diff --git a/thrust/thrust/iterator/transform_iterator.h b/thrust/thrust/iterator/transform_iterator.h
index 164fa1cf5199..1577a22fb2ea 100644
--- a/thrust/thrust/iterator/transform_iterator.h
+++ b/thrust/thrust/iterator/transform_iterator.h
@@ -48,12 +48,12 @@
 #include <thrust/iterator/iterator_traits.h>
 
 #include <cuda/std/__functional/identity.h>
+#include <cuda/std/__functional/invoke.h>
 #include <cuda/std/__memory/construct_at.h>
 #include <cuda/std/__type_traits/is_copy_assignable.h>
 #include <cuda/std/__type_traits/is_copy_constructible.h>
 #include <cuda/std/__type_traits/remove_cvref.h>
 #include <cuda/std/__type_traits/remove_reference.h>
-#include <cuda/std/__utility/declval.h>
 
 THRUST_NAMESPACE_BEGIN
 
@@ -66,7 +66,7 @@ template <class UnaryFunc, class Iterator>
 struct transform_iterator_reference
 {
   // by default, dereferencing the iterator yields the same as the function.
-  using type = decltype(::cuda::std::declval<UnaryFunc>()(::cuda::std::declval<it_value_t<Iterator>>()));
+  using type = ::cuda::std::invoke_result_t<UnaryFunc&, it_value_t<Iterator>>;
 };
 
 // for certain function objects, we need to tweak the reference type. Notably, identity functions must decay to values.
@@ -79,8 +79,8 @@ struct transform_iterator_reference<::cuda::std::identity, Iterator>
 template <typename Eval, class Iterator>
 struct transform_iterator_reference<functional::actor<Eval>, Iterator>
 {
-  using type = ::cuda::std::remove_reference_t<decltype(::cuda::std::declval<functional::actor<Eval>>()(
-    ::cuda::std::declval<it_value_t<Iterator>>()))>;
+  using type =
+    ::cuda::std::remove_reference_t<::cuda::std::invoke_result_t<functional::actor<Eval>&, it_value_t<Iterator>>>;
 };
 
 // Type function to compute the iterator_adaptor instantiation to be used for transform_iterator