Change AccCpuTbbBlocks object back to being local for each thread

alpaka-group · Nov 17, 2024 · c95b136 · c95b136
1 parent 215a292
commit c95b136
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 9 deletions.
diff --git a/include/alpaka/acc/AccCpuTbbBlocks.hpp b/include/alpaka/acc/AccCpuTbbBlocks.hpp
@@ -33,6 +33,7 @@
 
 // Implementation details.
 #include "alpaka/acc/Tag.hpp"
+#include "alpaka/core/BarrierTbb.h"
 #include "alpaka/core/ClipCast.hpp"
 #include "alpaka/core/Interface.hpp"
 #include "alpaka/dev/DevCpu.hpp"
@@ -90,13 +91,16 @@ namespace alpaka
 
     private:
         template<typename TWorkDiv>
-        ALPAKA_FN_HOST AccCpuTbbBlocks(TWorkDiv const& workDiv, std::size_t const& blockSharedMemDynSizeBytes)
+        ALPAKA_FN_HOST AccCpuTbbBlocks(
+            TWorkDiv const& workDiv,
+            std::size_t const& blockSharedMemDynSizeBytes,
+            core::tbb::BarrierThread<TIdx>& barrier)
             : WorkDivMembers<TDim, TIdx>(workDiv)
             , gb::IdxGbRef<TDim, TIdx>(m_gridBlockIdx)
             , BlockSharedMemDynMember<>(blockSharedMemDynSizeBytes)
             , BlockSharedMemStMember<>(staticMemBegin(), staticMemCapacity())
+            , GridSyncBarrierTbb<TIdx>(barrier)
             , m_gridBlockIdx(Vec<TDim, TIdx>::zeros())
-            , GridSyncBarrierTbb<TIdx>(getWorkDiv<Grid, Threads>(workDiv).prod())
         {
         }
 

diff --git a/include/alpaka/grid/GridSyncCpuTbbBlocks.hpp b/include/alpaka/grid/GridSyncCpuTbbBlocks.hpp
@@ -4,11 +4,11 @@
 
 #pragma once
 
+#include "alpaka/core/BarrierTbb.h"
 #include "alpaka/core/Common.hpp"
 #include "alpaka/grid/Traits.hpp"
 
 #ifdef ALPAKA_ACC_CPU_B_TBB_T_SEQ_ENABLED
-#    include "alpaka/core/BarrierTbb.h"
 
 namespace alpaka
 {
@@ -19,11 +19,13 @@ namespace alpaka
     public:
         using Barrier = core::tbb::BarrierThread<TIdx>;
 
-        ALPAKA_FN_HOST explicit GridSyncBarrierTbb(TIdx const& gridThreadCount) : m_barrier(gridThreadCount)
+        // Get reference to the barrier from the outside because we need it to be shared between blocks
+        ALPAKA_FN_HOST explicit GridSyncBarrierTbb(Barrier& barrier)
         {
+            m_barrier = &barrier;
         }
 
-        Barrier mutable m_barrier;
+        Barrier* m_barrier;
     };
 
     namespace trait
@@ -33,7 +35,7 @@ namespace alpaka
         {
             ALPAKA_FN_HOST static auto syncGridThreads(GridSyncBarrierTbb<TIdx> const& gridSync) -> void
             {
-                gridSync.m_barrier.wait();
+                gridSync.m_barrier->wait();
             }
         };
 

diff --git a/include/alpaka/kernel/TaskKernelCpuTbbBlocks.hpp b/include/alpaka/kernel/TaskKernelCpuTbbBlocks.hpp
@@ -86,15 +86,20 @@ namespace alpaka
             tbb::this_task_arena::isolate(
                 [&]
                 {
-                    AccCpuTbbBlocks<TDim, TIdx> acc(
-                        *static_cast<WorkDivMembers<TDim, TIdx> const*>(this),
-                        blockSharedMemDynSizeBytes);
+                    // Create a shared barrier for grid sync, which will be passed by reference to each thread to
+                    // achieve shared state
+                    core::tbb::BarrierThread<TIdx> barrier(numBlocksInGrid);
 
                     tbb::parallel_for(
                         static_cast<TIdx>(0),
                         static_cast<TIdx>(numBlocksInGrid),
                         [&](TIdx i)
                         {
+                            AccCpuTbbBlocks<TDim, TIdx> acc(
+                                *static_cast<WorkDivMembers<TDim, TIdx> const*>(this),
+                                blockSharedMemDynSizeBytes,
+                                barrier);
+
                             acc.m_gridBlockIdx
                                 = mapIdx<TDim::value>(Vec<DimInt<1u>, TIdx>(static_cast<TIdx>(i)), gridBlockExtent);