GPU: Replace OpenMP parallization with TBB

Author: davidrohr

GPU/GPUTracking/Base/GPUReconstruction.cxx

@@ -23,12 +23,9 @@
 #include <condition_variable>
 #include <array>
 
-#ifdef WITH_OPENMP
-#include <omp.h>
-#endif
-
 #include "GPUReconstruction.h"
 #include "GPUReconstructionIncludes.h"
+#include "GPUReconstructionThreading.h"
 #include "GPUROOTDumpCore.h"
 #include "GPUConfigDump.h"
 #include "GPUChainTracking.h"
@@ -121,17 +118,18 @@ void GPUReconstruction::GetITSTraits(std::unique_ptr<o2::its::TrackerTraits>* tr
   }
 }
 
-int32_t GPUReconstruction::SetNOMPThreads(int32_t n)
+void GPUReconstruction::SetNActiveThreads(int32_t n)
 {
-#ifdef WITH_OPENMP
-  omp_set_num_threads(mProcessingSettings.ompThreads = std::max(1, n < 0 ? mMaxOMPThreads : std::min(n, mMaxOMPThreads)));
+  mProcessingSettings.ompThreads = std::max(1, n < 0 ? mMaxOMPThreads : std::min(n, mMaxOMPThreads));
+  mThreading->activeThreads = std::make_unique<tbb::task_arena>(mProcessingSettings.ompThreads);
   if (mProcessingSettings.debugLevel >= 3) {
-    GPUInfo("Set number of OpenMP threads to %d (%d requested)", mProcessingSettings.ompThreads, n);
+    GPUInfo("Set number of active parallel threads to %d (%d requested)", mProcessingSettings.ompThreads, n);
   }
-  return n > mMaxOMPThreads;
-#else
-  return 1;
-#endif
+}
+
+int32_t GPUReconstruction::getOMPThreadNum()
+{
+  return std::max<int32_t>(0, tbb::this_task_arena::current_thread_index());
 }
 
 int32_t GPUReconstruction::Init()
@@ -299,23 +297,16 @@ int32_t GPUReconstruction::InitPhaseBeforeDevice()
     mMemoryScalers->rescaleMaxMem(mProcessingSettings.forceMaxMemScalers);
   }
 
-#ifdef WITH_OPENMP
   if (mProcessingSettings.ompThreads <= 0) {
-    mProcessingSettings.ompThreads = omp_get_max_threads();
+    mProcessingSettings.ompThreads = tbb::info::default_concurrency();
   } else {
     mProcessingSettings.ompAutoNThreads = false;
-    omp_set_num_threads(mProcessingSettings.ompThreads);
-  }
-  if (mProcessingSettings.ompKernels) {
-    if (omp_get_max_active_levels() < 2) {
-      omp_set_max_active_levels(2);
-    }
   }
-#else
-  mProcessingSettings.ompThreads = 1;
-#endif
   mMaxOMPThreads = mProcessingSettings.ompThreads;
-  mMaxThreads = std::max(mMaxThreads, mProcessingSettings.ompThreads);
+  mThreading->control = std::make_unique<tbb::global_control>(tbb::global_control::max_allowed_parallelism, mMaxOMPThreads);
+  mThreading->allThreads = std::make_unique<tbb::task_arena>(mMaxOMPThreads);
+  mThreading->activeThreads = std::make_unique<tbb::task_arena>(mProcessingSettings.ompThreads);
+  mMaxThreads = std::max(mMaxThreads, mMaxOMPThreads);
   if (IsGPU()) {
     mNStreams = std::max<int32_t>(mProcessingSettings.nStreams, 3);
   }

GPU/GPUTracking/Base/GPUReconstruction.h

@@ -51,6 +51,7 @@ namespace gpu
 class GPUChain;
 struct GPUMemorySizeScalers;
 struct GPUReconstructionPipelineContext;
+struct GPUReconstructionThreading;
 class GPUROOTDumpCore;
 
 namespace gpu_reconstruction_kernels
@@ -207,7 +208,7 @@ class GPUReconstruction
   void SetInputControl(void* ptr, size_t size);
   GPUOutputControl& OutputControl() { return mOutputControl; }
   int32_t GetMaxThreads() const { return mMaxThreads; }
