Optimization of __serial_merge function

include/oneapi/dpl/pstl/hetero/dpcpp/parallel_backend_sycl_merge.h

@@ -80,52 +80,35 @@ __find_start_point(const _Rng1& __rng1, const _Rng2& __rng2, const _Index __i_el
 
 // Do serial merge of the data from rng1 (starting from start1) and rng2 (starting from start2) and writing
 // to rng3 (starting from start3) in 'chunk' steps, but do not exceed the total size of the sequences (n1 and n2)
-template <typename _Rng1, typename _Rng2, typename _Rng3, typename _Index, typename _Compare>
+template <unsigned int _UnrollFactor = 4, typename _Rng1, typename _Rng2, typename _Rng3, typename _Index, typename _Compare>
 void
 __serial_merge(const _Rng1& __rng1, const _Rng2& __rng2, _Rng3& __rng3, _Index __start1, _Index __start2,
                const _Index __start3, const std::uint8_t __chunk, const _Index __n1, const _Index __n2, _Compare __comp)
 {
-    if (__start1 >= __n1)
+    const _Index __rng1_size = std::min<_Index>(__n1 > __start1 ? __n1 - __start1 : (_Index)0, (_Index)__chunk);
+    const _Index __rng2_size = std::min<_Index>(__n2 > __start2 ? __n2 - __start2 : (_Index)0, (_Index)__chunk);
+    const _Index __rng3_size = std::min<_Index>(__rng1_size + __rng2_size, (_Index)__chunk);
+
+    const auto __rng1_idx_end = __start1 + __rng1_size;
+    const auto __rng2_idx_end = __start2 + __rng2_size;
+    const auto __rng3_idx_end = __start3 + __rng3_size;
+
+    _Index __rng1_idx = __start1;
+    _Index __rng2_idx = __start2;
+    _Index __rng3_idx = __start3;
+
+    bool __rng1_idx_less__n1, __rng2_idx_less__n2;
+
+    #pragma unroll _UnrollFactor
+    for (_Index __rng3_idx = __start3; __rng3_idx < __rng3_idx_end; ++__rng3_idx)
     {
-        //copying a residual of the second seq
-        const _Index __n = std::min<_Index>(__n2 - __start2, __chunk);
-        for (std::uint8_t __i = 0; __i < __n; ++__i)
-            __rng3[__start3 + __i] = __rng2[__start2 + __i];
-    }
-    else if (__start2 >= __n2)
-    {
-        //copying a residual of the first seq
-        const _Index __n = std::min<_Index>(__n1 - __start1, __chunk);
-        for (std::uint8_t __i = 0; __i < __n; ++__i)
-            __rng3[__start3 + __i] = __rng1[__start1 + __i];
-    }
-    else
-    {
-        for (std::uint8_t __i = 0; __i < __chunk && __start1 < __n1 && __start2 < __n2; ++__i)
-        {
-            const auto& __val1 = __rng1[__start1];
-            const auto& __val2 = __rng2[__start2];
-            if (__comp(__val2, __val1))
-            {
-                __rng3[__start3 + __i] = __val2;
-                if (++__start2 == __n2)
-                {
-                    //copying a residual of the first seq
-                    for (++__i; __i < __chunk && __start1 < __n1; ++__i, ++__start1)
-                        __rng3[__start3 + __i] = __rng1[__start1];
-                }
-            }
-            else
-            {
-                __rng3[__start3 + __i] = __val1;
-                if (++__start1 == __n1)
-                {
-                    //copying a residual of the second seq
-                    for (++__i; __i < __chunk && __start2 < __n2; ++__i, ++__start2)
-                        __rng3[__start3 + __i] = __rng2[__start2];
-                }
-            }
-        }
+        __rng1_idx_less__n1 = __rng1_idx < __rng1_idx_end;
+        __rng2_idx_less__n2 = __rng2_idx < __rng2_idx_end;
+
+        __rng3[__rng3_idx] =
+            __rng1_idx_less__n1 && __rng2_idx_less__n2
+                ? (__comp(__rng2[__rng2_idx], __rng1[__rng1_idx]) ? __rng2[__rng2_idx++] : __rng1[__rng1_idx++])
+                : (__rng1_idx_less__n1 ? __rng1[__rng1_idx++] : __rng2[__rng2_idx++]);
     }
 }
 
@@ -136,6 +119,8 @@ struct __parallel_merge_submitter;
 template <typename _IdType, typename... _Name>
 struct __parallel_merge_submitter<_IdType, __internal::__optional_kernel_name<_Name...>>
 {
+    static constexpr std::uint32_t __gpu_chunk = 4;
+
     template <typename _ExecutionPolicy, typename _Range1, typename _Range2, typename _Range3, typename _Compare>
     auto
     operator()(_ExecutionPolicy&& __exec, _Range1&& __rng1, _Range2&& __rng2, _Range3&& __rng3, _Compare __comp) const
@@ -149,7 +134,7 @@ struct __parallel_merge_submitter<_IdType, __internal::__optional_kernel_name<_N
         _PRINT_INFO_IN_DEBUG_MODE(__exec);
 
         // Empirical number of values to process per work-item
-        const std::uint8_t __chunk = __exec.queue().get_device().is_cpu() ? 128 : 4;
+        const std::uint8_t __chunk = __exec.queue().get_device().is_cpu() ? 128 : __gpu_chunk;
 
         const _IdType __steps = oneapi::dpl::__internal::__dpl_ceiling_div(__n, __chunk);
 
@@ -158,8 +143,8 @@ struct __parallel_merge_submitter<_IdType, __internal::__optional_kernel_name<_N
             __cgh.parallel_for<_Name...>(sycl::range</*dim=*/1>(__steps), [=](sycl::item</*dim=*/1> __item_id) {
                 const _IdType __i_elem = __item_id.get_linear_id() * __chunk;
                 const auto __start = __find_start_point(__rng1, __rng2, __i_elem, __n1, __n2, __comp);
-                __serial_merge(__rng1, __rng2, __rng3, __start.first, __start.second, __i_elem, __chunk, __n1, __n2,
-                               __comp);
+                __serial_merge<__gpu_chunk>(__rng1, __rng2, __rng3, __start.first, __start.second, __i_elem, __chunk,
+                                            __n1, __n2, __comp);
             });
         });
         return __future(__event);

