Improve kernel speed by ~15%. Next test: for-loop in pad direction for coallesced access

ChSonnabend · ChSonnabend · commit eecb8a8e3f94 · 2025-07-16T00:57:13.000+02:00
diff --git a/GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx b/GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx
@@ -706,8 +706,8 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput)
         nnApplications[lane].initClusterizer(nn_settings, clustererNNShadow);
       }
       AllocateRegisteredMemory(clustererNN.mMemoryId);
-      nnApplications[lane].createBoundary(clustererNNShadow);
-      nnApplications[lane].createIndexLookup(clustererNNShadow);
+      // nnApplications[lane].createBoundary(clustererNNShadow);
+      // nnApplications[lane].createIndexLookup(clustererNNShadow);
     });
     if (doGPU) {
       WriteToConstantMemory(RecoStep::TPCClusterFinding, (char*)&processors()->tpcNNClusterer - (char*)processors(), &processorsShadow()->tpcNNClusterer, sizeof(GPUTPCNNClusterizer) * NSECTORS, mRec->NStreams() - 1, &mEvents->init);
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.cxx
@@ -61,8 +61,6 @@ void* GPUTPCNNClusterizer::setIOPointers(void* mem)
   }
   if (mNnClusterizerTotalClusters > 0) {
     computePointerWithAlignment(mem, mOutputDataClass, mNnClusterizerTotalClusters);
-    computePointerWithAlignment(mem, mIsBoundary, mBoundaryMapSize);
-    computePointerWithAlignment(mem, mIndexLookup, mIndexLookupSize);
   }
   return mem;
 }
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h
@@ -59,16 +59,22 @@ class GPUTPCNNClusterizer : public GPUProcessor
   int32_t mISector = -1;
   int32_t mDeviceId = -1;
 
+  // GPU optimizations
+  uint32_t mNnClusterizerFullRowSize = 0;
+  uint32_t mNnClusterizerFullPadSize = 0;
+  uint32_t mNnClusterizerFullTimeSize = 0;
+  uint32_t mNnClusterizerPadTimeSize = 0;
+
   // Boundary lookup table
-  int32_t mBoundaryMapSizeRow = 0;
-  int32_t mBoundaryMapSizePadsPerRow = 0;
-  int32_t mBoundaryMapSize = 0;
-  int32_t mBoundaryPadding = 11; // Padding on each side of the boundary map to account for pad_offset
-  int8_t* mIsBoundary = nullptr;
+  // int32_t mBoundaryMapSizeRow = 0;
+  // int32_t mBoundaryMapSizePadsPerRow = 0;
+  // int32_t mBoundaryMapSize = 0;
+  // int32_t mBoundaryPadding = 11; // Padding on each side of the boundary map to account for pad_offset
+  // int8_t* mIsBoundary = nullptr;
 
   // Index lookup table
-  int32_t mIndexLookupSize = 0;
-  int32_t* mIndexLookup = nullptr;
+  // int32_t mIndexLookupSize = 0;
+  // int32_t* mIndexLookup = nullptr;
 
   // Memory allocation for neural network
 
@@ -83,6 +89,7 @@ class GPUTPCNNClusterizer : public GPUProcessor
 
   // FP16
   OrtDataType::Float16_t* mInputData_16 = nullptr;
