From 046c3aa1abc75a3487cf31a21b9e62637f63cfd9 Mon Sep 17 00:00:00 2001 From: Alper Han <45140028+SirAlper@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:33:41 +0300 Subject: [PATCH] [tmva][sofie] Add support for the ONNX PRelu operator Implement ROperator_PRelu and its ONNX parser. The slope is unidirectionally broadcast to the shape of X as required since opset 7: instead of expanding the slope to the size of X, the generated loop computes the slope index from the strides of X, skipping the broadcast axes, e.g. tensor_slope[(id / 4) % 3] for a [1,3,2,2] input with the per-channel [3,1,1] slope exported by PyTorch's nn.PReLU. A slope that cannot be broadcast to X is rejected in Initialize(). Only float is supported for now. Add four model tests (same shape, last-axis, per-channel and single slope) and two parser tests checking the error for a non-broadcastable slope and that a per-channel slope is not expanded. --- tmva/sofie/README.md | 1 + tmva/sofie/inc/TMVA/ROperator_PRelu.hxx | 102 +++++++++++++++++++ tmva/sofie/test/TestCustomModelsFromONNX.cxx | 36 +++++++ tmva/sofie/test/TestSofieParser.cxx | 54 ++++++++++ tmva/sofie/test/generate_input_models.py | 49 +++++++++ tmva/sofie_parsers/CMakeLists.txt | 1 + tmva/sofie_parsers/src/ParsePRelu.cxx | 52 ++++++++++ tmva/sofie_parsers/src/RModelParser_ONNX.cxx | 2 + 8 files changed, 297 insertions(+) create mode 100644 tmva/sofie/inc/TMVA/ROperator_PRelu.hxx create mode 100644 tmva/sofie_parsers/src/ParsePRelu.cxx diff --git a/tmva/sofie/README.md b/tmva/sofie/README.md index 3a992fb32e888..c571d2e948c23 100644 --- a/tmva/sofie/README.md +++ b/tmva/sofie/README.md @@ -130,6 +130,7 @@ std::vector supportedOperators = parser.GetRegisteredOperators(); - [x] Min - [x] Mul - [x] Neg +- [x] PRelu - [x] Pad - [x] Pow - [x] RNN diff --git a/tmva/sofie/inc/TMVA/ROperator_PRelu.hxx b/tmva/sofie/inc/TMVA/ROperator_PRelu.hxx new file mode 100644 index 0000000000000..c49a29b89aee7 --- /dev/null +++ b/tmva/sofie/inc/TMVA/ROperator_PRelu.hxx @@ -0,0 +1,102 @@ +#ifndef TMVA_SOFIE_ROPERATOR_PRelu +#define TMVA_SOFIE_ROPERATOR_PRelu + +#include "TMVA/SOFIE_common.hxx" +#include "TMVA/ROperator.hxx" +#include "TMVA/RModel.hxx" + +#include + +namespace TMVA { +namespace Experimental { +namespace SOFIE { + +template +class ROperator_PRelu final : public ROperator { +private: + /* Attributes */ + std::string fNX; + std::string fNSlope; + std::string fNY; + std::vector fShapeX; + std::vector fShapeSlope; + std::string fType; + +public: + ROperator_PRelu() {} + ROperator_PRelu(std::string nameX, std::string nameSlope, std::string nameY) + : fNX(UTILITY::Clean_name(nameX)), fNSlope(UTILITY::Clean_name(nameSlope)), fNY(UTILITY::Clean_name(nameY)) + { + if (std::is_same::value) { + fType = "float"; + } else { + throw std::runtime_error("TMVA SOFIE Encountered unsupported type parsing a PRelu operator"); + } + fInputTensorNames = {fNX, fNSlope}; + fOutputTensorNames = {fNY}; + } + void Initialize(RModel &model) override + { + if (model.CheckIfTensorAlreadyExist(fNX) == false || model.CheckIfTensorAlreadyExist(fNSlope) == false) { + throw std::runtime_error("TMVA SOFIE PRelu Op Input Tensor is not found in model"); + } + fShapeX = model.GetTensorShape(fNX); + fShapeSlope = model.GetTensorShape(fNSlope); + + // check that slope is unidirectionally broadcastable to X (throws otherwise); + // this also prepends 1s to fShapeSlope so that it has the same rank as X + try { + UTILITY::UnidirectionalBroadcastShape(fShapeSlope, fShapeX); + } catch (const std::runtime_error &e) { + throw std::runtime_error("TMVA SOFIE PRelu Op slope tensor " + fNSlope + + " cannot be broadcast to input tensor " + fNX + ": " + e.what()); + } + + model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeX); + } + std::string Generate(std::string OpName) override + { + OpName = "op_" + OpName; + + if (fShapeX.empty() || fShapeSlope.empty()) { + throw std::runtime_error("TMVA SOFIE Operator PRelu called to Generate without being initialized first"); + } + + auto stridesX = UTILITY::ComputeStrideFromShape(fShapeX); + auto stridesSlope = UTILITY::ComputeStrideFromShape(fShapeSlope); + + std::string slopeIdx; + if (UTILITY::AreSameShape(fShapeX, fShapeSlope)) { + slopeIdx = "id"; + } else { + for (size_t k = 0; k < fShapeSlope.size(); k++) { + if (fShapeSlope[k] == 1) + continue; + if (!slopeIdx.empty()) + slopeIdx += " + "; + if (stridesX[k] == 1) + slopeIdx += "id % " + std::to_string(fShapeX[k]); + else + slopeIdx += "(id / " + std::to_string(stridesX[k]) + ") % " + std::to_string(fShapeX[k]); + if (stridesSlope[k] != 1) + slopeIdx += " * " + std::to_string(stridesSlope[k]); + } + if (slopeIdx.empty()) + slopeIdx = "0"; + } + std::stringstream out; + size_t length = ConvertShapeToLength(fShapeX); + out << "\n//------ PRELU\n"; + out << SP << "for (int id = 0; id < " << length << " ; id++){\n"; + out << SP << SP << "tensor_" << fNY << "[id] = ((tensor_" << fNX << "[id] >= 0 )? tensor_" << fNX + << "[id] : tensor_" << fNSlope << "[" << slopeIdx << "] * tensor_" << fNX << "[id]);\n"; + out << SP << "}\n"; + return out.str(); + } +}; + +} // namespace SOFIE +} // namespace Experimental +} // namespace TMVA + +#endif // TMVA_SOFIE_ROPERATOR_PRelu diff --git a/tmva/sofie/test/TestCustomModelsFromONNX.cxx b/tmva/sofie/test/TestCustomModelsFromONNX.cxx index f39859ca9fba2..550ff7ed0a952 100644 --- a/tmva/sofie/test/TestCustomModelsFromONNX.cxx +++ b/tmva/sofie/test/TestCustomModelsFromONNX.cxx @@ -647,6 +647,42 @@ TEST(ONNX, ArgMax_NegAxis_SelectLast) expectEqual(output, correct_output); } +TEST(ONNX, PRelu) +{ + SofieReference ref = readReference("PRelu"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "PRelu", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + +TEST(ONNX, PRelu_Broadcast) +{ + SofieReference ref = readReference("PRelu_Broadcast"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "PRelu_Broadcast", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + +TEST(ONNX, PRelu_Channel) +{ + SofieReference ref = readReference("PRelu_Channel"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "PRelu_Channel", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + +TEST(ONNX, PRelu_Scalar) +{ + SofieReference ref = readReference("PRelu_Scalar"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "PRelu_Scalar", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + // Elu on a tensor whose first dimension is only known at run time. TEST(ONNX, EluDynShape) { diff --git a/tmva/sofie/test/TestSofieParser.cxx b/tmva/sofie/test/TestSofieParser.cxx index 837a3007474c3..0bda5144b1339 100644 --- a/tmva/sofie/test/TestSofieParser.cxx +++ b/tmva/sofie/test/TestSofieParser.cxx @@ -187,6 +187,42 @@ GenerateGemmCode(const std::string &fileName, const std::vector &inp return code.str(); } +// Write a model holding a single PRelu node and return the code SOFIE generates +// for it. The slope is a graph input like X, which keeps the model file tiny. +std::string GeneratePReluCode(const std::string &fileName, const TestTensor &x, const TestTensor &slope) +{ + std::string node; + AppendBytesField(node, 1, x.name); // input + AppendBytesField(node, 1, slope.name); // input + AppendBytesField(node, 2, "y"); // output + AppendBytesField(node, 3, "prelu_0"); // name + AppendBytesField(node, 4, "PRelu"); // op_type + + std::string graph; + AppendBytesField(graph, 1, node); // node + AppendBytesField(graph, 2, "test_graph"); // name + AppendBytesField(graph, 11, FloatValueInfo(x)); // input + AppendBytesField(graph, 11, FloatValueInfo(slope)); // input + AppendBytesField(graph, 12, FloatValueInfo({"y", x.dims})); // output + + std::string opset; + AppendVarintField(opset, 2, 16); // version + std::string model; + AppendVarintField(model, 1, 8); // ir_version + AppendBytesField(model, 7, graph); // graph + AppendBytesField(model, 8, opset); // opset_import + + std::ofstream file(fileName, std::ios::binary); + file.write(model.data(), model.size()); + file.close(); + + RModel rmodel = RModelParser_ONNX{}.Parse(fileName); + rmodel.Generate(Options::kNoWeightFile); + std::ostringstream code; + rmodel.PrintGenerated(code); + return code.str(); +} + } // namespace // The "location" key of a tensor's external_data names the data file relative @@ -312,3 +348,21 @@ TEST(SOFIEParser, InputOfUnsupportedTypeThrows) RModel rmodel = RModelParser_ONNX{}.Parse("float16_input.onnx"); EXPECT_THROW(rmodel.Generate(Options::kNoWeightFile), std::runtime_error); } + +// A slope that cannot be unidirectionally broadcast to the shape of X is an +// error, not a silent out-of-bounds read in the generated code. +TEST(SOFIEParser, PReluWithNonBroadcastableSlopeThrows) +{ + EXPECT_THROW(GeneratePReluCode("prelu_bad_slope.onnx", {"x", {"3", "4", "5"}}, {"slope", {"4"}}), + std::runtime_error); +} + +// The per-channel slope exported by PyTorch's nn.PReLU is indexed in place +// instead of being expanded to the shape of X. +TEST(SOFIEParser, PReluIndexesAPerChannelSlopeWithoutExpandingIt) +{ + const std::string code = + GeneratePReluCode("prelu_channel.onnx", {"x", {"1", "3", "2", "2"}}, {"slope", {"3", "1", "1"}}); + + EXPECT_NE(code.find("tensor_slope[(id / 4) % 3]"), std::string::npos); +} diff --git a/tmva/sofie/test/generate_input_models.py b/tmva/sofie/test/generate_input_models.py index 5dd39ba342e08..374b09f4cf426 100644 --- a/tmva/sofie/test/generate_input_models.py +++ b/tmva/sofie/test/generate_input_models.py @@ -278,6 +278,47 @@ def make_ArgMax_NegAxis_SelectLast(): return _argmax_model("argmax_select_last_test", [2, 1], axis=-1, keepdims=1, select_last_index=1) +def _prelu_model(name, x_shape, slope_shape, slope_vals): + """PRelu on a float tensor X with the slope given as an initializer.""" + nodes = [ + helper.make_node("PRelu", ["X", "Slope"], ["Y"]), + ] + graph = helper.make_graph( + nodes, + name, + inputs=[ + _vi("X", FLOAT, x_shape), + ], + outputs=[ + _vi("Y", FLOAT, x_shape), + ], + initializer=[ + _tensor("Slope", FLOAT, slope_shape, slope_vals), + ], + ) + return _model(graph, opset=16, ir_version=8) + + +def make_PRelu(): + """Ops: PRelu (slope with the same shape as X)""" + return _prelu_model("prelu_test", [3, 4, 5], [3, 4, 5], np.linspace(0.05, 0.6, 60).tolist()) + + +def make_PRelu_Broadcast(): + """Ops: PRelu (slope broadcast along the last axis)""" + return _prelu_model("prelu_broadcast_test", [3, 4, 5], [5], [0.1, 0.2, 0.3, 0.4, 0.5]) + + +def make_PRelu_Channel(): + """Ops: PRelu (per-channel slope, as exported by PyTorch nn.PReLU)""" + return _prelu_model("prelu_channel_test", [1, 3, 2, 2], [3, 1, 1], [0.1, 0.2, 0.3]) + + +def make_PRelu_Scalar(): + """Ops: PRelu (single slope value)""" + return _prelu_model("prelu_scalar_test", [2, 3], [1], [0.25]) + + def make_AveragePool1d_CeilMode_Overhang(): """Ops: AveragePool""" nodes = [ @@ -6336,6 +6377,10 @@ def make_AliasDynShapeAcrossNewTensor(): "NonZero_Constant": make_NonZero_Constant, "NonZeroTwice": make_NonZeroTwice, "NotIsNaN": make_NotIsNaN, + "PRelu": make_PRelu, + "PRelu_Broadcast": make_PRelu_Broadcast, + "PRelu_Channel": make_PRelu_Channel, + "PRelu_Scalar": make_PRelu_Scalar, "Pad": make_Pad, "Pow": make_Pow, "Pow_broadcast": make_Pow_broadcast, @@ -7462,6 +7507,10 @@ def rand_f32(seed, shape): (12,), ) ], + "PRelu": [f32(np.arange(-30, 30) / 10, (3, 4, 5))], + "PRelu_Broadcast": [f32(np.arange(-30, 30) / 10, (3, 4, 5))], + "PRelu_Channel": [f32([-2.0, 1.0, -0.5, 0.0, -1.0, 2.0, 0.0, -3.0, 3.0, -4.0, -1.5, 0.5], (1, 3, 2, 2))], + "PRelu_Scalar": [f32([-2.0, -1.0, 0.0, 0.5, 1.0, 3.0], (2, 3))], "Pow": [ f32([1.0, 2.0, 3.0], (3,)), f32([4.0, 5.0, 6.0], (3,)), diff --git a/tmva/sofie_parsers/CMakeLists.txt b/tmva/sofie_parsers/CMakeLists.txt index 2402b725dc50e..8c3942e066acc 100644 --- a/tmva/sofie_parsers/CMakeLists.txt +++ b/tmva/sofie_parsers/CMakeLists.txt @@ -79,6 +79,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTTMVASofieParser src/ParseHardSigmoid.cxx src/ParseHardSwish.cxx src/ParseArgMax.cxx + src/ParsePRelu.cxx DEPENDENCIES ROOTTMVASofie ) diff --git a/tmva/sofie_parsers/src/ParsePRelu.cxx b/tmva/sofie_parsers/src/ParsePRelu.cxx new file mode 100644 index 0000000000000..d307c512703a6 --- /dev/null +++ b/tmva/sofie_parsers/src/ParsePRelu.cxx @@ -0,0 +1,52 @@ +#include "TMVA/RModelParser_ONNX.hxx" +#include "TMVA/ROperator_PRelu.hxx" +#include "onnx.hxx" + +namespace TMVA { +namespace Experimental { +namespace SOFIE { + +ParserFuncSignature ParsePRelu = [](RModelParser_ONNX &parser, const onnx::NodeProto &nodeproto) { + ETensorType input_type; + ETensorType slope_type; + + auto input_name = nodeproto.input(0); + if (parser.IsRegisteredTensorType(input_name)) { + input_type = parser.GetTensorType(input_name); + } else { + throw std::runtime_error("TMVA::SOFIE ONNX Parser PRelu op has input tensor " + input_name + + " but its type is not yet registered"); + } + auto slope_name = nodeproto.input(1); + if (parser.IsRegisteredTensorType(slope_name)) { + slope_type = parser.GetTensorType(slope_name); + } else { + throw std::runtime_error("TMVA::SOFIE ONNX Parser PRelu op has slope tensor " + slope_name + + " but its type is not yet registered"); + } + if (slope_type != input_type) { + throw std::runtime_error("TMVA::SOFIE ONNX Parser PRelu op has slope tensor " + slope_name + " of type " + + ConvertTypeToString(slope_type) + " but input tensor " + input_name + " is of type " + + ConvertTypeToString(input_type) + "; they must be the same"); + } + + std::unique_ptr op; + std::string output_name = nodeproto.output(0); + + // PRelu has no attributes (the legacy opset-1 'consumed_inputs' is ignored) + switch (input_type) { + case ETensorType::FLOAT: op.reset(new ROperator_PRelu(input_name, slope_name, output_name)); break; + default: + throw std::runtime_error("TMVA::SOFIE - Unsupported - Operator PRelu does not yet support input type " + + ConvertTypeToString(input_type)); + } + if (!parser.IsRegisteredTensorType(output_name)) { + parser.RegisterTensorType(output_name, input_type); + } + + return op; +}; + +} // namespace SOFIE +} // namespace Experimental +} // namespace TMVA diff --git a/tmva/sofie_parsers/src/RModelParser_ONNX.cxx b/tmva/sofie_parsers/src/RModelParser_ONNX.cxx index 7ec8dbfbfd1d9..eebd094650409 100644 --- a/tmva/sofie_parsers/src/RModelParser_ONNX.cxx +++ b/tmva/sofie_parsers/src/RModelParser_ONNX.cxx @@ -85,6 +85,7 @@ extern ParserFuncSignature ParseScatterND; extern ParserFuncSignature ParseNonZero; extern ParserFuncSignature ParseClip; extern ParserFuncSignature ParseArgMax; +extern ParserFuncSignature ParsePRelu; // Declaration of fused operators extern ParserFuseFuncSignature ParseFuseConvAdd; extern ParserFuseFuncSignature ParseFuseGemmRelu; @@ -359,6 +360,7 @@ RModelParser_ONNX::RModelParser_ONNX() noexcept : fOperatorsMapImpl(std::make_un RegisterOperator("NonZero", ParseNonZero); RegisterOperator("Clip", ParseClip); RegisterOperator("ArgMax", ParseArgMax); + RegisterOperator("PRelu", ParsePRelu); } // Destructor of the parser