GPU Interface: Support a callback for providing final set of inputs and output buffer late, to be propagated to GPUReconstruciton only after running TPC Clusterization

Author: davidrohr

GPU/GPUTracking/Interface/GPUO2Interface.cxx

@@ -163,23 +163,50 @@ void GPUO2Interface::DumpSettings()
   mCtx[0].mRec->DumpSettings();
 }
 
-int GPUO2Interface::RunTracking(GPUTrackingInOutPointers* data, GPUInterfaceOutputs* outputs, unsigned int iThread)
+int GPUO2Interface::RunTracking(GPUTrackingInOutPointers* data, GPUInterfaceOutputs* outputs, unsigned int iThread, GPUInterfaceInputUpdate* inputUpdateCallback)
 {
   if (mNContexts <= iThread) {
     return (1);
   }
 
   mCtx[iThread].mChain->mIOPtrs = *data;
-  if (mConfig->configInterface.outputToExternalBuffers) {
-    for (unsigned int i = 0; i < mCtx[iThread].mOutputRegions->count(); i++) {
-      if (outputs->asArray()[i].allocator) {
-        mCtx[iThread].mOutputRegions->asArray()[i].set(outputs->asArray()[i].allocator);
-      } else if (outputs->asArray()[i].ptrBase) {
-        mCtx[iThread].mOutputRegions->asArray()[i].set(outputs->asArray()[i].ptrBase, outputs->asArray()[i].size);
-      } else {
-        mCtx[iThread].mOutputRegions->asArray()[i].reset();
+
+  auto setOutputs = [this, iThread](GPUInterfaceOutputs* outputs) {
+    if (mConfig->configInterface.outputToExternalBuffers) {
+      for (unsigned int i = 0; i < mCtx[iThread].mOutputRegions->count(); i++) {
+        if (outputs->asArray()[i].allocator) {
+          mCtx[iThread].mOutputRegions->asArray()[i].set(outputs->asArray()[i].allocator);
+        } else if (outputs->asArray()[i].ptrBase) {
+          mCtx[iThread].mOutputRegions->asArray()[i].set(outputs->asArray()[i].ptrBase, outputs->asArray()[i].size);
+        } else {
+          mCtx[iThread].mOutputRegions->asArray()[i].reset();
+        }
       }
     }
+  };
+
+  auto inputWaitCallback = [this, iThread, inputUpdateCallback, &data, &outputs, &setOutputs]() {
+    GPUTrackingInOutPointers* updatedData;
+    GPUInterfaceOutputs* updatedOutputs;
+    if (inputUpdateCallback->callback) {
+      inputUpdateCallback->callback(updatedData, updatedOutputs);
+      mCtx[iThread].mChain->mIOPtrs = *updatedData;
+      outputs = updatedOutputs;
+      data = updatedData;
+      setOutputs(outputs);
+    }
+    if (inputUpdateCallback->notifyCallback) {
+      inputUpdateCallback->notifyCallback();
+    }
+  };
+
+  if (inputUpdateCallback) {
+    mCtx[iThread].mChain->SetFinalInputCallback(inputWaitCallback);
+  } else {
+    mCtx[iThread].mChain->SetFinalInputCallback(nullptr);
+  }
+  if (!inputUpdateCallback || !inputUpdateCallback->callback) {
+    setOutputs(outputs);
   }
 
   int retVal = mCtx[iThread].mRec->RunChains();

GPU/GPUTracking/Interface/GPUO2Interface.h

@@ -54,6 +54,7 @@ class GPUChainTracking;
 class GPUChainITS;
 struct GPUO2InterfaceConfiguration;
 struct GPUInterfaceOutputs;
+struct GPUInterfaceInputUpdate;
 struct GPUTrackingOutputs;
 struct GPUConstantMem;
 struct GPUNewCalibValues;
@@ -70,7 +71,7 @@ class GPUO2Interface
   int Initialize(const GPUO2InterfaceConfiguration& config);
   void Deinitialize();
 
-  int RunTracking(GPUTrackingInOutPointers* data, GPUInterfaceOutputs* outputs = nullptr, unsigned int iThread = 0);
+  int RunTracking(GPUTrackingInOutPointers* data, GPUInterfaceOutputs* outputs = nullptr, unsigned int iThread = 0, GPUInterfaceInputUpdate* inputUpdateCallback = nullptr);
   void Clear(bool clearOutputs, unsigned int iThread = 0);
   void DumpEvent(int nEvent, GPUTrackingInOutPointers* data);
   void DumpSettings();

GPU/GPUTracking/Interface/GPUO2InterfaceConfiguration.h

@@ -66,6 +66,11 @@ struct GPUInterfaceOutputs : public GPUTrackingOutputs {
   GPUInterfaceQAOutputs qa;
 };
 
+struct GPUInterfaceInputUpdate {
+  std::function<void(GPUTrackingInOutPointers*& data, GPUInterfaceOutputs*& outputs)> callback; // Callback which provides final data ptrs / outputRegions after Clusterization stage
+  std::function<void()> notifyCallback;                                                         // Callback called to notify that Clusterization state has finished without update
+};
+
 // Full configuration structure with all available settings of GPU...
 struct GPUO2InterfaceConfiguration {
   GPUO2InterfaceConfiguration() = default;