[alpaka] Refactor prefixScan implementation

Author: antoniopetre

src/alpaka/AlpakaCore/HistoContainer.h

@@ -41,6 +41,7 @@ namespace cms {
                                     T const *__restrict__ v,
                                     uint32_t const *__restrict__ offsets) const {
         const uint32_t nt = offsets[nh];
+
         cms::alpakatools::for_each_element_in_grid_strided(acc, nt, [&](uint32_t i) {
           auto off = alpaka_std::upper_bound(offsets, offsets + nh + 1, i);
           ALPAKA_ASSERT_OFFLOAD((*off) > 0);
@@ -74,17 +75,15 @@ namespace cms {
       const unsigned int nblocks = (num_items + nthreads - 1) / nthreads;
       const Vec1 blocksPerGrid(nblocks);
 
+      auto d_pc = cms::alpakatools::allocDeviceBuf<int32_t>(1u);
+      int32_t *pc = alpaka::getPtrNative(d_pc);
+      alpaka::memset(queue, d_pc, 0, 1u);
+
       const WorkDiv1 &workDiv = cms::alpakatools::make_workdiv(blocksPerGrid, threadsPerBlockOrElementsPerThread);
       alpaka::enqueue(queue,
                       alpaka::createTaskKernel<ALPAKA_ACCELERATOR_NAMESPACE::Acc1>(
-                          workDiv, multiBlockPrefixScanFirstStep<uint32_t>(), poff, poff, num_items));
-
-      const WorkDiv1 &workDivWith1Block =
-          cms::alpakatools::make_workdiv(Vec1::all(1), threadsPerBlockOrElementsPerThread);
-      alpaka::enqueue(
-          queue,
-          alpaka::createTaskKernel<ALPAKA_ACCELERATOR_NAMESPACE::Acc1>(
-              workDivWith1Block, multiBlockPrefixScanSecondStep<uint32_t>(), poff, poff, num_items, nblocks));
+                          workDiv, multiBlockPrefixScan<uint32_t>(), poff, poff, num_items, pc));
+      alpaka::wait(queue);
     }
 
     template <typename Histo, typename T>
@@ -318,4 +317,4 @@ namespace cms {
   }  // namespace alpakatools
 }  // namespace cms
 
-#endif  // HeterogeneousCore_CUDAUtilities_interface_HistoContainer_h
+#endif  // HeterogeneousCore_CUDAUtilities_interface_HistoContainer_h

src/alpaka/AlpakaCore/prefixScan.h

@@ -4,6 +4,7 @@
 #include <cstdint>
 #include "CUDACore/CMSUnrollLoop.h"
 #include "AlpakaCore/alpakaConfig.h"
+#include "AlpakaCore/threadfence.h"
 
 #ifdef ALPAKA_ACC_GPU_CUDA_ENABLED
 
