Skip to content

[tmva][sofie] Add support for optimal memory allocation of dynamic tensors #20434

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
lmoneta wants to merge 10 commits into
base: master
Choose a base branch
from
Open

[tmva][sofie] Add support for optimal memory allocation of dynamic tensors #20434

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
9edf0fa
[tmva][sofie] Apply fixes for supporting Dynamic tensors
lmoneta Nov 10, 2025
aa4d008
[tmva][sofie] Remove special case handling bool outputs
lmoneta Nov 10, 2025
6b5c35b
[tmva][sofie] Add support for greedy memory allocation for dynammic t…
lmoneta Nov 12, 2025
556f0d7
[tmva][sofie] Fix an issue in genereting code for dynamic tensor when…
lmoneta Nov 14, 2025
3e6691e
[tmva][sofie] Fix stacked MatMul and speedup LayerNorm
lmoneta Dec 17, 2025
2fba9e5
[tmva][sofie] Fix issue with order execution of tensors in GNN models
lmoneta Dec 20, 2025
55fb3e7
[tmva][sofie] Apply fixes for the TestCustiomModelsFrom ROOT
lmoneta Jan 5, 2026
4b87a64
[tmva][sofie] Sort in alphabetical order the shape parameter in Sessi…
lmoneta Jan 7, 2026
b95b9a7
[tmva][sofie] Do not permor broadcast of bias tensor in Gemm in Sessi…
lmoneta Jan 8, 2026
bdd7fcc
[tmva][sofie] Add alias tensors
lmoneta Jan 8, 2026
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
6 changes: 3 additions & 3 deletions tmva/sofie/inc/TMVA/RFunction.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ public:

class RFunction_Update: public RFunction {
protected:
std::shared_ptr<RModel> function_block;
std::shared_ptr<RModel> fFunction_block;
FunctionTarget fTarget;
GraphType fGraphType;
std::vector<std::string> fInputTensors;
Expand All @@ -50,9 +50,9 @@ public:
void AddInputTensors(const std::vector<std::vector<std::size_t>>& inputShapes);
void AddInputTensors(const std::vector<std::vector<Dim>>& inputShapes);
std::shared_ptr<RModel> GetFunctionBlock() {
return function_block;
return fFunction_block;
}
std::string GenerateModel(const std::string& filename, long read_pos = 0, long block_size = -1);
std::string GenerateModel(const std::string& filename, long read_pos = 0, long block_size = -1, bool verbose = false);
std::string Generate(const std::vector<std::string>& inputPtrs);
FunctionTarget GetFunctionTarget() {
return fTarget;
Expand Down
15 changes: 10 additions & 5 deletions tmva/sofie/inc/TMVA/RModel.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ private:
std::unordered_map<std::string, DynamicTensorInfo> fDynamicTensorInfos;
std::unordered_map<std::string, std::pair<std::vector<Dim>, bool>> fShapeTensors; // constant tensors describing a shape
std::unordered_map<std::string, std::string> fShapeParams; // parameters defining the dynamic shape (e.g. batch size), store also its default value
std::unordered_map<std::string, std::string> fAliasTensors; // list of alias tensors
std::vector<std::string> fDimShapeNames; // parameter names used to define the shapes
std::vector<std::string> fOutputTensorNames;
std::vector<std::string> fInputTensorNames; // input tensor names using ONNX order
Expand Down Expand Up @@ -82,6 +83,8 @@ public:
void AddConstantTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape,
std::shared_ptr<void> data);

void AddAliasTensor(const std::string & tensor_name, const std::string & orig_tensor_name);


template<class T>
void AddConstantTensor(const std::string & name, const std::vector<size_t> & shape, const T * data) {
Expand Down Expand Up @@ -130,6 +133,8 @@ public:
bool IsReadyInputTensor(const std::string &name) const;
/// check if a tensor is a shape tensor
bool IsShapeTensor(const std::string & name) const;
/// check if a tensor is a alias tensor
bool IsAliasTensor(const std::string & name) const;

// Add intermediate tensor
void AddIntermediateTensor(std::string tensor_name, ETensorType type, std::vector<Dim> dim_shape);
Expand Down Expand Up @@ -205,8 +210,8 @@ public:
void ReadInitializedTensorsFromFile(long);
long WriteInitializedTensorsToFile(std::string filename = "");

void PrintIntermediateTensors();
void PrintOutputTensors();
void PrintIntermediateTensors() const;
void PrintOutputTensors() const;
void OutputGenerated(std::string filename = "", bool append = false);
std::vector<std::string> GetOutputTensorNames() { return fOutputTensorNames; }
void SetFilename(std::string filename) { fName = filename; }
Expand All @@ -224,9 +229,9 @@ public:
}
*/

void PrintRequiredInputTensors();
void PrintInitializedTensors();
void PrintDynamicTensors();
void PrintRequiredInputTensors() const;
void PrintInitializedTensors() const;
void PrintDynamicTensors() const;
void HeadInitializedTensors(std::string name, int n_print = 50);

bool UseSession() const { return fUseSession; }
Expand Down
2 changes: 1 addition & 1 deletion tmva/sofie/inc/TMVA/ROperator_BasicBinary.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,7 +192,7 @@ public:
dataY[i] = BinaryOperatorTrait<T, Op>::Func(dataA[i], dataB[i]);
}
model.AddConstantTensor<T>(fNY, fShapeY, dataY.data());
// flag tensors to not be written in the weight file
// flag tensors to not be written in the generated code or weight file
model.SetNotWritableInitializedTensor(nameA);
model.SetNotWritableInitializedTensor(nameB);
fIsOutputConstant = true;
Expand Down
189 changes: 146 additions & 43 deletions tmva/sofie/inc/TMVA/ROperator_BasicNary.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,11 @@ struct NaryOperatorTraits<T, EBasicNaryOperator::Max> {
static const std::string Name() {return "Max";}
static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
std::stringstream out;
out << "\t" << "\t" << res << " = " << inputs[0] << ";\n";
out << res << " = std::max({ " << inputs[0];
for (size_t i = 1; i < inputs.size(); i++) {
out << "\t" << "\t" << res << " = std::max(" << res << ", " << inputs[i] << ");\n";
out << ", " << inputs[i];
}
out << "});\n";
return out.str();
}
};
Expand All @@ -36,10 +37,11 @@ struct NaryOperatorTraits<T, EBasicNaryOperator::Min> {
static const std::string Name() {return "Min";}
static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
std::stringstream out;
out << "\t" << "\t" << res << " = " << inputs[0] << ";\n";
out << res << " = std::min({ " << inputs[0];
for (size_t i = 1; i < inputs.size(); i++) {
out << "\t" << "\t" << res << " = std::min(" << res << ", " << inputs[i] << ");\n";
out << ", " << inputs[i];
}
out << "});\n";
return out.str();
}
};
Expand All @@ -52,7 +54,7 @@ struct NaryOperatorTraits<float, EBasicNaryOperator::Mean> {
static const std::string Name() {return "Mean";}
static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
std::stringstream out;
out << "\t" << "\t" << res << " = (" << inputs[0];
out << res << " = (" << inputs[0];
for (size_t i = 1; i < inputs.size(); i++) {
out << " + " << inputs[i];
}
Expand All @@ -66,7 +68,7 @@ struct NaryOperatorTraits<T, EBasicNaryOperator::Sum> {
static const std::string Name() {return "Sum";}
static std::string Op(const std::string& res, std::vector<std::string>& inputs) {
std::stringstream out;
out << "\t" << "\t" << res << " = " << inputs[0];
out << res << " = " << inputs[0];
for (size_t i = 1; i < inputs.size(); i++) {
out << " + " << inputs[i];
}
Expand All @@ -83,10 +85,11 @@ private:

std::vector<std::string> fNInputs;
std::string fNY;
std::vector<std::vector<size_t>> fShapeInputs;
std::vector<std::vector<Dim>> fShapeInputs;

std::vector<std::string> fNBroadcastedInputs;
std::vector<size_t> fShapeY;
std::vector<Dim> fDimShapeY;

bool fBroadcast = false;

Expand Down Expand Up @@ -119,64 +122,164 @@ public:
}

void Initialize(RModel& model) override {
std::vector<std::vector<size_t>> inputShapes;
for (auto &it : fNInputs) {
if (!model.CheckIfTensorAlreadyExist(it)) {
throw std::runtime_error("TMVA SOFIE BasicNary Op Input Tensor " + it + " is not found in model");
}
fShapeInputs.push_back(model.GetTensorShape(it));
fShapeInputs.push_back(model.GetDimTensorShape(it));
if (fNInputs.size()> 2) {
if (model.IsDimInputTensor(it))
throw std::runtime_error("TMVA SOFIE BasicNary : supports only 2 inputs for dynamic tensors");
else
inputShapes.push_back(model.GetTensorShape(it));
}
}
// Find the common shape of the input tensors
fShapeY = UTILITY::MultidirectionalBroadcastShape(fShapeInputs);
model.AddIntermediateTensor(fNY, model.GetTensorType(fNInputs[0]), fShapeY);
// Broadcasting
size_t N = fNInputs.size();
fNBroadcastedInputs.reserve(N);
for (size_t i = 0; i < N; i++) {
if (!UTILITY::AreSameShape(model.GetTensorShape(fNInputs[i]), fShapeY)) {
fBroadcast = true;
std::string name = "Broadcasted" + fNInputs[i];
model.AddIntermediateTensor(name, model.GetTensorType(fNInputs[0]), fShapeY);
fNBroadcastedInputs.emplace_back("tensor_" + name);
} else {
fNBroadcastedInputs.emplace_back("tensor_" + fNInputs[i]);
if (fShapeInputs.size() > 2 ) {
// support dynamic tensors now for input list of size=2
auto shapeY = UTILITY::MultidirectionalBroadcastShape(inputShapes);
fDimShapeY = ConvertShapeToDim(shapeY);
} else if (fShapeInputs.size() == 2 ) {
auto ret = UTILITY::MultidirectionalBroadcastShape(fShapeInputs[0], fShapeInputs[1]);
// use same code as in BinaryOperator (need to extend for input sizes > 2)
fBroadcast = ret.first;
fDimShapeY = ret.second;
// case of all parametric shapes and MultiDirectionalBroadcastShape return the max of the 2
// need to do before we declare the output tensor shape and the broadcasted ones
if (ret.first & 4) {
// check if one of the parameter is an input dimension
// define function to find this
auto IsInputDimParam = [&](const std::string &p) {
auto inputNames = model.GetInputTensorNames();
for (auto &input : inputNames) {
for (auto &i_s : model.GetDimTensorShape(input)) {
if (i_s.isParam && i_s.param == p)
return true;
}
}
return false;
};
auto & shapeA = fShapeInputs[0];
auto & shapeB = fShapeInputs[1];
for (size_t i = 0; i < fDimShapeY.size(); i++) {
auto &s = fDimShapeY[i];
if (s.isParam && s.param.find("std::max") != std::string::npos) {
if (IsInputDimParam(shapeA[i].param)) {
// case dim is 1 we indicate that the input parameter is equal to 1
if (shapeA[i].dim != 1)
s = shapeA[i];
else
s = shapeB[i];
} else if (IsInputDimParam(shapeB[i].param)) {
if (shapeB[i].dim != 1)
s = shapeB[i];
else
s = shapeA[i];
}
}
}
}
} else if (fShapeInputs.size() == 1 ) {
fDimShapeY = fShapeInputs[0];
}
if (!fShapeY.empty())
model.AddIntermediateTensor(fNY, model.GetTensorType(fNInputs[0]), fShapeY);
else
model.AddIntermediateTensor(fNY, model.GetTensorType(fNInputs[0]), fDimShapeY);


fType = ConvertTypeToString(model.GetTensorType(fNInputs[0]));

if (model.Verbose()) {
std::cout << NaryOperatorTraits<T, Op>::Name() << " : ";
if (fNInputs.size() == 2)
std::cout << ConvertShapeToString(fShapeInputs[0]) << " , "
<< ConvertShapeToString(fShapeInputs[1]);
std::cout << " --> " << ConvertShapeToString(fDimShapeY) << std::endl;
}
}

std::string Generate(std::string OpName) override {
OpName = "op_" + OpName;
if (fShapeY.empty()) {
if (fDimShapeY.empty()) {
throw std::runtime_error("TMVA SOFIE BasicNary called to Generate without being initialized first");
}
std::stringstream out;
size_t length = ConvertShapeToLength(fShapeY);
auto length = ConvertDimShapeToLength(fDimShapeY);
out << SP << "\n//------ BasicNary operator\n";
if (fBroadcast) {
for (size_t i = 0; i < fNInputs.size(); i++) {
if (fNBroadcastedInputs[i] != fNInputs[i]) {
out << SP << SP << "// Broadcasting " << fNInputs[i] << " to " << ConvertShapeToString(fShapeY) << "\n";
out << SP << SP << "{\n";
out << SP << SP << SP << fType << "* data = TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<" << fType << ">(tensor_" + fNInputs[i] << ", " << ConvertShapeToString(fShapeInputs[i]);
out << ", " << ConvertShapeToString(fShapeY) << ");\n";
out << SP << SP << SP << "std::copy(data, data + " << length << ", " << fNBroadcastedInputs[i] << ");\n";
out << SP << SP << SP << "delete[] data;\n";
out << SP << SP << "}\n";
}
}
}

if (fNInputs.size() == 1) {
int nInputs = fNInputs.size();

if (nInputs == 1) {
out << SP << "std::copy(tensor_" << fNInputs[0] << ", tensor_" << fNInputs[0] << " + ";
out << length << ", tensor_" << fNY << ");\n";
} else {
std::vector<std::string> inputs(fNBroadcastedInputs.size());
for (size_t i = 0; i < fNBroadcastedInputs.size(); i++) {
inputs[i] = fNBroadcastedInputs[i] + "[id]";

// implement operator without broadcasting, but using loos on all indices
Copy link
Contributor

@sanjibansg sanjibansg Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
// implement operator without broadcasting, but using loos on all indices
// implement operator without broadcasting, but using loops on all indices

std::vector<std::vector<Dim>> inputStrides(nInputs);
for (int i = 0; i < nInputs; i++)
inputStrides[i] = UTILITY::ComputeStrideFromShape(fShapeInputs[i]);

auto stridesY = UTILITY::ComputeStrideFromShape(fDimShapeY);

// make loop on output indices
std::string compute_idx_Y;
int nloop = 0;
if (fDimShapeY.empty() ||
std::all_of(fDimShapeY.begin(), fDimShapeY.end(), [](Dim d) { return d.dim == 1 || d.GetVal() == "1"; })) {
compute_idx_Y = "0";
} else {
for (size_t i = 0; i < fDimShapeY.size(); ++i) {
if (fDimShapeY[i].dim != 1 && fDimShapeY[i].GetVal() != "1") {
nloop++;
for (int j = 0; j < nloop; j++) out << SP;
out << "for (size_t idx_" << i << " = 0; idx_" << i << " < " << fDimShapeY[i]
<< "; ++idx_" << i << "){\n";
compute_idx_Y += "idx_" + std::to_string(i);
if (stridesY[i].GetVal() != "1")
compute_idx_Y += " * " + stridesY[i].GetVal();
compute_idx_Y += " + ";
}
}
// remove last 3 characters " + "
for (int j = 0; j < 3; j++)
compute_idx_Y.pop_back();
}
// find indices for input tensors
std::vector<std::string> inputs(nInputs);
for (int ipt = 0; ipt < nInputs; ipt++ ) {
std::string compute_idx_X;
auto & shape = fShapeInputs[ipt];
auto & stride = inputStrides[ipt];
if (shape.empty() ||
std::all_of(shape.begin(), shape.end(), [](Dim d) { return d.dim == 1 || d.GetVal() == "1"; })) {
compute_idx_X = "0";
} else {
for (size_t i = 0; i < shape.size(); ++i) {
if (shape[i].dim == 1 || shape[i].GetVal() == "1")
continue;
compute_idx_X += "idx_" + std::to_string(i + (fDimShapeY.size() - shape.size()));
if (stride[i].GetVal() != "1")
compute_idx_X += " * " + stride[i].GetVal();
compute_idx_X += " + ";
}
// remove last 3 character " + "
for (int j = 0; j < 3; j++)
compute_idx_X.pop_back();
}
inputs[ipt] = "tensor_" + fNInputs[ipt] + "[" + compute_idx_X + "]";
}

// perform the operation
for (int j = 0; j < nloop + 1; j++) out << SP;
std::string output = "tensor_" + fNY + "[" + compute_idx_Y + "]";
out << NaryOperatorTraits<T,Op>::Op(output, inputs);

for (int i = nloop; i > 0; i--) {
for (int j = 0; j < i; j++) out << SP;
out << "}\n";
}
out << SP << "for (size_t id = 0; id < " << length << "; id++) {\n";
out << NaryOperatorTraits<T,Op>::Op("tensor_" + fNY + "[id]", inputs);
out << SP << "}\n";
}
return out.str();
}
Expand Down
6 changes: 3 additions & 3 deletions tmva/sofie/inc/TMVA/ROperator_BatchNormalization.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,8 +141,8 @@ public:
}
}

std::string Generate(std::string OpName) override {
OpName = "op_" + OpName;
std::string Generate(std::string opName) override {
opName = "op_" + opName;
if (fShapeX.empty()){
throw std::runtime_error("TMVA SOFIE Batch Normalization called to Generate without being initialized first");
}
Expand All @@ -158,7 +158,7 @@ public:
spatial_dim = ConvertDimShapeToLength( spatialShape);
}

out << "\n\n//---- BatchNorm" << (fActivation == EActivationType::RELU ? " + ReLU" : "") << "\n";
out << "\n\n//---- BatchNorm" << (fActivation == EActivationType::RELU ? " + ReLU " : " ") << opName << "\n";
out << SP << "{\n";
out << SP << " size_t i = 0;\n";
out << SP << " for (size_t n = 0; n < " << batchSize << "; ++n) {\n";
Expand Down
Loading
Loading