Compiles on EPNs. Need to add shadow processors next. But for this, I will merge #14069 to have the changes in GPUChainTrackingClusterizer.

ChSonnabend · ChSonnabend · commit 70320c3afce4 · 2025-03-29T13:39:27.000+01:00
diff --git a/Common/ML/include/ML/OrtInterface.h b/Common/ML/include/ML/OrtInterface.h
@@ -26,6 +26,11 @@
 // O2 includes
 #include "Framework/Logger.h"
 
+namespace Ort {
+  struct SessionOptions;
+  struct MemoryInfo;
+}
+
 namespace o2
 {
 
@@ -42,6 +47,8 @@ class OrtModel
   void init(std::unordered_map<std::string, std::string> optionsMap) { reset(optionsMap); }
   void reset(std::unordered_map<std::string, std::string>);
   bool isInitialized() { return mInitialized; }
+  Ort::SessionOptions* updateSessionOptions();
+  Ort::MemoryInfo* updateMemoryInfo();
 
   virtual ~OrtModel() = default;
 
diff --git a/Common/ML/src/OrtInterface.cxx b/Common/ML/src/OrtInterface.cxx
@@ -35,6 +35,16 @@ struct OrtModel::OrtVariables { // The actual implementation is hidden in the .c
   Ort::MemoryInfo memoryInfo = Ort::MemoryInfo("Cpu", OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemType::OrtMemTypeDefault);
 };
 
+Ort::SessionOptions* OrtModel::updateSessionOptions()
+{
+  return &(pImplOrt->sessionOptions);
+}
+
+Ort::MemoryInfo* OrtModel::updateMemoryInfo()
+{
+  return &(pImplOrt->memoryInfo);
+}
+
 void OrtModel::reset(std::unordered_map<std::string, std::string> optionsMap)
 {
 
@@ -56,39 +66,41 @@ void OrtModel::reset(std::unordered_map<std::string, std::string> optionsMap)
     enableProfiling = (optionsMap.contains("enable-profiling") ? std::stoi(optionsMap["enable-profiling"]) : 0);
     enableOptimizations = (optionsMap.contains("enable-optimizations") ? std::stoi(optionsMap["enable-optimizations"]) : 0);
 
-    std::string dev_mem_str = "Hip";
-#if defined(ORT_ROCM_BUILD)
-#if ORT_ROCM_BUILD == 1
-  if (device == "ROCM") {
-    // Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_ROCM(pImplOrt->sessionOptions, streamId));
-    o2::gpu::SetONNXGPUStream(pImplOrt->sessionOptions, streamId);
-    LOG(info) << "(ORT) ROCM execution provider set";
-  }
-#endif
-#endif
-#if defined(ORT_MIGRAPHX_BUILD)
-#if ORT_MIGRAPHX_BUILD == 1
-  if (device == "MIGRAPHX") {
-    Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_MIGraphX(pImplOrt->sessionOptions, streamId));
-    LOG(info) << "(ORT) MIGraphX execution provider set";
-  }
-#endif
-#endif
-#if defined(ORT_CUDA_BUILD)
-#if ORT_CUDA_BUILD == 1
-  if (device == "CUDA") {
-    // Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(pImplOrt->sessionOptions, streamId));
-    o2::gpu::SetONNXGPUStream(pImplOrt->sessionOptions, streamId);
-    LOG(info) << "(ORT) CUDA execution provider set";
-    dev_mem_str = "Cuda";
-  }
-#endif
-#endif
-
+// #if defined(ORT_ROCM_BUILD) && ORT_ROCM_BUILD == 1
+//   if (device == "ROCM") {
+//     // Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_ROCM(pImplOrt->sessionOptions, streamId));
+//     SetONNXGPUStream(pImplOrt->sessionOptions, streamId);
+//     LOG(info) << "(ORT) ROCM execution provider set";
+//   }
+// #endif
+// #if defined(ORT_MIGRAPHX_BUILD) && ORT_MIGRAPHX_BUILD == 1
+//   if (device == "MIGRAPHX") {
+//     Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_MIGraphX(pImplOrt->sessionOptions, streamId));
+//     LOG(info) << "(ORT) MIGraphX execution provider set";
+//   }
+// #endif
+// #if defined(ORT_CUDA_BUILD) && ORT_CUDA_BUILD == 1
+//   if (device == "CUDA") {
+//     // Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(pImplOrt->sessionOptions, streamId));
+//     SetONNXGPUStream(pImplOrt->sessionOptions, streamId);
+//     LOG(info) << "(ORT) CUDA execution provider set";
+//     dev_mem_str = "Cuda";
+//   }
+// #endif
+
+#if (defined(ORT_ROCM_BUILD) && ORT_ROCM_BUILD == 1) || (defined(ORT_MIGRAPHX_BUILD) && ORT_MIGRAPHX_BUILD == 1) || (defined(ORT_CUDA_BUILD) && ORT_CUDA_BUILD == 1)
   if (allocateDeviceMemory) {
+    std::string dev_mem_str = "";
+    if (device == "ROCM") {
+      dev_mem_str = "Hip";
+    }
+    if (device == "CUDA") {
+      dev_mem_str = "Cuda";
+    }
     pImplOrt->memoryInfo = Ort::MemoryInfo(dev_mem_str.c_str(), OrtAllocatorType::OrtDeviceAllocator, streamId, OrtMemType::OrtMemTypeDefault);
     LOG(info) << "(ORT) Memory info set to on-device memory";
   }
+#endif
 
   if (device == "CPU") {
     (pImplOrt->sessionOptions).SetIntraOpNumThreads(intraOpNumThreads);
diff --git a/GPU/GPUTracking/Base/GPUReconstructionCPU.h b/GPU/GPUTracking/Base/GPUReconstructionCPU.h
@@ -24,6 +24,10 @@
 #include "GPUReconstructionKernelIncludes.h"
 #include "GPUReconstructionKernels.h"
 
+namespace Ort {
+  struct SessionOptions;
+}
+
 namespace o2::gpu
 {
 
@@ -111,6 +115,9 @@ class GPUReconstructionCPU : public GPUReconstructionKernels<GPUReconstructionCP
   size_t WriteToConstantMemory(size_t offset, const void* src, size_t size, int32_t stream = -1, deviceEvent* ev = nullptr) override;
   virtual size_t TransferMemoryInternal(GPUMemoryResource* res, int32_t stream, deviceEvent* ev, deviceEvent* evList, int32_t nEvents, bool toGPU, const void* src, void* dst);
 
+  // ONNX runtime
+  virtual void SetONNXGPUStream(Ort::SessionOptions*, int32_t) {}
+
   int32_t InitDevice() override;
   int32_t ExitDevice() override;
   int32_t GetThread();
diff --git a/GPU/GPUTracking/Base/GPUReconstructionProcessing.h b/GPU/GPUTracking/Base/GPUReconstructionProcessing.h
@@ -90,7 +90,7 @@ class GPUReconstructionProcessing : public GPUReconstruction
   void AddGPUEvents(T*& events);
 
   virtual std::unique_ptr<gpu_reconstruction_kernels::threadContext> GetThreadContext() override;
-  virtual int32_t SetONNXGPUStream(OrtSessionOptions* session_options, int32_t stream) { return 0; }
+  virtual void SetONNXGPUStream(OrtSessionOptions* session_options, int32_t stream) {}
 
   struct RecoStepTimerMeta {
     HighResTimer timerToGPU;
diff --git a/GPU/GPUTracking/Base/cuda/GPUReconstructionCUDA.cu b/GPU/GPUTracking/Base/cuda/GPUReconstructionCUDA.cu
@@ -661,24 +661,23 @@ void GPUReconstructionCUDA::endGPUProfiling()
   GPUChkErr(cudaProfilerStop());
 }
 
-#ifdef GPUCA_HAS_ONNX
-int32_t GPUReconstructionCUDA::SetONNXGPUStream(OrtSessionOptions* session_options, int32_t stream)
+#if defined(ORT_CUDA_BUILD) && ORT_CUDA_BUILD == 1
+void GPUReconstructionCUDA::SetONNXGPUStream(Ort::SessionOptions* session_options, int32_t stream)
 {
   OrtCUDAProviderOptionsV2* cuda_options = nullptr;
   CreateCUDAProviderOptions(&cuda_options);
+  OrtSessionOptions* raw_options = session_options->operator OrtSessionOptions*();
 
   // std::vector<const char*> keys{"device_id", "gpu_mem_limit", "arena_extend_strategy", "cudnn_conv_algo_search", "do_copy_in_default_stream", "cudnn_conv_use_max_workspace", "cudnn_conv1d_pad_to_nc1d"};
   // std::vector<const char*> values{"0", "2147483648", "kSameAsRequested", "DEFAULT", "1", "1", "1"};
   // UpdateCUDAProviderOptions(cuda_options, keys.data(), values.data(), keys.size());
 
   // this implicitly sets "has_user_compute_stream"
   UpdateCUDAProviderOptionsWithValue(cuda_options, "user_compute_stream", &mInternals->Streams[stream]);
-  Ort::ThrowOnError(SessionOptionsAppendExecutionProvider_CUDA_V2(session_options, cuda_options));
+  Ort::ThrowOnError(SessionOptionsAppendExecutionProvider_CUDA_V2(raw_options, cuda_options));
 
   // Finally, don't forget to release the provider options
   ReleaseCUDAProviderOptions(cuda_options);
-
-  return 0;
 }
 #endif // GPUCA_HAS_ONNX
 
@@ -690,19 +689,22 @@ void* GPUReconstructionHIP::getGPUPointer(void* ptr)
   return retVal;
 }
 
-#ifdef GPUCA_HAS_ONNX
-int32_t GPUReconstructionHIP::SetONNXGPUStream(OrtSessionOptions* session_options, int32_t stream)
+#if defined(ORT_ROCM_BUILD) && ORT_ROCM_BUILD == 1
+void GPUReconstructionHIP::SetONNXGPUStream(Ort::SessionOptions* session_options, int32_t stream)
 {
   // Create ROCm provider options
   const auto& api = Ort::GetApi();
   OrtROCMProviderOptions rocm_options{};
   rocm_options.has_user_compute_stream = 1; // Indicate that we are passing a user stream
   rocm_options.user_compute_stream = &mInternals->Streams[stream];
 
+  // Get the raw OrtSessionOptions pointer from the Ort::SessionOptions wrapper
+  OrtSessionOptions* raw_options = session_options->operator OrtSessionOptions*();
+
   // Append the ROCm execution provider with the custom HIP stream
-  Ort::ThrowOnError(api.SessionOptionsAppendExecutionProvider_ROCM(session_options, &rocm_options));
-  return 0;
+  Ort::ThrowOnError(api.SessionOptionsAppendExecutionProvider_ROCM(raw_options, &rocm_options));
 }
+
 #endif // GPUCA_HAS_ONNX
 #endif // __HIPCC__
 
diff --git a/GPU/GPUTracking/Base/cuda/GPUReconstructionCUDA.h b/GPU/GPUTracking/Base/cuda/GPUReconstructionCUDA.h
@@ -25,6 +25,10 @@ extern "C" __declspec(dllexport) o2::gpu::GPUReconstruction* GPUReconstruction_C
 extern "C" o2::gpu::GPUReconstruction* GPUReconstruction_Create_CUDA(const o2::gpu::GPUSettingsDeviceBackend& cfg);
 #endif
 
+namespace Ort {
+  struct SessionOptions;
+}
+
 namespace o2::gpu
 {
 struct GPUReconstructionCUDAInternals;
@@ -79,7 +83,7 @@ class GPUReconstructionCUDA : public GPUReconstructionKernels<GPUReconstructionC
   size_t GPUMemCpy(void* dst, const void* src, size_t size, int32_t stream, int32_t toGPU, deviceEvent* ev = nullptr, deviceEvent* evList = nullptr, int32_t nEvents = 1) override;
   void ReleaseEvent(deviceEvent ev) override;
   void RecordMarker(deviceEvent* ev, int32_t stream) override;
-  int32_t SetONNXGPUStream(OrtSessionOptions* session_options, int32_t stream) override;
+  void SetONNXGPUStream(Ort::SessionOptions* session_options, int32_t stream) override;
 
   void GetITSTraits(std::unique_ptr<o2::its::TrackerTraits>* trackerTraits, std::unique_ptr<o2::its::VertexerTraits>* vertexerTraits, std::unique_ptr<o2::its::TimeFrame>* timeFrame) override;
 
diff --git a/GPU/GPUTracking/CMakeLists.txt b/GPU/GPUTracking/CMakeLists.txt
@@ -14,6 +14,20 @@ set(MODULE GPUTracking)
 # set(CMAKE_CXX_FLAGS_${CMAKE_BUILD_TYPE_UPPER} "${CMAKE_CXX_FLAGS_${CMAKE_BUILD_TYPE_UPPER}} -O0") # to uncomment if needed, tired of typing this...
 # set(GPUCA_BUILD_DEBUG 1)
 
+# Pass ORT variables as a preprocessor definition
+if(DEFINED ENV{ORT_ROCM_BUILD})
+    add_compile_definitions(ORT_ROCM_BUILD=$ENV{ORT_ROCM_BUILD})
+endif()
+if(DEFINED ENV{ORT_CUDA_BUILD})
+    add_compile_definitions(ORT_CUDA_BUILD=$ENV{ORT_CUDA_BUILD})
+endif()
+if(DEFINED ENV{ORT_MIGRAPHX_BUILD})
+    add_compile_definitions(ORT_MIGRAPHX_BUILD=$ENV{ORT_MIGRAPHX_BUILD})
+endif()
+if(DEFINED ENV{ORT_TENSORRT_BUILD})
+    add_compile_definitions(ORT_TENSORRT_BUILD=$ENV{ORT_TENSORRT_BUILD})
+endif()
+
 if(GPUCA_DETERMINISTIC_MODE GREATER_EQUAL ${GPUCA_DETERMINISTIC_MODE_MAP_NO_FAST_MATH})
   set(CMAKE_CXX_FLAGS_${CMAKE_BUILD_TYPE_UPPER} "${CMAKE_CXX_FLAGS_${CMAKE_BUILD_TYPE_UPPER}} ${GPUCA_CXX_NO_FAST_MATH_FLAGS}")
   if(GPUCA_DETERMINISTIC_MODE GREATER_EQUAL ${GPUCA_DETERMINISTIC_MODE_MAP_OPTO2})
diff --git a/GPU/GPUTracking/Global/GPUChain.h b/GPU/GPUTracking/Global/GPUChain.h
@@ -83,6 +83,7 @@ class GPUChain
   inline GPUParam& param() { return mRec->param(); }
   inline const GPUConstantMem* processors() const { return mRec->processors(); }
   inline void SynchronizeStream(int32_t stream) { mRec->SynchronizeStream(stream); }
+  inline void SetONNXGPUStream(Ort::SessionOptions* opt, int32_t stream) { mRec->SetONNXGPUStream(opt, stream); }
   inline void SynchronizeEvents(deviceEvent* evList, int32_t nEvents = 1) { mRec->SynchronizeEvents(evList, nEvents); }
   inline void SynchronizeEventAndRelease(deviceEvent& ev, bool doGPU = true)
   {
diff --git a/GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx b/GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx
@@ -917,6 +917,9 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
           GPUTPCNNClusterizer& clustererNN = processors()->tpcNNClusterer[iSector];
           const GPUSettingsProcessingNNclusterizer& nn_settings = GetProcessingSettings().nn;
           GPUTPCNNClusterizerHost nnApplication(nn_settings, clustererNN, lane);
+          SetONNXGPUStream(nnApplication.model_class.updateSessionOptions(), lane);
+          SetONNXGPUStream(nnApplication.model_reg_1.updateSessionOptions(), lane);
+          SetONNXGPUStream(nnApplication.model_reg_2.updateSessionOptions(), lane);
 
           if (clustererNN.nnClusterizerUseCfRegression || (int)(nn_settings.nnClusterizerApplyCfDeconvolution)) {
             runKernel<GPUTPCCFDeconvolution>({GetGrid(clusterer.mPmemory->counters.nPositions, lane), {iSector}});

Original file line number	Diff line number	Diff line change
`@@ -83,6 +83,7 @@ class GPUChain`
`83`	`83`	`inline GPUParam& param() { return mRec->param(); }`
`84`	`84`	`inline const GPUConstantMem* processors() const { return mRec->processors(); }`
`85`	`85`	`inline void SynchronizeStream(int32_t stream) { mRec->SynchronizeStream(stream); }`
	`86`	`+ inline void SetONNXGPUStream(Ort::SessionOptions* opt, int32_t stream) { mRec->SetONNXGPUStream(opt, stream); }`
`86`	`87`	`inline void SynchronizeEvents(deviceEvent* evList, int32_t nEvents = 1) { mRec->SynchronizeEvents(evList, nEvents); }`
`87`	`88`	`inline void SynchronizeEventAndRelease(deviceEvent& ev, bool doGPU = true)`
`88`	`89`	`{`