Move address stability into libcu++

libcudacxx/include/cuda/std/__type_traits/address_stability.h

@@ -0,0 +1,64 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, 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) 2024 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef _LIBCUDACXX___TYPE_TRAITS_ADDRESS_STABILITY_H
+#define _LIBCUDACXX___TYPE_TRAITS_ADDRESS_STABILITY_H
+
+#include <cuda/std/detail/__config>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cuda/std/__type_traits/integral_constant.h>
+#include <cuda/std/__type_traits/void_t.h>
+#include <cuda/std/__utility/move.h>
+
+_LIBCUDACXX_BEGIN_NAMESPACE_STD
+
+// need a separate implementation trait because we SFINAE with a type parameter before the variadic pack
+template <typename F, typename SFINAE, typename... Args>
+struct __can_copy_arguments_impl : false_type
+{};
+
+template <typename F, typename... Args>
+struct __can_copy_arguments_impl<F, void_t<decltype(F::can_copy_arguments)>, Args...>
+{
+  static constexpr bool value = F::can_copy_arguments;
+};
+
+//! Trait telling whether a function object relies on the memory address of its arguments when called with the given set
+//! of types. The nested value is true when the addresses of the arguments do not matter and arguments can be provided
+//! from arbitrary copies of the respective sources.
+template <typename F, typename... Args>
+using can_copy_arguments = __can_copy_arguments_impl<F, void, Args...>;
+
+//! Wrapper for a callable to mark it as allowing copied arguments
+template <typename F>
+struct copied_arguments_allowing_wrapper : F
+{
+  using F::operator();
+  static constexpr bool can_copy_arguments = true;
+};
+
+//! Creates a new function object from an existing one, allowing its arguments to be copies of whatever source they come
+//! from. This implies that the addresses of the arguments are irrelevant to the function object.
+template <typename F>
+_CCCL_HOST_DEVICE constexpr auto allow_copied_arguments(F f) -> copied_arguments_allowing_wrapper<F>
+{
+  return copied_arguments_allowing_wrapper<F>{_CUDA_VSTD::move(f)};
+}
+
+_LIBCUDACXX_END_NAMESPACE_STD
+
+#endif // _LIBCUDACXX___TYPE_TRAITS_ADDRESS_STABILITY_H

libcudacxx/include/cuda/std/detail/libcxx/include/type_traits

@@ -435,6 +435,7 @@ namespace std
 #include <cuda/std/__type_traits/add_pointer.h>
 #include <cuda/std/__type_traits/add_rvalue_reference.h>
 #include <cuda/std/__type_traits/add_volatile.h>
+#include <cuda/std/__type_traits/address_stability.h>
 #include <cuda/std/__type_traits/aligned_storage.h>
 #include <cuda/std/__type_traits/aligned_union.h>
 #include <cuda/std/__type_traits/alignment_of.h>

thrust/benchmarks/bench/transform/basic.cu

@@ -25,7 +25,6 @@
  *
  ******************************************************************************/
 
-#include <thrust/address_stability.h>
 #include <thrust/copy.h>
 #include <thrust/count.h>
 #include <thrust/device_vector.h>
@@ -35,6 +34,7 @@
 #include <thrust/zip_function.h>
 
 #include <cuda/functional>
+#include <cuda/std/__type_traits/address_stability.h>
 
 #include <nvbench_helper.cuh>
 
@@ -125,7 +125,7 @@ static void mul(nvbench::state& state, nvbench::type_list<T>)
   state.exec(nvbench::exec_tag::no_batch | nvbench::exec_tag::sync, [&](nvbench::launch&) {
     const T scalar = startScalar;
     thrust::transform(
-      c.begin(), c.end(), b.begin(), thrust::allow_argument_copies([=] __device__ __host__(const T& ci) {
+      c.begin(), c.end(), b.begin(), cuda::std::allow_copied_arguments([=] __device__ __host__(const T& ci) {
         return ci * scalar;
       }));
   });
@@ -154,7 +154,7 @@ static void add(nvbench::state& state, nvbench::type_list<T>)
       a.end(),
       b.begin(),
       c.begin(),
-      thrust::allow_argument_copies([] _CCCL_DEVICE(const T& ai, const T& bi) -> T {
+      cuda::std::allow_copied_arguments([] _CCCL_DEVICE(const T& ai, const T& bi) -> T {
         return ai + bi;
       }));
   });
@@ -184,7 +184,7 @@ static void triad(nvbench::state& state, nvbench::type_list<T>)
       b.end(),
       c.begin(),
       a.begin(),
-      thrust::allow_argument_copies([=] _CCCL_DEVICE(const T& bi, const T& ci) {
+      cuda::std::allow_copied_arguments([=] _CCCL_DEVICE(const T& bi, const T& ci) {
         return bi + scalar * ci;
       }));
   });