-  int32_t SetNOMPThreads(int32_t n);
+  void SetNActiveThreads(int32_t n);
   int32_t NStreams() const { return mNStreams; }
   const void* DeviceMemoryBase() const { return mDeviceMemoryBase; }
 
@@ -234,6 +235,9 @@ class GPUReconstruction
   double GetStatKernelTime() { return mStatKernelTime; }
   double GetStatWallTime() { return mStatWallTime; }
 
+  std::unique_ptr<GPUReconstructionThreading> mThreading;
+  static int32_t getOMPThreadNum();
+
  protected:
   void AllocateRegisteredMemoryInternal(GPUMemoryResource* res, GPUOutputControl* control, GPUReconstruction* recPool);
   void FreeRegisteredMemory(GPUMemoryResource* res);

GPU/GPUTracking/Base/GPUReconstructionCPU.cxx

@@ -14,6 +14,7 @@
 
 #include "GPUReconstructionCPU.h"
 #include "GPUReconstructionIncludes.h"
+#include "GPUReconstructionThreading.h"
 #include "GPUChain.h"
 
 #include "GPUTPCClusterData.h"
@@ -40,13 +41,6 @@
 #include <unistd.h>
 #endif
 
-#if defined(WITH_OPENMP) || defined(_OPENMP)
-#include <omp.h>
-#else
-static inline int32_t omp_get_thread_num() { return 0; }
-static inline int32_t omp_get_max_threads() { return 1; }
-#endif
-
 using namespace o2::gpu;
 using namespace o2::gpu::gpu_reconstruction_kernels;
 
@@ -111,24 +105,20 @@ inline int32_t GPUReconstructionCPUBackend::runKernelBackendInternal(const krnlS
 template <>
 inline int32_t GPUReconstructionCPUBackend::runKernelBackendInternal<GPUMemClean16, 0>(const krnlSetupTime& _xyz, void* const& ptr, uint64_t const& size)
 {
-#ifdef WITH_OPENMP
   int32_t nOMPThreads = std::max<int32_t>(1, std::min<int32_t>(size / (16 * 1024 * 1024), getNOMPThreads()));
   if (nOMPThreads > 1) {
-    GPUCA_OPENMP(parallel num_threads(nOMPThreads))
-    {
-      size_t threadSize = size / omp_get_num_threads();
+    tbb::parallel_for(0, nOMPThreads, [&](int iThread) {
+      size_t threadSize = size / nOMPThreads;
       if (threadSize % 4096) {
         threadSize += 4096 - threadSize % 4096;
       }
-      size_t offset = threadSize * omp_get_thread_num();
+      size_t offset = threadSize * iThread;
       size_t mySize = std::min<size_t>(threadSize, size - offset);
       if (mySize) {
         memset((char*)ptr + offset, 0, mySize);
-      }
-    }
-  } else
-#endif
-  {
+      } // clang-format off
+    }, tbb::static_partitioner()); // clang-format on
+  } else {
     memset(ptr, 0, size);
   }
   return 0;
@@ -214,7 +204,7 @@ int32_t GPUReconstructionCPU::InitDevice()
     ClearAllocatedMemory();
   }
   if (mProcessingSettings.ompKernels) {
-    mBlockCount = getOMPMaxThreads();
+    mBlockCount = mMaxOMPThreads;
   }
   mThreadId = GetThread();
   mProcShadow.mProcessorsProc = processors();
@@ -351,16 +341,6 @@ void GPUReconstructionCPU::ResetDeviceProcessorTypes()
   }
 }
 
-int32_t GPUReconstructionCPUBackend::getOMPThreadNum()
-{
-  return omp_get_thread_num();
-}
-
-int32_t GPUReconstructionCPUBackend::getOMPMaxThreads()
-{
-  return omp_get_max_threads();
-}
-
 static std::atomic_flag timerFlag = ATOMIC_FLAG_INIT; // TODO: Should be a class member not global, but cannot be moved to header due to ROOT limitation
 
 GPUReconstructionCPU::timerMeta* GPUReconstructionCPU::insertTimer(uint32_t id, std::string&& name, int32_t J, int32_t num, int32_t type, RecoStep step)

