[SYCLomatic] Fix for in/out types segmented sort (#483)

Signed-off-by: Dan Hoeflinger <[email protected]>

Author: danhoeflinger

clang/runtime/dpct-rt/include/dpl_extras/algorithm.h.inc

@@ -1414,11 +1414,11 @@ sort_pairs_impl(_ExecutionPolicy &&policy, key_t keys_in, key_out_t keys_out,
 // DplExtrasAlgorithm|sort_pairs
 // DPCT_DEPENDENCY_END
 // DPCT_CODE
-template <typename _ExecutionPolicy, typename key_t, typename value_t,
-          typename OffsetIteratorT>
+template <typename _ExecutionPolicy, typename key_t, typename key_out_t,
+          typename value_t, typename value_out_t, typename OffsetIteratorT>
 inline void segmented_sort_pairs_by_parallel_sorts(
-    _ExecutionPolicy &&policy, key_t keys_in, key_t keys_out, value_t values_in,
-    value_t values_out, int64_t n, int64_t nsegments,
+    _ExecutionPolicy &&policy, key_t keys_in, key_out_t keys_out,
+    value_out_t values_in, value_t values_out, int64_t n, int64_t nsegments,
     OffsetIteratorT begin_offsets, OffsetIteratorT end_offsets,
     bool descending = false, int begin_bit = 0,
     int end_bit = sizeof(typename ::std::iterator_traits<key_t>::value_type) *
@@ -1457,11 +1457,11 @@ inline void segmented_sort_pairs_by_parallel_sorts(
 // DplExtrasAlgorithm|sort_pairs
 // DPCT_DEPENDENCY_END
 // DPCT_CODE
-template <typename _ExecutionPolicy, typename key_t, typename value_t,
-          typename OffsetIteratorT>
+template <typename _ExecutionPolicy, typename key_t, typename key_out_t,
+          typename value_t, typename value_out_t, typename OffsetIteratorT>
 inline void segmented_sort_pairs_by_parallel_for_of_sorts(
-    _ExecutionPolicy &&policy, key_t keys_in, key_t keys_out, value_t values_in,
-    value_t values_out, int64_t n, int64_t nsegments,
+    _ExecutionPolicy &&policy, key_t keys_in, key_out_t keys_out,
+    value_t values_in, value_out_t values_out, int64_t n, int64_t nsegments,
     OffsetIteratorT begin_offsets, OffsetIteratorT end_offsets,
     bool descending = false, int begin_bit = 0,
     int end_bit = sizeof(typename ::std::iterator_traits<key_t>::value_type) *
@@ -1489,11 +1489,11 @@ inline void segmented_sort_pairs_by_parallel_for_of_sorts(
 // DplExtrasAlgorithm|sort_pairs
 // DPCT_DEPENDENCY_END
 // DPCT_CODE
-template <typename _ExecutionPolicy, typename key_t, typename value_t,
-          typename OffsetIteratorT>
+template <typename _ExecutionPolicy, typename key_t, typename key_out_t,
+          typename value_t, typename value_out_t, typename OffsetIteratorT>
 inline void segmented_sort_pairs_by_two_pair_sorts(
-    _ExecutionPolicy &&policy, key_t keys_in, key_t keys_out, value_t values_in,
-    value_t values_out, int64_t n, int64_t nsegments,
+    _ExecutionPolicy &&policy, key_t keys_in, key_out_t keys_out,
+    value_out_t values_in, value_t values_out, int64_t n, int64_t nsegments,
     OffsetIteratorT begin_offsets, OffsetIteratorT end_offsets,
     bool descending = false, int begin_bit = 0,
     int end_bit = sizeof(typename ::std::iterator_traits<key_t>::value_type) *
@@ -1540,29 +1540,26 @@ inline void segmented_sort_pairs_by_two_pair_sorts(
     // coordinate to mark segments
     policy.queue()
         .submit([&](sycl::handler &h) {
-          h.parallel_for(sycl::nd_range<1>{work_group_size, work_group_size},
-                         ([=](sycl::nd_item<1> item) {
-                           auto sub_group = item.get_sub_group();
-                           ::std::size_t num_subgroups =
-                               sub_group.get_group_range().size();
-                           ::std::size_t local_size =
-                               sub_group.get_local_range().size();
-
-                           ::std::size_t sub_group_id =
-                               sub_group.get_group_id();
-                           while (sub_group_id < nsegments) {
-                             ::std::size_t subgroup_local_id =
-                                 sub_group.get_local_id();
-                             std::size_t i = begin_offsets[sub_group_id];
-                             std::size_t end = end_offsets[sub_group_id];
-                             while (i + subgroup_local_id < end) {
-                               segments[i + subgroup_local_id] =
-                                   sub_group_id;
-                               i += local_size;
-                             }
-                             sub_group_id += num_subgroups;
-                           }
-                         }));
+          h.parallel_for(
+              sycl::nd_range<1>{work_group_size, work_group_size},
+              ([=](sycl::nd_item<1> item) {
+                auto sub_group = item.get_sub_group();
+                ::std::size_t num_subgroups =
+                    sub_group.get_group_range().size();
+                ::std::size_t local_size = sub_group.get_local_range().size();
+
+                ::std::size_t sub_group_id = sub_group.get_group_id();
+                while (sub_group_id < nsegments) {
+                  ::std::size_t subgroup_local_id = sub_group.get_local_id();
+                  std::size_t i = begin_offsets[sub_group_id];
+                  std::size_t end = end_offsets[sub_group_id];
+                  while (i + subgroup_local_id < end) {
+                    segments[i + subgroup_local_id] = sub_group_id;
+                    i += local_size;
+                  }
+                  sub_group_id += num_subgroups;
+                }
+              }));
         })
         .wait();
   } else {
@@ -1716,13 +1713,18 @@ inline void sort_keys(
 // DplExtrasAlgorithm|segmented_sort_pairs_by_parallel_sorts
 // DplExtrasAlgorithm|segmented_sort_pairs_by_parallel_for_of_sorts
 // DplExtrasAlgorithm|segmented_sort_pairs_by_two_pair_sorts
+// DplExtrasVector|is_iterator
 // DPCT_DEPENDENCY_END
 // DPCT_CODE
-template <typename _ExecutionPolicy, typename key_t, typename value_t,
-          typename OffsetIteratorT>
-inline void segmented_sort_pairs(
-    _ExecutionPolicy &&policy, key_t keys_in, key_t keys_out, value_t values_in,
-    value_t values_out, int64_t n, int64_t nsegments,
+template <typename _ExecutionPolicy, typename key_t, typename key_out_t,
+          typename value_t, typename value_out_t, typename OffsetIteratorT>
+inline ::std::enable_if_t<dpct::internal::is_iterator<key_t>::value &&
+                          dpct::internal::is_iterator<key_out_t>::value &&
+                          dpct::internal::is_iterator<value_t>::value &&
+                          dpct::internal::is_iterator<value_out_t>::value>
+segmented_sort_pairs(
+    _ExecutionPolicy &&policy, key_t keys_in, key_out_t keys_out,
+    value_t values_in, value_out_t values_out, int64_t n, int64_t nsegments,
     OffsetIteratorT begin_offsets, OffsetIteratorT end_offsets,
     bool descending = false, int begin_bit = 0,
     int end_bit = sizeof(typename ::std::iterator_traits<key_t>::value_type) *

clang/test/dpct/helper_files_ref/include/dpl_extras/algorithm.h

@@ -1180,11 +1180,11 @@ sort_pairs_impl(_ExecutionPolicy &&policy, key_t keys_in, key_out_t keys_out,
   sycl::free(temp_keys_out, policy.queue());
 }
 
-template <typename _ExecutionPolicy, typename key_t, typename value_t,
-          typename OffsetIteratorT>
+template <typename _ExecutionPolicy, typename key_t, typename key_out_t,
+          typename value_t, typename value_out_t, typename OffsetIteratorT>
 inline void segmented_sort_pairs_by_parallel_sorts(
-    _ExecutionPolicy &&policy, key_t keys_in, key_t keys_out, value_t values_in,
-    value_t values_out, int64_t n, int64_t nsegments,
+    _ExecutionPolicy &&policy, key_t keys_in, key_out_t keys_out,
+    value_out_t values_in, value_t values_out, int64_t n, int64_t nsegments,
     OffsetIteratorT begin_offsets, OffsetIteratorT end_offsets,
     bool descending = false, int begin_bit = 0,
     int end_bit = sizeof(typename ::std::iterator_traits<key_t>::value_type) *
@@ -1217,11 +1217,11 @@ inline void segmented_sort_pairs_by_parallel_sorts(
   sycl::free(host_accessible_offset_ends, policy.queue());
 }
 
-template <typename _ExecutionPolicy, typename key_t, typename value_t,
-          typename OffsetIteratorT>
+template <typename _ExecutionPolicy, typename key_t, typename key_out_t,
+          typename value_t, typename value_out_t, typename OffsetIteratorT>
 inline void segmented_sort_pairs_by_parallel_for_of_sorts(
-    _ExecutionPolicy &&policy, key_t keys_in, key_t keys_out, value_t values_in,
-    value_t values_out, int64_t n, int64_t nsegments,
+    _ExecutionPolicy &&policy, key_t keys_in, key_out_t keys_out,
+    value_t values_in, value_out_t values_out, int64_t n, int64_t nsegments,
     OffsetIteratorT begin_offsets, OffsetIteratorT end_offsets,
     bool descending = false, int begin_bit = 0,
     int end_bit = sizeof(typename ::std::iterator_traits<key_t>::value_type) *
@@ -1243,11 +1243,11 @@ inline void segmented_sort_pairs_by_parallel_for_of_sorts(
   policy.queue().wait();
 }
 
-template <typename _ExecutionPolicy, typename key_t, typename value_t,
-          typename OffsetIteratorT>
+template <typename _ExecutionPolicy, typename key_t, typename key_out_t,
+          typename value_t, typename value_out_t, typename OffsetIteratorT>
 inline void segmented_sort_pairs_by_two_pair_sorts(
-    _ExecutionPolicy &&policy, key_t keys_in, key_t keys_out, value_t values_in,
-    value_t values_out, int64_t n, int64_t nsegments,
+    _ExecutionPolicy &&policy, key_t keys_in, key_out_t keys_out,
+    value_out_t values_in, value_t values_out, int64_t n, int64_t nsegments,
     OffsetIteratorT begin_offsets, OffsetIteratorT end_offsets,
     bool descending = false, int begin_bit = 0,
     int end_bit = sizeof(typename ::std::iterator_traits<key_t>::value_type) *
@@ -1294,29 +1294,26 @@ inline void segmented_sort_pairs_by_two_pair_sorts(
     // coordinate to mark segments
     policy.queue()
         .submit([&](sycl::handler &h) {
-          h.parallel_for(sycl::nd_range<1>{work_group_size, work_group_size},
-                         ([=](sycl::nd_item<1> item) {
-                           auto sub_group = item.get_sub_group();
-                           ::std::size_t num_subgroups =
-                               sub_group.get_group_range().size();
-                           ::std::size_t local_size =
-                               sub_group.get_local_range().size();
-
-                           ::std::size_t sub_group_id =
-                               sub_group.get_group_id();
-                           while (sub_group_id < nsegments) {
-                             ::std::size_t subgroup_local_id =
-                                 sub_group.get_local_id();
-                             std::size_t i = begin_offsets[sub_group_id];
-                             std::size_t end = end_offsets[sub_group_id];
-                             while (i + subgroup_local_id < end) {
-                               segments[i + subgroup_local_id] =
-                                   sub_group_id;
-                               i += local_size;
-                             }
-                             sub_group_id += num_subgroups;
-                           }
-                         }));
+          h.parallel_for(
+              sycl::nd_range<1>{work_group_size, work_group_size},
+              ([=](sycl::nd_item<1> item) {
+                auto sub_group = item.get_sub_group();
+                ::std::size_t num_subgroups =
+                    sub_group.get_group_range().size();
+                ::std::size_t local_size = sub_group.get_local_range().size();
+
+                ::std::size_t sub_group_id = sub_group.get_group_id();
+                while (sub_group_id < nsegments) {
+                  ::std::size_t subgroup_local_id = sub_group.get_local_id();
+                  std::size_t i = begin_offsets[sub_group_id];
+                  std::size_t end = end_offsets[sub_group_id];
+                  while (i + subgroup_local_id < end) {
+                    segments[i + subgroup_local_id] = sub_group_id;
+                    i += local_size;
+                  }
+                  sub_group_id += num_subgroups;
+                }
+              }));
         })
         .wait();
   } else {
@@ -1434,11 +1431,15 @@ inline void sort_keys(
     keys.swap();
 }
 
-template <typename _ExecutionPolicy, typename key_t, typename value_t,
-          typename OffsetIteratorT>
-inline void segmented_sort_pairs(
-    _ExecutionPolicy &&policy, key_t keys_in, key_t keys_out, value_t values_in,
-    value_t values_out, int64_t n, int64_t nsegments,
+template <typename _ExecutionPolicy, typename key_t, typename key_out_t,
+          typename value_t, typename value_out_t, typename OffsetIteratorT>
+inline ::std::enable_if_t<dpct::internal::is_iterator<key_t>::value &&
+                          dpct::internal::is_iterator<key_out_t>::value &&
+                          dpct::internal::is_iterator<value_t>::value &&
+                          dpct::internal::is_iterator<value_out_t>::value>
+segmented_sort_pairs(
+    _ExecutionPolicy &&policy, key_t keys_in, key_out_t keys_out,
+    value_t values_in, value_out_t values_out, int64_t n, int64_t nsegments,
     OffsetIteratorT begin_offsets, OffsetIteratorT end_offsets,
     bool descending = false, int begin_bit = 0,
     int end_bit = sizeof(typename ::std::iterator_traits<key_t>::value_type) *