ALICE3: use track extension + allow configuring passes via JSON

Detectors/Upgrades/ALICE3/GlobalReconstruction/workflow/README.md

@@ -19,6 +19,7 @@ o2-alice3-global-reconstruction-reco-workflow --tracking-from-hits-config config
 - `--tracking-from-hits-config <file>`: Path to tracking-from-hits configuration JSON file
 - `--tracking-from-clusters-config <file>`: Path to tracking-from-clusters configuration JSON file
 - `--gpu-device <id>`: Tracking device type (`1` CPU, `2` CUDA, `3` HIP)
+- `--tracking-threads <n>`: Number of CPU threads used by TRK tracking
 - `-b`: Batch mode (no GUI)
 - `--disable-root-output`: Skip writing tracks to ROOT file
 - `--help`: Show all available options
@@ -80,12 +81,10 @@ The tracking configuration is provided via a JSON file that specifies:
     "SaveTimeBenchmarks": false,
     "DoUPCIteration": false,
     "FataliseUponFailure": true,
-    "UseTrackFollower": true,
-    "UseTrackFollowerTop": false,
-    "UseTrackFollowerBot": false,
-    "UseTrackFollowerMix": true,
+    "TrackFollower": "mix",
     "TrackFollowerNSigmaCutZ": 1.0,
     "TrackFollowerNSigmaCutPhi": 1.0,
+    "TrackFollowerMaxHypotheses": 1,
     "createArtefactLabels": false,
     "PrintMemory": false,
     "DropTFUponFailure": false

Detectors/Upgrades/ALICE3/GlobalReconstruction/workflow/include/ALICE3GlobalReconstructionWorkflow/RecoWorkflow.h

@@ -23,7 +23,8 @@ o2::framework::WorkflowSpec getWorkflow(bool useMC,
                                         const std::string& hitRecoConfig,
                                         const std::string& clusterRecoConfig,
                                         bool disableRootOutput = false,
-                                        o2::gpu::gpudatatypes::DeviceType dType = o2::gpu::gpudatatypes::DeviceType::CPU);
+                                        o2::gpu::gpudatatypes::DeviceType dType = o2::gpu::gpudatatypes::DeviceType::CPU,
+                                        int trackingThreads = 1);
 
 } // namespace o2::trk::global_reco_workflow
 

Detectors/Upgrades/ALICE3/GlobalReconstruction/workflow/include/ALICE3GlobalReconstructionWorkflow/TrackerSpec.h

@@ -45,7 +45,8 @@ class TrackerDPL : public framework::Task
              bool isMC,
              const std::string& hitRecoConfig,
              const std::string& clusterRecoConfig,
-             gpu::gpudatatypes::DeviceType dType = gpu::gpudatatypes::DeviceType::CPU);
+             gpu::gpudatatypes::DeviceType dType = gpu::gpudatatypes::DeviceType::CPU,
+             int trackingThreads = 1);
   ~TrackerDPL() override = default;
   void init(framework::InitContext& ic) final;
   void run(framework::ProcessingContext& pc) final;
@@ -67,6 +68,7 @@ class TrackerDPL : public framework::Task
   //   ITSTrackingInterface mITSTrackingInterface;
   bool mIsMC{true};
   gpu::gpudatatypes::DeviceType mDeviceType{gpu::gpudatatypes::DeviceType::CPU};
+  int mTrackingThreads{1};
   std::shared_ptr<its::BoundedMemoryResource> mMemoryPool;
   std::shared_ptr<its::ExternalAllocator> mGPUAllocator;
   std::shared_ptr<tbb::task_arena> mTaskArena;
@@ -79,7 +81,7 @@ class TrackerDPL : public framework::Task
 #endif
 };
 
-framework::DataProcessorSpec getTrackerSpec(bool useMC, const std::string& hitRecoConfig, const std::string& clusterRecoConfig, gpu::gpudatatypes::DeviceType dType = gpu::gpudatatypes::DeviceType::CPU);
+framework::DataProcessorSpec getTrackerSpec(bool useMC, const std::string& hitRecoConfig, const std::string& clusterRecoConfig, gpu::gpudatatypes::DeviceType dType = gpu::gpudatatypes::DeviceType::CPU, int trackingThreads = 1);
 
 } // namespace o2::trk
 #endif /* O2_TRK_TRACKERDPL */

Detectors/Upgrades/ALICE3/GlobalReconstruction/workflow/src/RecoWorkflow.cxx