GPU/GPUTracking/Base/GPUReconstructionCPU.h

@@ -44,8 +44,6 @@ class GPUReconstructionCPUBackend : public GPUReconstruction
   template <class T, int32_t I>
   gpu_reconstruction_kernels::krnlProperties getKernelPropertiesBackend();
   uint32_t mNestedLoopOmpFactor = 1;
-  static int32_t getOMPThreadNum();
-  static int32_t getOMPMaxThreads();
   int32_t getNOMPThreads();
 };
 
@@ -287,7 +285,7 @@ HighResTimer& GPUReconstructionCPU::getTimer(const char* name, int32_t num)
   static int32_t id = getNextTimerId();
   timerMeta* timer = getTimerById(id);
   if (timer == nullptr) {
-    int32_t max = std::max<int32_t>({getOMPMaxThreads(), mProcessingSettings.nStreams});
+    int32_t max = std::max<int32_t>({mMaxOMPThreads, mProcessingSettings.nStreams});
     timer = insertTimer(id, name, J, max, 1, RecoStep::NoRecoStep);
   }
   if (num == -1) {

GPU/GPUTracking/Base/GPUReconstructionConvert.cxx

@@ -40,6 +40,8 @@
 #include "TPCBase/CRU.h"
 #include "DetectorsRaw/RDHUtils.h"
 
+#include <oneapi/tbb.h>
+
 using namespace o2::gpu;
 using namespace o2::tpc;
 using namespace o2::tpc::constants;
@@ -1306,6 +1308,17 @@ size_t zsEncoderRun<T>::compare(std::vector<zsPage>* buffer, std::vector<o2::tpc
 } // anonymous namespace
 #endif // GPUCA_TPC_GEOMETRY_O2
 
+namespace o2::gpu::internal
+{
+struct tmpReductionResult {
+  uint32_t totalPages = 0;
+  size_t totalSize = 0;
+  size_t nErrors = 0;
+  size_t digitsInput = 0;
+  size_t digitsEncoded = 0;
+};
+} // namespace o2::gpu::internal
+
 template <class S>
 void GPUReconstructionConvert::RunZSEncoder(const S& in, std::unique_ptr<uint64_t[]>* outBuffer, uint32_t* outSizes, o2::raw::RawFileWriter* raw, const o2::InteractionRecord* ir, const GPUParam& param, int32_t version, bool verify, float threshold, bool padding, std::function<void(std::vector<o2::tpc::Digit>&)> digitsFilter)
 {
@@ -1316,67 +1329,68 @@ void GPUReconstructionConvert::RunZSEncoder(const S& in, std::unique_ptr<uint64_
   }
 #ifdef GPUCA_TPC_GEOMETRY_O2
   std::vector<zsPage> buffer[NSLICES][GPUTrackingInOutZS::NENDPOINTS];
-  uint32_t totalPages = 0;
-  size_t totalSize = 0;
-  size_t nErrors = 0;
-  size_t digitsInput = 0;
-  size_t digitsEncoded = 0;
-  // clang-format off
-  GPUCA_OPENMP(parallel for reduction(+ : totalPages, nErrors, totalSize, digitsInput, digitsEncoded))
-  // clang-format on
-  for (uint32_t i = 0; i < NSLICES; i++) {
-    std::vector<o2::tpc::Digit> tmpBuffer;
-    digitsInput += ZSEncoderGetNDigits(in, i);
-    tmpBuffer.resize(ZSEncoderGetNDigits(in, i));
-    if (threshold > 0.f && !digitsFilter) {
-      auto it = std::copy_if(ZSEncoderGetDigits(in, i), ZSEncoderGetDigits(in, i) + ZSEncoderGetNDigits(in, i), tmpBuffer.begin(), [threshold](auto& v) { return v.getChargeFloat() >= threshold; });
-      tmpBuffer.resize(std::distance(tmpBuffer.begin(), it));
-    } else {
-      std::copy(ZSEncoderGetDigits(in, i), ZSEncoderGetDigits(in, i) + ZSEncoderGetNDigits(in, i), tmpBuffer.begin());
-    }
-
-    if (digitsFilter) {
-      digitsFilter(tmpBuffer);
-      if (threshold > 0.f) {
-        std::vector<o2::tpc::Digit> tmpBuffer2 = std::move(tmpBuffer);
-        tmpBuffer = std::vector<o2::tpc::Digit>(tmpBuffer2.size());
-        auto it = std::copy_if(tmpBuffer2.begin(), tmpBuffer2.end(), tmpBuffer.begin(), [threshold](auto& v) { return v.getChargeFloat() >= threshold; });
+  auto reduced = tbb::parallel_reduce(tbb::blocked_range<uint32_t>(0, NSLICES), o2::gpu::internal::tmpReductionResult(), [&](const auto range, auto red) {
+    for (uint32_t i = range.begin(); i < range.end(); i++) {
+      std::vector<o2::tpc::Digit> tmpBuffer;
+      red.digitsInput += ZSEncoderGetNDigits(in, i);
+      tmpBuffer.resize(ZSEncoderGetNDigits(in, i));
+      if (threshold > 0.f && !digitsFilter) {
+        auto it = std::copy_if(ZSEncoderGetDigits(in, i), ZSEncoderGetDigits(in, i) + ZSEncoderGetNDigits(in, i), tmpBuffer.begin(), [threshold](auto& v) { return v.getChargeFloat() >= threshold; });
         tmpBuffer.resize(std::distance(tmpBuffer.begin(), it));
+      } else {
+        std::copy(ZSEncoderGetDigits(in, i), ZSEncoderGetDigits(in, i) + ZSEncoderGetNDigits(in, i), tmpBuffer.begin());
       }
-    }
-    digitsEncoded += tmpBuffer.size();
-
-    auto runZS = [&](auto& encoder) {
-      encoder.zsVersion = version;
-      encoder.init();
-      totalPages += encoder.run(buffer[i], tmpBuffer, &totalSize);
-      if (verify) {
-        nErrors += encoder.compare(buffer[i], tmpBuffer); // Verification
+
+      if (digitsFilter) {
+        digitsFilter(tmpBuffer);
+        if (threshold > 0.f) {
+          std::vector<o2::tpc::Digit> tmpBuffer2 = std::move(tmpBuffer);
+          tmpBuffer = std::vector<o2::tpc::Digit>(tmpBuffer2.size());
+          auto it = std::copy_if(tmpBuffer2.begin(), tmpBuffer2.end(), tmpBuffer.begin(), [threshold](auto& v) { return v.getChargeFloat() >= threshold; });
+          tmpBuffer.resize(std::distance(tmpBuffer.begin(), it));
+        }
       }
-    };
+      red.digitsEncoded += tmpBuffer.size();
+
+      auto runZS = [&](auto& encoder) {
+        encoder.zsVersion = version;
+        encoder.init();
+        red.totalPages += encoder.run(buffer[i], tmpBuffer, &red.totalSize);
+        if (verify) {
+          red.nErrors += encoder.compare(buffer[i], tmpBuffer); // Verification
+        }
+      };
 
-    if (version >= ZSVersion::ZSVersionRowBased10BitADC && version <= ZSVersion::ZSVersionRowBased12BitADC) {
-      zsEncoderRun<zsEncoderRow> enc{{{.iSector = i, .raw = raw, .ir = ir, .param = &param, .padding = padding}}};
-      runZS(enc);
-    } else if (version >= ZSVersion::ZSVersionLinkBasedWithMeta && version <= ZSVersion::ZSVersionDenseLinkBasedV2) {
-#ifdef GPUCA_O2_LIB
-      if (version == ZSVersion::ZSVersionLinkBasedWithMeta) {
-        zsEncoderRun<zsEncoderImprovedLinkBased> enc{{{{.iSector = i, .raw = raw, .ir = ir, .param = &param, .padding = padding}}}};
+      if (version >= ZSVersion::ZSVersionRowBased10BitADC && version <= ZSVersion::ZSVersionRowBased12BitADC) {
+        zsEncoderRun<zsEncoderRow> enc{{{.iSector = i, .raw = raw, .ir = ir, .param = &param, .padding = padding}}};
         runZS(enc);
-      } else if (version >= ZSVersion::ZSVersionDenseLinkBased && version <= ZSVersion::ZSVersionDenseLinkBasedV2) {
-        zsEncoderRun<zsEncoderDenseLinkBased> enc{{{{.iSector = i, .raw = raw, .ir = ir, .param = &param, .padding = padding}}}};
-        runZS(enc);
-      }
+      } else if (version >= ZSVersion::ZSVersionLinkBasedWithMeta && version <= ZSVersion::ZSVersionDenseLinkBasedV2) {
+#ifdef GPUCA_O2_LIB
+        if (version == ZSVersion::ZSVersionLinkBasedWithMeta) {
+          zsEncoderRun<zsEncoderImprovedLinkBased> enc{{{{.iSector = i, .raw = raw, .ir = ir, .param = &param, .padding = padding}}}};
+          runZS(enc);
+        } else if (version >= ZSVersion::ZSVersionDenseLinkBased && version <= ZSVersion::ZSVersionDenseLinkBasedV2) {
+          zsEncoderRun<zsEncoderDenseLinkBased> enc{{{{.iSector = i, .raw = raw, .ir = ir, .param = &param, .padding = padding}}}};
+          runZS(enc);
+        }
 #else
-      throw std::runtime_error("Link based ZS encoding not supported in standalone build");
+        throw std::runtime_error("Link based ZS encoding not supported in standalone build");
 #endif
-    } else {
-      throw std::runtime_error("Invalid ZS version "s + std::to_string(version) + ", cannot decode"s);
+      } else {
+        throw std::runtime_error("Invalid ZS version "s + std::to_string(version) + ", cannot decode"s);
+      }
     }
-  }
+    return red; }, [&](const auto& red1, const auto& red2) {
+    auto red = red1;
+    red.totalPages += red2.totalPages;
+    red.totalSize += red2.totalSize;
+    red.nErrors += red2.nErrors;
+    red.digitsInput += red2.digitsInput;
+    red.digitsEncoded += red2.digitsEncoded;
+    return red; });
 
   if (outBuffer) {
-    outBuffer->reset(new uint64_t[totalPages * TPCZSHDR::TPC_ZS_PAGE_SIZE / sizeof(uint64_t)]);
+    outBuffer->reset(new uint64_t[reduced.totalPages * TPCZSHDR::TPC_ZS_PAGE_SIZE / sizeof(uint64_t)]);
     uint64_t offset = 0;
     for (uint32_t i = 0; i < NSLICES; i++) {
       for (uint32_t j = 0; j < GPUTrackingInOutZS::NENDPOINTS; j++) {
@@ -1386,12 +1400,12 @@ void GPUReconstructionConvert::RunZSEncoder(const S& in, std::unique_ptr<uint64_
       }
     }
   }
-  if (nErrors) {
-    GPUError("ERROR: %ld INCORRECT SAMPLES DURING ZS ENCODING VERIFICATION!!!", (int64_t)nErrors);
+  if (reduced.nErrors) {
+    GPUError("ERROR: %lu INCORRECT SAMPLES DURING ZS ENCODING VERIFICATION!!!", reduced.nErrors);
   } else if (verify) {
     GPUInfo("ENCODING VERIFICATION PASSED");
   }
-  GPUInfo("TOTAL ENCODED SIZE: %lu (%lu of %lu digits encoded)", totalSize, digitsEncoded, digitsInput);
+  GPUInfo("TOTAL ENCODED SIZE: %lu (%lu of %lu digits encoded)", reduced.totalSize, reduced.digitsEncoded, reduced.digitsInput);
 #endif
 }
 

GPU/GPUTracking/Base/GPUReconstructionLibrary.cxx

@@ -22,10 +22,6 @@
 #include <unistd.h>
 #endif
 
-#ifdef WITH_OPENMP
-#include <omp.h>
-#endif
-
 #include "GPUReconstruction.h"
 #include "GPUReconstructionAvailableBackends.h"