Adding shadow instance. Not sure if this correctly allocates GPU memory using AllocateRegisteredMemory

ChSonnabend · ChSonnabend · commit 9d9267f6d9af · 2025-03-29T19:20:21.000+01:00
diff --git a/GPU/GPUTracking/Base/GPUReconstructionCPU.h b/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) {}
+  virtual void SetONNXGPUStream(Ort::SessionOptions*, int32_t, int32_t*) {}
 
   int32_t InitDevice() override;
   int32_t ExitDevice() override;
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 void SetONNXGPUStream(OrtSessionOptions* session_options, int32_t stream) {}
+  virtual void SetONNXGPUStream(OrtSessionOptions*, int32_t, int32_t*) {}
 
   struct RecoStepTimerMeta {
     HighResTimer timerToGPU;
diff --git a/GPU/GPUTracking/Base/cuda/GPUReconstructionCUDA.cu b/GPU/GPUTracking/Base/cuda/GPUReconstructionCUDA.cu
@@ -662,8 +662,9 @@ void GPUReconstructionCUDA::endGPUProfiling()
 }
 
 #if defined(ORT_CUDA_BUILD) && ORT_CUDA_BUILD == 1
-void GPUReconstructionCUDA::SetONNXGPUStream(Ort::SessionOptions* session_options, int32_t stream)
+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*();
@@ -690,9 +691,10 @@ void* GPUReconstructionHIP::getGPUPointer(void* ptr)
 }
 
 #if defined(ORT_ROCM_BUILD) && ORT_ROCM_BUILD == 1
-void GPUReconstructionHIP::SetONNXGPUStream(Ort::SessionOptions* session_options, int32_t stream)
+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
diff --git a/GPU/GPUTracking/Base/cuda/GPUReconstructionCUDA.h b/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) 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;
 
diff --git a/GPU/GPUTracking/Global/GPUChain.h b/GPU/GPUTracking/Global/GPUChain.h
@@ -83,7 +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 SetONNXGPUStream(Ort::SessionOptions* opt, int32_t stream, int32_t* deviceId) { mRec->SetONNXGPUStream(opt, stream, deviceId); }
   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
