improvement of for-loop in loading data into SLM cache

Signed-off-by: Sergey Kopienko <[email protected]>
uxlfoundation · Nov 19, 2024 · 8b5fc76 · 8b5fc76
1 parent 1dd0f40
commit 8b5fc76
Showing 1 changed file with 16 additions and 6 deletions.
diff --git a/include/oneapi/dpl/pstl/hetero/dpcpp/parallel_backend_sycl_merge.h b/include/oneapi/dpl/pstl/hetero/dpcpp/parallel_backend_sycl_merge.h
@@ -393,12 +393,22 @@ struct __parallel_merge_submitter_large<_IdType, _CustomName,
                     const _IdType __i_elem = __local_id * __chunk;
 
                     // Cooperative data load from __rng1 to __rng1_cache_slm, from __rng2 to __rng1_cache_slm
-                    _ONEDPL_PRAGMA_UNROLL
-                    for (_IdType __idx = __i_elem; __idx < __i_elem + __chunk && __sp_base_left_global.first + __idx < __sp_base_right_global.first; ++__idx)
-                        __rng1_cache_slm[__idx] = __rng1[__sp_base_left_global.first + __idx];
-                    _ONEDPL_PRAGMA_UNROLL
-                    for (_IdType __idx = __i_elem; __idx < __i_elem + __chunk && __sp_base_left_global.second + __idx < __sp_base_right_global.second; ++__idx)
-                        __rng2_cache_slm[__idx] = __rng2[__sp_base_left_global.second + __idx];
+                    {
+                        const _IdType __idx_end = std::min(__i_elem + __chunk, __rng1_wg_data_size);
+                        _IdType __idx_rng = __sp_base_left_global.first;
+
+                        _ONEDPL_PRAGMA_UNROLL
+                        for (_IdType __idx = __i_elem; __idx < __idx_end; ++__idx, ++__idx_rng)
+                            __rng1_cache_slm[__idx] = __rng1[__idx_rng];
+                    }
+                    {
+                        const _IdType __idx_end = std::min(__i_elem + __chunk, __rng2_wg_data_size);
+                        _IdType __idx_rng = __sp_base_left_global.second;
+
+                        _ONEDPL_PRAGMA_UNROLL
+                        for (_IdType __idx = __i_elem; __idx < __idx_end; ++__idx, ++__idx_rng)
+                            __rng2_cache_slm[__idx] = __rng2[__idx_rng];
+                    }
 
                     // Wait until all the data is loaded
                     __dpl_sycl::__group_barrier(__nd_item);