[NN Clusterizer] CCDB fetching within reco workflow (AliceO2Group#14841)

* Improve GPU filling kernel speed

* Adjusting parameter bounds and additional GPU kernel optimizations

* Adding back if statement for early exit

* const'ing + fixing CPU kernel

* Remiving print statements

* Fixing CI build issue

* Working version of NN CCDB fetching and loading to file

* Cleanup

* Please consider the following formatting changes

* Using char* buffer for model loading

* Please consider the following formatting changes

* Bug-fix

* Working version of CCDB fetching and loading into ROOT class of std::vector<char>

* Please consider the following formatting changes

* Disable dumpToFile by default

* Moving macro, adding o2-test

---------

Co-authored-by: ALICE Action Bot <[email protected]>

Author: ChSonnabend

Common/ML/include/ML/OrtInterface.h

@@ -51,6 +51,7 @@ class OrtModel
   void initOptions(std::unordered_map<std::string, std::string> optionsMap);
   void initEnvironment();
   void initSession();
+  void initSessionFromBuffer(const char* buffer, size_t bufferSize);
   void memoryOnDevice(int32_t = 0);
   bool isInitialized() { return mInitialized; }
   void resetSession();

Common/ML/src/OrtInterface.cxx

@@ -138,6 +138,24 @@ void OrtModel::initEnvironment()
   (mPImplOrt->env)->DisableTelemetryEvents(); // Disable telemetry events
 }
 