@@ -21,14 +21,15 @@ framework::WorkflowSpec getWorkflow(bool useMC,
                                     const std::string& hitRecoConfig,
                                     const std::string& clusterRecoConfig,
                                     bool disableRootOutput,
-                                    o2::gpu::gpudatatypes::DeviceType dtype)
+                                    o2::gpu::gpudatatypes::DeviceType dtype,
+                                    int trackingThreads)
 {
   framework::WorkflowSpec specs;
 
   if (!hitRecoConfig.empty() || !clusterRecoConfig.empty()) {
     LOG_IF(info, !hitRecoConfig.empty()) << "Using hit reco config from file " << hitRecoConfig;
     LOG_IF(info, !clusterRecoConfig.empty()) << "Using cluster reco config from file " << clusterRecoConfig;
-    specs.emplace_back(o2::trk::getTrackerSpec(useMC, hitRecoConfig, clusterRecoConfig, dtype));
+    specs.emplace_back(o2::trk::getTrackerSpec(useMC, hitRecoConfig, clusterRecoConfig, dtype, trackingThreads));
     if (!disableRootOutput) {
       specs.emplace_back(o2::trk::getTrackWriterSpec(useMC));
     }

Detectors/Upgrades/ALICE3/GlobalReconstruction/workflow/src/TrackerSpec.cxx

@@ -71,7 +71,8 @@ TrackerDPL::TrackerDPL(std::shared_ptr<o2::base::GRPGeomRequest> gr,
                        bool isMC,
                        const std::string& hitRecoConfigFileName,
                        const std::string& clusterRecoConfigFileName,
-                       o2::gpu::gpudatatypes::DeviceType dType)
+                       o2::gpu::gpudatatypes::DeviceType dType,
+                       int trackingThreads)
 {
   if (!hitRecoConfigFileName.empty()) {
     std::ifstream configFile(hitRecoConfigFileName);
@@ -83,6 +84,7 @@ TrackerDPL::TrackerDPL(std::shared_ptr<o2::base::GRPGeomRequest> gr,
   }
   mIsMC = isMC;
   mDeviceType = dType;
+  mTrackingThreads = std::max(1, trackingThreads);
 }
 
 void TrackerDPL::init(InitContext& ic)
@@ -103,6 +105,15 @@ std::vector<o2::its::TrackingParameters> TrackerDPL::createTrackingParamsFromCon
   auto loadTrackingParamsFromJson = [](std::vector<o2::its::TrackingParameters>& trackingParams, const nlohmann::json& paramConfigJson) {
     for (const auto& paramConfig : paramConfigJson) {
       o2::its::TrackingParameters params;
+      auto applyPassFlag = [&](const char* name, o2::its::IterationStep step) {
+        if (paramConfig.contains(name)) {
+          if (paramConfig[name].get<bool>()) {
+            params.PassFlags.set(step);
+          } else {
+            params.PassFlags.reset(step);
+          }
+        }
+      };
 
       if (paramConfig.contains("NLayers")) {
         params.NLayers = paramConfig["NLayers"].get<int>();
@@ -172,6 +183,37 @@ std::vector<o2::its::TrackingParameters> TrackerDPL::createTrackingParamsFromCon
       if (paramConfig.contains("CreateArtefactLabels")) {
         params.CreateArtefactLabels = paramConfig["CreateArtefactLabels"].get<bool>();
       }
+      if (paramConfig.contains("TrackFollower")) {
+        const auto mode = paramConfig["TrackFollower"].get<std::string>();
+        if (mode == "top" || mode == "outward") {
+          params.PassFlags.set(o2::its::IterationStep::TrackFollowerTop);
+        } else if (mode == "bot" || mode == "bottom" || mode == "inward") {
+          params.PassFlags.set(o2::its::IterationStep::TrackFollowerBot);
+        } else if (mode == "mix" || mode == "both") {
+          params.PassFlags.set(o2::its::IterationStep::TrackFollowerTop);
+          params.PassFlags.set(o2::its::IterationStep::TrackFollowerBot);
+        } else if (mode != "off") {
+          LOGP(fatal, "Invalid ALICE3 TRK tracking parameter TrackFollower: {}", mode);
+        }
+      }
+      if (paramConfig.contains("TrackFollowerNSigmaCutZ")) {
+        params.TrackFollowerNSigmaCutZ = paramConfig["TrackFollowerNSigmaCutZ"].get<float>();
+      }
+      if (paramConfig.contains("TrackFollowerNSigmaCutPhi")) {
+        params.TrackFollowerNSigmaCutPhi = paramConfig["TrackFollowerNSigmaCutPhi"].get<float>();
+      }
+      if (paramConfig.contains("TrackFollowerMaxHypotheses")) {
+        params.TrackFollowerMaxHypotheses = std::max(1, paramConfig["TrackFollowerMaxHypotheses"].get<int>());
+      }
+      applyPassFlag("FirstPass", o2::its::IterationStep::FirstPass);
+      applyPassFlag("RebuildClusterLUT", o2::its::IterationStep::RebuildClusterLUT);
+      applyPassFlag("UseUPCMask", o2::its::IterationStep::UseUPCMask);
+      applyPassFlag("SelectUPCVertices", o2::its::IterationStep::SelectUPCVertices);
+      applyPassFlag("ResetVertices", o2::its::IterationStep::ResetVertices);
+      applyPassFlag("SkipROFsAboveThreshold", o2::its::IterationStep::SkipROFsAboveThreshold);
+      applyPassFlag("MarkVerticesAsUPC", o2::its::IterationStep::MarkVerticesAsUPC);
+      applyPassFlag("TrackFollowerTop", o2::its::IterationStep::TrackFollowerTop);
+      applyPassFlag("TrackFollowerBot", o2::its::IterationStep::TrackFollowerBot);
       if (paramConfig.contains("PrintMemory")) {
         params.PrintMemory = paramConfig["PrintMemory"].get<bool>();
       }
@@ -255,7 +297,7 @@ void TrackerDPL::run(ProcessingContext& pc)
     mMemoryPool = std::make_shared<its::BoundedMemoryResource>();
   }
   if (mTaskArena.get() == nullptr) {
-    mTaskArena = std::make_shared<tbb::task_arena>(1); /// TODO: make it configurable
+    mTaskArena = std::make_shared<tbb::task_arena>(mTrackingThreads);
   }
 
   mTrackingParams = createTrackingParamsFromConfig();
@@ -309,7 +351,7 @@ void TrackerDPL::endOfStream(EndOfStreamContext& ec)
   LOGF(info, "TRK CA-Tracker total timing: Cpu: %.3e Real: %.3e s in %d slots", mTimer.CpuTime(), mTimer.RealTime(), mTimer.Counter() - 1);
 }
 
-DataProcessorSpec getTrackerSpec(bool useMC, const std::string& hitRecoConfig, const std::string& clusterRecoConfig, o2::gpu::gpudatatypes::DeviceType dType)
+DataProcessorSpec getTrackerSpec(bool useMC, const std::string& hitRecoConfig, const std::string& clusterRecoConfig, o2::gpu::gpudatatypes::DeviceType dType, int trackingThreads)
 {
   std::vector<InputSpec> inputs;
   std::vector<OutputSpec> outputs;
@@ -337,7 +379,8 @@ DataProcessorSpec getTrackerSpec(bool useMC, const std::string& hitRecoConfig, c
                                               useMC,
                                               hitRecoConfig,
                                               clusterRecoConfig,
-                                              dType)},
+                                              dType,
+                                              trackingThreads)},
       Options{ConfigParamSpec{"max-loops", VariantType::Int, 1, {"max number of loops"}}
 #ifdef O2_WITH_ACTS
               ,
@@ -373,7 +416,8 @@ DataProcessorSpec getTrackerSpec(bool useMC, const std::string& hitRecoConfig, c
                                             useMC,
                                             hitRecoConfig,
                                             clusterRecoConfig,
-                                            dType)},
+                                            dType,
+                                            trackingThreads)},
     Options{}};
 }
 

