MaskedPlaced kernel optimized

Use local_accessor to improve memory bandwidth of the work-group.
oleksandr-pavlyk · Dec 10, 2024 · Dec 6, 2024 · Dec 6, 2024 · Dec 6, 2024 · Dec 7, 2024
commit 5161b8e30b7ab5f99b8fbb0b3d58a5e85b748adb
@@ -134,51 +134,65 @@ template <typename OrthogIndexerT,
           typename MaskedDstIndexerT,
           typename MaskedRhsIndexerT,
           typename dataT,
-          typename indT>
+          typename indT,
+          typename LocalAccessorT>
 struct MaskedPlaceStridedFunctor
 {
     MaskedPlaceStridedFunctor(char *dst_data_p,
                               const char *cumsum_data_p,
                               const char *rhs_data_p,
-                              size_t orthog_iter_size,
                               size_t masked_iter_size,
                               const OrthogIndexerT &orthog_dst_rhs_indexer_,
                               const MaskedDstIndexerT &masked_dst_indexer_,
-                              const MaskedRhsIndexerT &masked_rhs_indexer_)
+                              const MaskedRhsIndexerT &masked_rhs_indexer_,
+                              const LocalAccessorT &lacc_)
         : dst_cp(dst_data_p), cumsum_cp(cumsum_data_p), rhs_cp(rhs_data_p),
-          orthog_nelems(orthog_iter_size), masked_nelems(masked_iter_size),
+          masked_nelems(masked_iter_size),
           orthog_dst_rhs_indexer(orthog_dst_rhs_indexer_),
           masked_dst_indexer(masked_dst_indexer_),
-          masked_rhs_indexer(masked_rhs_indexer_)
+          masked_rhs_indexer(masked_rhs_indexer_), lacc(lacc_)
     {
+        static_assert(
+            std::is_same_v<indT, typename LocalAccessorT::value_type>);
     }
 
