Skip to content

[tmva][sofie] Emit self-contained inference code from SOFIE #22920

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
guitargeek wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[tmva][sofie] Emit self-contained inference code from SOFIE #22920

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 0 additions & 94 deletions math/mathcore/inc/Math/CladDerivator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
1 change: 1 addition & 0 deletions tmva/sofie/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
27 changes: 27 additions & 0 deletions tmva/sofie/inc/TMVA/RModel_Base.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,10 +51,19 @@ protected:
WeightFileType fWeightFile = WeightFileType::Text;

std::unordered_set<std::string> 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<std::string> fNeededStdLib = {"vector"};
std::unordered_set<std::string> 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<std::string> fNeededHelperFunctions;

std::string fName = "UnnamedModel";
std::string fGC; // generated code
bool fUseWeightFile = true;
Expand Down Expand Up @@ -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<std::string> &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; }
Expand Down
4 changes: 3 additions & 1 deletion tmva/sofie/inc/TMVA/ROperator_Concat.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,8 @@
}

void Initialize(RModel& model) override {
// the generated code may use the Copy inference helper
model.AddNeededHelperFunction("Copy");
std::vector<std::vector<size_t>> inputIntShapes;
for (auto &it : fInputs) {
if (model.CheckIfTensorAlreadyExist(it) == false) {
Expand Down Expand Up @@ -325,7 +327,7 @@
std::string offset;
for(size_t i=0; i<fInputs.size(); ++i) {
auto length = ConvertDimShapeToLength(fInputShapes[i]);
out << SP << "TMVA::Experimental::SOFIE::Copy(tensor_" << fOutput;
out << SP << "Copy(tensor_" << fOutput;
if (i > 0)
out << offset;
offset += " + " + length;
Expand Down
23 changes: 16 additions & 7 deletions tmva/sofie/inc/TMVA/ROperator_Conv.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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";
Expand Down Expand Up @@ -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<float>(tensor_" << fNX
out << SP << SP << "UTILITY::Im2col<float>(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,
Expand All @@ -490,7 +499,7 @@ public:
out << "," << "tensor_" <<fNX << "_xcol);\n\n ";
} else {
// 3d im2col
out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::Im2col_3d<float>(tensor_" << fNX
out << SP << SP << "UTILITY::Im2col_3d<float>(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,
Expand All @@ -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, "
Expand Down Expand Up @@ -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<float>(tensor_" << fNX
out << SP << SP << "UTILITY::Im2col<float>(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,
Expand All @@ -550,7 +559,7 @@ public:
out << ", tensor_" << fNX << "_xcol);\n\n ";
} else {
// 3d im2col
out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::Im2col_3d<float>(tensor_" << fNX
out << SP << SP << "UTILITY::Im2col_3d<float>(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,
Expand All @@ -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, "
Expand Down
16 changes: 11 additions & 5 deletions tmva/sofie/inc/TMVA/ROperator_ConvTranspose.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,12 @@ void ROperator_ConvTranspose<T>::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 <typename T>
Expand All @@ -313,7 +319,7 @@ std::string ROperator_ConvTranspose<T>::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<float>(tensor_" << fNB
out << SP << SP << "float * data = UTILITY::BroadcastConvBias<float>(tensor_" << fNB
<< ", " << bsize << ", " << ConvertShapeToString(fShapeY) << ");\n";
out << SP << SP << "std::copy(data, data + " << ConvertShapeToLength(fShapeY) << ", tensor_" << fNBroadcastedB
<< ");\n";
Expand Down Expand Up @@ -484,7 +490,7 @@ std::string ROperator_ConvTranspose<T>::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<float>(tensor_" << fNX
out << SP << SP << "UTILITY::col2im<float>(tensor_" << fNX
<< "_xcol,"
// channels, height, width, kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w, dilation_h,
// dilation_w,
Expand All @@ -500,7 +506,7 @@ std::string ROperator_ConvTranspose<T>::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<float>(tensor_" << fNX
out << SP << SP << "UTILITY::Im2col_3d<float>(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,
Expand Down Expand Up @@ -537,7 +543,7 @@ std::string ROperator_ConvTranspose<T>::Generate(std::string OpName)
<< "_xcol , &" << OpName << "_m);\n";

if (fDim < 3) {
out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::col2im<float>(tensor_" << fNX
out << SP << SP << "UTILITY::col2im<float>(tensor_" << fNX
<< "_xcol,"
// channels, height, width, kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w, dilation_h,
// dilation_w,
Expand All @@ -554,7 +560,7 @@ std::string ROperator_ConvTranspose<T>::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<float>(tensor_" << fNX
out << SP << SP << "UTILITY::Im2col_3d<float>(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,
Expand Down
6 changes: 4 additions & 2 deletions tmva/sofie/inc/TMVA/ROperator_Expand.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down Expand Up @@ -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<size_t> shapeX;
Expand Down Expand Up @@ -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_"<<fNY<<");\n";
out << SP << "} else {\n";
out << SP << SP << "std::copy(tensor_" << fNX << ", " << "tensor_" << fNX << " + (" << lengthX << "), tensor_" << fNY << ");\n";
Expand Down
17 changes: 14 additions & 3 deletions tmva/sofie/inc/TMVA/ROperator_Gemm.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -273,6 +273,17 @@ namespace SOFIE{
}

model.AddNeededStdLib("algorithm");

// register the inference helper functions used by the generated code
if (fType == "float")
model.AddNeededHelperFunction("Gemm_Call");
// bias handling emits Copy / Fill, fused activation emits Relu
if (fNC != "") {
model.AddNeededHelperFunction("Copy");
model.AddNeededHelperFunction("Fill");
}
if (fActivation == EActivationType::RELU)
model.AddNeededHelperFunction("Relu");
}

std::string Generate(std::string opName) override {
Expand Down Expand Up @@ -447,7 +458,7 @@ namespace SOFIE{
else
out << "j;\n";

std::string prefix = SP2 + SP + "TMVA::Experimental::SOFIE::";
std::string prefix = SP2 + SP;
std::string target = "tensor_" + fNY;
if (sC.size() != 2) {
throw std::runtime_error("TMVA SOFIE Gemm Op - invalid rank for bias tensor " + ConvertDimShapeToString(fDimShapeC) + ConvertDimShapeToString(sC));
Expand Down Expand Up @@ -477,7 +488,7 @@ namespace SOFIE{

if (fType == "float"){

out << SP2 << "TMVA::Experimental::SOFIE::Gemm_Call(" << "tensor_" << fNY;
out << SP2 << "Gemm_Call(" << "tensor_" << fNY;
if (doStackMul) out << " + " << opName << "_y_offset";
out << ", "
<< (fAttrTransB ? "true, " : "false, ")
Expand Down Expand Up @@ -518,7 +529,7 @@ namespace SOFIE{
out << SP << "//--- applying RELU to output\n";
std::string tnsr = "tensor_" + fNY;
std::string reluSize = ConvertDimShapeToLength(fShapeY);
out << SP << "TMVA::Experimental::SOFIE::Relu(" << tnsr << ", " << tnsr << ", " << reluSize << ");\n";
out << SP << "Relu(" << tnsr << ", " << tnsr << ", " << reluSize << ");\n";
}

return out.str();
Expand Down
6 changes: 5 additions & 1 deletion tmva/sofie/inc/TMVA/ROperator_LayerNormalization.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,10 @@ public:
}

model.AddNeededStdLib("cmath");

// the generated init code may broadcast the bias with UnidirectionalBroadcast
if (!fNBroadcastedB.empty())
model.AddNeededHelperFunction("UnidirectionalBroadcast");
}

std::string GenerateInitCode() override
Expand All @@ -168,7 +172,7 @@ public:
if (!fNBroadcastedB.empty()) {
out << SP << "// Broadcasting the bias of LayerNormalization op\n";
out << SP << "{\n";
out << SP << SP << "float* data = TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast(tensor_";
out << SP << SP << "float* data = UTILITY::UnidirectionalBroadcast(tensor_";
out << fNB << ", " << ConvertDimShapeToString(fShapeB) << ", " << ConvertDimShapeToString(fShapeX) << ");\n";
out << SP << "std::copy(data, data + " << fLength << ", tensor_" << fNBroadcastedB << ");\n";
out << SP << "delete[] data;\n";
Expand Down
Loading
Loading