Detectors/Upgrades/ALICE3/GlobalReconstruction/workflow/src/alice3-global-reconstruction-workflow.cxx

@@ -39,7 +39,8 @@ void customize(std::vector<ConfigParamSpec>& workflowOptions)
     {"tracking-from-hits-config", VariantType::String, "", {"JSON file with tracking from hits configuration"}},
     {"tracking-from-clusters-config", VariantType::String, "", {"JSON file with tracking from clusters configuration"}},
     {"configKeyValues", VariantType::String, "", {"Semicolon separated key=value strings"}},
-    {"gpu-device", VariantType::Int, 1, {"use gpu device: CPU=1,CUDA=2,HIP=3 (default: CPU)"}}};
+    {"gpu-device", VariantType::Int, 1, {"use gpu device: CPU=1,CUDA=2,HIP=3 (default: CPU)"}},
+    {"tracking-threads", VariantType::Int, 1, {"number of CPU threads used by TRK tracking"}}};
   std::swap(workflowOptions, options);
 }
 
@@ -52,6 +53,7 @@ WorkflowSpec defineDataProcessing(ConfigContext const& configcontext)
   auto hitRecoConfig = configcontext.options().get<std::string>("tracking-from-hits-config");
   auto clusterRecoConfig = configcontext.options().get<std::string>("tracking-from-clusters-config");
   auto gpuDevice = static_cast<o2::gpu::gpudatatypes::DeviceType>(configcontext.options().get<int>("gpu-device"));
+  auto trackingThreads = configcontext.options().get<int>("tracking-threads");
   auto disableRootOutput = configcontext.options().get<bool>("disable-root-output");
   o2::conf::ConfigurableParam::updateFromString(configcontext.options().get<std::string>("configKeyValues"));
 
@@ -61,5 +63,5 @@ WorkflowSpec defineDataProcessing(ConfigContext const& configcontext)
 
   o2::conf::ConfigurableParam::writeINI("o2alice3globalrecoflow_configuration.ini");
 
-  return o2::trk::global_reco_workflow::getWorkflow(useMC, hitRecoConfig, clusterRecoConfig, disableRootOutput, gpuDevice);
+  return o2::trk::global_reco_workflow::getWorkflow(useMC, hitRecoConfig, clusterRecoConfig, disableRootOutput, gpuDevice, trackingThreads);
 }