This runs, but will eventually fill up the VRAM. Need to include a mem clean

Author: ChSonnabend

Common/ML/CMakeLists.txt

@@ -10,18 +10,10 @@
 # or submit itself to any jurisdiction.
 
 # 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()
+add_compile_definitions(ORT_ROCM_BUILD=${ORT_ROCM_BUILD})
+add_compile_definitions(ORT_CUDA_BUILD=${ORT_CUDA_BUILD})
+add_compile_definitions(ORT_MIGRAPHX_BUILD=${ORT_MIGRAPHX_BUILD})
+add_compile_definitions(ORT_TENSORRT_BUILD=${ORT_TENSORRT_BUILD})
 
 o2_add_library(ML
                SOURCES src/OrtInterface.cxx

Common/ML/include/ML/OrtInterface.h

@@ -43,12 +43,19 @@ class OrtModel
  public:
   // Constructor
   OrtModel() = default;
-  OrtModel(std::unordered_map<std::string, std::string> optionsMap) { reset(optionsMap); }
-  void init(std::unordered_map<std::string, std::string> optionsMap) { reset(optionsMap); }
-  void reset(std::unordered_map<std::string, std::string>);
+  OrtModel(std::unordered_map<std::string, std::string> optionsMap) {
+    initOptions(optionsMap);
+    initEnvironment();
+  }
+  void init(std::unordered_map<std::string, std::string> optionsMap) {
+    initOptions(optionsMap);
+    initEnvironment();
+  }
+  void initOptions(std::unordered_map<std::string, std::string> optionsMap);
+  void initEnvironment();
   bool isInitialized() { return mInitialized; }
-  Ort::SessionOptions* updateSessionOptions();
-  Ort::MemoryInfo* updateMemoryInfo();
+  Ort::SessionOptions& updateSessionOptions();
+  void setIO();
 
   virtual ~OrtModel() = default;
 
@@ -91,7 +98,7 @@ class OrtModel
 
   // Environment settings
   bool mInitialized = false;
-  std::string modelPath, device = "cpu", thread_affinity = ""; // device options should be cpu, rocm, migraphx, cuda
+  std::string modelPath, envName = "", device = "cpu", thread_affinity = ""; // device options should be cpu, rocm, migraphx, cuda
   int intraOpNumThreads = 1, interOpNumThreads = 1, deviceId = 0, enableProfiling = 0, loggingLevel = 0, allocateDeviceMemory = 0, enableOptimizations = 0;
 
   std::string printShape(const std::vector<int64_t>&);

Common/ML/src/OrtInterface.cxx

@@ -35,19 +35,13 @@ 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()
+Ort::SessionOptions& OrtModel::updateSessionOptions()
 {
-  return &(pImplOrt->sessionOptions);
+  return pImplOrt->sessionOptions;
 }
 
-Ort::MemoryInfo* OrtModel::updateMemoryInfo()
+void OrtModel::initOptions(std::unordered_map<std::string, std::string> optionsMap)
 {
-  return &(pImplOrt->memoryInfo);
-}
-
-void OrtModel::reset(std::unordered_map<std::string, std::string> optionsMap)
-{
-
   pImplOrt = new OrtVariables();
 
   // Load from options map
@@ -58,71 +52,57 @@ void OrtModel::reset(std::unordered_map<std::string, std::string> optionsMap)
   if (!optionsMap["model-path"].empty()) {
     modelPath = optionsMap["model-path"];
     device = (optionsMap.contains("device") ? optionsMap["device"] : "CPU");
-    deviceId = (optionsMap.contains("device-id") ? std::stoi(optionsMap["device-id"]) : 0);
     allocateDeviceMemory = (optionsMap.contains("allocate-device-memory") ? std::stoi(optionsMap["allocate-device-memory"]) : 0);
     intraOpNumThreads = (optionsMap.contains("intra-op-num-threads") ? std::stoi(optionsMap["intra-op-num-threads"]) : 0);
     interOpNumThreads = (optionsMap.contains("inter-op-num-threads") ? std::stoi(optionsMap["inter-op-num-threads"]) : 0);
     loggingLevel = (optionsMap.contains("logging-level") ? std::stoi(optionsMap["logging-level"]) : 0);
     enableProfiling = (optionsMap.contains("enable-profiling") ? std::stoi(optionsMap["enable-profiling"]) : 0);
     enableOptimizations = (optionsMap.contains("enable-optimizations") ? std::stoi(optionsMap["enable-optimizations"]) : 0);
-
-// #if defined(ORT_ROCM_BUILD) && ORT_ROCM_BUILD == 1
-//   if (device == "ROCM") {
-//     // Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_ROCM(pImplOrt->sessionOptions, deviceId));
-//     SetONNXGPUStream(pImplOrt->sessionOptions, deviceId);
-//     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, deviceId));
-//     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, deviceId));
-//     SetONNXGPUStream(pImplOrt->sessionOptions, deviceId);
-//     LOG(info) << "(ORT) CUDA execution provider set";
-//     dev_mem_str = "Cuda";
-//   }
-// #endif
-
-  if (device == "CPU") {
-    (pImplOrt->sessionOptions).SetIntraOpNumThreads(intraOpNumThreads);
-    (pImplOrt->sessionOptions).SetInterOpNumThreads(interOpNumThreads);
-    if (intraOpNumThreads > 1 || interOpNumThreads > 1) {
-      (pImplOrt->sessionOptions).SetExecutionMode(ExecutionMode::ORT_PARALLEL);
-    } else if (intraOpNumThreads == 1) {
-      (pImplOrt->sessionOptions).SetExecutionMode(ExecutionMode::ORT_SEQUENTIAL);
-    }
-    if (loggingLevel < 2) {
-      LOG(info) << "(ORT) CPU execution provider set with " << intraOpNumThreads << " (intraOpNumThreads) and " << interOpNumThreads << " (interOpNumThreads) threads";
+    envName = (optionsMap.contains("onnx-environment-name") ? optionsMap["onnx-environment-name"] : "onnx_model_inference");
+
+    if (device == "CPU") {
+      (pImplOrt->sessionOptions).SetIntraOpNumThreads(intraOpNumThreads);
+      (pImplOrt->sessionOptions).SetInterOpNumThreads(interOpNumThreads);
+      if (intraOpNumThreads > 1 || interOpNumThreads > 1) {
+        (pImplOrt->sessionOptions).SetExecutionMode(ExecutionMode::ORT_PARALLEL);
+      } else if (intraOpNumThreads == 1) {
+        (pImplOrt->sessionOptions).SetExecutionMode(ExecutionMode::ORT_SEQUENTIAL);
+      }
+      if (loggingLevel < 2) {
+        LOG(info) << "(ORT) CPU execution provider set with " << intraOpNumThreads << " (intraOpNumThreads) and " << interOpNumThreads << " (interOpNumThreads) threads";
+      }
     }
-  }
 
-  (pImplOrt->sessionOptions).DisableMemPattern();
-  (pImplOrt->sessionOptions).DisableCpuMemArena();
+    // OrtROCMProviderOptions rocm_options{};
+    // (pImplOrt->sessionOptions).AppendExecutionProvider_ROCM(rocm_options);
 
-  if (enableProfiling) {
-    if (optionsMap.contains("profiling-output-path")) {
-      (pImplOrt->sessionOptions).EnableProfiling((optionsMap["profiling-output-path"] + "/ORT_LOG_").c_str());
+    (pImplOrt->sessionOptions).DisableMemPattern();
+    (pImplOrt->sessionOptions).DisableCpuMemArena();
+
+    if (enableProfiling) {
+      if (optionsMap.contains("profiling-output-path")) {
+        (pImplOrt->sessionOptions).EnableProfiling((optionsMap["profiling-output-path"] + "/ORT_LOG_").c_str());
+      } else {
+        LOG(warning) << "(ORT) If profiling is enabled, optionsMap[\"profiling-output-path\"] should be set. Disabling profiling for now.";
+        (pImplOrt->sessionOptions).DisableProfiling();
+      }
     } else {
-      LOG(warning) << "(ORT) If profiling is enabled, optionsMap[\"profiling-output-path\"] should be set. Disabling profiling for now.";
       (pImplOrt->sessionOptions).DisableProfiling();
     }
+
+    (pImplOrt->sessionOptions).SetGraphOptimizationLevel(GraphOptimizationLevel(enableOptimizations));
+    (pImplOrt->sessionOptions).SetLogSeverityLevel(OrtLoggingLevel(loggingLevel));
   } else {
-    (pImplOrt->sessionOptions).DisableProfiling();
+    LOG(fatal) << "(ORT) Model path cannot be empty!";
   }
+}
 
+void OrtModel::initEnvironment()
+{
   mInitialized = true;
-
-  (pImplOrt->sessionOptions).SetGraphOptimizationLevel(GraphOptimizationLevel(enableOptimizations));
-  (pImplOrt->sessionOptions).SetLogSeverityLevel(OrtLoggingLevel(loggingLevel));
-
   pImplOrt->env = std::make_shared<Ort::Env>(
     OrtLoggingLevel(loggingLevel),
-    (optionsMap["onnx-environment-name"].empty() ? "onnx_model_inference" : optionsMap["onnx-environment-name"].c_str()),
+    (envName.empty() ? "ORT" : envName.c_str()),
     // Integrate ORT logging into Fairlogger
     [](void* param, OrtLoggingLevel severity, const char* category, const char* logid, const char* code_location, const char* message) {
       if (severity == ORT_LOGGING_LEVEL_VERBOSE) {
@@ -143,6 +123,10 @@ void OrtModel::reset(std::unordered_map<std::string, std::string> optionsMap)
   (pImplOrt->env)->DisableTelemetryEvents(); // Disable telemetry events
   pImplOrt->session = std::make_shared<Ort::Session>(*(pImplOrt->env), modelPath.c_str(), pImplOrt->sessionOptions);
 
+  setIO();
+}
+
+void OrtModel::setIO() {
   for (size_t i = 0; i < (pImplOrt->session)->GetInputCount(); ++i) {
     mInputNames.push_back((pImplOrt->session)->GetInputNameAllocated(i, pImplOrt->allocator).get());
   }
@@ -162,7 +146,6 @@ void OrtModel::reset(std::unordered_map<std::string, std::string> optionsMap)
   outputNamesChar.resize(mOutputNames.size(), nullptr);
   std::transform(std::begin(mOutputNames), std::end(mOutputNames), std::begin(outputNamesChar),
                  [&](const std::string& str) { return str.c_str(); });
-  }
   if (loggingLevel < 2) {
     LOG(info) << "(ORT) Model loaded successfully! (input: " << printShape(mInputShapes[0]) << ", output: " << printShape(mOutputShapes[0]) << ")";
   }
@@ -203,18 +186,15 @@ std::vector<O> OrtModel::inference(std::vector<I>& input, int32_t deviceIndex)
 {
 #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) {
-    if (deviceIndex >= 0) {
-      deviceId = deviceIndex;
-    }
     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, deviceId, OrtMemType::OrtMemTypeDefault);
-    LOG(info) << "(ORT) Memory info set to on-device memory";
+    pImplOrt->memoryInfo = Ort::MemoryInfo(dev_mem_str.c_str(), OrtAllocatorType::OrtDeviceAllocator, deviceIndex, OrtMemType::OrtMemTypeDefault);
+    LOG(info) << "(ORT) Memory info set to on-device memory for device " << device << " with ID "<< deviceIndex;
   }
 #endif
   std::vector<int64_t> inputShape{(int64_t)(input.size() / mInputShapes[0][1]), (int64_t)mInputShapes[0][1]};
@@ -241,20 +221,21 @@ template std::vector<OrtDataType::Float16_t> OrtModel::inference<OrtDataType::Fl
 template <class I, class O>
 void OrtModel::inference(I* input, size_t input_size, O* output, int32_t deviceIndex)
 {
+  // std::vector<std::string> providers = Ort::GetAvailableProviders();
+  // for (const auto& provider : providers) {
+  //     LOG(info) << "Available Execution Provider: " << provider;
+  // }
 #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) {
-    if (deviceIndex >= 0) {
-      deviceId = deviceIndex;
-    }
     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, deviceId, OrtMemType::OrtMemTypeDefault);
-    LOG(info) << "(ORT) Memory info set to on-device memory";
+    pImplOrt->memoryInfo = Ort::MemoryInfo(dev_mem_str.c_str(), OrtAllocatorType::OrtDeviceAllocator, deviceIndex, OrtMemType::OrtMemTypeDefault);
+    LOG(info) << "(ORT) Memory info set to on-device memory for device " << device << " with ID "<< deviceIndex;
   }
 #endif
   std::vector<int64_t> inputShape{input_size, (int64_t)mInputShapes[0][1]};
@@ -268,7 +249,7 @@ void OrtModel::inference(I* input, size_t input_size, O* output, int32_t deviceI
   std::vector<int64_t> outputShape{input_size, mOutputShapes[0][1]};
   Ort::Value outputTensor = Ort::Value::CreateTensor<O>(pImplOrt->memoryInfo, output, input_size * mOutputShapes[0][1], outputShape.data(), outputShape.size());
 
-  (pImplOrt->session)->Run(pImplOrt->runOptions, inputNamesChar.data(), &inputTensor, 1, outputNamesChar.data(), &outputTensor, outputNamesChar.size()); // TODO: Not sure if 1 is always correct here
+  (pImplOrt->session)->Run(pImplOrt->runOptions, inputNamesChar.data(), &inputTensor, 1, outputNamesChar.data(), &outputTensor, outputNamesChar.size());
 }
 
 template void OrtModel::inference<OrtDataType::Float16_t, float>(OrtDataType::Float16_t*, size_t, float*, int32_t);
@@ -280,18 +261,15 @@ std::vector<O> OrtModel::inference(std::vector<std::vector<I>>& input, int32_t d
 {
 #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) {
-    if (deviceIndex >= 0) {
-       deviceId = deviceIndex;
-    }
     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, deviceId, OrtMemType::OrtMemTypeDefault);
-    LOG(info) << "(ORT) Memory info set to on-device memory";
+    pImplOrt->memoryInfo = Ort::MemoryInfo(dev_mem_str.c_str(), OrtAllocatorType::OrtDeviceAllocator, deviceIndex, OrtMemType::OrtMemTypeDefault);
+    LOG(info) << "(ORT) Memory info set to on-device memory for device " << device << " with ID " << deviceIndex;
   }
 #endif
   std::vector<Ort::Value> inputTensor;

GPU/GPUTracking/Base/GPUReconstructionCPU.h

@@ -116,7 +116,7 @@ class GPUReconstructionCPU : public GPUReconstructionKernels<GPUReconstructionCP
   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*) {}
+  virtual void SetONNXGPUStream(Ort::SessionOptions&, int32_t, int32_t*) {}
 
   int32_t InitDevice() override;
   int32_t ExitDevice() override;

GPU/GPUTracking/Base/GPUReconstructionProcessing.h

@@ -22,7 +22,9 @@
 #include <functional>
 #include <atomic>
 
-struct OrtSessionOptions;
+namespace Ort {
+struct SessionOptions;
+}
 
 namespace o2::gpu
 {
@@ -90,7 +92,7 @@ class GPUReconstructionProcessing : public GPUReconstruction
   void AddGPUEvents(T*& events);
 
   virtual std::unique_ptr<gpu_reconstruction_kernels::threadContext> GetThreadContext() override;
-  virtual void SetONNXGPUStream(OrtSessionOptions*, int32_t, int32_t*) {}
+  virtual void SetONNXGPUStream(Ort::SessionOptions&, int32_t, int32_t*) {}
 
   struct RecoStepTimerMeta {
     HighResTimer timerToGPU;

GPU/GPUTracking/Base/cuda/GPUReconstructionCUDA.cu

@@ -662,20 +662,19 @@ void GPUReconstructionCUDA::endGPUProfiling()
 }
 
 #if defined(ORT_CUDA_BUILD) && ORT_CUDA_BUILD == 1
-void GPUReconstructionCUDA::SetONNXGPUStream(Ort::SessionOptions* session_options, int32_t stream, int32_t* deviceId)
+void GPUReconstructionCUDA::SetONNXGPUStream(Ort::SessionOptions& session_options, int32_t stream, int32_t* deviceId)
 {
   cudaGetDevice(deviceId);
   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(raw_options, cuda_options));
+  session_options.AppendExecutionProvider_CUDA_V2(cuda_options);
 
   // Finally, don't forget to release the provider options
   ReleaseCUDAProviderOptions(cuda_options);
@@ -691,20 +690,23 @@ void* GPUReconstructionHIP::getGPUPointer(void* ptr)
 }
 
 #if defined(ORT_ROCM_BUILD) && ORT_ROCM_BUILD == 1
-void GPUReconstructionHIP::SetONNXGPUStream(Ort::SessionOptions* session_options, int32_t stream, int32_t* deviceId)
+void GPUReconstructionHIP::SetONNXGPUStream(Ort::SessionOptions& session_options, int32_t stream, int32_t* deviceId)
 {
   // Create ROCm provider options
   cudaGetDevice(deviceId);
   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(raw_options, &rocm_options));
+  // api.GetCurrentGpuDeviceId(deviceId);
+  OrtROCMProviderOptions rocm_options;
+  LOG(info) << "Creating ROCm provider options";
+  // rocm_options.has_user_compute_stream = 1; // Indicate that we are passing a user stream
+  // LOG(info) << "Setting user compute stream";
+  // rocm_options.user_compute_stream = &(mInternals->Streams[stream]);
+  // LOG(info) << "Stream is set with streamId " << stream << " and reference " << &(mInternals->Streams[stream]);
+  session_options.AppendExecutionProvider_ROCM(rocm_options);
+  LOG(info) << "Appending ROCm provider options";
+  // OrtSessionOptionsAppendExecutionProvider_ROCM(session_options, *deviceId);
+  // api.ReleaseROCMProviderOptions(rocm_options);
+  LOG(info) << "Releasing ROCm provider options";
 }
 
 #endif // GPUCA_HAS_ONNX

GPU/GPUTracking/Base/cuda/GPUReconstructionCUDA.h

@@ -83,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;
-  void SetONNXGPUStream(Ort::SessionOptions* session_options, int32_t stream, int32_t* deviceId) override;
+  void SetONNXGPUStream(Ort::SessionOptions& session_options, int32_t stream, int32_t* deviceId) 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;