Skip to content

[tmva][sofie] Restructure emitted code to be differentiable with Clad #18332

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

Merged
guitargeek merged 2 commits into from
Mar 30, 2026
Merged

[tmva][sofie] Restructure emitted code to be differentiable with Clad #18332

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
855380c
Revert "[tmva][sofie] Fix NonZero to define max output shape values i…
guitargeek Mar 12, 2026
ac44395
[tmva][sofie] Restructure emitted code to be differentiable with Clad
guitargeek Apr 7, 2025
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
4 changes: 2 additions & 2 deletions tmva/sofie/inc/TMVA/RModel.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,6 @@ private:
std::vector<std::string> fOutputTensorNames;
std::vector<std::string> fInputTensorNames; // input tensor names using ONNX order



std::vector<std::unique_ptr<ROperator>> fOperators;

std::vector<std::shared_ptr<RModel>> fSubGraphs; ///<! sub-graph models (transient)
Expand Down Expand Up @@ -196,6 +194,8 @@ protected:
void GenerateIntermediateMemoryPool();
// Generate all session code
void GenerateSessionCode();
bool IsInputTensorShapeParam(std::string const &name) const;
std::vector<std::string> CollectTensorMemberNames(const std::string &input);

public:
const std::vector<std::string> & GetInputTensorNames() const { return fInputTensorNames; }
Expand Down
4 changes: 1 addition & 3 deletions tmva/sofie/inc/TMVA/ROperator.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,7 @@ public:
virtual std::vector<ETensorType> TypeInference(std::vector<ETensorType>) { return {}; };
virtual void Initialize(RModel&) = 0;
virtual std::string Generate(std::string OpName) = 0; //expect unique opName for each operator within the same RModel
// generate code for Session constructor before tensor allocation
virtual std::string GenerateSessionCtorCode() { return "";}
// generate initialization code for session constructor after tensor allocations
// generate initialization code for session constructor
virtual std::string GenerateInitCode() { return "";}
// generate some specific declaration code for Session
virtual std::string GenerateDeclCode() { return "";}
Expand Down
7 changes: 5 additions & 2 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 @@ -393,9 +393,12 @@ namespace SOFIE{
<< (fAttrTransB ? "true, " : "false, ")
<< (fAttrTransA ? "true, " : "false, ")
<< n << ", " << m << ", " << k << ", ";
out << std::setprecision(std::numeric_limits<float>::max_digits10) << fAttrAlpha << ", tensor_" << fNB;
// TODO: the cast to (float *) is not needed here from the C++ language perspective (the arguments to
// Gemm_Call are const already), but Clad bug https://github.com/vgvassilev/clad/issues/1721 is requiring
// us to do this cast to keep Clad working. Remove this hack once the Clad issue is fixed.
out << std::setprecision(std::numeric_limits<float>::max_digits10) << fAttrAlpha << ", (float*)tensor_" << fNB;
if (extraB) out << " + " << opName << "_B_offset";
out << ", tensor_" << fNA;
out << ", (float*)tensor_" << fNA; // TODO: same here
if (extraA) out << " + " << opName << "_A_offset";
out << ", " << std::setprecision(std::numeric_limits<float>::max_digits10) << fAttrBeta << ",";
// in the case of bias and no broadcasting needed
Expand Down
121 changes: 69 additions & 52 deletions tmva/sofie/inc/TMVA/ROperator_LSTM.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -390,45 +390,62 @@ std::string ROperator_LSTM<T>::GenerateSessionMembersCode(std::string opName)
size_t batch_size = (fAttrLayout == 0) ? fShapeX[1] : fShapeX[0];
size_t input_size = fShapeX[2];

struct Block {
std::string name;
size_t size;
};

std::vector<Block> blocks;

size_t ff_size = seq_length * batch_size * fAttrHiddenSize;
size_t hs_size = seq_length * num_directions * batch_size * fAttrHiddenSize;

// Layout-dependent buffers
if (fAttrLayout != 0) {
out << "std::vector<" << fType << "> fVec_" << opName << "_input = std::vector<" << fType << ">("
<< seq_length * batch_size * input_size << ");\n";
out << "std::vector<" << fType << "> fVec_" << opName << "_initial_hidden_state = std::vector<" << fType << ">("
<< num_directions * batch_size * fAttrHiddenSize << ");\n";
out << "std::vector<" << fType << "> fVec_" << opName << "_initial_cell_state = std::vector<" << fType << ">("
<< num_directions * batch_size * fAttrHiddenSize << ");\n";
blocks.push_back({"input", seq_length * batch_size * input_size});
blocks.push_back({"initial_hidden_state", num_directions * batch_size * fAttrHiddenSize});
blocks.push_back({"initial_cell_state", num_directions * batch_size * fAttrHiddenSize});
}
// Set the feedforward
size_t ff_size = seq_length * batch_size * fAttrHiddenSize;
out << "std::vector<" << fType << "> fVec_" << opName << "_ff_input_gate = std::vector<" << fType << ">(" << ff_size
<< ");\n";
out << "std::vector<" << fType << "> fVec_" << opName << "_ff_output_gate = std::vector<" << fType << ">(" << ff_size
<< ");\n";
out << "std::vector<" << fType << "> fVec_" << opName << "_ff_cell_gate = std::vector<" << fType << ">(" << ff_size
<< ");\n";

// Feedforward gates
blocks.push_back({"ff_input_gate", ff_size});
blocks.push_back({"ff_output_gate", ff_size});
blocks.push_back({"ff_cell_gate", ff_size});
if (fAttrInputForget == 0)
out << "std::vector<" << fType << "> fVec_" << opName << "_ff_forget_gate = std::vector<" << fType << ">("
<< ff_size << ");\n";
// gate results
size_t hs_size = seq_length * num_directions * batch_size * fAttrHiddenSize;
out << "std::vector<" << fType << "> fVec_" << opName << "_input_gate = std::vector<" << fType << ">(" << hs_size
<< ");\n";
out << "std::vector<" << fType << "> fVec_" << opName << "_output_gate = std::vector<" << fType << ">(" << hs_size
<< ");\n";
out << "std::vector<" << fType << "> fVec_" << opName << "_cell_gate = std::vector<" << fType << ">(" << hs_size
<< ");\n";
blocks.push_back({"ff_forget_gate", ff_size});

// Gate outputs
blocks.push_back({"input_gate", hs_size});
blocks.push_back({"output_gate", hs_size});
blocks.push_back({"cell_gate", hs_size});
if (fAttrInputForget == 0)
out << "std::vector<" << fType << "> fVec_" << opName << "_forget_gate = std::vector<" << fType << ">(" << hs_size
<< ");\n";
// cell state
out << "std::vector<" << fType << "> fVec_" << opName << "_cell_state = std::vector<" << fType << ">(" << hs_size
<< ");\n";
out << "std::vector<" << fType << "> fVec_" << opName << "_new_cell_state = std::vector<" << fType << ">(" << hs_size
<< ");\n";
// hiddden state
blocks.push_back({"forget_gate", hs_size});

// Cell state
blocks.push_back({"cell_state", hs_size});
blocks.push_back({"new_cell_state", hs_size});

// Hidden state (conditional)
if (fAttrLayout != 0 || fNY.empty()) {
out << "std::vector<" << fType << "> fVec_" << opName << "_hidden_state = std::vector<" << fType << ">("
<< hs_size << ");\n";
blocks.push_back({"hidden_state", hs_size});
}

// Compute total size
size_t total_size = 0;
for (const auto &b : blocks) {
total_size += b.size;
}

// Backing storage
out << "std::vector<" << fType << "> fVec_" << opName << "_buffer = std::vector<" << fType << ">(" << total_size
<< ");\n";

// Emit pointers
std::size_t offset = 0;
for (const auto &b : blocks) {
out << fType << "* fVec_" << opName << "_" << b.name << " = fVec_" << opName << "_buffer.data() + " << offset
<< ";\n";
offset += b.size;
}

out << "\n";
Expand All @@ -452,7 +469,7 @@ auto ROperator_LSTM<T>::Generate(std::string OpName) -> std::string
out << SP << fType << " const *" << OpName << "_input = tensor_" << fNX << ";\n";
} else {
if (fUseSession)
out << SP << fType << " * " << OpName << "_input = this->fVec_" << OpName << "_input.data();\n";
out << SP << fType << " * " << OpName << "_input = this->fVec_" << OpName << "_input;\n";
else
out << SP << fType << " " << OpName << "_input[" << seq_length * batch_size * input_size << "] = {0};\n";

Expand All @@ -470,11 +487,11 @@ auto ROperator_LSTM<T>::Generate(std::string OpName) -> std::string
// Set the initial hidden state
if (!fNInitial_h.empty()) {
if (fAttrLayout == 0) {
out << SP << fType << " *" << OpName << "_initial_hidden_state = " << " tensor_" << fNInitial_h << ";\n";
out << SP << fType << " const*" << OpName << "_initial_hidden_state = " << " tensor_" << fNInitial_h << ";\n";
} else {
if (fUseSession)
out << SP << fType << " * " << OpName << "_initial_hidden_state = this->fVec_" << OpName
<< "_initial_hidden_state.data();\n";
out << SP << fType << " const* " << OpName << "_initial_hidden_state = this->fVec_" << OpName
<< "_initial_hidden_state;\n";
else
out << SP << fType << " " << OpName << "_initial_hidden_state["
<< num_directions * batch_size * fAttrHiddenSize << "] = {0};\n";
Expand All @@ -494,11 +511,11 @@ auto ROperator_LSTM<T>::Generate(std::string OpName) -> std::string
// Set the initial cell state
if (!fNInitial_c.empty()) {
if (fAttrLayout == 0) {
out << SP << fType << " *" << OpName << "_initial_cell_state = " << " tensor_" << fNInitial_c << ";\n";
out << SP << fType << " const*" << OpName << "_initial_cell_state = " << " tensor_" << fNInitial_c << ";\n";
} else {
if (fUseSession)
out << SP << fType << " * " << OpName << "_initial_cell_state = this->fVec_" << OpName
<< "_initial_cell_state.data();\n";
out << SP << fType << " const* " << OpName << "_initial_cell_state = this->fVec_" << OpName
<< "_initial_cell_state;\n";
else
out << SP << fType << " " << OpName << "_initial_cell_state["
<< num_directions * batch_size * fAttrHiddenSize << "] = {0};\n";
Expand All @@ -518,12 +535,12 @@ auto ROperator_LSTM<T>::Generate(std::string OpName) -> std::string
// Set the feedforward
size_t ff_size = seq_length * batch_size * fAttrHiddenSize;
if (fUseSession) {
out << SP << fType << " * " << OpName << "_ff_input_gate = this->fVec_" << OpName << "_ff_input_gate.data();\n";
out << SP << fType << " * " << OpName << "_ff_output_gate = this->fVec_" << OpName << "_ff_output_gate.data();\n";
out << SP << fType << " * " << OpName << "_ff_cell_gate = this->fVec_" << OpName << "_ff_cell_gate.data();\n";
out << SP << fType << " * " << OpName << "_ff_input_gate = this->fVec_" << OpName << "_ff_input_gate;\n";
out << SP << fType << " * " << OpName << "_ff_output_gate = this->fVec_" << OpName << "_ff_output_gate;\n";
out << SP << fType << " * " << OpName << "_ff_cell_gate = this->fVec_" << OpName << "_ff_cell_gate;\n";
if (fAttrInputForget == 0) {
out << SP << fType << " * " << OpName << "_ff_forget_gate = this->fVec_" << OpName
<< "_ff_forget_gate.data();\n";
<< "_ff_forget_gate;\n";
}
} else {
out << SP << fType << " " << OpName << "_ff_input_gate[" << ff_size << "] = {0};\n";
Expand All @@ -536,11 +553,11 @@ auto ROperator_LSTM<T>::Generate(std::string OpName) -> std::string
// Set the gates
size_t hidden_state_size = seq_length * num_directions * batch_size * fAttrHiddenSize;
if (fUseSession) {
out << SP << fType << " * " << OpName << "_input_gate = this->fVec_" << OpName << "_input_gate.data();\n";
out << SP << fType << " * " << OpName << "_output_gate = this->fVec_" << OpName << "_output_gate.data();\n";
out << SP << fType << " * " << OpName << "_cell_gate = this->fVec_" << OpName << "_cell_gate.data();\n";
out << SP << fType << " * " << OpName << "_input_gate = this->fVec_" << OpName << "_input_gate;\n";
out << SP << fType << " * " << OpName << "_output_gate = this->fVec_" << OpName << "_output_gate;\n";
out << SP << fType << " * " << OpName << "_cell_gate = this->fVec_" << OpName << "_cell_gate;\n";
if (fAttrInputForget == 0) {
out << SP << fType << " * " << OpName << "_forget_gate = this->fVec_" << OpName << "_forget_gate.data();\n";
out << SP << fType << " * " << OpName << "_forget_gate = this->fVec_" << OpName << "_forget_gate;\n";
}
} else {
out << SP << fType << " " << OpName << "_input_gate[" << hidden_state_size << "] = {0};\n";
Expand All @@ -552,8 +569,8 @@ auto ROperator_LSTM<T>::Generate(std::string OpName) -> std::string
}
// Set the cell state and the new cell state = h(cell state)
if (fUseSession) {
out << SP << fType << " * " << OpName << "_cell_state = this->fVec_" << OpName << "_cell_state.data();\n";
out << SP << fType << " * " << OpName << "_new_cell_state = this->fVec_" << OpName << "_new_cell_state.data();\n";
out << SP << fType << " * " << OpName << "_cell_state = this->fVec_" << OpName << "_cell_state;\n";
out << SP << fType << " * " << OpName << "_new_cell_state = this->fVec_" << OpName << "_new_cell_state;\n";
} else {
out << SP << fType << " " << OpName << "_cell_state[" << hidden_state_size << "] = {0};\n";
out << SP << fType << " " << OpName << "_new_cell_state[" << hidden_state_size << "] = {0};\n";
Expand All @@ -564,7 +581,7 @@ auto ROperator_LSTM<T>::Generate(std::string OpName) -> std::string
out << SP << fType << " *" << OpName << "_hidden_state = tensor_" << fNY << ";\n";
} else {
if (fUseSession) {
out << SP << fType << " * " << OpName << "_hidden_state = this->fVec_" << OpName << "_hidden_state.data();\n";
out << SP << fType << " * " << OpName << "_hidden_state = this->fVec_" << OpName << "_hidden_state;\n";
} else {
out << SP << fType << " " << OpName << "_hidden_state[" << hidden_state_size << "] = {0};\n";
}
Expand Down
10 changes: 5 additions & 5 deletions tmva/sofie/inc/TMVA/ROperator_NonZero.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,25 +92,25 @@ public:
fShapeY.resize(2);
fShapeY[0] = fShapeX.size();

// flag -1 to define the shape variable in the constructor code and not in the constructor signature
fShapeY[1] = Dim{std::string("v_NonZero_") + fNX, static_cast<size_t>(-1) };
// identify as -1 since we will declare maximum as size of input
fShapeY[1] = Dim{std::string("v_NonZero_") + fNX, static_cast<size_t>(-1)};

model.AddIntermediateTensor(fNY, ETensorType::INT64, fShapeY);
if (model.Verbose()) {
std::cout << "NonZero : " << fNX << " -> " << fNY << " " << ConvertDimShapeToString(fShapeY) << std::endl;
}
}
}
std::string GenerateSessionCtorCode() override {

std::string GenerateSessionMembersCode(std::string /*opName*/) override {
if (fIsOutputConstant) return "";
// define output value used as max non zero with max size = input shape * N
auto inputLength = ConvertDimShapeToLength(fShapeX);
std::stringstream out;
out << SP << "size_t v_NonZero_" << fNX << " = " << inputLength << ";\n";
out << SP << "size_t fV_NonZero_" << fNX << " = " << inputLength << ";\n";
return out.str();
}


std::string Generate(std::string opName) override {
if (fIsOutputConstant) {
return "";
Expand Down
46 changes: 33 additions & 13 deletions tmva/sofie/inc/TMVA/ROperator_RNN.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -308,18 +308,38 @@ std::string ROperator_RNN<T>::GenerateSessionMembersCode(std::string opName)
size_t batch_size = (fAttrLayout == 0) ? fShapeX[1] : fShapeX[0];
size_t input_size = fShapeX[2];

struct Block {
std::string name;
size_t size;
};

std::vector<Block> blocks;

if (fAttrLayout != 0) {
out << "std::vector<" << fType << "> fVec_" << opName << "_input = std::vector<" << fType << ">("
<< seq_length * batch_size * input_size << ");\n";
out << "std::vector<" << fType << "> fVec_" << opName << "_initial_hidden_state = std::vector<" << fType << ">("
<< num_directions * batch_size * fAttrHiddenSize << ");\n";
blocks.push_back({"input", seq_length * batch_size * input_size});
blocks.push_back({"initial_hidden_state", num_directions * batch_size * fAttrHiddenSize});
}
out << "std::vector<" << fType << "> fVec_" << opName << "_feedforward = std::vector<" << fType << ">("
<< seq_length * batch_size * fAttrHiddenSize << ");\n";

blocks.push_back({"feedforward", seq_length * batch_size * fAttrHiddenSize});
if (fAttrLayout != 0 || fNY.empty()) {
out << "std::vector<" << fType << "> fVec_" << opName << "_hidden_state = std::vector<" << fType << ">("
<< seq_length * num_directions * batch_size * fAttrHiddenSize << ");\n";
blocks.push_back({"hidden_state", seq_length * num_directions * batch_size * fAttrHiddenSize});
}

// Compute total size
size_t total_size = 0;
for (const auto &b : blocks) {
total_size += b.size;
}

// Emit backing storage
out << "std::vector<" << fType << "> fVec_" << opName << "_buffer = std::vector<" << fType << ">(" << total_size
<< ");\n";

// Emit pointers
std::size_t offset = 0;
for (const auto &b : blocks) {
out << fType << "* fVec_" << opName << "_" << b.name << " = fVec_" << opName << "_buffer.data() + " << offset
<< ";\n";
offset += b.size;
}

out << "\n";
Expand All @@ -346,7 +366,7 @@ auto ROperator_RNN<T>::Generate(std::string OpName) -> std::string
}
} else {
if (fUseSession)
out << SP << fType << " * " << OpName << "_input = this->fVec_" << OpName << "_input.data();\n";
out << SP << fType << " * " << OpName << "_input = this->fVec_" << OpName << "_input;\n";
else
out << SP << fType << " " << OpName << "_input[" << seq_length * batch_size * input_size << "];\n";
out << SP << "for(size_t seq = 0; seq < " << seq_length << "; seq++) {\n";
Expand All @@ -367,7 +387,7 @@ auto ROperator_RNN<T>::Generate(std::string OpName) -> std::string
} else {
if (fUseSession)
out << SP << fType << " * " << OpName << "_initial_hidden_state = this->fVec_" << OpName
<< "_initial_hidden_state.data();\n";
<< "_initial_hidden_state;\n";
else
out << fType << " " << OpName << "_initial_hidden_state[" << num_directions * batch_size * fAttrHiddenSize
<< "] = {0};\n";
Expand All @@ -385,7 +405,7 @@ auto ROperator_RNN<T>::Generate(std::string OpName) -> std::string
}

if (fUseSession)
out << SP << fType << " * " << OpName << "_feedforward = this->fVec_" << OpName << "_feedforward.data();\n";
out << SP << fType << " * " << OpName << "_feedforward = this->fVec_" << OpName << "_feedforward;\n";
else
out << SP << fType << " " << OpName << "_feedforward[" << seq_length * batch_size * fAttrHiddenSize
<< "] = {0};\n";
Expand All @@ -395,7 +415,7 @@ auto ROperator_RNN<T>::Generate(std::string OpName) -> std::string
out << SP << fType << " *" << OpName << "_hidden_state = tensor_" << fNY << ";\n";
} else {
if (fUseSession)
out << SP << fType << " * " << OpName << "_hidden_state = this->fVec_" << OpName << "_hidden_state.data();\n";
out << SP << fType << " * " << OpName << "_hidden_state = this->fVec_" << OpName << "_hidden_state;\n";
else
out << SP << fType << " " << OpName << "_hidden_state["
<< seq_length * num_directions * batch_size * fAttrHiddenSize << "] = {0};\n";
Expand Down
10 changes: 0 additions & 10 deletions tmva/sofie/inc/TMVA/SOFIE_common.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -681,16 +681,6 @@ void col2im(const Dtype* data_col, const int channels,
//std::cout << "finishing col2imp" << std::endl;
}

// Used at the end of infer() to fill the return object.
template <class T>
void FillOutput(T const *arr, std::vector<T> &out, std::size_t n)
{
out.resize(n);
for (std::size_t i = 0; i < n; ++i) {
out[i] = arr[i];
}
}

} // end namespace UTILITY

namespace BLAS{
Expand Down
Loading
Loading