core: Add semantic ROOT::EnableImplicitMT (#18694)

pcanal · pcanal · commit b9235eff15f4 · 2025-05-16T20:21:30.000-05:00
Instead of specificy the number of core, the new overload can be specified to select a configuration type. Currently the following 2 configuration are supported: ROOT::EIMTConfig::kWholeMachine ROOT::EIMTConfig::kExistingTBBArena calling ``` ROOT::EnableImplicitMT(ROOT::EIMTConfig::kExistingTBBArena); ``` will attach to an existing TBB Arena instead of creating a new one. (cherry picked from commit eb0b903)
diff --git a/core/base/inc/TROOT.h b/core/base/inc/TROOT.h
@@ -80,12 +80,20 @@ namespace Internal {
 } } // End ROOT::Internal
 
 namespace ROOT {
+   enum class EIMTConfig {
+      kWholeMachine = 0,     ///< Default configuration
+      kExistingTBBArena = 1, ///< Use the existing TBB arena
+      kNumConfigs = 2        ///< Number of support IMT semantic configurations
+   };
    /// \brief Enable support for multi-threading within the ROOT code
    /// in particular, enables the global mutex to make ROOT thread safe/aware.
    void EnableThreadSafety();
    /// \brief Enable ROOT's implicit multi-threading for all objects and methods that provide an internal
    /// parallelisation mechanism.
    void EnableImplicitMT(UInt_t numthreads = 0);
+   /// \brief Enable ROOT's implicit multi-threading for all objects and methods that provide an internal
+   /// parallelisation mechanism.
+   void EnableImplicitMT(ROOT::EIMTConfig config);
    void DisableImplicitMT();
    Bool_t IsImplicitMTEnabled();
    UInt_t GetThreadPoolSize();
diff --git a/core/base/src/TROOT.cxx b/core/base/src/TROOT.cxx
@@ -551,6 +551,31 @@ namespace Internal {
 #endif
    }
 
+   ////////////////////////////////////////////////////////////////////////////////
+   /// @param[in] config Configuration to use. The default is kWholeMachine, which
+   ///                   will create a thread pool that spans the whole machine.
+   ///
+   /// EnableImplicitMT calls in turn EnableThreadSafety.
+   /// If ImplicitMT is already enabled, this function does nothing.
+
+   void EnableImplicitMT(ROOT::EIMTConfig config)
+   {
+#ifdef R__USE_IMT
+      if (ROOT::Internal::IsImplicitMTEnabledImpl())
+         return;
+      EnableThreadSafety();
+      static void (*sym)(ROOT::EIMTConfig) =
+         (void (*)(ROOT::EIMTConfig))Internal::GetSymInLibImt("ROOT_TImplicitMT_EnableImplicitMT_Config");
+      if (sym)
+         sym(config);
+      ROOT::Internal::IsImplicitMTEnabledImpl() = true;
+#else
+      ::Warning("EnableImplicitMT",
+                "Cannot enable implicit multi-threading with config %d, please build ROOT with -Dimt=ON",
+                static_cast<int>(config));
+#endif
+   }
+
    ////////////////////////////////////////////////////////////////////////////////
    /// Disables the implicit multi-threading in ROOT (see EnableImplicitMT).
    void DisableImplicitMT()
diff --git a/core/imt/inc/ROOT/RTaskArena.hxx b/core/imt/inc/ROOT/RTaskArena.hxx
@@ -24,6 +24,7 @@
 #define ROOT_RTaskArena
 
 #include "RConfigure.h"
+#include "TROOT.h" // For ROOT::EIMTConfig
 #include <memory>
 
 // exclude in case ROOT does not have IMT support
@@ -65,9 +66,13 @@ public:
    ~RTaskArenaWrapper(); // necessary to set size back to zero
    static unsigned TaskArenaSize(); // A static getter lets us check for RTaskArenaWrapper's existence
    ROOT::ROpaqueTaskArena &Access();
-private:
+   struct Attach {}; ///< Marker for attaching to an existing tbb::task_arena
+
    RTaskArenaWrapper(unsigned maxConcurrency = 0);
-   friend std::shared_ptr<ROOT::Internal::RTaskArenaWrapper> GetGlobalTaskArena(unsigned maxConcurrency);
+   RTaskArenaWrapper(Attach);
+
+private:
+   friend std::shared_ptr<ROOT::Internal::RTaskArenaWrapper> GetGlobalTaskArena(unsigned, ROOT::EIMTConfig);
    std::unique_ptr<ROOT::ROpaqueTaskArena> fTBBArena;
    static unsigned fNWorkers;
 };