+  OrtDataType::Float16_t* mInputData_16_Test = nullptr;
   OrtDataType::Float16_t* mModelProbabilities_16 = nullptr;
   OrtDataType::Float16_t* mOutputDataReg1_16 = nullptr;
   OrtDataType::Float16_t* mOutputDataReg2_16 = nullptr;
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx
@@ -89,13 +89,17 @@ void GPUTPCNNClusterizerHost::initClusterizer(const GPUSettingsProcessingNNclust
   clustererNN.mNnClusterizerSizeInputRow = settings.nnClusterizerSizeInputRow;
   clustererNN.mNnClusterizerSizeInputPad = settings.nnClusterizerSizeInputPad;
   clustererNN.mNnClusterizerSizeInputTime = settings.nnClusterizerSizeInputTime;
-  clustererNN.mNnClusterizerChargeArraySize = ((2 * settings.nnClusterizerSizeInputRow + 1) * (2 * settings.nnClusterizerSizeInputPad + 1) * (2 * settings.nnClusterizerSizeInputTime + 1));
+  clustererNN.mNnClusterizerFullRowSize = 2 * settings.nnClusterizerSizeInputRow + 1;
+  clustererNN.mNnClusterizerFullPadSize = 2 * settings.nnClusterizerSizeInputPad + 1;
+  clustererNN.mNnClusterizerFullTimeSize = 2 * settings.nnClusterizerSizeInputTime + 1;
+  clustererNN.mNnClusterizerChargeArraySize = clustererNN.mNnClusterizerFullRowSize * clustererNN.mNnClusterizerFullPadSize * clustererNN.mNnClusterizerFullTimeSize;
+  clustererNN.mNnClusterizerPadTimeSize = clustererNN.mNnClusterizerFullPadSize * clustererNN.mNnClusterizerFullTimeSize;
   clustererNN.mNnClusterizerElementSize = clustererNN.mNnClusterizerChargeArraySize + (settings.nnClusterizerAddIndexData ? 3 : 0);
-  clustererNN.mBoundaryMapSizeRow = 3 * clustererNN.mNnClusterizerSizeInputRow + o2::tpc::constants::MAXGLOBALPADROW;
-  clustererNN.mBoundaryPadding = 11; // padding on each side to account for pad_offset. N=11 since then mIsBoundary = 24320 ~< (1.5 x 2^14 = 24576) && N must be bigger than (NPads[row(end_iroc + 1)] - NPads[row(end_iroc)])/2 (=6) for pad_offset to work
-  clustererNN.mBoundaryMapSizePadsPerRow = GPUTPCGeometry::NPads(o2::tpc::constants::MAXGLOBALPADROW - 1) + 2 * clustererNN.mBoundaryPadding;
-  clustererNN.mBoundaryMapSize = clustererNN.mBoundaryMapSizeRow * clustererNN.mBoundaryMapSizePadsPerRow;
-  clustererNN.mIndexLookupSize = 3 * clustererNN.mNnClusterizerChargeArraySize; // local row, pad, time shift from flat index
+  // clustererNN.mBoundaryMapSizeRow = 3 * clustererNN.mNnClusterizerSizeInputRow + o2::tpc::constants::MAXGLOBALPADROW;
+  // clustererNN.mBoundaryPadding = 11; // padding on each side to account for pad_offset. N=11 since then mIsBoundary = 24320 ~< (1.5 x 2^14 = 24576) && N must be bigger than (NPads[row(end_iroc + 1)] - NPads[row(end_iroc)])/2 (=6) for pad_offset to work
+  // clustererNN.mBoundaryMapSizePadsPerRow = GPUTPCGeometry::NPads(o2::tpc::constants::MAXGLOBALPADROW - 1) + 2 * clustererNN.mBoundaryPadding;
+  // clustererNN.mBoundaryMapSize = clustererNN.mBoundaryMapSizeRow * clustererNN.mBoundaryMapSizePadsPerRow;
+  // clustererNN.mIndexLookupSize = 3 * clustererNN.mNnClusterizerChargeArraySize; // local row, pad, time shift from flat index
   clustererNN.mNnClusterizerAddIndexData = settings.nnClusterizerAddIndexData;
   clustererNN.mNnClusterizerBatchedMode = settings.nnClusterizerBatchedMode;
   clustererNN.mNnClusterizerBoundaryFillValue = settings.nnClusterizerBoundaryFillValue;
@@ -124,38 +128,38 @@ void GPUTPCNNClusterizerHost::initClusterizer(const GPUSettingsProcessingNNclust
   }
 }
 
