From c659362daef40159071ec9b02c7c8f2218cb61ca Mon Sep 17 00:00:00 2001
From: Michael Schellenberger Costa <miscco@nvidia.com>
Date: Wed, 12 Jul 2023 17:15:42 +0200
Subject: [PATCH] Enable use of `cudaMemcpyAsync` for `thrust::copy`

In case of contigous ranges of trivially relocatable types we can directly utilize `cudaMemcpyAsync` instead of going through transform.

Fixes #210
---
 .../detail/internal/copy_device_to_device.h   | 48 +++++++++++++++++--
 1 file changed, 45 insertions(+), 3 deletions(-)

diff --git a/thrust/thrust/system/cuda/detail/internal/copy_device_to_device.h b/thrust/thrust/system/cuda/detail/internal/copy_device_to_device.h
index 69c4e20dfdf..a0d9d492dac 100644
--- a/thrust/thrust/system/cuda/detail/internal/copy_device_to_device.h
+++ b/thrust/thrust/system/cuda/detail/internal/copy_device_to_device.h
@@ -33,12 +33,38 @@
 #include <thrust/system/cuda/config.h>
 #include <thrust/system/cuda/detail/execution_policy.h>
 #include <thrust/system/cuda/detail/transform.h>
+#include <thrust/system/cuda/detail/util.h>
+#include <thrust/distance.h>
 #include <thrust/functional.h>
+#include <thrust/type_traits/is_trivially_relocatable.h>
 
 THRUST_NAMESPACE_BEGIN
 namespace cuda_cub {
 
 namespace __copy {
+  template <class Derived,
+            class InputIt,
+            class OutputIt>
+  OutputIt THRUST_RUNTIME_FUNCTION
+  device_to_device(execution_policy<Derived>& policy,
+                   InputIt                    first,
+                   InputIt                    last,
+                   OutputIt                   result,
+                   thrust::detail::true_type)
+  {
+    typedef typename thrust::iterator_traits<InputIt>::value_type InputTy;
+    const auto n = thrust::distance(first, last);
+    if (n > 0) {
+      cudaError status;
+      status = trivial_copy_device_to_device(policy,
+                                             reinterpret_cast<InputTy*>(thrust::raw_pointer_cast(&*result)),
+                                             reinterpret_cast<InputTy const*>(thrust::raw_pointer_cast(&*first)),
+                                             n);
+      cuda_cub::throw_on_error(status, "__copy:: D->D: failed");
+    }
+
+    return result + n;
+  }
 
   template <class Derived,
             class InputIt,
@@ -47,16 +73,32 @@ namespace __copy {
   device_to_device(execution_policy<Derived>& policy,
                    InputIt                    first,
                    InputIt                    last,
-                   OutputIt                   result)
+                   OutputIt                   result,
+                   thrust::detail::false_type)
   {
     typedef typename thrust::iterator_traits<InputIt>::value_type InputTy;
     return cuda_cub::transform(policy,
+                              first,
+                              last,
+                              result,
+                              thrust::identity<InputTy>());
+  }
+
+  template <class Derived,
+            class InputIt,
+            class OutputIt>
+  OutputIt THRUST_RUNTIME_FUNCTION
+  device_to_device(execution_policy<Derived>& policy,
+                   InputIt                    first,
+                   InputIt                    last,
+                   OutputIt                   result)
+  {
+    return device_to_device(policy,
                             first,
                             last,
                             result,
-                            thrust::identity<InputTy>());
+                            typename is_indirectly_trivially_relocatable_to<InputIt, OutputIt>::type());
   }
-
 }    // namespace __copy
 
 }    // namespace cuda_cub