@@ -214,9 +214,10 @@ static void nstream(nvbench::state& state, nvbench::type_list<T>)
       thrust::make_zip_iterator(a.end(), b.end(), c.end()),
       a.begin(),
 
-      thrust::make_zip_function(thrust::allow_argument_copies([=] _CCCL_DEVICE(const T& ai, const T& bi, const T& ci) {
-        return ai + bi + scalar * ci;
-      })));
+      thrust::make_zip_function(
+        cuda::std::allow_copied_arguments([=] _CCCL_DEVICE(const T& ai, const T& bi, const T& ci) {
+          return ai + bi + scalar * ci;
+        })));
   });
 }
 

thrust/testing/address_stability.cu

@@ -1,7 +1,9 @@
-#include <thrust/address_stability.h>
+#include <cuda/std/__type_traits/address_stability.h>
 
 #include <unittest/unittest.h>
 
+// TODO(bgruber): move this test into libcu++
+
 struct MyPlus
 {
   _CCCL_HOST_DEVICE auto operator()(int a, int b) const -> int
@@ -12,12 +14,13 @@ struct MyPlus
 
 void TestAddressStability()
 {
-  using thrust::can_copy_arguments;
+  using ::cuda::std::allow_copied_arguments;
+  using ::cuda::std::can_copy_arguments;
 
   static_assert(!can_copy_arguments<thrust::plus<int>, int, int>::value, "");
-  static_assert(can_copy_arguments<decltype(thrust::allow_copied_arguments(thrust::plus<int>{})), int, int>::value, "");
+  static_assert(can_copy_arguments<decltype(allow_copied_arguments(thrust::plus<int>{})), int, int>::value, "");
 
   static_assert(!can_copy_arguments<MyPlus, int, int>::value, "");
-  static_assert(can_copy_arguments<decltype(thrust::allow_copied_arguments(MyPlus{})), int, int>::value, "");
+  static_assert(can_copy_arguments<decltype(allow_copied_arguments(MyPlus{})), int, int>::value, "");
 }
 DECLARE_UNITTEST(TestAddressStability);

thrust/thrust/system/cuda/detail/transform.h

@@ -41,7 +41,6 @@
 
 #  include <cub/device/device_transform.cuh>
 
-#  include <thrust/address_stability.h>
 #  include <thrust/detail/temporary_array.h>
 #  include <thrust/distance.h>
 #  include <thrust/iterator/zip_iterator.h>
@@ -50,6 +49,8 @@
 #  include <thrust/system/cuda/detail/util.h>
 #  include <thrust/zip_function.h>
 
+#  include <cuda/std/__type_traits/address_stability.h>
+
 #  include <cstdint>
 
 THRUST_NAMESPACE_BEGIN
@@ -252,8 +253,9 @@ OutputIt THRUST_FUNCTION cub_transform_many(
   }
 
   // TODO(bgruber): iterator_reference_t or iterator_value_type?
-  constexpr auto requires_stable_address = !can_copy_arguments<TransformOp, iterator_reference_t<InputIts>...>::value;
-  using dispatch32_t                     = cub::detail::transform::
+  constexpr auto requires_stable_address =
+    !::cuda::std::can_copy_arguments<TransformOp, iterator_reference_t<InputIts>...>::value;
+  using dispatch32_t = cub::detail::transform::
     dispatch_t<requires_stable_address, std::int32_t, ::cuda::std::tuple<InputIts...>, OutputIt, TransformOp>;
   using dispatch64_t = cub::detail::transform::
     dispatch_t<requires_stable_address, std::int64_t, ::cuda::std::tuple<InputIts...>, OutputIt, TransformOp>;