Modification for new SOFIE repo

Author: ChSonnabend

GPU/GPUTracking/CMakeLists.txt

@@ -16,11 +16,7 @@ if(NOT GPUCA_BUILD_ORT)
   set(onnxruntime_FOUND OFF)
 endif()
 if(GPUCA_BUILD_SOFIE)
-  find_package(ROOT CONFIG REQUIRED COMPONENTS ROOTTMVASofie ROOTTMVASofieParser)
-  find_path(GPUCA_SOFIE_GPU_INCLUDE TMVA/RGPUModel.hxx HINTS ${ROOT_INCLUDE_DIRS} NO_DEFAULT_PATH)
-  if(NOT GPUCA_SOFIE_GPU_INCLUDE)
-    message(FATAL_ERROR "SOFIE GPU requires the ROOT development package with RGPUModel support")
-  endif()
+  find_package(SOFIE CONFIG REQUIRED)
 endif()
 
 # set(CMAKE_CXX_FLAGS_${CMAKE_BUILD_TYPE_UPPER} "${CMAKE_CXX_FLAGS_${CMAKE_BUILD_TYPE_UPPER}} -O0") # to uncomment if needed, tired of typing this...
@@ -468,7 +464,7 @@ endif()
 target_link_libraries(${targetName} PRIVATE TBB::tbb)
 if(GPUCA_BUILD_SOFIE)
   target_compile_definitions(${targetName} PRIVATE GPUCA_HAS_SOFIE=1)
-  target_link_libraries(${targetName} PRIVATE ROOT::ROOTTMVASofie ROOT::ROOTTMVASofieParser)
+  target_link_libraries(${targetName} PRIVATE SOFIE::SOFIE_core SOFIE::SOFIE_parsers)
 endif()
 
 target_compile_options(${targetName} PRIVATE -Wno-instantiation-after-specialization)

GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx

@@ -12,10 +12,6 @@
 /// \file GPUTPCNNClusterizerHost.cxx
 /// \author Christian Sonnabend
 
-#ifdef GPUCA_HAS_SOFIE
-#include "Rtypes.h"
-#endif
-
 #include <CommonUtils/StringUtils.h>
 
 #include "GPUTPCNNClusterizerHost.h"
@@ -33,8 +29,8 @@
 #endif
 
 #ifdef GPUCA_HAS_SOFIE
-#include <TMVA/RGPUModel.hxx>
-#include <TMVA/RModelParser_ONNX.hxx>
+#include <SOFIE/RGPUModel.hxx>
+#include <SOFIE/RModelParser_ONNX.hxx>
 #endif
 #include <array>
 #include <limits>