@@ -81,6 +86,7 @@ private:
 /// references to the previous one are gone and the object destroyed.
 ////////////////////////////////////////////////////////////////////////////////
 std::shared_ptr<ROOT::Internal::RTaskArenaWrapper> GetGlobalTaskArena(unsigned maxConcurrency = 0);
+std::shared_ptr<ROOT::Internal::RTaskArenaWrapper> GetGlobalTaskArena(ROOT::EIMTConfig config);
 
 } // namespace Internal
 } // namespace ROOT
diff --git a/core/imt/src/ROpaqueTaskArena.hxx b/core/imt/src/ROpaqueTaskArena.hxx
@@ -1,5 +1,8 @@
 #include "tbb/task_arena.h"
 
 namespace ROOT {
-class ROpaqueTaskArena: public tbb::task_arena {};
+class ROpaqueTaskArena : public tbb::task_arena {
+public:
+   using tbb::task_arena::task_arena;
+};
 }
diff --git a/core/imt/src/RTaskArena.cxx b/core/imt/src/RTaskArena.cxx
@@ -108,6 +108,20 @@ RTaskArenaWrapper::RTaskArenaWrapper(unsigned maxConcurrency) : fTBBArena(new RO
    ROOT::EnableThreadSafety();
 }
 
+////////////////////////////////////////////////////////////////////////////////
+/// Initializes the tbb::task_arena within RTaskArenaWrapper by attaching to an
+/// existing arena.
+///
+/// * Can't be reinitialized
+////////////////////////////////////////////////////////////////////////////////
+RTaskArenaWrapper::RTaskArenaWrapper(RTaskArenaWrapper::Attach)
+   : fTBBArena(new ROpaqueTaskArena{tbb::task_arena::attach{}})
+{
+   fTBBArena->initialize(tbb::task_arena::attach{});
+   fNWorkers = fTBBArena->max_concurrency();
+   ROOT::EnableThreadSafety();
+}
+
 RTaskArenaWrapper::~RTaskArenaWrapper()
 {
    fNWorkers = 0u;
@@ -127,7 +141,8 @@ ROOT::ROpaqueTaskArena &RTaskArenaWrapper::Access()
    return *fTBBArena;
 }
 
-std::shared_ptr<ROOT::Internal::RTaskArenaWrapper> GetGlobalTaskArena(unsigned maxConcurrency)
+std::shared_ptr<ROOT::Internal::RTaskArenaWrapper>
+GetGlobalTaskArena(unsigned maxConcurrency, ROOT::EIMTConfig config)
 {
    static std::weak_ptr<ROOT::Internal::RTaskArenaWrapper> weak_GTAWrapper;
 
@@ -140,10 +155,31 @@ std::shared_ptr<ROOT::Internal::RTaskArenaWrapper> GetGlobalTaskArena(unsigned m
       }
       return sp;
    }
-   std::shared_ptr<ROOT::Internal::RTaskArenaWrapper> sp(new ROOT::Internal::RTaskArenaWrapper(maxConcurrency));
+   std::shared_ptr<ROOT::Internal::RTaskArenaWrapper> sp;
+   if (config == ROOT::EIMTConfig::kExistingTBBArena) {
+      sp = std::make_shared<ROOT::Internal::RTaskArenaWrapper>(ROOT::Internal::RTaskArenaWrapper::Attach{});
+   } else {
+      if (config == ROOT::EIMTConfig::kWholeMachine) {
+         maxConcurrency = 0;
+      }
+      sp = std::make_shared<ROOT::Internal::RTaskArenaWrapper>(maxConcurrency);
+   }
    weak_GTAWrapper = sp;
    return sp;
 }
 
+std::shared_ptr<ROOT::Internal::RTaskArenaWrapper> GetGlobalTaskArena(ROOT::EIMTConfig config)
+{
+   if (config >= ROOT::EIMTConfig::kNumConfigs)
+      ::Fatal("ROOT::Internal::GetGlobalTaskArena",
+              "Unsupported enum value %d", (int)config);
+   return GetGlobalTaskArena(0, config);
+}
+
+std::shared_ptr<ROOT::Internal::RTaskArenaWrapper> GetGlobalTaskArena(unsigned maxConcurrency)
+{
+   return GetGlobalTaskArena(maxConcurrency, ROOT::EIMTConfig::kNumConfigs);
+}
+
 } // namespace Internal
 } // namespace ROOT
