use custom stream for thrust policy

ginkgo-project · Aug 16, 2023 · 09bfa7f · 09bfa7f
1 parent 5285e41
commit 09bfa7f
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/common/cuda_hip/distributed/partition_helpers_kernels.hpp.inc b/common/cuda_hip/distributed/partition_helpers_kernels.hpp.inc
@@ -46,8 +46,8 @@ void sort_by_range_start(
         range_start_ends.get_data() + 1, strided_indices);
     auto zip_it = thrust::make_zip_iterator(
         thrust::make_tuple(end_it, part_ids.get_data()));
-    thrust::stable_sort_by_key(thrust::device, start_it, start_it + num_ranges,
-                               zip_it);
+    thrust::stable_sort_by_key(thrust_policy(exec), start_it,
+                               start_it + num_ranges, zip_it);
 }
 
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(

diff --git a/cuda/distributed/partition_helpers_kernels.cu b/cuda/distributed/partition_helpers_kernels.cu
@@ -39,6 +39,9 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <thrust/sort.h>
 
 
+#include "cuda/base/thrust.cuh"
+
+
 namespace gko {
 namespace kernels {
 namespace cuda {

diff --git a/hip/distributed/partition_helpers_kernels.hip.cpp b/hip/distributed/partition_helpers_kernels.hip.cpp
@@ -39,6 +39,9 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <thrust/sort.h>
 
 
+#include "hip/base/thrust.hip.hpp"
+
+
 namespace gko {
 namespace kernels {
 namespace hip {