arrayfire · umar456 · Mar 25, 2020 · Mar 24, 2020 · Mar 24, 2020
diff --git a/src/backend/cuda/CMakeLists.txt b/src/backend/cuda/CMakeLists.txt
@@ -377,6 +377,9 @@ cuda_add_library(afcuda
     Array.hpp
     LookupTable1D.hpp
     Param.hpp
+    ThrustAllocator.cuh
+    ThrustArrayFirePolicy.hpp
+    ThrustArrayFirePolicy.cpp
     anisotropic_diffusion.hpp
     approx.hpp
     arith.hpp
@@ -411,7 +414,7 @@ cuda_add_library(afcuda
     device_manager.cpp
     device_manager.hpp
     debug_cuda.hpp
-    debug_thrust.hpp
+    thrust_utils.hpp
     diagonal.cpp
     diagonal.hpp
     diff.cpp

diff --git a/src/backend/cuda/ThrustArrayFirePolicy.cpp b/src/backend/cuda/ThrustArrayFirePolicy.cpp
@@ -0,0 +1,20 @@
+/*******************************************************
+ * Copyright (c) 2020, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#include <ThrustArrayFirePolicy.hpp>
+
+namespace cuda {
+
+cudaStream_t get_stream(ThrustArrayFirePolicy) { return getActiveStream(); }
+
+cudaError_t synchronize_stream(ThrustArrayFirePolicy) {
+    return cudaStreamSynchronize(getActiveStream());
+}
+
+}  // namespace cuda
diff --git a/src/backend/cuda/ThrustArrayFirePolicy.hpp b/src/backend/cuda/ThrustArrayFirePolicy.hpp
@@ -0,0 +1,41 @@
+/*******************************************************
+ * Copyright (c) 2020, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#pragma once
+
+#include <backend.hpp>
+#include <memory.hpp>
+#include <platform.hpp>
+#include <thrust/execution_policy.h>
+
+namespace cuda {
+struct ThrustArrayFirePolicy
+    : thrust::device_execution_policy<ThrustArrayFirePolicy> {};
+
+__DH__
+cudaStream_t get_stream(ThrustArrayFirePolicy);
+
+__DH__
+cudaError_t synchronize_stream(ThrustArrayFirePolicy);
+
+template<typename T>
+thrust::pair<thrust::pointer<T, ThrustArrayFirePolicy>, std::ptrdiff_t>
+get_temporary_buffer(ThrustArrayFirePolicy, std::ptrdiff_t n) {
+    thrust::pointer<T, ThrustArrayFirePolicy> result(
+        cuda::memAlloc<T>(n / sizeof(T)).release());
+
+    return thrust::make_pair(result, n);
+}
+
+template<typename Pointer>
+void return_temporary_buffer(ThrustArrayFirePolicy, Pointer p) {
+    memFree(p.get());
+}
+
+}  // namespace cuda
diff --git a/src/backend/cuda/kernel/regions.hpp b/src/backend/cuda/kernel/regions.hpp
@@ -9,11 +9,11 @@
 
 #include <common/dispatch.hpp>
 #include <debug_cuda.hpp>
-#include <debug_thrust.hpp>
 #include <err_cuda.hpp>
 #include <math.hpp>
 #include <memory.hpp>
-#include <stdio.h>
+#include <thrust_utils.hpp>
+
 #include <thrust/adjacent_difference.h>
 #include <thrust/binary_search.h>
 #include <thrust/device_vector.h>

diff --git a/src/backend/cuda/kernel/sift_nonfree.hpp b/src/backend/cuda/kernel/sift_nonfree.hpp
@@ -74,9 +74,9 @@
 
 #include <common/dispatch.hpp>
 #include <debug_cuda.hpp>
-#include <debug_thrust.hpp>
 #include <err_cuda.hpp>
 #include <memory.hpp>
+#include <thrust_utils.hpp>
 #include <af/defines.h>
 #include "shared.hpp"
 

diff --git a/src/backend/cuda/kernel/sort.hpp b/src/backend/cuda/kernel/sort.hpp
@@ -10,13 +10,13 @@
 #include <Param.hpp>
 #include <common/dispatch.hpp>
 #include <debug_cuda.hpp>
-#include <debug_thrust.hpp>
 #include <err_cuda.hpp>
 #include <handle.hpp>
 #include <iota.hpp>
 #include <kernel/thrust_sort_by_key.hpp>
 #include <math.hpp>
 #include <thrust/sort.h>
+#include <thrust_utils.hpp>
 
 namespace cuda {
 namespace kernel {

diff --git a/src/backend/cuda/kernel/thrust_sort_by_key_impl.hpp b/src/backend/cuda/kernel/thrust_sort_by_key_impl.hpp
@@ -8,9 +8,9 @@
  ********************************************************/
 
 #include <debug_cuda.hpp>
-#include <debug_thrust.hpp>
 #include <kernel/thrust_sort_by_key.hpp>
 #include <thrust/sort.h>
+#include <thrust_utils.hpp>
 #include <types.hpp>
 
 namespace cuda {

diff --git a/src/backend/cuda/set.cu b/src/backend/cuda/set.cu
@@ -10,18 +10,18 @@
 #include <Array.hpp>
 #include <copy.hpp>
 #include <debug_cuda.hpp>
-#include <debug_thrust.hpp>
+#include <thrust_utils.hpp>
 #include <set.hpp>
 #include <sort.hpp>
 #include <af/dim4.hpp>
 
-#include <algorithm>
-
 #include <thrust/device_ptr.h>
 #include <thrust/set_operations.h>
 #include <thrust/sort.h>
 #include <thrust/unique.h>
 
+#include <algorithm>
+
 namespace cuda {
 using af::dim4;
 

diff --git a/src/backend/cuda/debug_thrust.hpp → src/backend/cuda/thrust_utils.hpp b/src/backend/cuda/debug_thrust.hpp → src/backend/cuda/thrust_utils.hpp
@@ -7,6 +7,8 @@
  * http://arrayfire.com/licenses/BSD-3-Clause
  ********************************************************/
 
+#pragma once
+#include <ThrustArrayFirePolicy.hpp>
 #include <thrust/system/cuda/detail/par.h>
 #include <thrust/version.h>
 #include <ThrustAllocator.cuh>
@@ -16,12 +18,11 @@ template<typename T>
 using ThrustVector = thrust::device_vector<T, cuda::ThrustAllocator<T>>;
 }
 
-#define THRUST_STREAM thrust::cuda::par.on(cuda::getActiveStream())
-
 #if THRUST_MAJOR_VERSION >= 1 && THRUST_MINOR_VERSION >= 8
 
-#define THRUST_SELECT(fn, ...) fn(THRUST_STREAM, __VA_ARGS__)
-#define THRUST_SELECT_OUT(res, fn, ...) res = fn(THRUST_STREAM, __VA_ARGS__)
+#define THRUST_SELECT(fn, ...) fn(cuda::ThrustArrayFirePolicy(), __VA_ARGS__)
+#define THRUST_SELECT_OUT(res, fn, ...) \
+    res = fn(cuda::ThrustArrayFirePolicy(), __VA_ARGS__)
 
 #else