diff --git a/core/imt/src/TImplicitMT.cxx b/core/imt/src/TImplicitMT.cxx
@@ -18,6 +18,7 @@
 //                                                                      //
 //////////////////////////////////////////////////////////////////////////
 
+#include "TROOT.h"
 #include "TError.h"
 #include "ROOT/RTaskArena.hxx"
 #include <atomic>
@@ -55,6 +56,22 @@ extern "C" void ROOT_TImplicitMT_EnableImplicitMT(UInt_t numthreads)
    }
 };
 
+extern "C" void ROOT_TImplicitMT_EnableImplicitMT_Config(ROOT::EIMTConfig config)
+{
+   if (!GetImplicitMTFlag()) {
+      if (config < ROOT::EIMTConfig::kNumConfigs) {
+         R__GetTaskArena4IMT() = ROOT::Internal::GetGlobalTaskArena(config);
+      } else {
+         ::Warning("ROOT_TImplicitMT_EnableImplicitMT_Config",
+                   "Unknown enum value %d defaulting to EIMTCconfig::kWholeMachine", (int)config);
+         R__GetTaskArena4IMT() = ROOT::Internal::GetGlobalTaskArena(0);
+      }
+      GetImplicitMTFlag() = true;
+   } else {
+      ::Warning("ROOT_TImplicitMT_EnableImplicitMT_Config", "Implicit multi-threading is already enabled");
+   }
+};
+
 extern "C" void ROOT_TImplicitMT_DisableImplicitMT()
 {
    if (GetImplicitMTFlag()) {
diff --git a/core/imt/test/CMakeLists.txt b/core/imt/test/CMakeLists.txt
@@ -6,3 +6,4 @@
 
 ROOT_ADD_GTEST(testTaskArena testRTaskArena.cxx LIBRARIES Imt ${TBB_LIBRARIES} FAILREGEX "")
 ROOT_ADD_GTEST(testTBBGlobalControl testTBBGlobalControl.cxx LIBRARIES Imt ${TBB_LIBRARIES})
+ROOT_ADD_GTEST(testEnableImt testEnableImt.cxx LIBRARIES Imt ${TBB_LIBRARIES})
diff --git a/core/imt/test/testEnableImt.cxx b/core/imt/test/testEnableImt.cxx
@@ -0,0 +1,26 @@
+#include "TROOT.h"
+#include "ROOT/RTaskArena.hxx"
+
+#include "tbb/task_arena.h"
+
+#include "ROOT/TestSupport.hxx"
+#include "gtest/gtest.h"
+
+#ifdef R__USE_IMT
+
+static const unsigned gMaxConcurrency = ROOT::Internal::LogicalCPUBandwidthControl();
+
+TEST(EnableImt, TBBAttach)
+{
+   tbb::task_arena main_arena{2};
+
+   main_arena.execute([&]() { ROOT::EnableImplicitMT(ROOT::EIMTConfig::kExistingTBBArena); });
+
+   auto psize = ROOT::GetThreadPoolSize();
+
+   EXPECT_TRUE(psize > 1);
+   EXPECT_EQ(main_arena.max_concurrency(), 2);
+   EXPECT_EQ(psize, 2);
+}
+
+#endif

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,8 @@`
`1`	`1`	`#include "tbb/task_arena.h"`
`2`	`2`
`3`	`3`	`namespace ROOT {`
`4`		`-class ROpaqueTaskArena: public tbb::task_arena {};`
	`4`	`+class ROpaqueTaskArena : public tbb::task_arena {`
	`5`	`+public:`
	`6`	`+ using tbb::task_arena::task_arena;`
	`7`	`+};`
`5`	`8`	`}`
Original file line number	Diff line number	Diff line change
`@@ -6,3 +6,4 @@`
`6`	`6`
`7`	`7`	`ROOT_ADD_GTEST(testTaskArena testRTaskArena.cxx LIBRARIES Imt ${TBB_LIBRARIES} FAILREGEX "")`
`8`	`8`	`ROOT_ADD_GTEST(testTBBGlobalControl testTBBGlobalControl.cxx LIBRARIES Imt ${TBB_LIBRARIES})`
	`9`	`+ROOT_ADD_GTEST(testEnableImt testEnableImt.cxx LIBRARIES Imt ${TBB_LIBRARIES})`