GPU: Template GPUChainTracking::DoTRDGPUTracking to support both o2track and gputrack types

davidrohr · davidrohr · commit 4fc2edec5b69 · 2022-01-30T19:16:48.000+01:00
diff --git a/GPU/GPUTracking/Base/GPUConstantMem.h b/GPU/GPUTracking/Base/GPUConstantMem.h
@@ -81,8 +81,32 @@ struct GPUConstantMem {
 #ifdef GPUCA_KERNEL_DEBUGGER_OUTPUT
   GPUKernelDebugOutput debugOutput;
 #endif
+
+#if defined(GPUCA_HAVE_O2HEADERS) && defined(GPUCA_NOCOMPAT)
+  template <int I>
+  GPUd() auto& getTRDTracker();
+#else  // GPUCA_HAVE_O2HEADERS
+  template <int I>
+  GPUdi() GPUTRDTrackerGPU& getTRDTracker()
+  {
+    return trdTrackerGPU;
+  }
+#endif // !GPUCA_HAVE_O2HEADERS
 };
 
+#if defined(GPUCA_HAVE_O2HEADERS) && defined(GPUCA_NOCOMPAT)
+template <>
+GPUdi() auto& GPUConstantMem::getTRDTracker<0>()
+{
+  return trdTrackerGPU;
+}
+template <>
+GPUdi() auto& GPUConstantMem::getTRDTracker<1>()
+{
+  return trdTrackerO2;
+}
+#endif
+
 #ifdef GPUCA_NOCOMPAT
 union GPUConstantMemCopyable {
   GPUConstantMemCopyable() {}  // NOLINT: We want an empty constructor, not a default one
diff --git a/GPU/GPUTracking/Global/GPUChainTracking.h b/GPU/GPUTracking/Global/GPUChainTracking.h
@@ -158,7 +158,8 @@ class GPUChainTracking : public GPUChain, GPUReconstructionHelpers::helperDelega
   int RunTPCTrackingSlices();
   int RunTPCTrackingMerger(bool synchronizeOutput = true);
   int RunTRDTracking();
-  int DoTRDGPUTracking();
+  template <int I>
+  int DoTRDGPUTracking(GPUTRDTracker* externalInstance = nullptr);
   int RunTPCCompression();
   int RunTPCDecompression();
   int RunRefit();
diff --git a/GPU/GPUTracking/Global/GPUChainTrackingTRD.cxx b/GPU/GPUTracking/Global/GPUChainTrackingTRD.cxx
@@ -55,7 +55,7 @@ int GPUChainTracking::RunTRDTracking()
     }
   }
 
-  Tracker.DoTracking(this);
+  DoTRDGPUTracking<GPUTRDTrackerKernels::gpuVersion>();
 
   mIOPtrs.nTRDTracks = Tracker.NTracks();
   mIOPtrs.trdTracks = Tracker.Tracks();
@@ -64,7 +64,8 @@ int GPUChainTracking::RunTRDTracking()
   return 0;
 }
 
-int GPUChainTracking::DoTRDGPUTracking()
+template <int I>
+int GPUChainTracking::DoTRDGPUTracking(GPUTRDTracker* externalInstance)
 {
 #ifdef GPUCA_HAVE_O2HEADERS
   bool doGPU = GetRecoStepsGPU() & RecoStep::TRDTracking;
@@ -88,3 +89,6 @@ int GPUChainTracking::DoTRDGPUTracking()
 #endif
   return (0);
 }