@@ -622,28 +622,45 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
     }
     for (uint32_t iSector = 0; iSector < NSECTORS; iSector++) {
       GPUTPCNNClusterizer& clustererNN = processors()->tpcNNClusterer[iSector];
-      clustererNN.nnClusterizerUseCfRegression = nn_settings.nnClusterizerUseCfRegression;
-      clustererNN.nnClusterizerSizeInputRow = nn_settings.nnClusterizerSizeInputRow;
-      clustererNN.nnClusterizerSizeInputPad = nn_settings.nnClusterizerSizeInputPad;
-      clustererNN.nnClusterizerSizeInputTime = nn_settings.nnClusterizerSizeInputTime;
-      clustererNN.nnClusterizerAddIndexData = nn_settings.nnClusterizerAddIndexData;
-      clustererNN.nnClusterizerElementSize = ((2 * nn_settings.nnClusterizerSizeInputRow + 1) * (2 * nn_settings.nnClusterizerSizeInputPad + 1) * (2 * nn_settings.nnClusterizerSizeInputTime + 1)) + (nn_settings.nnClusterizerAddIndexData ? 3 : 0);
-      clustererNN.nnClusterizerBatchedMode = nn_settings.nnClusterizerBatchedMode;
-      clustererNN.nnClusterizerBoundaryFillValue = nn_settings.nnClusterizerBoundaryFillValue;
-      clustererNN.nnClusterizerTotalClusters = maxClusters;
-      clustererNN.nnClassThreshold = nn_settings.nnClassThreshold;
-      clustererNN.nnSigmoidTrafoClassThreshold = nn_settings.nnSigmoidTrafoClassThreshold;
-      if (clustererNN.nnSigmoidTrafoClassThreshold) {
-        clustererNN.nnClassThreshold = (float)std::log(clustererNN.nnClassThreshold / (1.f - clustererNN.nnClassThreshold));
+      GPUTPCNNClusterizer& clustererNNShadow = doGPU ? processorsShadow()->tpcNNClusterer[iSector] : clustererNN;
+      clustererNNShadow.nnClusterizerUseCfRegression = nn_settings.nnClusterizerUseCfRegression;
+      clustererNNShadow.nnClusterizerSizeInputRow = nn_settings.nnClusterizerSizeInputRow;
+      clustererNNShadow.nnClusterizerSizeInputPad = nn_settings.nnClusterizerSizeInputPad;
+      clustererNNShadow.nnClusterizerSizeInputTime = nn_settings.nnClusterizerSizeInputTime;
+      clustererNNShadow.nnClusterizerAddIndexData = nn_settings.nnClusterizerAddIndexData;
+      clustererNNShadow.nnClusterizerElementSize = ((2 * nn_settings.nnClusterizerSizeInputRow + 1) * (2 * nn_settings.nnClusterizerSizeInputPad + 1) * (2 * nn_settings.nnClusterizerSizeInputTime + 1)) + (nn_settings.nnClusterizerAddIndexData ? 3 : 0);
+      clustererNNShadow.nnClusterizerBatchedMode = nn_settings.nnClusterizerBatchedMode;
+      clustererNNShadow.nnClusterizerBoundaryFillValue = nn_settings.nnClusterizerBoundaryFillValue;
+      clustererNNShadow.nnClusterizerTotalClusters = maxClusters;
+      clustererNNShadow.nnClassThreshold = nn_settings.nnClassThreshold;
+      clustererNNShadow.nnSigmoidTrafoClassThreshold = nn_settings.nnSigmoidTrafoClassThreshold;
+      if (clustererNNShadow.nnSigmoidTrafoClassThreshold) {
+        clustererNNShadow.nnClassThreshold = (float)std::log(clustererNNShadow.nnClassThreshold / (1.f - clustererNNShadow.nnClassThreshold));
       }
       if (nn_settings.nnClusterizerVerbosity < 0) {
-        clustererNN.nnClusterizerVerbosity = nn_settings.nnInferenceVerbosity;
+        clustererNNShadow.nnClusterizerVerbosity = nn_settings.nnInferenceVerbosity;
       } else {
-        clustererNN.nnClusterizerVerbosity = nn_settings.nnClusterizerVerbosity;
+        clustererNNShadow.nnClusterizerVerbosity = nn_settings.nnClusterizerVerbosity;
       }
-      clustererNN.nnInferenceInputDType = nn_settings.nnInferenceInputDType.find("32") != std::string::npos;
-      nnApplication.initClusterizer(nn_settings, clustererNN);
-      AllocateRegisteredMemory(clustererNN.mMemoryId);
+      clustererNNShadow.nnInferenceInputDType = nn_settings.nnInferenceInputDType.find("32") != std::string::npos;
+      nnApplication.initClusterizer(nn_settings, clustererNNShadow);
+      // if (doGPU) {
+      //   std::vector<int32_t> pointerSizes = clustererNNShadow.pointerSizes();
+      //   // FIXME: These are for sure not needed. The arrays are empty at this point, only the space needs to be reserved. Is this already handeled by computePointerWithAlignment?
+      //   // Once a GPU is available, everything should be done on the GPU for now.
+      //   GPUMemCpy(RecoStep::TPCClusterFinding, clustererNNShadow.inputData32, clustererNN.inputData32, pointerSizes[0], lane, true);
+      //   GPUMemCpy(RecoStep::TPCClusterFinding, clustererNNShadow.inputData16, clustererNN.inputData16, pointerSizes[1], lane, true);
+      //   GPUMemCpy(RecoStep::TPCClusterFinding, clustererNNShadow.outputDataClass, clustererNN.outputDataClass, pointerSizes[2], lane, true);
+      //   GPUMemCpy(RecoStep::TPCClusterFinding, clustererNNShadow.modelProbabilities, clustererNN.modelProbabilities, pointerSizes[3], lane, true);
+      //   GPUMemCpy(RecoStep::TPCClusterFinding, clustererNNShadow.outputDataReg1, clustererNN.outputDataReg1, pointerSizes[4], lane, true);
+      //   GPUMemCpy(RecoStep::TPCClusterFinding, clustererNNShadow.outputDataReg2, clustererNN.outputDataReg2, pointerSizes[5], lane, true);
+      //   GPUMemCpy(RecoStep::TPCClusterFinding, clustererNNShadow.peakPositions, clustererNN.peakPositions, pointerSizes[6], lane, true);
+      //   GPUMemCpy(RecoStep::TPCClusterFinding, clustererNNShadow.clusterFlags, clustererNN.clusterFlags, pointerSizes[7], lane, true);
+      //   GPUMemCpy(RecoStep::TPCClusterFinding, clustererNNShadow.centralCharges, clustererNN.centralCharges, pointerSizes[8], lane, true);
+      // } else {
+      //   AllocateRegisteredMemory(clustererNNShadow.mMemoryId);
+      // }
+      AllocateRegisteredMemory(clustererNNShadow.mMemoryId);
     }
   }
 #endif
