From 4551adfec39fe11f76ef39ce574ded60fe531296 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Sun, 26 Jul 2026 08:19:06 +0000 Subject: [PATCH] [tmva][sofie] Emit self-contained inference code from SOFIE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The C++ that TMVA SOFIE generates used to call helper functions defined in TMVA/SOFIE_common.hxx, so the generated header was not truly standalone: it required the ROOT/TMVA headers to compile. This makes the generated inference code self-contained. - New src/SOFIE_common_helpers.cxx provides GenerateHelperFunctionsCode(), which returns standalone copies of only the helpers a given model needs (Gemm_Call, Copy, Fill, Relu, Im2col/Im2col_3d/col2im, the broadcast helpers, ReadTensorFromStream, the input-dim and dynamic-memory helpers). RModel records the needed helpers via RModel_Base::AddNeededHelperFunction() and dumps them into the generated model namespace; the generated header no longer includes TMVA/SOFIE_common.hxx. - Clad autodiff: the custom-derivative pullbacks (Gemm_Call/Copy/Fill/Relu) are moved out of Math/CladDerivator.h and emitted alongside the generated helpers, into clad::custom_derivatives:: where Clad looks them up. The generated code can therefore be differentiated with Clad without depending on SOFIE_common.hxx or CladDerivator.h. - Remove the now generated-code-only helpers from SOFIE_common.hxx/.cxx and the SOFIE custom-derivative block from CladDerivator.h. Helpers still used at generation time (ConvertShapeTo*, the broadcast helpers, GNN_Data, ...) are kept. - TestGemmDerivative now exercises the emitted helper + pullback via GenerateHelperFunctionsCode instead of the library helper + CladDerivator.h. - GNN / GraphIndependent keep the shared TMVA::Experimental::SOFIE::GNN_Data as the inference boundary type (emit a using-alias, not a per-model struct) so the type stays compatible with the callers that build and chain it. - Avoid declaring the BLAS sgemm_ routine twice in generated headers (the fNeededBlasRoutines block and the Gemm_Call helper no longer both emit it). - Fix a latent double-broadcast in the buffer-filling UnidirectionalBroadcast overload (missing return after the shape-prepend branch). 🤖 Done with the help of AI. --- math/mathcore/inc/Math/CladDerivator.h | 94 -- tmva/sofie/CMakeLists.txt | 1 + tmva/sofie/inc/TMVA/RModel_Base.hxx | 27 + tmva/sofie/inc/TMVA/ROperator_Concat.hxx | 4 +- tmva/sofie/inc/TMVA/ROperator_Conv.hxx | 23 +- .../inc/TMVA/ROperator_ConvTranspose.hxx | 16 +- tmva/sofie/inc/TMVA/ROperator_Expand.hxx | 6 +- tmva/sofie/inc/TMVA/ROperator_Gemm.hxx | 17 +- .../inc/TMVA/ROperator_LayerNormalization.hxx | 6 +- tmva/sofie/inc/TMVA/SOFIE_common.hxx | 304 +------ tmva/sofie/src/RModel.cxx | 20 +- tmva/sofie/src/RModel_Base.cxx | 45 +- tmva/sofie/src/RModel_GNN.cxx | 16 +- tmva/sofie/src/RModel_GraphIndependent.cxx | 14 +- tmva/sofie/src/SOFIE_common.cxx | 125 --- tmva/sofie/src/SOFIE_common_helpers.cxx | 831 ++++++++++++++++++ tmva/sofie/test/CMakeLists.txt | 2 +- tmva/sofie/test/TestGemmDerivative.cxx | 24 +- 18 files changed, 1046 insertions(+), 529 deletions(-) create mode 100644 tmva/sofie/src/SOFIE_common_helpers.cxx diff --git a/math/mathcore/inc/Math/CladDerivator.h b/math/mathcore/inc/Math/CladDerivator.h index 9fc0eaa482883..0a317d474f2f2 100644 --- a/math/mathcore/inc/Math/CladDerivator.h +++ b/math/mathcore/inc/Math/CladDerivator.h @@ -1067,98 +1067,4 @@ inline void inc_gamma_c_pullback(double a, double x, double _d_y, double *_d_a, } // namespace custom_derivatives } // namespace clad -// Forward declare BLAS functions. -extern "C" void sgemm_(const char *transa, const char *transb, const int *m, const int *n, const int *k, - const float *alpha, const float *A, const int *lda, const float *B, const int *ldb, - const float *beta, float *C, const int *ldc); - -namespace clad::custom_derivatives { - -#ifdef R__HAS_TMVASOFIE -namespace TMVA::Experimental::SOFIE { - -inline void Gemm_Call_pullback(float *output, bool transa, bool transb, int m, int n, int k, float alpha, - const float *A, const float *B, float beta, const float *C, float *_d_output, bool *, - bool *, int *, int *, int *, float *_d_alpha, float *_d_A, float *_d_B, float *_d_beta, - float *_d_C) -{ - using ::TMVA::Experimental::SOFIE::Gemm_Call; - - // TODO: - // - fix and test the implementation for alpha != 1.0 - if (alpha != 1.0f) { - return; - } - - // beta needs to be one because we want to add to _d_A and _d_B instead of - // overwriting it. - float one = 1.; - - // ---- dA ---- - if (!transa) { - // dA += dY * op(B)^T - Gemm_Call(_d_A, false, !transb, m, k, n, one, _d_output, B, one, _d_A); - } else { - // dA += op(B) * dY^T - Gemm_Call(_d_A, transb, true, k, m, n, one, B, _d_output, one, _d_A); - } - - // ---- dB ---- - if (!transb) { - // dB += op(A)^T * dY - Gemm_Call(_d_B, !transa, false, k, n, m, one, A, _d_output, one, _d_B); - } else { - // dB += dY^T * op(A) - Gemm_Call(_d_B, true, transa, n, k, m, one, _d_output, A, one, _d_B); - } - - int sizeC = n * m; - - for (int i = 0; i < sizeC; ++i) { - if (C) { - *_d_alpha += _d_output[i] * (output[i] - beta * C[i]); - *_d_beta += _d_output[i] * C[i]; - } else { - *_d_alpha += _d_output[i] * output[i]; - } - if (_d_C) - _d_C[i] += _d_output[i] * beta; - } -} - -inline void Copy_pullback(float *output, const float *input, int size, float *_d_output, float *_d_input, int *) -{ - for (int i = 0; i < size; i++) { - output[i] = input[i]; - _d_input[i] += _d_output[i]; - _d_output[i] = 0.F; - } -} - -inline void Fill_pullback(float *output, float value, int size, float *_d_output, float *_d_value, int *) -{ - for (int i = 0; i < size; i++) { - output[i] = value; - *_d_value += _d_output[i]; - _d_output[i] = 0.F; - } -} - -inline void Relu_pullback(float *output, const float *input, int size, float *_d_output, float *_d_input, int *) -{ - for (int i = 0; i < size; i++) { - output[i] = input[i] > 0.F ? input[i] : 0.F; - float _r_d0 = _d_output[i]; - _d_output[i] = 0.F; - if (input[i] > 0.F) - _d_input[i] += _r_d0; - } -} - -} // namespace TMVA::Experimental::SOFIE - -#endif // R__HAS_TMVASOFIE - -} // namespace clad::custom_derivatives - #endif // CLAD_DERIVATOR diff --git a/tmva/sofie/CMakeLists.txt b/tmva/sofie/CMakeLists.txt index e6bded2e18470..b23733ce87319 100644 --- a/tmva/sofie/CMakeLists.txt +++ b/tmva/sofie/CMakeLists.txt @@ -94,6 +94,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTTMVASofie src/RFunction_Mean.cxx src/RFunction_Sum.cxx src/SOFIE_common.cxx + src/SOFIE_common_helpers.cxx DEPENDENCIES Core TMVA diff --git a/tmva/sofie/inc/TMVA/RModel_Base.hxx b/tmva/sofie/inc/TMVA/RModel_Base.hxx index 891f67044b405..c4bdadb1ded74 100644 --- a/tmva/sofie/inc/TMVA/RModel_Base.hxx +++ b/tmva/sofie/inc/TMVA/RModel_Base.hxx @@ -51,10 +51,19 @@ protected: WeightFileType fWeightFile = WeightFileType::Text; std::unordered_set fNeededBlasRoutines; + // Set to true once GenerateHeaderInfo has emitted the extern "C" declaration + // of the BLAS sgemm_ routine (from fNeededBlasRoutines). It lets the + // standalone Gemm_Call helper skip emitting a second, duplicate declaration. + bool fBlasSgemmDeclared = false; //! std::unordered_set fNeededStdLib = {"vector"}; std::unordered_set fCustomOpHeaders; + // Inference helper functions (from SOFIE_common) that the generated code + // needs. Their standalone definitions are emitted into the generated header + // so that it does not depend on including TMVA/SOFIE_common.hxx. + std::set fNeededHelperFunctions; + std::string fName = "UnnamedModel"; std::string fGC; // generated code bool fUseWeightFile = true; @@ -89,7 +98,25 @@ public: { fCustomOpHeaders.insert(filename); } + // Register an inference helper function that the generated code needs. See + // GenerateHelperFunctionsCode for the list of recognised keys. + void AddNeededHelperFunction(std::string name) + { + fNeededHelperFunctions.insert(std::move(name)); + } + const std::set &GetNeededHelperFunctions() const { return fNeededHelperFunctions; } + + // Placeholder tokens emitted by GenerateHeaderInfo and later replaced by + // EmitHelperFunctionsCode with the actual helper includes / definitions. + // This two-step approach is needed because the full set of required helpers + // is only known once all operators (and sub-graphs) have been generated. + static constexpr const char *kHelperIncludesMarker = "//@SOFIE_HELPER_INCLUDES@\n"; + static constexpr const char *kHelperFunctionsMarker = "//@SOFIE_HELPER_FUNCTIONS@\n"; + void GenerateHeaderInfo(std::string &hgname); + // Replace the helper markers in the generated code with the standalone + // definitions of the helper functions collected in fNeededHelperFunctions. + void EmitHelperFunctionsCode(); void PrintGenerated(std::ostream &os=std::cout) { os << fGC; } std::string ReturnGenerated() { return fGC; } diff --git a/tmva/sofie/inc/TMVA/ROperator_Concat.hxx b/tmva/sofie/inc/TMVA/ROperator_Concat.hxx index 75b764c3294b3..61ecd72027bd2 100644 --- a/tmva/sofie/inc/TMVA/ROperator_Concat.hxx +++ b/tmva/sofie/inc/TMVA/ROperator_Concat.hxx @@ -171,6 +171,8 @@ } void Initialize(RModel& model) override { + // the generated code may use the Copy inference helper + model.AddNeededHelperFunction("Copy"); std::vector> inputIntShapes; for (auto &it : fInputs) { if (model.CheckIfTensorAlreadyExist(it) == false) { @@ -325,7 +327,7 @@ std::string offset; for(size_t i=0; i 0) out << offset; offset += " + " + length; diff --git a/tmva/sofie/inc/TMVA/ROperator_Conv.hxx b/tmva/sofie/inc/TMVA/ROperator_Conv.hxx index a306121fb0714..0ad780ec35171 100644 --- a/tmva/sofie/inc/TMVA/ROperator_Conv.hxx +++ b/tmva/sofie/inc/TMVA/ROperator_Conv.hxx @@ -330,6 +330,15 @@ public: std::cout << "Conv - " << fDim << " " << fNX << " : " << ConvertDimShapeToString(fShapeX) << " --> " << fNY << " : " << ConvertDimShapeToString(fShapeY) << std::endl; } + + // register the inference helper functions used by the generated code + if (fDim < 3) + model.AddNeededHelperFunction("Im2col"); + else + model.AddNeededHelperFunction("Im2col_3d"); + model.AddNeededHelperFunction("Gemm_Call"); + if (fBroadcastBias) + model.AddNeededHelperFunction("UnidirectionalBroadcast"); } std::string GenerateInitCode() override { @@ -349,7 +358,7 @@ public: out << SP << "if (" << length << " > " << ConvertShapeToLength(shape) << ") {\n"; else out << SP << "{\n"; - out << SP << SP << "float * data = TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast(tensor_" + out << SP << SP << "float * data = UTILITY::UnidirectionalBroadcast(tensor_" << fNB << ", " << ConvertShapeToString(shape) << ", " << ConvertDimShapeToString(fShapeY) << ");\n"; out << SP << SP << "fTensor_" << fNB << ".resize(" << length << ");\n"; out << SP << SP << "std::copy(data, data + " << length << ", fTensor_" << fNB << ".begin());\n"; @@ -473,7 +482,7 @@ public: // when using im2col - resulting matrix is transposed, the dimension is (input_c * filter_h * filter_y, output_h * // output_w) if (fDim < 3) { - out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::Im2col(tensor_" << fNX + out << SP << SP << "UTILITY::Im2col(tensor_" << fNX << " + x_offset," // channels, height, width, kernel_h, kernel_w, pad_h_begin, pad_h_end, pad_w_begin, // pad_w_end, stride_h, stride_w, dilation_h, dilation_w, @@ -490,7 +499,7 @@ public: out << "," << "tensor_" <(tensor_" << fNX + out << SP << SP << "UTILITY::Im2col_3d(tensor_" << fNX << " + x_offset," // channels, d, h, w, k_d, k_h, k_w, pad_d_begin, pad_d_end, pad_h_begin, pad_h_end, // pad_w_begin, pad_w_end, stride_d, stride_h, stride_w, dilation_d, dilation_h, dilation_w, @@ -504,7 +513,7 @@ public: << "tensor_" << fNX << "_xcol);\n\n "; } // BLAS - out << SP << "TMVA::Experimental::SOFIE::Gemm_Call(" + out << SP << "Gemm_Call(" << "tensor_" << fNY << " + out_offset, false, false, " << OpName << "_m, " << OpName << "_n, " << OpName << "_k, " << OpName << "_alpha, " << "tensor_" << fNX << "_xcol, tensor_" << fNX << "_f, " @@ -533,7 +542,7 @@ public: out << SP << SP << "size_t out_offset = n * " << outputBatchStride << " + g_offset;\n"; if (fDim < 3) { - out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::Im2col(tensor_" << fNX + out << SP << SP << "UTILITY::Im2col(tensor_" << fNX << " + x_offset," // channels, height, width, kernel_h, kernel_w, pad_h_begin, pad_h_end, pad_w_begin, // pad_w_end, stride_h, stride_w, dilation_h, dilation_w, @@ -550,7 +559,7 @@ public: out << ", tensor_" << fNX << "_xcol);\n\n "; } else { // 3d im2col - out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::Im2col_3d(tensor_" << fNX + out << SP << SP << "UTILITY::Im2col_3d(tensor_" << fNX << " + x_offset," // channels, d, h, w, k_d, k_h, k_w, pad_d_begin, pad_d_end, pad_h_begin, pad_h_end, // pad_w_begin, pad_w_end, stride_d, stride_h, stride_w, dilation_d, dilation_h, dilation_w, @@ -571,7 +580,7 @@ public: << fShapeW[0] * fShapeW[1] * fAttrKernelShape[0] * fAttrKernelShape[1] * fAttrKernelShape[2] / fAttrGroup << ";\n"; - out << SP << "TMVA::Experimental::SOFIE::Gemm_Call(" + out << SP << "Gemm_Call(" << "tensor_" << fNY << " + out_offset, false, false, " << OpName << "_m, " << OpName << "_n, " << OpName << "_k, " << OpName << "_alpha, " << "tensor_" << fNX << "_xcol, tensor_" << fNX << "_f + offset_f, " diff --git a/tmva/sofie/inc/TMVA/ROperator_ConvTranspose.hxx b/tmva/sofie/inc/TMVA/ROperator_ConvTranspose.hxx index fd307ec350bf2..4fd149f747e06 100644 --- a/tmva/sofie/inc/TMVA/ROperator_ConvTranspose.hxx +++ b/tmva/sofie/inc/TMVA/ROperator_ConvTranspose.hxx @@ -301,6 +301,12 @@ void ROperator_ConvTranspose::Initialize(RModel &model) fImcol = fNX + "_xcol"; fOutputTensorNames.emplace_back(fConvK); fOutputTensorNames.emplace_back(fImcol); + + // register the inference helper functions used by the generated code + // (only the <3D case is supported, which uses col2im) + model.AddNeededHelperFunction("col2im"); + if (!fNB.empty()) + model.AddNeededHelperFunction("BroadcastConvBias"); } template @@ -313,7 +319,7 @@ std::string ROperator_ConvTranspose::GenerateInitCode() if (bsize != ysize && !fNBroadcastedB.empty()) { // include a separate scope to avoid defining unique operator temp variables out << SP << "{\n"; - out << SP << SP << "float * data = TMVA::Experimental::SOFIE::UTILITY::BroadcastConvBias(tensor_" << fNB + out << SP << SP << "float * data = UTILITY::BroadcastConvBias(tensor_" << fNB << ", " << bsize << ", " << ConvertShapeToString(fShapeY) << ");\n"; out << SP << SP << "std::copy(data, data + " << ConvertShapeToLength(fShapeY) << ", tensor_" << fNBroadcastedB << ");\n"; @@ -484,7 +490,7 @@ std::string ROperator_ConvTranspose::Generate(std::string OpName) // output_w) // before using col2im I need to transpose matrix if (fDim < 3) { - out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::col2im(tensor_" << fNX + out << SP << SP << "UTILITY::col2im(tensor_" << fNX << "_xcol," // channels, height, width, kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w, dilation_h, // dilation_w, @@ -500,7 +506,7 @@ std::string ROperator_ConvTranspose::Generate(std::string OpName) } else { // 3d : needs a col2im for 3d throw std::runtime_error("TMVA SOFIE 3D Conv Transpose not yet supported"); - out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::Im2col_3d(tensor_" << fNX + out << SP << SP << "UTILITY::Im2col_3d(tensor_" << fNX << " + x_offset," // channels, d, h, w, k_d, k_h, k_w, pad_d_begin, pad_d_end, pad_h_begin, pad_h_end, // pad_w_begin, pad_w_end, stride_d, stride_h, stride_w, dilation_d, dilation_h, dilation_w, @@ -537,7 +543,7 @@ std::string ROperator_ConvTranspose::Generate(std::string OpName) << "_xcol , &" << OpName << "_m);\n"; if (fDim < 3) { - out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::col2im(tensor_" << fNX + out << SP << SP << "UTILITY::col2im(tensor_" << fNX << "_xcol," // channels, height, width, kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w, dilation_h, // dilation_w, @@ -554,7 +560,7 @@ std::string ROperator_ConvTranspose::Generate(std::string OpName) // 3d im2col throw std::runtime_error("TMVA SOFIE 3D Conv Transpose not yet supported"); - out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::Im2col_3d(tensor_" << fNX + out << SP << SP << "UTILITY::Im2col_3d(tensor_" << fNX << " + x_offset," // channels, d, h, w, k_d, k_h, k_w, pad_d_begin, pad_d_end, pad_h_begin, pad_h_end, // pad_w_begin, pad_w_end, stride_d, stride_h, stride_w, dilation_d, dilation_h, dilation_w, diff --git a/tmva/sofie/inc/TMVA/ROperator_Expand.hxx b/tmva/sofie/inc/TMVA/ROperator_Expand.hxx index a3a0c589a333d..48a236ea9aa39 100644 --- a/tmva/sofie/inc/TMVA/ROperator_Expand.hxx +++ b/tmva/sofie/inc/TMVA/ROperator_Expand.hxx @@ -40,6 +40,8 @@ public: void Initialize(RModel& model) override { + // the generated code may use the UnidirectionalBroadcast inference helper + model.AddNeededHelperFunction("UnidirectionalBroadcast"); // input must be a graph input, or already initialized intermediate tensor if (!model.CheckIfTensorAlreadyExist(fNX)) { throw std::runtime_error("TMVA SOFIE Expand Op Input Tensor " + fNX + " is not found in model"); @@ -75,7 +77,7 @@ public: } } // Y is the common shape of fShapeX and shape - auto ret = TMVA::Experimental::SOFIE::UTILITY::MultidirectionalBroadcastShape(fShapeX, fShapeDim); + auto ret = UTILITY::MultidirectionalBroadcastShape(fShapeX, fShapeDim); fShapeY = ret.second; fInitialized = model.IsInitializedTensor(fNX) && fInitializedShape; std::vector shapeX; @@ -153,7 +155,7 @@ public: if (lengthX != lengthY) { out << SP << "if ( (" << lengthX << ") < (" << lengthY << ") ) {\n"; out << SP << SP << "// Broadcasting uninitialized tensor " << fNX << "\n"; - out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast(tensor_" << fNX << ", " << ConvertDimShapeToString(fShapeX) << ", " << ConvertDimShapeToString(fShapeY) + out << SP << SP << "UTILITY::UnidirectionalBroadcast(tensor_" << fNX << ", " << ConvertDimShapeToString(fShapeX) << ", " << ConvertDimShapeToString(fShapeY) << ", tensor_"< #include #include +#include #include #include #include @@ -503,6 +504,7 @@ void UnidirectionalBroadcast(const T* data, const std::vector& shape, co size_t offset = targetSize - shape.size(); std::copy(shape.begin(), shape.end(), newShape.begin() + offset); BroadcastTensor(inData, newShape, targetShape, broadcastedData); + return; } BroadcastTensor(inData, shape, targetShape, broadcastedData); } @@ -511,173 +513,6 @@ void UnidirectionalBroadcast(const T* data, const std::vector& shape, co std::vector ComputeStrideFromShape(const std::vector & shape); std::vector ComputeStrideFromShape(const std::vector & shape); -/// function to check if a >> 0 and a < MAX using a single comparison -//// use trick casting to unsigned values so it becomes a single comparison -inline bool is_a_ge_zero_and_a_lt_b(int a, int b) { - return static_cast(a) < static_cast(b); -} - - -/// im2col : efficient function to re-arrange input data of convolution to a matrix -/// that can be used by BLAS -/// Use trick to loop on each element of filtered region first and follow input data layout -/// By doing this reads and writes are of consecutive data in memory and one gains in efficiency -/// The resulting matrix will be already transposed and can be used directly in BLAS -/// since output will be a matrix : (channels*kernel_h*kernel_w , output_h*output_w) -/// Example: with an input matrix -/// a1 a2 a3 -/// b1 b2 b3 and a 2x2 kernel (k1,k2,k3,k4) and padding 1 : -/// c1 c2 c3 -/// outpout will be a matrix (4 x 16) -/// the routine will follow output order : -// first all elements which will be operated by k1 then k2 then k3 -/// -> ( 0 0 0 0 0 a1 a2 a3 0 b1 b2 b3 0 c1 c2 c3 ) all elements for k1 -/// ( 0 0 0 0 a1 a2 a3 0 b1 b2 b3 0 c1 c2 c3 0 ) for k2 -/// ( 0 a1 a2 a3 0 b1 b2 b3 0 c1 c2 c3 0 0 0 0 ) for k3 -/// ( a1 a2 a3 0 b1 b2 b3 0 c1 c2 c3 0 0 0 0 0 ) for k4 -/// - -/// The padding at the beginning and at the end of an axis can differ, as ONNX -/// allows for the "pads" attribute and as the SAME_UPPER / SAME_LOWER autopad -/// modes produce whenever the total padding is odd. -template -void Im2col(const T *data_im, const int channels, const int height, const int width, const int kernel_h, - const int kernel_w, const int pad_h_begin, const int pad_h_end, const int pad_w_begin, - const int pad_w_end, const int stride_h, const int stride_w, - const int dilation_h, const int dilation_w, T *data_col) -{ - const int output_h = (height + pad_h_begin + pad_h_end - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1; - const int output_w = (width + pad_w_begin + pad_w_end - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1; - const int channel_size = height * width; - for (int channel = channels; channel--; data_im += channel_size) { - for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) { - for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) { - int input_row = -pad_h_begin + kernel_row * dilation_h; - for (int output_rows = output_h; output_rows; output_rows--) { - if (!is_a_ge_zero_and_a_lt_b(input_row, height)) { - for (int output_cols = output_w; output_cols; output_cols--) { - *(data_col++) = 0; - } - } else { - int input_col = -pad_w_begin + kernel_col * dilation_w; - for (int output_col = output_w; output_col; output_col--) { - if (is_a_ge_zero_and_a_lt_b(input_col, width)) { - *(data_col++) = data_im[input_row * width + input_col]; - } else { - *(data_col++) = 0; - } - input_col += stride_w; - } - } - input_row += stride_h; - } - } - } - } -} - -/// 3d implementation -template -void Im2col_3d(const T *data_im, const int channels, - const int depth, const int height, const int width, - const int kernel_d, const int kernel_h, const int kernel_w, - const int pad_d_begin, const int pad_d_end, const int pad_h_begin, const int pad_h_end, - const int pad_w_begin, const int pad_w_end, - const int stride_d, const int stride_h, const int stride_w, - const int dilation_d, const int dilation_h, const int dilation_w, T *data_col) -{ - const int output_h = (height + pad_h_begin + pad_h_end - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1; - const int output_w = (width + pad_w_begin + pad_w_end - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1; - const int output_d = (depth + pad_d_begin + pad_d_end - (dilation_d * (kernel_d - 1) + 1)) / stride_d + 1; - const int channel_size = height * width * depth; - // assume data are c x d x h x w - for (int channel = channels; channel--; data_im += channel_size) { - for (int kernel_depth = 0; kernel_depth < kernel_d; kernel_depth++) { - for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) { - for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) { - int input_dep = -pad_d_begin + kernel_depth * dilation_d; - for (int output_dep = output_d; output_dep; output_dep--) { - if (!is_a_ge_zero_and_a_lt_b(input_dep, depth)) { - for (int output_rows = output_h; output_rows; output_rows--) { - for (int output_cols = output_w; output_cols; output_cols--) { - *(data_col++) = 0; - } - } - } else { - int input_row = -pad_h_begin + kernel_row * dilation_h; - for (int output_rows = output_h; output_rows; output_rows--) { - if (!is_a_ge_zero_and_a_lt_b(input_row, height)) { - for (int output_cols = output_w; output_cols; output_cols--) { - *(data_col++) = 0; - } - } else { - int input_col = -pad_w_begin + kernel_col * dilation_w; - for (int output_col = output_w; output_col; output_col--) { - if (is_a_ge_zero_and_a_lt_b(input_col, width)) { - *(data_col++) = data_im[input_dep * width * height + input_row * width + input_col]; - } else { - *(data_col++) = 0; - } - input_col += stride_w; - } - } - input_row += stride_h; - } - } - input_dep += stride_d; - } - } - } - } - } -} - -template -void col2im(const Dtype* data_col, const int channels, - const int height, const int width, const int kernel_h, const int kernel_w, - const int pad_h, const int pad_w, - const int stride_h, const int stride_w, - const int dilation_h, const int dilation_w, - Dtype* data_im) { - // note that output data_im needs to be set to zero value!!!! - std::fill(data_im, data_im + height * width * channels, 0.); - //caffe_set(height * width * channels, Dtype(0), data_im); - // data_im must be a zero vector - //const Dtype * data_col_0 = data_col; - const int output_h = (height + 2 * pad_h - - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1; - const int output_w = (width + 2 * pad_w - - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1; - const int channel_size = height * width; - for (int channel = channels; channel--; data_im += channel_size) { - for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) { - for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) { - int input_row = -pad_h + kernel_row * dilation_h; - for (int output_rows = output_h; output_rows; output_rows--) { - if (!is_a_ge_zero_and_a_lt_b(input_row, height)) { - data_col += output_w; - } else { - int input_col = -pad_w + kernel_col * dilation_w; - for (int output_col = output_w; output_col; output_col--) { - if (is_a_ge_zero_and_a_lt_b(input_col, width)) { - //assert(input_row*width+input_col < height * width * channels); - //assert(data_col - data_col_0 < output_h*output_w*channels); - // std::cout << "COL2IM: input_row" << " " << input_row << " " << input_col - // << " <---- " << data_col - data_col_0 << " values: " - // << data_im[input_row * width + input_col] << " <--- " << *data_col << std::endl; - data_im[input_row * width + input_col] += *data_col; - } - data_col++; - input_col += stride_w; - } - } - input_row += stride_h; - } - } - } - } - //std::cout << "finishing col2imp" << std::endl; -} } // end namespace UTILITY @@ -762,127 +597,40 @@ inline GNN_Data Copy(const GNN_Data & data) { return out; } -inline void Gemm_Call(float *output, bool transa, bool transb, int m, int n, int k, float alpha, const float *A, - const float *B, float beta, const float *C) -{ - char ct = 't'; - char cn = 'n'; - const int *lda = transa ? &k : &m; - const int *ldb = transb ? &n : &k; - const int *ldc = &m; - if (C != nullptr) { - std::copy(C, C + m * n, output); - } - TMVA::Experimental::SOFIE::BLAS::sgemm_(transa ? &ct : &cn, transb ? &ct : &cn, &m, &n, &k, &alpha, A, lda, B, ldb, - &beta, output, ldc); -} - -inline void Fill(float *output, float value, int size) -{ - std::fill(output, output + size, value); -} - -template -inline void Copy(T *output, T const *input, int size) -{ - std::copy(input, input + size, output); -} - -inline void Relu(float *output, float const *input, int size) -{ - for (int i = 0; i < size; i++) { - output[i] = (input[i] > 0.0f) ? input[i] : 0.0f; - } -} -// function to read float from the file dealing with inf and nan values -inline float ParseFloatToken (const std::string & s) { - if (s == "inf") return std::numeric_limits::infinity(); - if (s == "-inf") return -std::numeric_limits::infinity(); - if (s == "nan") return std::numeric_limits::quiet_NaN(); - return std::stof(s); -} - -template -void ReadTensorFromStream(std::istream &is, T &target, std::string const &expectedName, std::size_t expectedLength) -{ - std::string name; - std::size_t length; - is >> name >> length; - if (name != expectedName) { - std::string err_msg = - "TMVA-SOFIE failed to read the correct tensor name; expected name is " + expectedName + " , read " + name; - throw std::runtime_error(err_msg); - } - if (length != expectedLength) { - std::string err_msg = "TMVA-SOFIE failed to read the correct tensor size; expected size is " + - std::to_string(expectedLength) + " , read " + std::to_string(length); - throw std::runtime_error(err_msg); - } - std::string token; - for (size_t i = 0; i < length; ++i) { - is >> token; - target[i] = ParseFloatToken(token); - } - if (is.fail()) { - throw std::runtime_error("TMVA-SOFIE failed to read the values for tensor " + expectedName); - } -} //Utility functions to generate code void EmitNestedLoops(std::stringstream &out, size_t loopRank, const std::vector shape); void CloseNestedLoops(std::stringstream &out, size_t loopRank); -// code for the memory greeding allocations -struct TensorLifeInfo { - int begin; // start time (op index) lifetime - int end; // end time lifetime - size_t size; // size of tensors in bytes -}; - -struct MemoryResult { - std::size_t total_bytes = 0; // total memory needed - std::vector offsets; // resulted offsets for each tensor -}; - -/// Greedy best-fit planner with coalescing free list. -MemoryResult OrganizeMemory(const std::vector & tensorsInfo ); - -// Simple Dimension classes ans helpers to add constexpr meta info on input -// tensors to the emitted code. -struct SingleDim { - enum class Kind { - Static, - Symbolic - }; - Kind kind; - std::size_t dim; - std::string_view name; - - constexpr SingleDim(std::size_t v) : kind(Kind::Static), dim(v), name() {} - constexpr SingleDim(const char *v) : kind(Kind::Symbolic), dim(0), name(v) {} +/// Source code of the inference helper functions to embed in generated code so +/// that it is standalone and does not need to include TMVA/SOFIE_common.hxx. +struct HelperFunctionsCode { + std::string includes; ///< #include directives to place in the header preamble + std::string definitions; ///< function/type definitions to place inside the generated model namespace + std::string cladDefinitions; ///< Clad custom-derivative definitions to place at file scope (outside the model + ///< namespace) so that Clad discovers them; empty when none are needed }; -struct TensorDims { - const SingleDim *data; - std::size_t size; - - constexpr std::size_t total_size() const - { - std::size_t result = 1; - for (std::size_t i = 0; i < size; ++i) { - result *= data[i].dim; - } - return result; - } -}; +/// Return the standalone C++ source of the inference helper functions requested +/// in `neededHelpers` (see RModel_Base::AddNeededHelperFunction), resolving +/// their inter-dependencies. Recognised keys are: "Im2col", "Im2col_3d", +/// "col2im", "UnidirectionalBroadcast", "BroadcastConvBias", "Gemm_Call", +/// "Relu", "Fill", "Copy", "ReadTensorFromStream", "InputTensorDims", +/// "DynamicMemory" and "GNN_Data". +/// +/// `modelNamespace` (e.g. "TMVA_SOFIE_MyModel") is the generated model namespace; +/// the Clad pullbacks are emitted into clad::custom_derivatives:: +/// so the model stays differentiable without SOFIE_common.hxx / CladDerivator.h. +/// +/// `sgemmAlreadyDeclared`: set true if the caller already emitted the `extern "C"` +/// sgemm_ declaration (fNeededBlasRoutines block), so Gemm_Call skips its own and +/// avoids a duplicate. Default false emits it, keeping the returned code self-contained. +HelperFunctionsCode GenerateHelperFunctionsCode(const std::set & neededHelpers, + const std::string & modelNamespace, + bool sgemmAlreadyDeclared = false); -template -constexpr TensorDims makeDims(Arr const &arr) -{ - return TensorDims{arr.data(), arr.size()}; -} } // namespace SOFIE } // namespace Experimental diff --git a/tmva/sofie/src/RModel.cxx b/tmva/sofie/src/RModel.cxx index 624914d1cbbee..62a8c1c06d665 100644 --- a/tmva/sofie/src/RModel.cxx +++ b/tmva/sofie/src/RModel.cxx @@ -751,6 +751,10 @@ void RModel::InitializeSubGraph(std::shared_ptr graph) { AddBlasRoutines(blasRoutines); for (auto e : graph->fNeededStdLib) AddNeededStdLib(e); + // helper functions used by the subgraph must be emitted in the top-level + // header, so propagate them to the parent model + for (auto const &h : graph->GetNeededHelperFunctions()) + AddNeededHelperFunction(h); // add parent input tensors to current graph for (auto & name : fInputTensorNames) @@ -940,9 +944,12 @@ void RModel::GenerateDynamicTensorInfo() PrintDynamicTensors(); } + // the generated code uses the TensorLifeInfo / OrganizeMemory inference helpers + AddNeededHelperFunction("DynamicMemory"); + std::stringstream out; out << "// dynamic tensor memory management\n"; - out << SP << "std::vector dynamicTensorInfos;\n"; + out << SP << "std::vector dynamicTensorInfos;\n"; out << SP << "dynamicTensorInfos.reserve(" << fDynamicTensorInfos.size() << ");\n"; // loop on all the operators to find begin/end life of the tensors @@ -1536,6 +1543,9 @@ void RModel::Generate(std::underlying_type_t options, int batchSize, lo if (!fIsGNNComponent && !fIsSubGraph) { fGC += ("} //TMVA_SOFIE_" + fName + "\n"); fGC += "\n#endif // " + hgname + "\n"; + // dump the standalone definitions of the helper functions this model uses + // so that the generated header does not depend on TMVA/SOFIE_common.hxx + EmitHelperFunctionsCode(); } } @@ -1556,7 +1566,8 @@ void RModel::ReadInitializedTensorsFromFile(long pos) { fGC += " f.seekg(" + std::to_string(pos) + ");\n"; } - fGC += " using TMVA::Experimental::SOFIE::ReadTensorFromStream;\n"; + // ReadTensorFromStream is emitted as a standalone helper in the header + AddNeededHelperFunction("ReadTensorFromStream"); // loop on tensors and parse the file for (auto& i: fInitializedTensors) { @@ -1751,9 +1762,8 @@ void RModel::PrintSummary() const { void RModel::GenerateRequiredInputTensorInfo() { fGC += "\n// Input tensor dimensions\n"; - fGC += "using TMVA::Experimental::SOFIE::SingleDim;\n"; - fGC += "using TMVA::Experimental::SOFIE::TensorDims;\n"; - fGC += "using TMVA::Experimental::SOFIE::makeDims;\n\n"; + // SingleDim / TensorDims / makeDims are emitted as standalone helpers + AddNeededHelperFunction("InputTensorDims"); bool hasDynamicInputTensors = false; for (std::size_t iInput = 0; iInput < fInputTensorNames.size(); ++iInput) { diff --git a/tmva/sofie/src/RModel_Base.cxx b/tmva/sofie/src/RModel_Base.cxx index de4e080358fac..d1a305df4e1ed 100644 --- a/tmva/sofie/src/RModel_Base.cxx +++ b/tmva/sofie/src/RModel_Base.cxx @@ -22,15 +22,21 @@ void RModel_Base::GenerateHeaderInfo(std::string& hgname) { hgname = "ROOT_TMVA_SOFIE_" + hgname; fGC += "\n#ifndef " + hgname + "\n"; fGC += "#define " + hgname + "\n\n"; + // Standard library headers the generated code relies on. Listed explicitly + // now that they are no longer pulled in transitively via SOFIE_common.hxx. + for (const char *h : {"cstdint", "cstring", "string", "vector", "map", "memory", "sstream", "iostream", "iomanip", + "limits", "stdexcept", "algorithm", "cmath", "cassert"}) { + fNeededStdLib.insert(h); + } for (auto& i: fNeededStdLib) { fGC += "#include <" + i + ">\n"; } for (auto& i: fCustomOpHeaders) { fGC += "#include \"" + i + "\"\n"; } - // for the session we need to include SOFIE_Common functions - //needed for convolution operator (need to add a flag) - fGC += "#include \"TMVA/SOFIE_common.hxx\"\n"; + // Placeholder for the #include directives needed by the embedded helper + // functions (filled in by EmitHelperFunctionsCode). + fGC += kHelperIncludesMarker; if (fUseWeightFile) fGC += "#include \n"; // Include TFile when saving the weights in a binary ROOT file @@ -45,6 +51,8 @@ void RModel_Base::GenerateHeaderInfo(std::string& hgname) { fGC += ("\textern \"C\" void sgemm_(const char * transa, const char * transb, const int * m, const int * n, const int * k,\n" "\t const float * alpha, const float * A, const int * lda, const float * B, const int * ldb,\n" "\t const float * beta, float * C, const int * ldc);\n"); + // sgemm_ now declared; the standalone Gemm_Call helper will skip its copy. + fBlasSgemmDeclared = true; } else if (routine == "Gemv") { fGC += ("\textern \"C\" void sgemv_(const char * trans, const int * m, const int * n, const float * alpha, const float * A,\n" "\t const int * lda, const float * X, const int * incx, const float * beta, const float * Y, const int * incy);\n"); @@ -57,6 +65,37 @@ void RModel_Base::GenerateHeaderInfo(std::string& hgname) { } fGC += ("}//BLAS\n"); } + // Placeholder for the standalone definitions of the inference helper + // functions used by this model (filled in by EmitHelperFunctionsCode). It + // sits inside the generated model namespace, right before the session code. + fGC += kHelperFunctionsMarker; +} + +void RModel_Base::EmitHelperFunctionsCode() +{ + HelperFunctionsCode code = + GenerateHelperFunctionsCode(fNeededHelperFunctions, "TMVA_SOFIE_" + fName, fBlasSgemmDeclared); + + auto replaceMarker = [this](const std::string &marker, const std::string &replacement) { + auto pos = fGC.find(marker); + if (pos != std::string::npos) { + fGC.replace(pos, marker.size(), replacement); + } + }; + + replaceMarker(kHelperIncludesMarker, code.includes); + replaceMarker(kHelperFunctionsMarker, code.definitions); + + // Clad derivatives live at file scope and reference the model's helpers, so + // insert them after the whole model namespace, just before the include guard. + if (!code.cladDefinitions.empty()) { + auto pos = fGC.rfind("#endif"); + if (pos != std::string::npos) { + fGC.insert(pos, code.cladDefinitions + "\n"); + } else { + fGC += code.cladDefinitions; + } + } } void RModel_Base::OutputGenerated(std::string filename, bool append) { diff --git a/tmva/sofie/src/RModel_GNN.cxx b/tmva/sofie/src/RModel_GNN.cxx index 57df6dadf8ac1..1975227b72a2a 100644 --- a/tmva/sofie/src/RModel_GNN.cxx +++ b/tmva/sofie/src/RModel_GNN.cxx @@ -35,6 +35,8 @@ RModel_GNN::RModel_GNN(GNN_Init& graph_input_struct) { void RModel_GNN::Generate() { std::string hgname; + // the inference interface uses the GNN_Data helper type + AddNeededHelperFunction("GNN_Data"); GenerateHeaderInfo(hgname); std::ofstream f; @@ -143,7 +145,7 @@ void RModel_GNN::Generate() { fGC += "std::vector fNodeEdgeAggregate = std::vector(" + n_num + "*" + n_size_input + ", 0);\n"; fGC += "std::vector fNodeAggregateTemp;\n"; - fGC += "\nvoid infer(TMVA::Experimental::SOFIE::GNN_Data& input_graph){\n"; + fGC += "\nvoid infer(GNN_Data& input_graph){\n"; // computing updated edge attributes fGC += "\n// --- Edge Update ---\n"; @@ -265,8 +267,20 @@ void RModel_GNN::Generate() { fGC+="\n}\n"; fGC+="};\n"; + // propagate the helper functions needed by the update-function components + // (they are generated as GNN components and share the top-level namespace) + for (auto *block : {edges_update_block.get(), nodes_update_block.get(), globals_update_block.get()}) { + if (block && block->GetFunctionBlock()) { + for (auto const &h : block->GetFunctionBlock()->GetNeededHelperFunctions()) + AddNeededHelperFunction(h); + } + } + fGC += ("} //TMVA_SOFIE_" + fName + "\n"); fGC += "\n#endif // TMVA_SOFIE_" + hgname + "\n"; + + // dump the standalone helper-function definitions into the generated header + EmitHelperFunctionsCode(); } }//SOFIE diff --git a/tmva/sofie/src/RModel_GraphIndependent.cxx b/tmva/sofie/src/RModel_GraphIndependent.cxx index 9ef9ad5b6617d..0b8bfa6b3fe12 100644 --- a/tmva/sofie/src/RModel_GraphIndependent.cxx +++ b/tmva/sofie/src/RModel_GraphIndependent.cxx @@ -29,6 +29,8 @@ RModel_GraphIndependent::RModel_GraphIndependent(GraphIndependent_Init& graph_in void RModel_GraphIndependent::Generate() { std::string hgname; + // the inference interface uses the GNN_Data helper type + AddNeededHelperFunction("GNN_Data"); GenerateHeaderInfo(hgname); std::ofstream f; @@ -123,7 +125,7 @@ void RModel_GraphIndependent::Generate() { //fGC += "std::vector fGlobalUpdates {" + std::to_string(num_global_features) + "};"; } - fGC += "\nvoid infer(TMVA::Experimental::SOFIE::GNN_Data& input_graph){\n"; + fGC += "\nvoid infer(GNN_Data& input_graph){\n"; // computing updated edge attributes // could use std::span @@ -200,9 +202,19 @@ void RModel_GraphIndependent::Generate() { fGC += "\n"; } + // propagate helper functions needed by the update-function components + for (auto *block : {edges_update_block.get(), nodes_update_block.get(), globals_update_block.get()}) { + if (block && block->GetFunctionBlock()) { + for (auto const &h : block->GetFunctionBlock()->GetNeededHelperFunctions()) + AddNeededHelperFunction(h); + } + } + fGC += ("}\n};\n} //TMVA_SOFIE_" + fName + "\n"); fGC += "\n#endif // TMVA_SOFIE_" + hgname + "\n"; + // dump the standalone helper-function definitions into the generated header + EmitHelperFunctionsCode(); } }//SOFIE diff --git a/tmva/sofie/src/SOFIE_common.cxx b/tmva/sofie/src/SOFIE_common.cxx index abafff23652e3..672d77048ef38 100644 --- a/tmva/sofie/src/SOFIE_common.cxx +++ b/tmva/sofie/src/SOFIE_common.cxx @@ -571,131 +571,6 @@ void CloseNestedLoops(std::stringstream &out, size_t loopRank) { } - -struct FreeBlock { - std::size_t offset; - std::size_t size; - bool operator<(const FreeBlock& other) const { - // order by offset for deterministic coalescing - return offset < other.offset; - } -}; - -struct MemoryEvent { - int t; // time (i.e. operator index) - int type; // 0 = END first, 1 = START - int idx; // tensor index - bool operator<(const MemoryEvent& o) const { - if (t != o.t) return t < o.t; - return type < o.type; // END before START at the same time - } -}; - -/// Greedy best-fit planner with coalescing free list. -MemoryResult OrganizeMemory(const std::vector & tensorsInfo ) -{ - // Basic validation - for (const auto &t : tensorsInfo) { - if (!(t.end > t.begin)) { - throw std::runtime_error("Each tensor must have end > begin."); - } - } - - // Build events: free before allocate at equal times. - std::vector events; - events.reserve(tensorsInfo.size() * 2); - for (int i = 0; i < (int)tensorsInfo.size(); ++i) { - events.push_back({tensorsInfo[i].end, 0, i}); // END - events.push_back({tensorsInfo[i].begin, 1, i}); // START - } - std::sort(events.begin(), events.end()); - - std::vector tensorsOffset(tensorsInfo.size()); - - // Free list ordered by offset (for O(log n) coalescing) - // and faster insert/erase with respect to a vector - std::set free_list; - - // Bookkeeping: size/offset map for frees. - std::unordered_map live_size; - std::unordered_map live_offset; - - std::size_t total_bytes = 0; - - auto allocate_best_fit = [&](std::size_t need) -> std::size_t { - // Find the *smallest* block whose size >= need (best-fit). - // Since free_list is ordered by offset, we scan to find best by size. - // (For very large sets you could maintain a multimap by size as well.) - auto best = free_list.end(); - for (auto it = free_list.begin(); it != free_list.end(); ++it) { - if (it->size >= need) { - if (best == free_list.end() || it->size < best->size) - best = it; - } - } - if (best != free_list.end()) { - std::size_t off = best->offset; - if (best->size == need) { - free_list.erase(best); - } else { - FreeBlock updated{best->offset + need, best->size - need}; - free_list.erase(best); - free_list.insert(updated); - } - return off; - } - // No free block large enough; grow the heap. - std::size_t off = total_bytes; - total_bytes += need; - return off; - }; - - auto try_coalesce = [&](std::set::iterator it) { - // Coalesce with previous - if (it != free_list.begin()) { - auto prev = std::prev(it); - if (prev->offset + prev->size == it->offset) { - FreeBlock merged{prev->offset, prev->size + it->size}; - free_list.erase(prev); - it = free_list.erase(it); - it = free_list.insert(merged).first; - } - } - // Coalesce with next - auto next = std::next(it); - if (next != free_list.end() && it->offset + it->size == next->offset) { - FreeBlock merged{it->offset, it->size + next->size}; - free_list.erase(next); - it = free_list.erase(it); - free_list.insert(merged); - } - }; - - // Sweep through time. - for (const auto &e : events) { - if (e.type == 0) { // END: free - auto it_sz = live_size.find(e.idx); - auto it_off = live_offset.find(e.idx); - if (it_sz != live_size.end() && it_off != live_offset.end()) { - FreeBlock fb{it_off->second, it_sz->second}; - // Insert and coalesce with neighbors - auto it = free_list.insert(fb).first; - try_coalesce(it); - live_size.erase(it_sz); - live_offset.erase(it_off); - } - } else { // START: allocate - auto &t = tensorsInfo[e.idx]; - std::size_t off = allocate_best_fit(t.size); - tensorsOffset[e.idx] = off; - live_size[e.idx] = t.size; - live_offset[e.idx] = off; - } - } - - return MemoryResult{total_bytes, std::move(tensorsOffset)}; -} - } // namespace SOFIE } // namespace Experimental } // namespace TMVA diff --git a/tmva/sofie/src/SOFIE_common_helpers.cxx b/tmva/sofie/src/SOFIE_common_helpers.cxx new file mode 100644 index 0000000000000..7b4b64397450d --- /dev/null +++ b/tmva/sofie/src/SOFIE_common_helpers.cxx @@ -0,0 +1,831 @@ +/// \file SOFIE_common_helpers.cxx +/// Standalone definitions of the SOFIE inference helpers (Im2col, Gemm_Call, ...). +/// RModel records which helpers a model needs (RModel_Base::AddNeededHelperFunction) +/// and dumps only those into the generated namespace, so the emitted header is +/// self-contained (no TMVA/SOFIE_common.hxx include). These are dependency-free +/// copies of the SOFIE_common.hxx originals: keep them in sync when those change. + +#include "TMVA/SOFIE_common.hxx" + +namespace TMVA { +namespace Experimental { +namespace SOFIE { + +namespace { + +// Standalone helper snippets, each emitted verbatim inside the generated model +// namespace. They depend only on the C++ stdlib (headers collected in +// GenerateHelperFunctionsCode) and, for Gemm, on the local BLAS::sgemm_ below. + +// extern "C" declaration of the BLAS routine Gemm_Call needs, in a nested BLAS +// namespace. Skipped when the caller already declared sgemm_ (see the +// sgemmAlreadyDeclared parameter of GenerateHelperFunctionsCode). +constexpr const char *kBlasSgemm = R"SOFIE( +namespace BLAS { +extern "C" void sgemm_(const char *transa, const char *transb, const int *m, const int *n, const int *k, + const float *alpha, const float *A, const int *lda, const float *B, const int *ldb, + const float *beta, float *C, const int *ldc); +} // namespace BLAS +)SOFIE"; + +constexpr const char *kConvertShapeToLength = R"SOFIE( +inline std::size_t ConvertShapeToLength(const std::vector &shape) +{ + std::size_t length = 1; + for (auto &dim : shape) + length *= dim; + return length; +} +)SOFIE"; + +constexpr const char *kConvertShapeToString = R"SOFIE( +inline std::string ConvertShapeToString(const std::vector &shape) +{ + std::stringstream out; + out << "{ "; + for (std::size_t i = 0; i < shape.size(); i++) { + out << shape[i]; + if (i < shape.size() - 1) + out << " , "; + } + out << " }"; + return out.str(); +} +)SOFIE"; + +// Branchless bounds check `0 <= a < b`: casting to unsigned collapses it to a +// single comparison (a negative `a` wraps to a large value that fails `< b`). +constexpr const char *kIsAGeZero = R"SOFIE( +inline bool is_a_ge_zero_and_a_lt_b(int a, int b) +{ + return static_cast(a) < static_cast(b); +} +)SOFIE"; + +// im2col: re-arrange convolution input into a matrix usable directly by BLAS. +// It loops over each element of the filtered region first, following the input +// layout, so reads/writes stay consecutive in memory; the result is already +// transposed -- a (channels*kernel_h*kernel_w , output_h*output_w) matrix. +// Example: input a1 a2 a3 +// b1 b2 b3 with a 2x2 kernel (k1,k2,k3,k4) and padding 1 +// c1 c2 c3 +// gives a 4x16 matrix, output-ordered (all elements for k1, then k2, ...): +// ( 0 0 0 0 0 a1 a2 a3 0 b1 b2 b3 0 c1 c2 c3 ) k1 +// ( 0 0 0 0 a1 a2 a3 0 b1 b2 b3 0 c1 c2 c3 0 ) k2 +// ( 0 a1 a2 a3 0 b1 b2 b3 0 c1 c2 c3 0 0 0 0 ) k3 +// ( a1 a2 a3 0 b1 b2 b3 0 c1 c2 c3 0 0 0 0 0 ) k4 +// Per-axis begin/end padding can differ (ONNX "pads" attribute, and odd total +// padding from the SAME_UPPER / SAME_LOWER autopad modes). +constexpr const char *kIm2col = R"SOFIE( +template +void Im2col(const T *data_im, const int channels, const int height, const int width, const int kernel_h, + const int kernel_w, const int pad_h_begin, const int pad_h_end, const int pad_w_begin, + const int pad_w_end, const int stride_h, const int stride_w, + const int dilation_h, const int dilation_w, T *data_col) +{ + const int output_h = (height + pad_h_begin + pad_h_end - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1; + const int output_w = (width + pad_w_begin + pad_w_end - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1; + const int channel_size = height * width; + for (int channel = channels; channel--; data_im += channel_size) { + for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) { + for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) { + int input_row = -pad_h_begin + kernel_row * dilation_h; + for (int output_rows = output_h; output_rows; output_rows--) { + if (!is_a_ge_zero_and_a_lt_b(input_row, height)) { + for (int output_cols = output_w; output_cols; output_cols--) { + *(data_col++) = 0; + } + } else { + int input_col = -pad_w_begin + kernel_col * dilation_w; + for (int output_col = output_w; output_col; output_col--) { + if (is_a_ge_zero_and_a_lt_b(input_col, width)) { + *(data_col++) = data_im[input_row * width + input_col]; + } else { + *(data_col++) = 0; + } + input_col += stride_w; + } + } + input_row += stride_h; + } + } + } + } +} +)SOFIE"; + +constexpr const char *kIm2col3d = R"SOFIE( +template +void Im2col_3d(const T *data_im, const int channels, + const int depth, const int height, const int width, + const int kernel_d, const int kernel_h, const int kernel_w, + const int pad_d_begin, const int pad_d_end, const int pad_h_begin, const int pad_h_end, + const int pad_w_begin, const int pad_w_end, + const int stride_d, const int stride_h, const int stride_w, + const int dilation_d, const int dilation_h, const int dilation_w, T *data_col) +{ + const int output_h = (height + pad_h_begin + pad_h_end - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1; + const int output_w = (width + pad_w_begin + pad_w_end - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1; + const int output_d = (depth + pad_d_begin + pad_d_end - (dilation_d * (kernel_d - 1) + 1)) / stride_d + 1; + const int channel_size = height * width * depth; + for (int channel = channels; channel--; data_im += channel_size) { + for (int kernel_depth = 0; kernel_depth < kernel_d; kernel_depth++) { + for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) { + for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) { + int input_dep = -pad_d_begin + kernel_depth * dilation_d; + for (int output_dep = output_d; output_dep; output_dep--) { + if (!is_a_ge_zero_and_a_lt_b(input_dep, depth)) { + for (int output_rows = output_h; output_rows; output_rows--) { + for (int output_cols = output_w; output_cols; output_cols--) { + *(data_col++) = 0; + } + } + } else { + int input_row = -pad_h_begin + kernel_row * dilation_h; + for (int output_rows = output_h; output_rows; output_rows--) { + if (!is_a_ge_zero_and_a_lt_b(input_row, height)) { + for (int output_cols = output_w; output_cols; output_cols--) { + *(data_col++) = 0; + } + } else { + int input_col = -pad_w_begin + kernel_col * dilation_w; + for (int output_col = output_w; output_col; output_col--) { + if (is_a_ge_zero_and_a_lt_b(input_col, width)) { + *(data_col++) = data_im[input_dep * width * height + input_row * width + input_col]; + } else { + *(data_col++) = 0; + } + input_col += stride_w; + } + } + input_row += stride_h; + } + } + input_dep += stride_d; + } + } + } + } + } +} +)SOFIE"; + +constexpr const char *kCol2im = R"SOFIE( +template +void col2im(const Dtype *data_col, const int channels, + const int height, const int width, const int kernel_h, const int kernel_w, + const int pad_h, const int pad_w, + const int stride_h, const int stride_w, + const int dilation_h, const int dilation_w, + Dtype *data_im) +{ + // output must start zeroed: col2im scatters with += so overlapping columns accumulate + std::fill(data_im, data_im + height * width * channels, 0.); + const int output_h = (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1; + const int output_w = (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1; + const int channel_size = height * width; + for (int channel = channels; channel--; data_im += channel_size) { + for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) { + for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) { + int input_row = -pad_h + kernel_row * dilation_h; + for (int output_rows = output_h; output_rows; output_rows--) { + if (!is_a_ge_zero_and_a_lt_b(input_row, height)) { + data_col += output_w; + } else { + int input_col = -pad_w + kernel_col * dilation_w; + for (int output_col = output_w; output_col; output_col--) { + if (is_a_ge_zero_and_a_lt_b(input_col, width)) { + data_im[input_row * width + input_col] += *data_col; + } + data_col++; + input_col += stride_w; + } + } + input_row += stride_h; + } + } + } + } +} +)SOFIE"; + +// Broadcast helpers implementing numpy broadcasting rules (see +// https://numpy.org/doc/stable/user/basics.broadcasting.html and +// https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md). Unidirectional +// broadcast: only the input shape is stretched to targetShape, not vice versa. +// These are rewritten with respect to TMVA/SOFIE_common.hxx to avoid the +// dependency on std::span (which would require C++20 or ROOT's RSpan.hxx in the +// generated code): the input is passed as a raw pointer plus a length instead. +constexpr const char *kBroadcastTensor = R"SOFIE( +template +void BroadcastTensor(const T *data, std::size_t curLength, const std::vector &shape, + const std::vector &targetShape, T *broadcastedData) +{ + std::size_t size = shape.size(); + if (size > 1 && shape.front() == targetShape.front() && shape.back() == 1) { + std::size_t bsize = targetShape.back(); + for (int k = int(size) - 2; k >= 0; k--) { + if (shape[k] != 1) + break; + bsize *= targetShape[k]; + } + for (std::size_t i = 0; i < curLength; i++) { + std::fill(broadcastedData + i * bsize, broadcastedData + (i + 1) * bsize, data[i]); + } + return; + } + + std::copy(data, data + curLength, broadcastedData); + std::size_t arrayNum = 1; + std::vector newData(ConvertShapeToLength(targetShape)); + + for (std::size_t idx = 0; idx < size; idx++) { + std::size_t dim = shape[idx]; + std::size_t targetDim = targetShape[idx]; + if (dim == 1 && targetDim > 1) { + std::size_t newLength = curLength * targetDim; + std::size_t arrayLength = curLength / arrayNum; + if (arrayLength > 1) { + for (std::size_t arrayIdx = 0; arrayIdx < arrayNum; arrayIdx++) { + for (std::size_t targetIdx = 0; targetIdx < targetDim; targetIdx++) { + std::size_t offset = arrayIdx * arrayLength * targetDim + targetIdx * arrayLength; + std::copy(broadcastedData + arrayIdx * arrayLength, + broadcastedData + (arrayIdx + 1) * arrayLength, + newData.begin() + offset); + } + } + } else { + for (std::size_t arrayIdx = 0; arrayIdx < arrayNum; arrayIdx++) { + std::fill(newData.begin() + arrayIdx * targetDim, + newData.begin() + (arrayIdx + 1) * targetDim, broadcastedData[arrayIdx]); + } + } + curLength = newLength; + std::copy(newData.begin(), newData.begin() + newLength, broadcastedData); + } + arrayNum *= targetDim; + } +} + +template +T *CreateBroadcastTensor(const T *data, const std::vector &shape, + const std::vector &targetShape, std::size_t targetLength) +{ + T *broadcastedData = new T[targetLength]; + std::size_t curLength = ConvertShapeToLength(shape); + BroadcastTensor(data, curLength, shape, targetShape, broadcastedData); + return broadcastedData; +} + +template +T *UnidirectionalBroadcast(const T *data, const std::vector &shape, + const std::vector &targetShape) +{ + if (shape.size() < targetShape.size()) { + std::size_t targetSize = targetShape.size(); + std::vector newShape(targetSize, 1); + std::size_t offset = targetSize - shape.size(); + std::copy(shape.begin(), shape.end(), newShape.begin() + offset); + return CreateBroadcastTensor(data, newShape, targetShape, ConvertShapeToLength(targetShape)); + } + return CreateBroadcastTensor(data, shape, targetShape, ConvertShapeToLength(targetShape)); +} + +template +void UnidirectionalBroadcast(const T *data, const std::vector &shape, + const std::vector &targetShape, T *broadcastedData) +{ + std::size_t curLength = ConvertShapeToLength(shape); + if (shape.size() < targetShape.size()) { + std::size_t targetSize = targetShape.size(); + std::vector newShape(targetSize, 1); + std::size_t offset = targetSize - shape.size(); + std::copy(shape.begin(), shape.end(), newShape.begin() + offset); + BroadcastTensor(data, curLength, newShape, targetShape, broadcastedData); + return; + } + BroadcastTensor(data, curLength, shape, targetShape, broadcastedData); +} +)SOFIE"; + +constexpr const char *kBroadcastConvBias = R"SOFIE( +template +T *BroadcastConvBias(const T *data, const std::size_t channel, const std::vector &targetShape) +{ + std::size_t size = targetShape.size(); + if (targetShape[1] != channel) { + std::stringstream ss; + ss << "TMVA::SOFIE - Error broadcasting Conv Bias of shape {"; + ss << std::to_string(channel); + ss << "} to "; + ss << ConvertShapeToString(targetShape); + throw std::runtime_error(ss.str()); + } + + std::size_t targetLength = ConvertShapeToLength(targetShape); + T *newData = new T[targetLength]; + + if (targetLength == channel) { + std::copy(data, data + channel, newData); + return newData; + } + + std::size_t cStride = 1; + for (std::size_t i = 2; i < size; i++) + cStride *= targetShape[i]; + for (std::size_t i = 0; i < channel; i++) { + std::fill(newData + i * cStride, newData + (i + 1) * cStride, data[i]); + } + std::size_t batch = targetShape[0]; + std::size_t bStride = channel * cStride; + for (std::size_t i = 1; i < batch; i++) { + std::copy(newData, newData + bStride, newData + i * bStride); + } + return newData; +} +)SOFIE"; + +constexpr const char *kGemmCall = R"SOFIE( +inline void Gemm_Call(float *output, bool transa, bool transb, int m, int n, int k, float alpha, const float *A, + const float *B, float beta, const float *C) +{ + char ct = 't'; + char cn = 'n'; + const int *lda = transa ? &k : &m; + const int *ldb = transb ? &n : &k; + const int *ldc = &m; + if (C != nullptr) { + std::copy(C, C + m * n, output); + } + BLAS::sgemm_(transa ? &ct : &cn, transb ? &ct : &cn, &m, &n, &k, &alpha, A, lda, B, ldb, &beta, output, ldc); +} +)SOFIE"; + +// Custom Clad reverse-mode pullbacks for the helpers (they used to live in +// Math/CladDerivator.h). Gemm_Call and Copy need a hand-written pullback because +// their bodies call bodyless routines (sgemm_, std::copy -> memmove) that Clad +// cannot differentiate; Fill and Relu are included for completeness so a +// differentiated model needs neither SOFIE_common.hxx nor CladDerivator.h. How +// they are placed and found is explained at the emission site +// (GenerateHelperFunctionsCode). +constexpr const char *kGemmCallPullback = R"SOFIE( +inline void Gemm_Call_pullback(float *output, bool transa, bool transb, int m, int n, int k, float alpha, + const float *A, const float *B, float beta, const float *C, float *_d_output, bool *, + bool *, int *, int *, int *, float *_d_alpha, float *_d_A, float *_d_B, float *_d_beta, + float *_d_C) +{ + // TODO: + // - fix and test the implementation for alpha != 1.0 + if (alpha != 1.0f) { + return; + } + + // beta needs to be one because we want to add to _d_A and _d_B instead of + // overwriting it. + float one = 1.; + + // ---- dA ---- + if (!transa) { + // dA += dY * op(B)^T + Gemm_Call(_d_A, false, !transb, m, k, n, one, _d_output, B, one, _d_A); + } else { + // dA += op(B) * dY^T + Gemm_Call(_d_A, transb, true, k, m, n, one, B, _d_output, one, _d_A); + } + + // ---- dB ---- + if (!transb) { + // dB += op(A)^T * dY + Gemm_Call(_d_B, !transa, false, k, n, m, one, A, _d_output, one, _d_B); + } else { + // dB += dY^T * op(A) + Gemm_Call(_d_B, true, transa, n, k, m, one, _d_output, A, one, _d_B); + } + + int sizeC = n * m; + + for (int i = 0; i < sizeC; ++i) { + if (C) { + *_d_alpha += _d_output[i] * (output[i] - beta * C[i]); + *_d_beta += _d_output[i] * C[i]; + } else { + *_d_alpha += _d_output[i] * output[i]; + } + if (_d_C) + _d_C[i] += _d_output[i] * beta; + } +} +)SOFIE"; + +// Pullback for the Copy helper. Copy's body uses std::copy, which lowers to the +// bodyless __builtin_memmove that Clad cannot differentiate, so a hand-written +// pullback is required. The generated Copy is a template; this float overload +// matches its float instantiation (the only one used by inference code). +constexpr const char *kCopyPullback = R"SOFIE( +inline void Copy_pullback(float *output, const float *input, int size, float *_d_output, float *_d_input, int *) +{ + for (int i = 0; i < size; i++) { + output[i] = input[i]; + _d_input[i] += _d_output[i]; + _d_output[i] = 0.F; + } +} +)SOFIE"; + +// Pullback for the Fill helper (std::fill -> bodyless builtin). +constexpr const char *kFillPullback = R"SOFIE( +inline void Fill_pullback(float *output, float value, int size, float *_d_output, float *_d_value, int *) +{ + for (int i = 0; i < size; i++) { + output[i] = value; + *_d_value += _d_output[i]; + _d_output[i] = 0.F; + } +} +)SOFIE"; + +// Pullback for the Relu helper. Relu's body is differentiable by Clad on its +// own, but providing the pullback keeps the derivative identical to the one +// previously supplied by Math/CladDerivator.h. +constexpr const char *kReluPullback = R"SOFIE( +inline void Relu_pullback(float *output, const float *input, int size, float *_d_output, float *_d_input, int *) +{ + for (int i = 0; i < size; i++) { + output[i] = input[i] > 0.F ? input[i] : 0.F; + float _r_d0 = _d_output[i]; + _d_output[i] = 0.F; + if (input[i] > 0.F) + _d_input[i] += _r_d0; + } +} +)SOFIE"; + +constexpr const char *kRelu = R"SOFIE( +inline void Relu(float *output, float const *input, int size) +{ + for (int i = 0; i < size; i++) { + output[i] = (input[i] > 0.0f) ? input[i] : 0.0f; + } +} +)SOFIE"; + +constexpr const char *kFill = R"SOFIE( +inline void Fill(float *output, float value, int size) +{ + std::fill(output, output + size, value); +} +)SOFIE"; + +constexpr const char *kCopy = R"SOFIE( +template +inline void Copy(T *output, T const *input, int size) +{ + std::copy(input, input + size, output); +} +)SOFIE"; + +constexpr const char *kReadTensorFromStream = R"SOFIE( +inline float ParseFloatToken(const std::string &s) +{ + if (s == "inf") + return std::numeric_limits::infinity(); + if (s == "-inf") + return -std::numeric_limits::infinity(); + if (s == "nan") + return std::numeric_limits::quiet_NaN(); + return std::stof(s); +} + +template +void ReadTensorFromStream(std::istream &is, T &target, std::string const &expectedName, std::size_t expectedLength) +{ + std::string name; + std::size_t length; + is >> name >> length; + if (name != expectedName) { + std::string err_msg = + "TMVA-SOFIE failed to read the correct tensor name; expected name is " + expectedName + " , read " + name; + throw std::runtime_error(err_msg); + } + if (length != expectedLength) { + std::string err_msg = "TMVA-SOFIE failed to read the correct tensor size; expected size is " + + std::to_string(expectedLength) + " , read " + std::to_string(length); + throw std::runtime_error(err_msg); + } + std::string token; + for (std::size_t i = 0; i < length; ++i) { + is >> token; + target[i] = ParseFloatToken(token); + } + if (is.fail()) { + throw std::runtime_error("TMVA-SOFIE failed to read the values for tensor " + expectedName); + } +} +)SOFIE"; + +// Constexpr helpers carrying static/symbolic shape metadata for the model's +// input tensors into the emitted code. +constexpr const char *kInputTensorDims = R"SOFIE( +struct SingleDim { + enum class Kind { Static, Symbolic }; + Kind kind; + std::size_t dim; + std::string_view name; + constexpr SingleDim(std::size_t v) : kind(Kind::Static), dim(v), name() {} + constexpr SingleDim(const char *v) : kind(Kind::Symbolic), dim(0), name(v) {} +}; + +struct TensorDims { + const SingleDim *data; + std::size_t size; + constexpr std::size_t total_size() const + { + std::size_t result = 1; + for (std::size_t i = 0; i < size; ++i) { + result *= data[i].dim; + } + return result; + } +}; + +template +constexpr TensorDims makeDims(Arr const &arr) +{ + return TensorDims{arr.data(), arr.size()}; +} +)SOFIE"; + +constexpr const char *kDynamicMemory = R"SOFIE( +struct TensorLifeInfo { + int begin; // start time (operator index) of the tensor's lifetime + int end; // end time (operator index) + std::size_t size; // size in bytes +}; + +struct MemoryResult { + std::size_t total_bytes = 0; // total memory needed + std::vector offsets; // resulting offset for each tensor +}; + +namespace memory_detail { +struct FreeBlock { + std::size_t offset; + std::size_t size; + // order by offset for deterministic coalescing + bool operator<(const FreeBlock &other) const { return offset < other.offset; } +}; +struct MemoryEvent { + int t; // time (operator index) + int type; // 0 = END, 1 = START + int idx; // tensor index + bool operator<(const MemoryEvent &o) const + { + if (t != o.t) + return t < o.t; + return type < o.type; // END before START at the same time + } +}; +} // namespace memory_detail + +// Greedy best-fit planner with a coalescing free list. +inline MemoryResult OrganizeMemory(const std::vector &tensorsInfo) +{ + using memory_detail::FreeBlock; + using memory_detail::MemoryEvent; + for (const auto &t : tensorsInfo) { + if (!(t.end > t.begin)) { + throw std::runtime_error("Each tensor must have end > begin."); + } + } + + std::vector events; + events.reserve(tensorsInfo.size() * 2); + for (int i = 0; i < (int)tensorsInfo.size(); ++i) { + events.push_back({tensorsInfo[i].end, 0, i}); + events.push_back({tensorsInfo[i].begin, 1, i}); + } + std::sort(events.begin(), events.end()); + + std::vector tensorsOffset(tensorsInfo.size()); + std::set free_list; + std::unordered_map live_size; + std::unordered_map live_offset; + std::size_t total_bytes = 0; + + auto allocate_best_fit = [&](std::size_t need) -> std::size_t { + // Smallest free block with size >= need. free_list is ordered by offset, so + // this scans linearly; for very large tensor sets a size-keyed multimap + // would avoid the O(n) scan. + auto best = free_list.end(); + for (auto it = free_list.begin(); it != free_list.end(); ++it) { + if (it->size >= need) { + if (best == free_list.end() || it->size < best->size) + best = it; + } + } + if (best != free_list.end()) { + std::size_t off = best->offset; + if (best->size == need) { + free_list.erase(best); + } else { + FreeBlock updated{best->offset + need, best->size - need}; + free_list.erase(best); + free_list.insert(updated); + } + return off; + } + std::size_t off = total_bytes; + total_bytes += need; + return off; + }; + + auto try_coalesce = [&](std::set::iterator it) { + if (it != free_list.begin()) { + auto prev = std::prev(it); + if (prev->offset + prev->size == it->offset) { + FreeBlock merged{prev->offset, prev->size + it->size}; + free_list.erase(prev); + it = free_list.erase(it); + it = free_list.insert(merged).first; + } + } + auto next = std::next(it); + if (next != free_list.end() && it->offset + it->size == next->offset) { + FreeBlock merged{it->offset, it->size + next->size}; + free_list.erase(next); + it = free_list.erase(it); + free_list.insert(merged); + } + }; + + for (const auto &e : events) { + if (e.type == 0) { + auto it_sz = live_size.find(e.idx); + auto it_off = live_offset.find(e.idx); + if (it_sz != live_size.end() && it_off != live_offset.end()) { + FreeBlock fb{it_off->second, it_sz->second}; + auto it = free_list.insert(fb).first; + try_coalesce(it); + live_size.erase(it_sz); + live_offset.erase(it_off); + } + } else { + auto &t = tensorsInfo[e.idx]; + std::size_t off = allocate_best_fit(t.size); + tensorsOffset[e.idx] = off; + live_size[e.idx] = t.size; + live_offset[e.idx] = off; + } + } + + return MemoryResult{total_bytes, std::move(tensorsOffset)}; +} +)SOFIE"; + +// GNN_Data is the shared input/output type passed between GNN inference +// sessions (callers build a TMVA::Experimental::SOFIE::GNN_Data and hand it to +// infer()), so unlike the other helpers it must NOT be re-defined per model: +// each model aliases the one shared type. The definition (and RTensor) comes +// from TMVA/SOFIE_common.hxx, which the generated header includes in this case. +constexpr const char *kGNNData = R"SOFIE( +using GNN_Data = TMVA::Experimental::SOFIE::GNN_Data; +)SOFIE"; + +} // anonymous namespace + +HelperFunctionsCode GenerateHelperFunctionsCode(const std::set &neededHelpers, + const std::string &modelNamespace, bool sgemmAlreadyDeclared) +{ + auto need = [&](const char *key) { return neededHelpers.count(key) > 0; }; + + const bool im2col = need("Im2col"); + const bool im2col3d = need("Im2col_3d"); + const bool col2im = need("col2im"); + const bool uniBroadcast = need("UnidirectionalBroadcast"); + const bool convBias = need("BroadcastConvBias"); + const bool gemm = need("Gemm_Call"); + const bool relu = need("Relu"); + const bool fill = need("Fill"); + const bool copy = need("Copy"); + const bool readTensor = need("ReadTensorFromStream"); + const bool inputDims = need("InputTensorDims"); + const bool dynMemory = need("DynamicMemory"); + const bool gnnData = need("GNN_Data"); + + const bool im2colFamily = im2col || im2col3d || col2im; + const bool needConvertLength = uniBroadcast || convBias; + const bool needConvertString = convBias; + + // ---- collect the required standard headers ----------------------------- + std::set stdHeaders; + std::set otherHeaders; + auto addStd = [&](std::initializer_list hs) { + for (auto h : hs) + stdHeaders.insert(h); + }; + + if (im2colFamily || uniBroadcast || convBias || gemm || fill || copy || dynMemory) + addStd({"algorithm"}); + if (needConvertLength || uniBroadcast || dynMemory) + addStd({"vector", "cstddef"}); + if (needConvertString || convBias) + addStd({"sstream", "string", "stdexcept"}); + if (readTensor) + addStd({"string", "istream", "stdexcept", "limits"}); + if (inputDims) + addStd({"array", "string_view", "cstddef"}); + if (dynMemory) + addStd({"set", "unordered_map", "stdexcept", "iterator"}); + if (gnnData) + // GNN_Data (and, transitively, RTensor) is provided by SOFIE_common.hxx; + // see kGNNData for why GNN keeps using the shared type. + otherHeaders.insert("TMVA/SOFIE_common.hxx"); + + std::string includes; + for (auto const &h : stdHeaders) + includes += "#include <" + h + ">\n"; + for (auto const &h : otherHeaders) + includes += "#include \"" + h + "\"\n"; + + // ---- assemble the definitions ------------------------------------------ + // The order matters: a definition must precede any non-dependent use of it. + std::string defs; + defs += "\n// --- Standalone SOFIE inference helper functions ---\n"; + + // sgemm_ declaration for Gemm_Call, unless the caller already emitted one. + if (gemm && !sgemmAlreadyDeclared) + defs += kBlasSgemm; + + if (needConvertLength) + defs += kConvertShapeToLength; + if (needConvertString) + defs += kConvertShapeToString; + + if (im2colFamily || uniBroadcast || convBias) { + defs += "\nnamespace UTILITY {\n"; + if (im2colFamily) + defs += kIsAGeZero; + if (im2col) + defs += kIm2col; + if (im2col3d) + defs += kIm2col3d; + if (col2im) + defs += kCol2im; + if (uniBroadcast) + defs += kBroadcastTensor; + if (convBias) + defs += kBroadcastConvBias; + defs += "} // namespace UTILITY\n"; + } + + if (gemm) + defs += kGemmCall; + if (relu) + defs += kRelu; + if (fill) + defs += kFill; + if (copy) + defs += kCopy; + if (readTensor) + defs += kReadTensorFromStream; + if (inputDims) + defs += kInputTensorDims; + if (dynMemory) + defs += kDynamicMemory; + if (gnnData) + defs += kGNNData; + + defs += "// --- End of SOFIE inference helper functions ---\n\n"; + + // ---- Clad custom derivatives (pullbacks) ------------------------------- + // Some helpers cannot be differentiated automatically by Clad (Gemm_Call + // calls the bodyless BLAS routine sgemm_). For those we emit a hand-written + // pullback. Clad looks up custom derivatives in + // clad::custom_derivatives::, so the pullback is placed + // there (mirroring the generated model namespace) rather than next to the + // function. The definitions reference the model's own helpers via a using + // declaration, so no Clad header is pulled in and a user who never + // differentiates the model simply carries an unused inline function. + std::string cladDefs; + if (gemm || copy || fill || relu) { + cladDefs += "\nnamespace clad {\nnamespace custom_derivatives {\nnamespace " + modelNamespace + " {\n"; + if (gemm) { + // Gemm_Call_pullback calls Gemm_Call, so bring it into scope. + cladDefs += "using ::" + modelNamespace + "::Gemm_Call;\n"; + cladDefs += kGemmCallPullback; + } + if (copy) + cladDefs += kCopyPullback; + if (fill) + cladDefs += kFillPullback; + if (relu) + cladDefs += kReluPullback; + cladDefs += "} // namespace " + modelNamespace + "\n} // namespace custom_derivatives\n} // namespace clad\n"; + } + + return HelperFunctionsCode{std::move(includes), std::move(defs), std::move(cladDefs)}; +} + +} // namespace SOFIE +} // namespace Experimental +} // namespace TMVA diff --git a/tmva/sofie/test/CMakeLists.txt b/tmva/sofie/test/CMakeLists.txt index 3c3224cb4325b..6d26ffe1e7ad2 100644 --- a/tmva/sofie/test/CMakeLists.txt +++ b/tmva/sofie/test/CMakeLists.txt @@ -117,7 +117,7 @@ endif() if (BLAS_FOUND) if (clad) # Creating a Google Test for the automatic differentiation of Gemm_Call - ROOT_ADD_GTEST(TestGemmDerivative TestGemmDerivative.cxx LIBRARIES Core BLAS::BLAS) + ROOT_ADD_GTEST(TestGemmDerivative TestGemmDerivative.cxx LIBRARIES Core BLAS::BLAS ROOTTMVASofie) endif() endif() diff --git a/tmva/sofie/test/TestGemmDerivative.cxx b/tmva/sofie/test/TestGemmDerivative.cxx index 2c011436617e2..c3ccbd6f2f6d9 100644 --- a/tmva/sofie/test/TestGemmDerivative.cxx +++ b/tmva/sofie/test/TestGemmDerivative.cxx @@ -1,8 +1,11 @@ #include +#include "TMVA/SOFIE_common.hxx" + #include "gtest/gtest.h" #include +#include class GemmTest : public testing::TestWithParam> { public: @@ -71,8 +74,25 @@ TEST_P(GemmTest, GemmTestDerivative) { static bool declared = false; if (declared == false) { + // Test the *emitted* helper code: ask SOFIE for the standalone Gemm_Call + // source (and its Clad pullback) exactly as it is dumped into generated + // inference headers, then differentiate that with Clad. This exercises the + // helper and the pullback that ship inside generated models, without + // relying on TMVA/SOFIE_common.hxx or Math/CladDerivator.h at inference + // time. + using namespace TMVA::Experimental::SOFIE; + HelperFunctionsCode code = GenerateHelperFunctionsCode({"Gemm_Call"}, "TMVA_SOFIE_GemmTest"); + std::string src = code.includes + "namespace TMVA_SOFIE_GemmTest {\n" + code.definitions + "}\n" + + code.cladDefinitions; + + if (!gInterpreter->Declare(src.c_str())) { + throw std::runtime_error("TestGemmDerivative: failed to declare emitted SOFIE helper code"); + } + gInterpreter->Declare(R"cpp( - #include + // Clad differentiation machinery only (no SOFIE custom derivatives): + // the pullback used below comes from the emitted code declared above. + #include float gemm_function(float *variables, int m, int n, int k) { // variable is assumed to pack Amk and Bkn, @@ -89,7 +109,7 @@ TEST_P(GemmTest, GemmTestDerivative) float output[n_out]; - TMVA::Experimental::SOFIE::Gemm_Call(output, false, false, m, n, k, alpha, matA, matB, beta, matC); + TMVA_SOFIE_GemmTest::Gemm_Call(output, false, false, m, n, k, alpha, matA, matB, beta, matC); float ret = 0; for (int i = 0; i < m*n; ++i) {