@@ -44,8 +40,8 @@ using namespace o2::gpu;
 
 struct GPUTPCNNClusterizerHost::SofieState {
 #ifdef GPUCA_HAS_SOFIE
-  std::array<std::shared_ptr<TMVA::Experimental::SOFIE::RGPUModel>, 3> models;
-  std::array<std::unique_ptr<TMVA::Experimental::SOFIE::RGPUModel::Session>, 3> sessions;
+  std::array<std::shared_ptr<SOFIE::RGPUModel>, 3> models;
+  std::array<std::unique_ptr<SOFIE::RGPUModel::Session>, 3> sessions;
 #endif
   std::array<std::string, 3> buffers;
   size_t workspaceSize = 0;
@@ -168,7 +164,7 @@ void GPUTPCNNClusterizerHost::initClusterizer(const GPUSettingsProcessingNNclust
     }
     clustererNN.mSofieWorkspaceSize = mSofie->workspaceSize;
     for (const auto& model : mSofie->models) {
-      if (model && ((model->GetPrecision() == TMVA::Experimental::SOFIE::RGPUModel::Precision::Float16) != (clustererNN.mNnInferenceInputDType == 1) || clustererNN.mNnInferenceInputDType != clustererNN.mNnInferenceOutputDType)) {
+      if (model && ((model->GetPrecision() == SOFIE::RGPUModel::Precision::Float16) != (clustererNN.mNnInferenceInputDType == 1) || clustererNN.mNnInferenceInputDType != clustererNN.mNnInferenceOutputDType)) {
         throw std::runtime_error("Changing SOFIE precision requires chain reinitialization");
       }
       if (model && model->InputSize() != static_cast<size_t>(clustererNN.mNnClusterizerElementSize)) {
@@ -328,7 +324,7 @@ MockedOrtAllocator* GPUTPCNNClusterizerHost::getMockedAllocator()
 void GPUTPCNNClusterizerHost::initSofie(const GPUSettingsProcessingNNclusterizer& settings, void* stream, int32_t device, bool hip, const GPUTPCNNClusterizerHost* source, const std::array<std::string_view, 3>& buffers, const GPUTPCNNClusterizerHost* previous)
 {
 #ifdef GPUCA_HAS_SOFIE
-  using Model = TMVA::Experimental::SOFIE::RGPUModel;
+  using Model = SOFIE::RGPUModel;
   if (!settings.nnClusterizerBatchedMode || settings.nnClusterizerBatchedMode > static_cast<unsigned>(std::numeric_limits<int32_t>::max())) {
     throw std::runtime_error("SOFIE requires a positive batch capacity within int32 range");
   }
@@ -339,7 +335,7 @@ void GPUTPCNNClusterizerHost::initSofie(const GPUSettingsProcessingNNclusterizer
     if (value == "FP16" || value == "fp16") {
       return Model::Precision::Float16;
     }
-    throw std::runtime_error("SOFIE precision must be FP32 or FP16");
+    throw std::runtime_error("SOFIE supports FP32 or FP16 model precision");
   };
   const auto inputType = precision(settings.nnInferenceInputDType);
   if (inputType != precision(settings.nnInferenceOutputDType)) {
@@ -383,7 +379,7 @@ void GPUTPCNNClusterizerHost::initSofie(const GPUSettingsProcessingNNclusterizer
       paths[i + 1] = regression[i];
     }
   }
-  TMVA::Experimental::SOFIE::RModelParser_ONNX parser;
+  SOFIE::RModelParser_ONNX parser;
   std::unordered_map<std::string, std::shared_ptr<Model>> compiled;
   if (settings.nnLoadFromCCDB && previous && previous->mSofie) {
     for (size_t i = 0; i < previous->mSofie->models.size(); i++) {
@@ -413,17 +409,16 @@ void GPUTPCNNClusterizerHost::initSofie(const GPUSettingsProcessingNNclusterizer
                 inputType == Model::Precision::Float16 ? "FP16" : "FP32");
         if (settings.nnLoadFromCCDB) {
           std::istringstream input(state->buffers[i], std::ios::in | std::ios::binary);
-          model = std::make_shared<Model>(parser.ParseGPU(input));
+          model = std::make_shared<Model>(parser.Parse(input, "clusterizer"));
         } else {
-          model = std::make_shared<Model>(parser.ParseGPU(paths[i]));
+          model = std::make_shared<Model>(parser.Parse(paths[i]));
         }
         if (model->GetPrecision() != inputType) {
           throw std::runtime_error("SOFIE model precision differs from nnInferenceInputDType: " + paths[i]);
         }
         if (model->InputSize() != inputSize || model->OutputSize() > static_cast<size_t>(std::numeric_limits<int32_t>::max()) / settings.nnClusterizerBatchedMode) {
           throw std::runtime_error("SOFIE model shape does not match the clusterizer window or index range: " + paths[i]);
         }
-        model->WorkspaceSize(settings.nnClusterizerBatchedMode);
         model->Compile({hip ? Model::Backend::HIP : Model::Backend::CUDA, settings.sofieCompiler, settings.sofieArchitecture});
         GPUInfo("SOFIE: model %zu compilation succeeded, input=%zu, output=%zu",
                 i, model->InputSize(), model->OutputSize());
@@ -471,14 +466,16 @@ void GPUTPCNNClusterizerHost::bindSofieWorkspace(void* workspace, size_t bytes)
   if (!mSofie) {
     return;
   }
-  if (!workspace || bytes < mSofie->workspaceSize) {
+  if ((!workspace && mSofie->workspaceSize) || bytes < mSofie->workspaceSize) {
     throw std::runtime_error("O2 supplied insufficient SOFIE workspace");
   }
   auto* next = static_cast<char*>(workspace);
   for (const auto& session : mSofie->sessions) {
     if (session) {
       session->SetWorkspace(next, session->WorkspaceSize());
-      next += session->WorkspaceSize();
+      if (session->WorkspaceSize()) {
+        next += session->WorkspaceSize();
+      }
     }
   }
 #endif

GPU/README.md

@@ -13,7 +13,9 @@ This module contains the following submodules:
 
 ## Experimental SOFIE NN clusterizer
 
-Build against the ROOT development package containing `TMVA/RGPUModel.hxx`:
+Build against standalone SOFIE on its `alice-custom` branch, with the
+`alpaka` and `sofieBLAS` dependencies supplied by alidist. A custom ROOT
+SOFIE implementation is no longer required:
 
 - `GPUCA_BUILD_SOFIE=ON` enables the SOFIE host integration.
 - `GPUCA_BUILD_ORT=ON` (default) retains ORT; set it to `OFF` for a SOFIE-only
@@ -27,14 +29,18 @@ Select `ml-framework=SOFIE` (the `mlFramework` field in the NN settings), or
 CUDA or `gfx90a` for HIP. `sofie-compiler` (`sofieCompiler`) optionally supplies
 an explicit compiler path; otherwise `nvcc`/`hipcc` is found on PATH at runtime.
 The compiler and GPU toolkit must be available on the worker at startup.
+CUDA requires cuBLAS/cuBLASLt; HIP requires hipBLAS/hipBLASLt. Load the SOFIE
+package environment so `SOFIE_ROOT`, `ALPAKA_ROOT`, and `SOFIEBLAS_ROOT` locate
+the installed headers when models are compiled at runtime.
 
 Use GPU TPC cluster finding, `nnInferenceDevice=cuda` or `rocm` as appropriate,
 local classification/regression ONNX paths with `nnLoadFromCCDB=0`, or
 `nnLoadFromCCDB=1` to parse the ONNX buffers fetched by the existing CCDB workflow.
 CCDB model slots follow `nnEvalMode` (classification, regression 1, optional regression 2). Set both
-`nnInferenceInputDType` and `nnInferenceOutputDType` to `FP32` or `FP16` to match
-the model. Mixed input/output precision, unsupported graphs, CPU SOFIE inference
-are rejected explicitly.
+`nnInferenceInputDType` and `nnInferenceOutputDType` to matching `FP32` or `FP16`
+and use ONNX models of that precision. FP16 inputs, weights, intermediate tensors
+and outputs stay in binary16, with FP32 GEMM/reduction accumulation. Dense models
+use hipBLAS/rocBLAS on HIP, including gfx906, without requiring hipBLASLt kernels. Unsupported graphs and CPU SOFIE inference fail explicitly.
 
 Each distinct local model path is parsed and compiled once during chain initialization.
 For CCDB, initialization is deferred until the first event has loaded its calibration
@@ -47,23 +53,27 @@ reload propagates an error without replacing the active sessions. O2 sizes the
 workspace for the new models before inference.
 Missing or empty required buffers fail explicitly. No temporary ONNX files are needed.
 Each lane then gets a session bound to its existing native stream. O2 allocates
-input, output, weights and intermediate storage through its registered clusterizer
-memory. When O2 recycles the scratch arena, weights are uploaded again on the
-lane's stream; this does not recompile the model. Sessions persist until a model
+input, output and intermediate tensor storage through its registered clusterizer
+memory. When O2 recycles the scratch arena, sessions rebind non-owning views.
+Weights, persistent state, shape metadata and BLAS resources remain session-owned;
+they are not uploaded again when scratch memory changes. Sessions persist until a model
 reload or chain finalization, with GPU synchronization before unloading their libraries.
 Changing local model paths, batch capacity, precision or backend requires reinitializing the chain.
 
-The supplied classification/regression networks have six Gemm layers, five Relu
-layers, 246 input values and respectively 1/5 output values. The clusterizer's
+Model layouts vary between exports, including dense and attention/preprocessing
+graphs. Check the deployed model input/output dimensions rather than assuming a
+particular layer count. The clusterizer's
 input-window and index-data settings must reproduce the training layout;
 the default window does not have 246 values. `nnClusterizerBatchedMode` sets
 the maximum batch; smaller final batches reuse the same code and workspace.
 
-Validation to run after building: ROOT's `TestSofieGPU` and CUDA/HIP
-`TestSofieGPUDevice`, then compare SOFIE against ORT using both supplied
-precisions and multiple batch sizes. Check numerical tolerances and downstream
-cluster decisions; the scalar dense kernels are not performance tuned. For CCDB,
+Validation to run after building: standalone SOFIE's `SOFIEWorkspace` and
+`SOFIERuntimeWorkspace` tests (see `SOFIE/core/README.md`), then compare SOFIE
+against ORT using FP32 and FP16 models and multiple batch sizes. Check numerical
+tolerances and downstream cluster decisions. Inference uses SOFIE's operator
+implementations and Alpaka/BLAS backends. For CCDB,
 check identical bytes in a new buffer, one changed model, changed hidden-layer
 sizes, and invalid replacement data. Only changed contents should compile; an
 invalid replacement must fail before switching sessions.
-No ROOT/O2 build or GPU test was run while implementing this change.
+No ROOT/O2/SOFIE build or generated-model GPU test was run while implementing
+this change. Small direct hipBLAS capability probes are not end-to-end validation.