[tmva][sofie] Add support for the ONNX PRelu operator

tmva/sofie/README.md

@@ -130,6 +130,7 @@ std::vector<std::string> supportedOperators = parser.GetRegisteredOperators();
 - [x] Min
 - [x] Mul
 - [x] Neg
+- [x] PRelu
 - [x] Pad
 - [x] Pow
 - [x] RNN

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 <sstream>
+
+namespace TMVA {
+namespace Experimental {
+namespace SOFIE {
+
+template <typename T>
+class ROperator_PRelu final : public ROperator {
+private:
+   /* Attributes */
+   std::string fNX;
+   std::string fNSlope;
+   std::string fNY;
+   std::vector<size_t> fShapeX;
+   std::vector<size_t> 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<T, float>::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

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<float>, "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<float>, "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<float>, "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<float>, "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)
 {

tmva/sofie/test/TestSofieParser.cxx

@@ -187,6 +187,42 @@ GenerateGemmCode(const std::string &fileName, const std::vector<TestTensor> &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);
+}

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,)),

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
 )

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<ROperator> 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<float>(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

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