@@ -51,7 +52,7 @@ namespace cms {
       uint32_t const blockDimension(alpaka::getWorkDiv<alpaka::Block, alpaka::Threads>(acc)[0u]);
       uint32_t const gridBlockIdx(alpaka::getIdx<alpaka::Grid, alpaka::Blocks>(acc)[0u]);
       uint32_t const blockThreadIdx(alpaka::getIdx<alpaka::Block, alpaka::Threads>(acc)[0u]);
-      assert(ws);
+      ALPAKA_ASSERT_OFFLOAD(ws);
       ALPAKA_ASSERT_OFFLOAD(size <= 1024);
       ALPAKA_ASSERT_OFFLOAD(0 == blockDimension % 32);
       auto first = blockThreadIdx;
@@ -99,7 +100,7 @@ namespace cms {
       uint32_t const blockDimension(alpaka::getWorkDiv<alpaka::Block, alpaka::Threads>(acc)[0u]);
       uint32_t const gridBlockIdx(alpaka::getIdx<alpaka::Grid, alpaka::Blocks>(acc)[0u]);
       uint32_t const blockThreadIdx(alpaka::getIdx<alpaka::Block, alpaka::Threads>(acc)[0u]);
-      assert(ws);
+      ALPAKA_ASSERT_OFFLOAD(ws);
       ALPAKA_ASSERT_OFFLOAD(size <= 1024);
       ALPAKA_ASSERT_OFFLOAD(0 == blockDimension % 32);
       auto first = blockThreadIdx;
@@ -134,43 +135,46 @@ namespace cms {
 
     // limited to 1024*1024 elements....
     template <typename T>
-    struct multiBlockPrefixScanFirstStep {
+    struct multiBlockPrefixScan {
       template <typename T_Acc>
-      ALPAKA_FN_ACC void operator()(const T_Acc& acc, T const* ci, T* co, int32_t size) const {
+      ALPAKA_FN_ACC void operator()(const T_Acc& acc, T const* ci, T* co, int32_t size, int32_t* pc) const {
         uint32_t const blockDimension(alpaka::getWorkDiv<alpaka::Block, alpaka::Threads>(acc)[0u]);
         uint32_t const threadDimension(alpaka::getWorkDiv<alpaka::Thread, alpaka::Elems>(acc)[0u]);
         uint32_t const blockIdx(alpaka::getIdx<alpaka::Grid, alpaka::Blocks>(acc)[0u]);
+        uint32_t const threadIdx(alpaka::getIdx<alpaka::Block, alpaka::Threads>(acc)[0u]);
 
-        auto& ws = alpaka::declareSharedVar<T[32], __COUNTER__>(acc);
         // first each block does a scan of size 1024; (better be enough blocks....)
 #ifndef NDEBUG
         uint32_t const gridDimension(alpaka::getWorkDiv<alpaka::Grid, alpaka::Blocks>(acc)[0u]);
         ALPAKA_ASSERT_OFFLOAD(gridDimension / threadDimension <= 1024);
 #endif
         int off = blockDimension * blockIdx * threadDimension;
+        auto& ws = alpaka::declareSharedVar<T[32], __COUNTER__>(acc);
         if (size - off > 0)
           blockPrefixScan(acc, ci + off, co + off, std::min(int(blockDimension * threadDimension), size - off), ws);
-      }
-    };
 
-    // limited to 1024*1024 elements....
-    template <typename T>
-    struct multiBlockPrefixScanSecondStep {
-      template <typename T_Acc>
-      ALPAKA_FN_ACC void operator()(const T_Acc& acc, T const* ci, T* co, int32_t size, int32_t numBlocks) const {
-        uint32_t const blockDimension(alpaka::getWorkDiv<alpaka::Block, alpaka::Threads>(acc)[0u]);
-        uint32_t const threadDimension(alpaka::getWorkDiv<alpaka::Thread, alpaka::Elems>(acc)[0u]);
+        auto& isLastBlockDone = alpaka::declareSharedVar<bool, __COUNTER__>(acc);
+        if (0 == threadIdx) {
+          cms::alpakatools::threadfence(acc);
+          auto value = alpaka::atomicAdd(acc, pc, 1, alpaka::hierarchy::Blocks{});  // block counter
+          isLastBlockDone = (value == (int(gridDimension) - 1));
+        }
 
-        uint32_t const threadIdx(alpaka::getIdx<alpaka::Block, alpaka::Threads>(acc)[0u]);
+        alpaka::syncBlockThreads(acc);
 
-        auto* const psum(alpaka::getDynSharedMem<T>(acc));
+        if (!isLastBlockDone)
+          return;
+
+        ALPAKA_ASSERT_OFFLOAD(int(gridDimension) == *pc);
+
+        auto& psum = alpaka::declareSharedVar<T[1024], __COUNTER__>(acc);
+
+        ALPAKA_ASSERT_OFFLOAD(static_cast<int32_t>(blockDimension * threadDimension) >= gridDimension);
 
-        // first each block does a scan of size 1024; (better be enough blocks....)
-        ALPAKA_ASSERT_OFFLOAD(static_cast<int32_t>(blockDimension * threadDimension) >= numBlocks);
         for (int elemId = 0; elemId < static_cast<int>(threadDimension); ++elemId) {
           int index = +threadIdx * threadDimension + elemId;
 
-          if (index < numBlocks) {
+          if (index < gridDimension) {
             int lastElementOfPreviousBlockId = index * blockDimension * threadDimension - 1;
             psum[index] = (lastElementOfPreviousBlockId < size and lastElementOfPreviousBlockId >= 0)
                               ? co[lastElementOfPreviousBlockId]
@@ -179,9 +183,7 @@ namespace cms {
         }
 
         alpaka::syncBlockThreads(acc);
-
-        auto& ws = alpaka::declareSharedVar<T[32], __COUNTER__>(acc);
-        blockPrefixScan(acc, psum, psum, numBlocks, ws);
+        blockPrefixScan(acc, psum, psum, gridDimension, ws);
 
         for (int elemId = 0; elemId < static_cast<int>(threadDimension); ++elemId) {
           int first = threadIdx * threadDimension + elemId;
@@ -192,40 +194,7 @@ namespace cms {
         }
       }
     };
-
   }  // namespace alpakatools
 }  // namespace cms
 
-namespace alpaka {
-  namespace traits {
-
-    //#############################################################################
-    //! The trait for getting the size of the block shared dynamic memory for a kernel.
-    template <typename T, typename TAcc>
-    struct BlockSharedMemDynSizeBytes<cms::alpakatools::multiBlockPrefixScanSecondStep<T>, TAcc> {
-      //-----------------------------------------------------------------------------
-      //! \return The size of the shared memory allocated for a block.
-      template <typename TVec>
-      ALPAKA_FN_HOST_ACC static auto getBlockSharedMemDynSizeBytes(
-          cms::alpakatools::multiBlockPrefixScanSecondStep<T> const& myKernel,
-          TVec const& blockThreadExtent,
-          TVec const& threadElemExtent,
-          T const* ci,
-          T* co,
-          int32_t size,
-          int32_t numBlocks) -> T {
-        alpaka::ignore_unused(myKernel);
-        alpaka::ignore_unused(blockThreadExtent);
-        alpaka::ignore_unused(threadElemExtent);
-        alpaka::ignore_unused(ci);
-        alpaka::ignore_unused(co);
-        alpaka::ignore_unused(size);
-
-        return static_cast<size_t>(numBlocks) * sizeof(T);
-      }
-    };
-
-  }  // namespace traits
-}  // namespace alpaka
-
 #endif  // HeterogeneousCore_AlpakaUtilities_interface_prefixScan_h

src/alpaka/test/alpaka/prefixScan_t.cc

@@ -176,28 +176,19 @@ int main() {
         cms::alpakatools::make_workdiv(blocksPerGrid4, threadsPerBlockOrElementsPerThread4);
 
     std::cout << "launch multiBlockPrefixScan " << num_items << ' ' << nBlocks << std::endl;
-    alpaka::enqueue(queue,
-                    alpaka::createTaskKernel<Acc1>(workDivMultiBlock,
-                                                   cms::alpakatools::multiBlockPrefixScanFirstStep<uint32_t>(),
-                                                   input_d,
-                                                   output1_d,
-                                                   num_items));
-
-    const Vec1 blocksPerGridSecondStep(Vec1::all(1));
-    const WorkDiv1& workDivMultiBlockSecondStep =
-        cms::alpakatools::make_workdiv(blocksPerGridSecondStep, threadsPerBlockOrElementsPerThread4);
-    alpaka::enqueue(queue,
-                    alpaka::createTaskKernel<Acc1>(workDivMultiBlockSecondStep,
-                                                   cms::alpakatools::multiBlockPrefixScanSecondStep<uint32_t>(),
-                                                   input_d,
-                                                   output1_d,
-                                                   num_items,
-                                                   nBlocks));
+    auto d_pc(alpaka::allocBuf<int32_t, Idx>(device, size));
+    int32_t* pc = alpaka::getPtrNative(d_pc);
+
+    alpaka::memset(queue, d_pc, 0, size);
+    alpaka::enqueue(
+        queue,
+        alpaka::createTaskKernel<Acc1>(
+            workDivMultiBlock, cms::alpakatools::multiBlockPrefixScan<uint32_t>(), input_d, output1_d, num_items, pc));
 
     alpaka::enqueue(queue, alpaka::createTaskKernel<Acc1>(workDivMultiBlock, verify(), output1_d, num_items));
 
     alpaka::wait(queue);  // input_dBuf and output1_dBuf end of scope
   }                       // ksize
 
   return 0;
-}
+}