+
+template int GPUChainTracking::DoTRDGPUTracking<0>(GPUTRDTracker*);
+template int GPUChainTracking::DoTRDGPUTracking<1>(GPUTRDTracker*);
diff --git a/GPU/GPUTracking/TRDTracking/GPUTRDTracker.cxx b/GPU/GPUTracking/TRDTracking/GPUTRDTracker.cxx
@@ -260,7 +260,7 @@ void GPUTRDTracker_t<TRDTRK, PROP>::DoTracking(GPUChainTracking* chainTracking)
   auto timeStart = std::chrono::high_resolution_clock::now();
 
   if (mRec->GetRecoStepsGPU() & GPUDataTypes::RecoStep::TRDTracking) {
-    chainTracking->DoTRDGPUTracking();
+    chainTracking->DoTRDGPUTracking<0>();
   } else {
 #ifdef WITH_OPENMP
 #pragma omp parallel for num_threads(mRec->GetProcessingSettings().ompThreads)
diff --git a/GPU/GPUTracking/TRDTracking/GPUTRDTrackerKernels.cxx b/GPU/GPUTracking/TRDTracking/GPUTRDTrackerKernels.cxx
@@ -21,31 +21,10 @@
 
 using namespace GPUCA_NAMESPACE::gpu;
 
-#ifdef GPUCA_HAVE_O2HEADERS
-template <int I>
-GPUd() auto& getTracker(GPUTRDTrackerKernels::processorType& processors);
-template <>
-GPUdi() auto& getTracker<0>(GPUTRDTrackerKernels::processorType& processors)
-{
-  return processors.trdTrackerGPU;
-}
-template <>
-GPUdi() auto& getTracker<1>(GPUTRDTrackerKernels::processorType& processors)
-{
-  return processors.trdTrackerO2;
-}
-#else
-template <int I>
-GPUdi() GPUTRDTrackerGPU& getTracker(GPUTRDTrackerKernels::processorType& processors)
-{
-  return processors.trdTrackerGPU;
-}
-#endif
-
 template <int I>
 GPUdii() void GPUTRDTrackerKernels::Thread(int nBlocks, int nThreads, int iBlock, int iThread, GPUsharedref() GPUSharedMemory& smem, processorType& processors)
 {
-  auto& trdTracker = getTracker<I>(processors);
+  auto& trdTracker = processors.getTRDTracker<I>();
   GPUCA_OPENMP(parallel for if(!trdTracker.GetRec().GetProcessingSettings().ompKernels) num_threads(trdTracker.GetRec().GetProcessingSettings().ompThreads))
   for (int i = get_global_id(0); i < trdTracker.NTracks(); i += get_global_size(0)) {
     trdTracker.DoTrackingThread(i, get_global_id(0));
diff --git a/GPU/GPUTracking/TRDTracking/GPUTRDTrackerKernels.h b/GPU/GPUTracking/TRDTracking/GPUTRDTrackerKernels.h
@@ -25,6 +25,9 @@ namespace gpu
 class GPUTRDTrackerKernels : public GPUKernelTemplate
 {
  public:
+  enum K { defaultKernel = 0,
+           gpuVersion = 0,
+           o2Version = 1 };
   GPUhdi() CONSTEXPR static GPUDataTypes::RecoStep GetRecoStep() { return GPUCA_RECO_STEP::TRDTracking; }
   template <int iKernel = defaultKernel>
   GPUd() static void Thread(int nBlocks, int nThreads, int iBlock, int iThread, GPUsharedref() GPUSharedMemory& smem, processorType& processors);

Original file line number	Diff line number	Diff line change
`@@ -55,7 +55,7 @@ int GPUChainTracking::RunTRDTracking()`
`55`	`55`	`}`
`56`	`56`	`}`
`57`	`57`
`58`		`- Tracker.DoTracking(this);`
	`58`	`+ DoTRDGPUTracking<GPUTRDTrackerKernels::gpuVersion>();`
`59`	`59`
`60`	`60`	`mIOPtrs.nTRDTracks = Tracker.NTracks();`
`61`	`61`	`mIOPtrs.trdTracks = Tracker.Tracks();`
`@@ -64,7 +64,8 @@ int GPUChainTracking::RunTRDTracking()`
`64`	`64`	`return 0;`
`65`	`65`	`}`
`66`	`66`
`67`		`-int GPUChainTracking::DoTRDGPUTracking()`
	`67`	`+template <int I>`
	`68`	`+int GPUChainTracking::DoTRDGPUTracking(GPUTRDTracker* externalInstance)`
`68`	`69`	`{`
`69`	`70`	`#ifdef GPUCA_HAVE_O2HEADERS`
`70`	`71`	`bool doGPU = GetRecoStepsGPU() & RecoStep::TRDTracking;`
`@@ -88,3 +89,6 @@ int GPUChainTracking::DoTRDGPUTracking()`
`88`	`89`	`#endif`
`89`	`90`	`return (0);`
`90`	`91`	`}`
	`92`	`+`
	`93`	`+template int GPUChainTracking::DoTRDGPUTracking<0>(GPUTRDTracker*);`
	`94`	`+template int GPUChainTracking::DoTRDGPUTracking<1>(GPUTRDTracker*);`