-    void operator()(sycl::id<1> idx) const
+    void operator()(sycl::nd_item<2> ndit) const
     {
         dataT *dst_data = reinterpret_cast<dataT *>(dst_cp);
         const indT *cumsum_data = reinterpret_cast<const indT *>(cumsum_cp);
         const dataT *rhs_data = reinterpret_cast<const dataT *>(rhs_cp);
 
-        size_t global_i = idx[0];
-        size_t orthog_i = global_i / masked_nelems;
-        size_t masked_i = global_i - masked_nelems * orthog_i;
-
-        indT current_running_count = cumsum_data[masked_i];
-        bool mask_set =
-            (masked_i == 0)
-                ? (current_running_count == 1)
-                : (current_running_count == cumsum_data[masked_i - 1] + 1);
-
-        // src[i, j] = rhs[cumsum[i] - 1, j] if cumsum[i] == ((i > 0) ?
-        // cumsum[i-1]
-        // + 1 : 1)
-        if (mask_set) {
-            auto orthog_offsets =
-                orthog_dst_rhs_indexer(static_cast<ssize_t>(orthog_i));
-
-            size_t total_dst_offset = masked_dst_indexer(masked_i) +
-                                      orthog_offsets.get_first_offset();
-            size_t total_rhs_offset =
+        const std::size_t orthog_i = ndit.get_global_id(0);
+        const std::size_t group_i = ndit.get_group(1);
+        const std::uint32_t l_i = ndit.get_local_id(1);
+        const std::uint32_t lws = ndit.get_local_range(1);
+
+        const size_t masked_block_start = group_i * lws;
+        const size_t masked_i = masked_block_start + l_i;
+
+        for (std::uint32_t i = l_i; i < lacc.size(); i += lws) {
+            const size_t offset = masked_block_start + i;
+            lacc[i] = (offset == 0) ? indT(0)
+                      : (offset - 1 < masked_nelems)
+                          ? cumsum_data[offset - 1]
+                          : cumsum_data[masked_nelems - 1] + 1;
+        }
+
+        sycl::group_barrier(ndit.get_group());
+
+        const indT current_running_count = lacc[l_i + 1];
+        const bool mask_set = (masked_i == 0)
+                                  ? (current_running_count == 1)
+                                  : (current_running_count == lacc[l_i] + 1);
+
+        // src[i, j] = rhs[cumsum[i] - 1, j]
+        // if cumsum[i] == ((i > 0) ? cumsum[i-1] + 1 : 1)
+        if (mask_set && (masked_i < masked_nelems)) {
+            const auto &orthog_offsets = orthog_dst_rhs_indexer(orthog_i);
+
+            const size_t total_dst_offset = masked_dst_indexer(masked_i) +
+                                            orthog_offsets.get_first_offset();
+            const size_t total_rhs_offset =
                 masked_rhs_indexer(current_running_count - 1) +
                 orthog_offsets.get_second_offset();
 
@@ -190,8 +204,7 @@ struct MaskedPlaceStridedFunctor
     char *dst_cp = nullptr;
     const char *cumsum_cp = nullptr;
     const char *rhs_cp = nullptr;
-    size_t orthog_nelems = 0;
-    size_t masked_nelems = 0;
+    const size_t masked_nelems = 0;
     // has nd, shape, dst_strides, rhs_strides for
     // dimensions that ARE NOT masked
     const OrthogIndexerT orthog_dst_rhs_indexer;
@@ -200,6 +213,7 @@ struct MaskedPlaceStridedFunctor
     const MaskedDstIndexerT masked_dst_indexer;
     // has 1, rhs_strides for dimensions that ARE masked
     const MaskedRhsIndexerT masked_rhs_indexer;
+    LocalAccessorT lacc;
 };
 
 // ======= Masked extraction ================================
@@ -537,18 +551,35 @@ sycl::event masked_place_all_slices_strided_impl(
     const StridedIndexer masked_dst_indexer(nd, 0, packed_dst_shape_strides);
     const Strided1DCyclicIndexer masked_rhs_indexer(0, rhs_size, rhs_stride);
 
+    using KernelName = class masked_place_all_slices_strided_impl_krn<
+        TwoZeroOffsets_Indexer, StridedIndexer, Strided1DCyclicIndexer, dataT,
+        indT>;
+
+    constexpr std::size_t nominal_lws = 256;
+    const std::size_t masked_extent = iteration_size;
+    const std::size_t lws = std::min(masked_extent, nominal_lws);
+
+    const std::size_t n_groups = (masked_extent + lws - 1) / lws;
+
+    sycl::range<2> gRange{1, n_groups * lws};
+    sycl::range<2> lRange{1, lws};
+    sycl::nd_range<2> ndRange{gRange, lRange};
+
+    using LocalAccessorT = sycl::local_accessor<indT, 1>;
+
     sycl::event comp_ev = exec_q.submit([&](sycl::handler &cgh) {
         cgh.depends_on(depends);
 
-        cgh.parallel_for<class masked_place_all_slices_strided_impl_krn<
-            TwoZeroOffsets_Indexer, StridedIndexer, Strided1DCyclicIndexer,
-            dataT, indT>>(
-            sycl::range<1>(static_cast<size_t>(iteration_size)),
+        const std::size_t lacc_size = std::min(masked_extent, lws) + 1;
+        LocalAccessorT lacc(lacc_size, cgh);
+
+        cgh.parallel_for<KernelName>(
+            ndRange,
             MaskedPlaceStridedFunctor<TwoZeroOffsets_Indexer, StridedIndexer,
-                                      Strided1DCyclicIndexer, dataT, indT>(
-                dst_p, cumsum_p, rhs_p, 1, iteration_size,
-                orthog_dst_rhs_indexer, masked_dst_indexer,
-                masked_rhs_indexer));
+                                      Strided1DCyclicIndexer, dataT, indT,
+                                      LocalAccessorT>(
+                dst_p, cumsum_p, rhs_p, iteration_size, orthog_dst_rhs_indexer,
+                masked_dst_indexer, masked_rhs_indexer, lacc));
     });
 
     return comp_ev;
@@ -615,18 +646,32 @@ sycl::event masked_place_some_slices_strided_impl(
         TwoOffsets_StridedIndexer, StridedIndexer, Strided1DCyclicIndexer,
         dataT, indT>;
 
-    sycl::range<1> gRange(static_cast<size_t>(orthog_nelems * masked_nelems));
+    constexpr std::size_t nominal_lws = 256;
+    const std::size_t orthog_extent = orthog_nelems;
+    const std::size_t masked_extent = masked_nelems;
+    const std::size_t lws = std::min(masked_extent, nominal_lws);
+
+    const std::size_t n_groups = (masked_extent + lws - 1) / lws;
+
+    sycl::range<2> gRange{orthog_extent, n_groups * lws};
+    sycl::range<2> lRange{1, lws};
+    sycl::nd_range<2> ndRange{gRange, lRange};
+
+    using LocalAccessorT = sycl::local_accessor<indT, 1>;
 
     sycl::event comp_ev = exec_q.submit([&](sycl::handler &cgh) {
         cgh.depends_on(depends);
 
+        const std::size_t lacc_size = std::min(masked_extent, lws) + 1;
+        LocalAccessorT lacc(lacc_size, cgh);
+
         cgh.parallel_for<KernelName>(
-            gRange,
+            ndRange,
             MaskedPlaceStridedFunctor<TwoOffsets_StridedIndexer, StridedIndexer,
-                                      Strided1DCyclicIndexer, dataT, indT>(
-                dst_p, cumsum_p, rhs_p, orthog_nelems, masked_nelems,
-                orthog_dst_rhs_indexer, masked_dst_indexer,
-                masked_rhs_indexer));
+                                      Strided1DCyclicIndexer, dataT, indT,
+                                      LocalAccessorT>(
+                dst_p, cumsum_p, rhs_p, masked_nelems, orthog_dst_rhs_indexer,
+                masked_dst_indexer, masked_rhs_indexer, lacc));
     });
 
     return comp_ev;