@@ -917,41 +934,43 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
         if (GetProcessingSettings().nn.applyNNclusterizer) {
 #ifdef GPUCA_HAS_ONNX
           GPUTPCNNClusterizer& clustererNN = processors()->tpcNNClusterer[iSector];
+          GPUTPCNNClusterizer& clustererNNShadow = doGPU ? processorsShadow()->tpcNNClusterer[iSector] : clustererNN;
           const GPUSettingsProcessingNNclusterizer& nn_settings = GetProcessingSettings().nn;
-          GPUTPCNNClusterizerHost nnApplication(nn_settings, lane);
-          SetONNXGPUStream(nnApplication.model_class.updateSessionOptions(), lane);
-          SetONNXGPUStream(nnApplication.model_reg_1.updateSessionOptions(), lane);
-          SetONNXGPUStream(nnApplication.model_reg_2.updateSessionOptions(), lane);
+          GPUTPCNNClusterizerHost nnApplication(nn_settings, lane); // FIXME: This needs to be the deviceID. If that is the lane, then this line is correct
+          int32_t deviceId = -1;
+          SetONNXGPUStream(nnApplication.model_class.updateSessionOptions(), lane, &deviceId);
+          SetONNXGPUStream(nnApplication.model_reg_1.updateSessionOptions(), lane, &deviceId);
+          SetONNXGPUStream(nnApplication.model_reg_2.updateSessionOptions(), lane, &deviceId);
           int withMC = (doGPU && propagateMCLabels);
 
-          if (clustererNN.nnClusterizerUseCfRegression || (int)(nn_settings.nnClusterizerApplyCfDeconvolution)) {
+          if (clustererNNShadow.nnClusterizerUseCfRegression || (int)(nn_settings.nnClusterizerApplyCfDeconvolution)) {
             runKernel<GPUTPCCFDeconvolution>({GetGrid(clusterer.mPmemory->counters.nPositions, lane), {iSector}});
             DoDebugAndDump(RecoStep::TPCClusterFinding, 262144 << 4, clusterer, &GPUTPCClusterFinder::DumpChargeMap, *mDebugFile, "Split Charges");
           }
 
           float time_clusterizer = 0, time_fill = 0;
-          for (int batch = 0; batch < std::ceil((float)clusterer.mPmemory->counters.nClusters / clustererNN.nnClusterizerBatchedMode); batch++) {
-            uint batchStart = batch * clustererNN.nnClusterizerBatchedMode;
-            size_t iSize = CAMath::Min((uint)clustererNN.nnClusterizerBatchedMode, (uint)(clusterer.mPmemory->counters.nClusters - batchStart));
+          for (int batch = 0; batch < std::ceil((float)clusterer.mPmemory->counters.nClusters / clustererNNShadow.nnClusterizerBatchedMode); batch++) {
+            uint batchStart = batch * clustererNNShadow.nnClusterizerBatchedMode;
+            size_t iSize = CAMath::Min((uint)clustererNNShadow.nnClusterizerBatchedMode, (uint)(clusterer.mPmemory->counters.nClusters - batchStart));
 
             auto start0 = std::chrono::high_resolution_clock::now();
-            runKernel<GPUTPCNNClusterizerKernels, GPUTPCNNClusterizerKernels::fillInputNN>({GetGrid(iSize, lane), krnlRunRangeNone}, iSector, clustererNN.nnInferenceInputDType, withMC, batchStart); // Filling the data
+            runKernel<GPUTPCNNClusterizerKernels, GPUTPCNNClusterizerKernels::fillInputNN>({GetGrid(iSize, lane), krnlRunRangeNone}, iSector, clustererNNShadow.nnInferenceInputDType, withMC, batchStart); // Filling the data
 
             auto stop0 = std::chrono::high_resolution_clock::now();
             auto start1 = std::chrono::high_resolution_clock::now();
-            nnApplication.networkInference(nnApplication.model_class, clustererNN, iSize, clustererNN.modelProbabilities, clustererNN.nnInferenceInputDType);
+            nnApplication.networkInference(nnApplication.model_class, clustererNN, iSize, clustererNNShadow.modelProbabilities, clustererNNShadow.nnInferenceInputDType, deviceId);
             if (nnApplication.model_class.getNumOutputNodes()[0][1] == 1) {
-              runKernel<GPUTPCNNClusterizerKernels, GPUTPCNNClusterizerKernels::determineClass1Labels>({GetGrid(iSize, lane), krnlRunRangeNone}, iSector, clustererNN.nnInferenceInputDType, withMC, batchStart); // Assigning class labels
+              runKernel<GPUTPCNNClusterizerKernels, GPUTPCNNClusterizerKernels::determineClass1Labels>({GetGrid(iSize, lane), krnlRunRangeNone}, iSector, clustererNNShadow.nnInferenceInputDType, withMC, batchStart); // Assigning class labels
             } else {
-              runKernel<GPUTPCNNClusterizerKernels, GPUTPCNNClusterizerKernels::determineClass2Labels>({GetGrid(iSize, lane), krnlRunRangeNone}, iSector, clustererNN.nnInferenceInputDType, withMC, batchStart); // Assigning class labels
+              runKernel<GPUTPCNNClusterizerKernels, GPUTPCNNClusterizerKernels::determineClass2Labels>({GetGrid(iSize, lane), krnlRunRangeNone}, iSector, clustererNNShadow.nnInferenceInputDType, withMC, batchStart); // Assigning class labels
             }
 
-            if (!clustererNN.nnClusterizerUseCfRegression) {
-              nnApplication.networkInference(nnApplication.model_reg_1, clustererNN, iSize, clustererNN.outputDataReg1, clustererNN.nnInferenceInputDType);
-              runKernel<GPUTPCNNClusterizerKernels, GPUTPCNNClusterizerKernels::publishClass1Regression>({GetGrid(iSize, lane), krnlRunRangeNone}, iSector, clustererNN.nnInferenceInputDType, withMC, batchStart); // Running the NN for regression class 1
+            if (!clustererNNShadow.nnClusterizerUseCfRegression) {
+              nnApplication.networkInference(nnApplication.model_reg_1, clustererNN, iSize, clustererNNShadow.outputDataReg1, clustererNNShadow.nnInferenceInputDType, deviceId);
+              runKernel<GPUTPCNNClusterizerKernels, GPUTPCNNClusterizerKernels::publishClass1Regression>({GetGrid(iSize, lane), krnlRunRangeNone}, iSector, clustererNNShadow.nnInferenceInputDType, withMC, batchStart); // Running the NN for regression class 1
               if (nnApplication.model_class.getNumOutputNodes()[0][1] > 1 && nnApplication.model_reg_2.isInitialized()) {
-                nnApplication.networkInference(nnApplication.model_reg_2, clustererNN, iSize, clustererNN.outputDataReg2, clustererNN.nnInferenceInputDType);
-                runKernel<GPUTPCNNClusterizerKernels, GPUTPCNNClusterizerKernels::publishClass2Regression>({GetGrid(iSize, lane), krnlRunRangeNone}, iSector, clustererNN.nnInferenceInputDType, withMC, batchStart); // Running the NN for regression class 2
+                nnApplication.networkInference(nnApplication.model_reg_2, clustererNN, iSize, clustererNNShadow.outputDataReg2, clustererNNShadow.nnInferenceInputDType, deviceId);
+                runKernel<GPUTPCNNClusterizerKernels, GPUTPCNNClusterizerKernels::publishClass2Regression>({GetGrid(iSize, lane), krnlRunRangeNone}, iSector, clustererNNShadow.nnInferenceInputDType, withMC, batchStart); // Running the NN for regression class 2
               }
             }
             auto stop1 = std::chrono::high_resolution_clock::now();
@@ -960,15 +979,15 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
             time_fill += std::chrono::duration_cast<std::chrono::nanoseconds>(stop0 - start0).count() / 1e9;
           }
           auto start1 = std::chrono::high_resolution_clock::now();
-          if (clustererNN.nnClusterizerUseCfRegression) {
-            runKernel<GPUTPCNNClusterizerKernels, GPUTPCNNClusterizerKernels::runCfClusterizer>({GetGrid(clusterer.mPmemory->counters.nClusters, lane), krnlRunRangeNone}, iSector, clustererNN.nnInferenceInputDType, withMC, 0); // Running the CF regression kernel - no batching needed: batchStart = 0
+          if (clustererNNShadow.nnClusterizerUseCfRegression) {
+            runKernel<GPUTPCNNClusterizerKernels, GPUTPCNNClusterizerKernels::runCfClusterizer>({GetGrid(clusterer.mPmemory->counters.nClusters, lane), krnlRunRangeNone}, iSector, clustererNNShadow.nnInferenceInputDType, withMC, 0); // Running the CF regression kernel - no batching needed: batchStart = 0
           }
           auto stop1 = std::chrono::high_resolution_clock::now();
           time_clusterizer += std::chrono::duration_cast<std::chrono::nanoseconds>(stop1 - start1).count() / 1e9;
-          if (clustererNN.nnClusterizerVerbosity < 3) {
+          if (clustererNNShadow.nnClusterizerVerbosity < 3) {
             int acceptedClusters = 0;
             for (size_t i = 0; i < clusterer.mPmemory->counters.nClusters; ++i) {
-              acceptedClusters += clustererNN.outputDataClass[i];
+              acceptedClusters += clustererNNShadow.outputDataClass[i];
             }
             LOG(info) << "[NN CF] Apply NN (fragment " << fragment.index << ", lane: " << lane << ", sector: " << iSector << "): filling data " << time_fill << "s ; clusterizer: " << time_clusterizer << "s ; " << clusterer.mPmemory->counters.nClusters << " clusters, " << acceptedClusters << " accepted. --> " << clusterer.mPmemory->counters.nClusters / (time_fill + time_clusterizer) << " clusters/s";
           }
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.cxx
@@ -51,6 +51,35 @@ void* GPUTPCNNClusterizer::setIOPointers(void* mem)
   return mem;
 }
 
+std::vector<int32_t> GPUTPCNNClusterizer::pointerSizes() {
+  std::vector<int32_t> sizes(9, -1);
+  if (nnClusterizerBatchedMode > 0) {
+    if (nnInferenceInputDType == 0 && nnClusterizerElementSize > 0) {
+      sizes[0] = nnClusterizerBatchedMode * nnClusterizerElementSize; // inputData16
+    } else if (nnInferenceInputDType == 1 && nnClusterizerElementSize > 0) {
+      sizes[1] = nnClusterizerBatchedMode * nnClusterizerElementSize; // inputData32
+    }
+    sizes[2] = nnClusterizerBatchedMode; // peakPositions
+    sizes[3] = 2 * nnClusterizerBatchedMode; // clusterFlags
+    sizes[4] = nnClusterizerBatchedMode; // centralCharges
+    if (nnClusterizerModelClassNumOutputNodes > 0) {
+      sizes[5] = nnClusterizerBatchedMode * nnClusterizerModelClassNumOutputNodes; // modelProbabilities
+    }
+    if (!nnClusterizerUseCfRegression) {
+      if (nnClusterizerModelReg1NumOutputNodes > 0) {
+        sizes[6] = nnClusterizerBatchedMode * nnClusterizerModelReg1NumOutputNodes; // outputDataReg1
+      }
+      if (nnClusterizerModelReg2NumOutputNodes > 0) {
+        sizes[7] = nnClusterizerBatchedMode * nnClusterizerModelReg2NumOutputNodes; // outputDataReg2
+      }
+    }
+  }
+  if (nnClusterizerTotalClusters > 0) {
+    sizes[8] = nnClusterizerTotalClusters; // outputDataClass
+  }
+  return sizes;
+}
+
 void GPUTPCNNClusterizer::RegisterMemoryAllocation()
 {
   AllocateAndInitializeLate();
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h
@@ -34,6 +34,7 @@ class GPUTPCNNClusterizer : public GPUProcessor
   void RegisterMemoryAllocation();
   void InitializeProcessor();
   void SetMaxData(const GPUTrackingInOutPointers&);
+  std::vector<int32_t> pointerSizes();
 
   // Neural network clusterization
 
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.h b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.h

Original file line number	Diff line number	Diff line change
`@@ -662,8 +662,9 @@ void GPUReconstructionCUDA::endGPUProfiling()`
`662`	`662`	`}`
`663`	`663`
`664`	`664`	`#if defined(ORT_CUDA_BUILD) && ORT_CUDA_BUILD == 1`
`665`		`-void GPUReconstructionCUDA::SetONNXGPUStream(Ort::SessionOptions* session_options, int32_t stream)`
	`665`	`+void GPUReconstructionCUDA::SetONNXGPUStream(Ort::SessionOptions* session_options, int32_t stream, int32_t* deviceId)`
`666`	`666`	`{`
	`667`	`+ cudaGetDevice(deviceId);`
`667`	`668`	`OrtCUDAProviderOptionsV2* cuda_options = nullptr;`
`668`	`669`	`CreateCUDAProviderOptions(&cuda_options);`
`669`	`670`	`OrtSessionOptions* raw_options = session_options->operator OrtSessionOptions*();`
`@@ -690,9 +691,10 @@ void* GPUReconstructionHIP::getGPUPointer(void* ptr)`
`690`	`691`	`}`
`691`	`692`
`692`	`693`	`#if defined(ORT_ROCM_BUILD) && ORT_ROCM_BUILD == 1`
`693`		`-void GPUReconstructionHIP::SetONNXGPUStream(Ort::SessionOptions* session_options, int32_t stream)`
	`694`	`+void GPUReconstructionHIP::SetONNXGPUStream(Ort::SessionOptions* session_options, int32_t stream, int32_t* deviceId)`
`694`	`695`	`{`
`695`	`696`	`// Create ROCm provider options`
	`697`	`+ cudaGetDevice(deviceId);`
`696`	`698`	`const auto& api = Ort::GetApi();`
`697`	`699`	`OrtROCMProviderOptions rocm_options{};`
`698`	`700`	`rocm_options.has_user_compute_stream = 1; // Indicate that we are passing a user stream`
Original file line number	Diff line number	Diff line change
`@@ -83,7 +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`	`+ inline void SetONNXGPUStream(Ort::SessionOptions* opt, int32_t stream, int32_t* deviceId) { mRec->SetONNXGPUStream(opt, stream, deviceId); }`
`87`	`87`	`inline void SynchronizeEvents(deviceEvent* evList, int32_t nEvents = 1) { mRec->SynchronizeEvents(evList, nEvents); }`
`88`	`88`	`inline void SynchronizeEventAndRelease(deviceEvent& ev, bool doGPU = true)`
`89`	`89`	`{`