+void OrtModel::initSessionFromBuffer(const char* buffer, size_t bufferSize)
+{
+  mPImplOrt->sessionOptions.AddConfigEntry("session.load_model_format", "ONNX");
+  mPImplOrt->sessionOptions.AddConfigEntry("session.use_ort_model_bytes_directly", "1");
+
+  mPImplOrt->session = std::make_unique<Ort::Session>(*mPImplOrt->env,
+                                                      buffer,
+                                                      bufferSize,
+                                                      mPImplOrt->sessionOptions);
+  mPImplOrt->ioBinding = std::make_unique<Ort::IoBinding>(*mPImplOrt->session);
+
+  setIO();
+
+  if (mLoggingLevel < 2) {
+    LOG(info) << "(ORT) Model loaded successfully from buffer! (inputs: " << printShape(mInputShapes, mInputNames) << ", outputs: " << printShape(mOutputShapes, mInputNames) << ")";
+  }
+}
+
 void OrtModel::initSession()
 {
   if (mAllocateDeviceMemory) {

Detectors/TPC/base/test/testTPCCDBInterface.cxx

@@ -22,7 +22,6 @@
 
 // o2 includes
 #include "TPCBase/CDBInterface.h"
-#include "TPCBase/CDBInterface.h"
 #include "TPCBase/CalArray.h"
 #include "TPCBase/CalDet.h"
 #include "TPCBase/Mapper.h"

Detectors/TPC/calibration/CMakeLists.txt

@@ -25,7 +25,6 @@ o2_add_library(TPCCalibration
                        src/CalibPadGainTracksBase.cxx
                        src/CalibLaserTracks.cxx
                        src/LaserTracksCalibrator.cxx
-                       src/NeuralNetworkClusterizer.cxx
                        src/SACDecoder.cxx
                        src/IDCAverageGroup.cxx
                        src/IDCAverageGroupBase.cxx
@@ -84,7 +83,6 @@ o2_target_root_dictionary(TPCCalibration
                                   include/TPCCalibration/FastHisto.h
                                   include/TPCCalibration/CalibLaserTracks.h
                                   include/TPCCalibration/LaserTracksCalibrator.h
-                                  include/TPCCalibration/NeuralNetworkClusterizer.h
                                   include/TPCCalibration/SACDecoder.h
                                   include/TPCCalibration/IDCAverageGroup.h
                                   include/TPCCalibration/IDCAverageGroupBase.h

Detectors/TPC/calibration/include/TPCCalibration/NeuralNetworkClusterizer.h

@@ -209,6 +209,7 @@ set(SRCS_DATATYPES
     DataTypes/TPCPadBitMap.cxx
     DataTypes/TPCZSLinkMapping.cxx
     DataTypes/CalibdEdxContainer.cxx
+    DataTypes/ORTRootSerializer.cxx
     DataTypes/CalibdEdxTrackTopologyPol.cxx
     DataTypes/CalibdEdxTrackTopologySpline.cxx
     DataTypes/GPUTRDTrackO2.cxx)

Detectors/TPC/calibration/src/NeuralNetworkClusterizer.cxx

@@ -85,6 +85,7 @@ class Cluster;
 namespace tpc
 {
 class CalibdEdxContainer;
+class ORTRootSerializer;
 } // namespace tpc
 } // namespace o2
 
@@ -182,6 +183,9 @@ struct GPUCalibObjectsTemplate { // use only pointers on PODs or flat objects he
   typename S<o2::tpc::CalibdEdxContainer>::type* dEdxCalibContainer = nullptr;
   typename S<o2::base::PropagatorImpl<float>>::type* o2Propagator = nullptr;
   typename S<o2::itsmft::TopologyDictionary>::type* itsPatternDict = nullptr;
+
+  // NN clusterizer objects
+  typename S<o2::tpc::ORTRootSerializer>::type* nnClusterizerNetworks[3] = {nullptr, nullptr, nullptr};
 };
 typedef GPUCalibObjectsTemplate<DefaultPtr> GPUCalibObjects; // NOTE: These 2 must have identical layout since they are memcopied
 typedef GPUCalibObjectsTemplate<ConstPtr> GPUCalibObjectsConst;

GPU/GPUTracking/CMakeLists.txt

@@ -0,0 +1,25 @@
+// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
+// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
+// All rights not expressly granted are reserved.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// In applying this license CERN does not waive the privileges and immunities
+// granted to it by virtue of its status as an Intergovernmental Organization
+// or submit itself to any jurisdiction.
+
+/// \file   ORTRootSerializer.cxx
+/// \author Christian Sonnabend <[email protected]>
+
+#include "ORTRootSerializer.h"
+#include <cstring>
+
+using namespace o2::tpc;
+
+/// Initialize the serialization from a char* buffer containing the model
+void ORTRootSerializer::setOnnxModel(const char* onnxModel, uint32_t size)
+{
+  mModelBuffer.resize(size);
+  std::memcpy(mModelBuffer.data(), onnxModel, size);
+}

GPU/GPUTracking/DataTypes/GPUDataTypes.h

@@ -0,0 +1,43 @@
+// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
+// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
+// All rights not expressly granted are reserved.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// In applying this license CERN does not waive the privileges and immunities
+// granted to it by virtue of its status as an Intergovernmental Organization
+// or submit itself to any jurisdiction.
+
+/// \file   ORTRootSerializer.h
+/// \brief  Class to serialize ONNX objects for ROOT snapshots of CCDB objects at runtime
+/// \author Christian Sonnabend <[email protected]>
+
+#ifndef ALICEO2_TPC_ORTROOTSERIALIZER_H_
+#define ALICEO2_TPC_ORTROOTSERIALIZER_H_
+
+#include "GPUCommonRtypes.h"
+#include <vector>
+#include <string>
+
+namespace o2::tpc
+{
+
+class ORTRootSerializer
+{
+ public:
+  ORTRootSerializer() = default;
+  ~ORTRootSerializer() = default;
+
+  void setOnnxModel(const char* onnxModel, uint32_t size);
+  const char* getONNXModel() const { return mModelBuffer.data(); }
+  uint32_t getONNXModelSize() const { return static_cast<uint32_t>(mModelBuffer.size()); }
+
+ private:
+  std::vector<char> mModelBuffer; ///< buffer for serialization
+  ClassDefNV(ORTRootSerializer, 1);
+};
+
+} // namespace o2::tpc
+
+#endif // ALICEO2_TPC_ORTROOTSERIALIZER_H_