-void GPUTPCNNClusterizerHost::createBoundary(GPUTPCNNClusterizer& clustererNN)
-{
-  // Call after init of the clustererNN elements
-  for (int r = 0; r < clustererNN.mBoundaryMapSizeRow; r++) {
-    int8_t skipCheckInRow = 0;
-    for (int p = 0; p < clustererNN.mBoundaryMapSizePadsPerRow; p++) {
-      int32_t i = r * clustererNN.mBoundaryMapSizePadsPerRow + p;
-      clustererNN.mIsBoundary[i] = 1;
-      if (!skipCheckInRow && (p >= clustererNN.mBoundaryPadding || r >= clustererNN.mNnClusterizerSizeInputRow)) {
-        if (r < (GPUTPCGeometry::EndIROC() + clustererNN.mNnClusterizerSizeInputRow)) {
-          clustererNN.mIsBoundary[i] = (int32_t)((p - clustererNN.mBoundaryPadding) >= static_cast<int>(GPUTPCGeometry::NPads(r - clustererNN.mNnClusterizerSizeInputRow)));
-        } else if (r >= (GPUTPCGeometry::EndIROC() + 2 * clustererNN.mNnClusterizerSizeInputRow) && r < (o2::tpc::constants::MAXGLOBALPADROW + 2 * clustererNN.mNnClusterizerSizeInputRow)) {
-          clustererNN.mIsBoundary[i] = (int32_t)((p - clustererNN.mBoundaryPadding) >= static_cast<int>(GPUTPCGeometry::NPads(r - 2 * clustererNN.mNnClusterizerSizeInputRow)));
-        }
-        skipCheckInRow = (clustererNN.mIsBoundary[i] == 1); // No need to check further pads in this row
-      }
-    }
-  }
-}
-
-void GPUTPCNNClusterizerHost::createIndexLookup(GPUTPCNNClusterizer& clustererNN)
-{
-  for (int32_t i = 0; i < clustererNN.mNnClusterizerChargeArraySize; i++) {
-    int32_t r = CAMath::Floor(i / ((2 * clustererNN.mNnClusterizerSizeInputPad + 1) * (2 * clustererNN.mNnClusterizerSizeInputTime + 1))) - clustererNN.mNnClusterizerSizeInputRow;
-    int32_t rest_1 = i % ((2 * clustererNN.mNnClusterizerSizeInputPad + 1) * (2 * clustererNN.mNnClusterizerSizeInputTime + 1));
-    int32_t p = CAMath::Floor(rest_1 / (2 * clustererNN.mNnClusterizerSizeInputTime + 1)) - clustererNN.mNnClusterizerSizeInputPad;
-    int32_t t = (rest_1 % (2 * clustererNN.mNnClusterizerSizeInputTime + 1)) - clustererNN.mNnClusterizerSizeInputTime;
-    clustererNN.mIndexLookup[3 * i] = r;
-    clustererNN.mIndexLookup[3 * i + 1] = p;
-    clustererNN.mIndexLookup[3 * i + 2] = t;
-  }
-}
+// void GPUTPCNNClusterizerHost::createBoundary(GPUTPCNNClusterizer& clustererNN)
+// {
+//   // Call after init of the clustererNN elements
+//   for (int r = 0; r < clustererNN.mBoundaryMapSizeRow; r++) {
+//     int8_t skipCheckInRow = 0;
+//     for (int p = 0; p < clustererNN.mBoundaryMapSizePadsPerRow; p++) {
+//       int32_t i = r * clustererNN.mBoundaryMapSizePadsPerRow + p;
+//       clustererNN.mIsBoundary[i] = 1;
+//       if (!skipCheckInRow && (p >= clustererNN.mBoundaryPadding || r >= clustererNN.mNnClusterizerSizeInputRow)) {
+//         if (r < (GPUTPCGeometry::EndIROC() + clustererNN.mNnClusterizerSizeInputRow)) {
+//           clustererNN.mIsBoundary[i] = (int32_t)((p - clustererNN.mBoundaryPadding) >= static_cast<int>(GPUTPCGeometry::NPads(r - clustererNN.mNnClusterizerSizeInputRow)));
+//         } else if (r >= (GPUTPCGeometry::EndIROC() + 2 * clustererNN.mNnClusterizerSizeInputRow) && r < (o2::tpc::constants::MAXGLOBALPADROW + 2 * clustererNN.mNnClusterizerSizeInputRow)) {
+//           clustererNN.mIsBoundary[i] = (int32_t)((p - clustererNN.mBoundaryPadding) >= static_cast<int>(GPUTPCGeometry::NPads(r - 2 * clustererNN.mNnClusterizerSizeInputRow)));
+//         }
+//         skipCheckInRow = (clustererNN.mIsBoundary[i] == 1); // No need to check further pads in this row
+//       }
+//     }
+//   }
+// }
+
+// void GPUTPCNNClusterizerHost::createIndexLookup(GPUTPCNNClusterizer& clustererNN)
+// {
+//   for (int32_t i = 0; i < clustererNN.mNnClusterizerChargeArraySize; i++) {
+//     int32_t r = CAMath::Floor(i / ((2 * clustererNN.mNnClusterizerSizeInputPad + 1) * (2 * clustererNN.mNnClusterizerSizeInputTime + 1))) - clustererNN.mNnClusterizerSizeInputRow;
+//     int32_t rest_1 = i % ((2 * clustererNN.mNnClusterizerSizeInputPad + 1) * (2 * clustererNN.mNnClusterizerSizeInputTime + 1));
+//     int32_t p = CAMath::Floor(rest_1 / (2 * clustererNN.mNnClusterizerSizeInputTime + 1)) - clustererNN.mNnClusterizerSizeInputPad;
+//     int32_t t = (rest_1 % (2 * clustererNN.mNnClusterizerSizeInputTime + 1)) - clustererNN.mNnClusterizerSizeInputTime;
+//     clustererNN.mIndexLookup[3 * i] = r;
+//     clustererNN.mIndexLookup[3 * i + 1] = p;
+//     clustererNN.mIndexLookup[3 * i + 2] = t;
+//   }
+// }
 
 // MockedOrtAllocator implementation to be able to use volatile assignment
 struct MockedOrtAllocator : OrtAllocator {
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.cxx
diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.h b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.h

Original file line number	Diff line number	Diff line change
`@@ -61,8 +61,6 @@ void* GPUTPCNNClusterizer::setIOPointers(void* mem)`
`61`	`61`	`}`
`62`	`62`	`if (mNnClusterizerTotalClusters > 0) {`
`63`	`63`	`computePointerWithAlignment(mem, mOutputDataClass, mNnClusterizerTotalClusters);`
`64`		`- computePointerWithAlignment(mem, mIsBoundary, mBoundaryMapSize);`
`65`		`- computePointerWithAlignment(mem, mIndexLookup, mIndexLookupSize);`
`66`	`64`	`}`
`67`	`65`	`return mem;`
`68`	`66`	`}`