include/oneapi/dpl/pstl/hetero/dpcpp/parallel_backend_sycl_merge_sort.h

@@ -233,6 +233,8 @@ struct __merge_sort_global_submitter;
 template <typename _IndexT, typename... _GlobalSortName>
 struct __merge_sort_global_submitter<_IndexT, __internal::__optional_kernel_name<_GlobalSortName...>>
 {
+    static constexpr std::uint32_t __gpu_chunk = 4;
+
     template <typename _Range, typename _Compare, typename _TempBuf, typename _LeafSizeT>
     std::pair<sycl::event, bool>
     operator()(sycl::queue& __q, _Range& __rng, _Compare __comp, _LeafSizeT __leaf_size, _TempBuf& __temp_buf,
@@ -241,7 +243,7 @@ struct __merge_sort_global_submitter<_IndexT, __internal::__optional_kernel_name
         const _IndexT __n = __rng.size();
         _IndexT __n_sorted = __leaf_size;
         const bool __is_cpu = __q.get_device().is_cpu();
-        const std::uint32_t __chunk = __is_cpu ? 32 : 4;
+        const std::uint32_t __chunk = __is_cpu ? 32 : __gpu_chunk;
         const std::size_t __steps = oneapi::dpl::__internal::__dpl_ceiling_div(__n, __chunk);
         bool __data_in_temp = false;
 
@@ -272,17 +274,17 @@ struct __merge_sort_global_submitter<_IndexT, __internal::__optional_kernel_name
                             const oneapi::dpl::__ranges::drop_view_simple __rng2(__dst, __offset + __n1);
 
                             const auto start = __find_start_point(__rng1, __rng2, __i_elem_local, __n1, __n2, __comp);
-                            __serial_merge(__rng1, __rng2, __rng /*__rng3*/, start.first, start.second, __i_elem,
-                                           __chunk, __n1, __n2, __comp);
+                            __serial_merge<__gpu_chunk>(__rng1, __rng2, __rng /*__rng3*/, start.first, start.second,
+                                                        __i_elem, __chunk, __n1, __n2, __comp);
                         }
                         else
                         {
                             const oneapi::dpl::__ranges::drop_view_simple __rng1(__rng, __offset);
                             const oneapi::dpl::__ranges::drop_view_simple __rng2(__rng, __offset + __n1);
 
                             const auto start = __find_start_point(__rng1, __rng2, __i_elem_local, __n1, __n2, __comp);
-                            __serial_merge(__rng1, __rng2, __dst /*__rng3*/, start.first, start.second, __i_elem,
-                                           __chunk, __n1, __n2, __comp);
+                            __serial_merge<__gpu_chunk>(__rng1, __rng2, __dst /*__rng3*/, start.first, start.second,
+                                                        __i_elem, __chunk, __n1, __n2, __comp);
                         }
                     });
             });