[tmva][sofie] Add support for NonZero operator and fix handling of booleans

lmoneta · lmoneta · commit 8817a96aac0a · 2025-08-04T16:53:36.000+02:00
Correct handle booleans types when parsing from ONNX
diff --git a/tmva/sofie/CMakeLists.txt b/tmva/sofie/CMakeLists.txt
@@ -66,6 +66,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTTMVASofie
    TMVA/ROperator_ScatterElements.hxx
    TMVA/ROperator_Gather.hxx
    TMVA/ROperator_GatherND.hxx
+   TMVA/ROperator_NonZero.hxx
    TMVA/SOFIE_common.hxx
    TMVA/SOFIEHelpers.hxx
 
diff --git a/tmva/sofie/inc/TMVA/ROperator_Constant.hxx b/tmva/sofie/inc/TMVA/ROperator_Constant.hxx
@@ -122,9 +122,7 @@ public:
          model.AddConstantTensor(fNY, fShape, fValues);
          if (model.Verbose()) {
             std::cout << "adding constant tensor " << fNY << " with shape " << ConvertShapeToString(fShape)
-            << " and values [";
-            for (auto v : fValues) std::cout << " " << v;
-            std::cout << "]" << std::endl;
+            << " and values " << ConvertValuesToString(fValues) << std::endl;
          }
       } else {
          model.AddIntermediateTensor(fNY, ConvertStringToType(TensorType<T>::Name()), fDimOutputShape);
diff --git a/tmva/sofie/inc/TMVA/ROperator_GatherND.hxx b/tmva/sofie/inc/TMVA/ROperator_GatherND.hxx
@@ -207,8 +207,7 @@ public:
       std::string idIndex;
       for (size_t j = 0; j < fBatchDims; j++) {
          std::string index = "i_" + std::to_string(j);
-         for (size_t k = 0; k <= j; k++)
-            out << SP;
+         for (size_t k = 0; k <= j; k++) out << SP;
          out << "for (size_t " << index << " = 0; " << index << " < " << fShapeY[j] << "; " << index << "++) {\n";
          if (j > 0) {
             outIndex += " + ";
diff --git a/tmva/sofie/inc/TMVA/ROperator_NonZero.hxx b/tmva/sofie/inc/TMVA/ROperator_NonZero.hxx
@@ -0,0 +1,179 @@
+#ifndef TMVA_SOFIE_ROPERATOR_NONZERO
+#define TMVA_SOFIE_ROPERATOR_NONZERO
+
+#include "TMVA/SOFIE_common.hxx"
+#include "TMVA/ROperator.hxx"
+#include "TMVA/RModel.hxx"
+
+#include <sstream>
+
+namespace TMVA{
+namespace Experimental{
+namespace SOFIE{
+
+template<class T>
+class ROperator_NonZero final : public ROperator
+{
+
+private:
+
+   std::string fNX;
+   std::string fNY;
+   std::vector<Dim> fShapeX;
+   std::vector<Dim> fShapeY;
+
+public:
+   ROperator_NonZero(){}
+   ROperator_NonZero(std::string nameX, std::string nameY):
+      fNX(UTILITY::Clean_name(nameX)), fNY(UTILITY::Clean_name(nameY)){
+         fInputTensorNames = { fNX };
+         fOutputTensorNames = { fNY };
+      }
+
+
+
+   void Initialize(RModel& model) override {
+      if (model.CheckIfTensorAlreadyExist(fNX) == false){   //input must be a graph input, or already initialized intermediate tensor
+         throw std::runtime_error("TMVA SOFIE NonZero Op Input Tensor " + fNX + " is not found in model");
+      }
+
+
+      // case input is constant
+      if (model.IsConstantTensor(fNX)) {
+         // compute output directly
+         T * data = static_cast<T*>(model.GetInitializedTensorData(fNX).get());
+         // shape is fully known
+         auto shapeX = model.GetTensorShape(fNX);
+         std::vector<size_t> shapeY(2);
+         shapeY[0] = shapeX.size();
+         auto length = ConvertShapeToLength(shapeX);
+         auto strides = UTILITY::ComputeStrideFromShape(shapeX);
+         std::vector<std::vector<int64_t>> nonzero_indices;
+         for (size_t i = 0; i < length; i++) {
+            if (data[i] != 0) {
+               // get indices
+               size_t flat_index = i;
+               std::vector<int64_t> indices(shapeX.size());
+               for (size_t j = 0; j < shapeX.size(); ++j) {
+                  indices[j] = flat_index / strides[j];
+                  flat_index %= strides[j];
+               }
+               nonzero_indices.emplace_back(indices);
+            }
+         }
+         shapeY[1] = nonzero_indices.size();
+         std::vector<int64_t> dataY(shapeY[0]* shapeY[1]);
+         size_t k = 0;
+         for (size_t i = 0; i < shapeY[0]; i++) {
+            for (size_t j = 0; j < shapeY[1]; j++) {
+               dataY[k] = nonzero_indices[j][i];
+               k++;
+            }
+         }
+         if (dataY.empty()) {
+            // no zero elements found
+            dataY.resize(1);
+            shapeY.clear();  // use an empty shape
+         }
+
+         model.AddConstantTensor(fNY, shapeY, dataY);
+         if (model.Verbose()) {
+            std::cout << "NonZero : " << fNX << " -> " << fNY << " " << ConvertShapeToString(shapeY)
+                     << " : " << ConvertValuesToString(dataY) << std::endl;
+         }
+         fIsOutputConstant = true;
+
+      } else {
+
+         fShapeX = model.GetDimTensorShape(fNX);
+
+         // output shape(-1) depends on number of elements of non zero values
+         // first dim is rank of input
+         fShapeY.resize(2);
+         fShapeY[0] = fShapeX.size();
+
+         // 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 << " " << ConvertShapeToString(fShapeY) << std::endl;
+         }
+      }
+   }
+   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";
+      return out.str();
+   }
+
+
+   std::string Generate(std::string opName) override {
+      if (fIsOutputConstant) {
+         return "";
+      }
+      opName = "op_" + opName;
+      if (fShapeX.empty()) {
+         throw std::runtime_error("TMVA SOFIE Operator NonZero called to Generate without being initialized first");
+      }
+      std::stringstream out;
+      auto inputLength = ConvertDimShapeToLength(fShapeX);
+      auto maxStrideY = inputLength;
+      size_t dims = fShapeX.size();
+      out << "\n//------ NonZero\n";
+
+      std::string vnonzero = "v_NonZero_" + fNX;
+
+      // loop on input indices
+      out << "size_t offset_" << opName << " = 0;\n";
+      out << vnonzero << " = 0;\n";
+      for (size_t j = 0; j < dims; j++) {
+         std::string index = "i_" + std::to_string(j);
+         for (size_t k = 0; k <= j; k++) out << SP;
+         out << "for (size_t " << index << " = 0; " << index << " < " << fShapeX[j] << "; " << index << "++) {\n";
+      }
+      for (size_t k = 0; k <= dims; k++) out << SP;
+      out << "if (tensor_" << fNX << "[offset_" << opName << "]) {\n";
+      for (size_t k = 0; k <= dims+1; k++) out << SP;
+      out << vnonzero << "++;\n";
+      for (size_t j = 0; j < dims; j++) {
+         for (size_t k = 0; k <= dims+1; k++) out << SP;
+         out << "tensor_" << fNY << "[" << maxStrideY << " * " << j << " + " << vnonzero << "] = i_" << j << ";\n";
+      }
+      for (size_t k = 0; k <= dims; k++) out << SP;
+      out << "}\n";
+      //end loops
+      for (size_t j = dims; j > 0; j--) {
+         for (size_t k = 0; k <j; k++) out << SP;
+         out << "}\n";
+      }
+      // now we need to rearrange the vector if nonzero is less than length of input
+      out << SP << "if (" << vnonzero << " < " << inputLength << "){\n";
+      for (size_t j = 1; j < dims; j++) {
+         out << SP << SP << "std::copy(tensor_" << fNY;
+         if (j>0) out << " + " << maxStrideY;
+         if (j>1) out << " * " << j;
+         out << ", tensor_" << fNY;
+         if (j>0) out << " + " << maxStrideY;
+         if (j>1) out << " * " << j;
+         out << " + " << vnonzero << ", tensor_" <<  fNY;
+         if (j>0) out << " + " << vnonzero;
+         if (j>1) out << "* " << j;
+         out << ");\n";
+      }
+      out << SP << "}\n";
+
+      return out.str();
+   }
+
+};
+
+}//SOFIE
+}//Experimental
+}//TMVA
+
+
+#endif //TMVA_SOFIE_ROPERATOR_NonZero
diff --git a/tmva/sofie/inc/TMVA/SOFIE_common.hxx b/tmva/sofie/inc/TMVA/SOFIE_common.hxx
@@ -164,6 +164,10 @@ template<>
 struct TensorType<bool> {
    static const std::string Name() { return "bool"; }
 };
+template<>
+struct TensorType<int8_t> {
+   static const std::string Name() { return "int8_t"; }
+};
 
 struct TensorMemoryInfo {
    std::string_view tensor_name;
@@ -225,8 +229,11 @@ std::string ConvertValuesToString(size_t n, const T * data) {
    ret << "{ ";
    for (size_t i = 0; i < n; i++) {
       if (std::is_floating_point_v<T>)
-         ret << std::setprecision(std::numeric_limits<T>::max_digits10);
-      ret << data[i];
+         ret << std::setprecision(std::numeric_limits<T>::max_digits10) << data[i];
+      else
+         // cast in case of boolean (int8)
+         ret << (int64_t) data[i];
+
       if (i < n-1) ret << ", ";
    }
    ret << "}";
diff --git a/tmva/sofie/src/SOFIE_common.cxx b/tmva/sofie/src/SOFIE_common.cxx
@@ -143,21 +143,21 @@ std::string ConvertDimShapeToLength(const std::vector<Dim> & shape) {
    std::string length;
    // case of empty vectors return 1
    if (shape.empty()) return "1";
-   size_t int_length = 0;
+   int64_t int_length = -1;
    for (size_t i = 0; i < shape.size(); i++) {
       if (shape[i].isParam) {
          if (!length.empty()) length += " * ";
          length += shape[i].param;
       } else {
-         if (int_length == 0)
+         if (int_length == -1)
             int_length = shape[i].dim;
          else
             int_length *= shape[i].dim;
       }
    }
    // multiply the integer components to the parametric one
-   // if larger than 1
-   if (int_length > 0) {
+   // if larger than 1 - otherwise returns -1
+   if (int_length >= 0) {
       if (!length.empty() && int_length > 1) {
          length += " * ";
          length += std::to_string(int_length);
diff --git a/tmva/sofie/test/TestCustomModelsFromONNX.cxx b/tmva/sofie/test/TestCustomModelsFromONNX.cxx
@@ -327,6 +327,9 @@
 
 #include "ScatterElements_FromONNX.hxx"
 
+#include "NonZero_FromONNX.hxx"
+#include "NonZero_Constant_FromONNX.hxx"
+
 #include "gtest/gtest.h"
 
 constexpr float DEFAULT_TOLERANCE = 1e-3f;
@@ -3283,3 +3286,42 @@ TEST(ONNX, GatherND_3)
    }
 }
 
+TEST(ONNX, NonZero)
+{
+   // test GatherND elements using batch size as first dim (bs=2)
+   std::vector<int8_t> input = {0,1,0, 1,1,0, 0,0,1, 0,1,1 }; // shape is (2x2x3)
+   // output is tensor shape { 3, number of non zeros}
+   std::vector<int32_t> correct_output = { 0,0,0,1,1,1 ,   0,1,1,0,1,1 ,    1,0,1,2,1,2 };
+
+   TMVA_SOFIE_NonZero::Session s("NonZero_FromONNX.dat");
+
+   auto output = s.infer(input.data());
+
+   // Checking output size
+   EXPECT_EQ(output.size(), correct_output.size());
+   // Checking output
+   for (size_t i = 0; i < output.size(); ++i) {
+      EXPECT_EQ(output[i] , correct_output[i]);
+   }
+}
+
+TEST(ONNX, NonZero_Constant)
+{
+   // test GatherND elements using batch size as first dim (bs=2)
+   //std::vector<int8_t> input = {0,1,0, 1,1,0, 0,0,1, 0,1,1 }; // shape is (2x2x3)
+   // output is tensor shape { 3, number of non zeros}
+   std::vector<int32_t> correct_output = { 0,0,0,1,1,1 ,   0,1,1,0,1,1 ,    1,0,1,2,1,2 };
+
+   TMVA_SOFIE_NonZero_Constant::Session s("NonZero_Constant_FromONNX.dat");
+
+   auto output = s.infer();
+
+   // Checking output size
+   EXPECT_EQ(output.size(), correct_output.size());
+   // Checking output
+   for (size_t i = 0; i < output.size(); ++i) {
+      EXPECT_EQ(output[i] , correct_output[i]);
+   }
+}
+
+
diff --git a/tmva/sofie/test/input_models/NonZero.onnx b/tmva/sofie/test/input_models/NonZero.onnx
@@ -0,0 +1,13 @@
+
+onnx-example:Z
+
+dataoutput"NonZero	TestGraphZ
+data
+
+
+
+b
+output
+
+
+B
diff --git a/tmva/sofie/test/input_models/NonZero_Constant.onnx b/tmva/sofie/test/input_models/NonZero_Constant.onnx
diff --git a/tmva/sofie_parsers/CMakeLists.txt b/tmva/sofie_parsers/CMakeLists.txt
@@ -73,6 +73,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTTMVASofieParser
     src/ParseEinsum.cxx
     src/ParseRandom.cxx
     src/ParseScatterElements.cxx
+    src/ParseNonZero.cxx
     ${PROTO_SRCS}
   LIBRARIES PUBLIC
     protobuf::libprotobuf
diff --git a/tmva/sofie_parsers/src/ParseConstant.cxx b/tmva/sofie_parsers/src/ParseConstant.cxx
@@ -109,12 +109,19 @@ ParserFuncSignature ParseConstant = [](RModelParser_ONNX &parser, const onnx::No
             break;
          }
          case ETensorType::BOOL: {
-            std::vector<bool> values(length);
-            auto raw_data_ptr = reinterpret_cast<bool *>(const_cast<char *>(t.raw_data().c_str()));
-            // cannot use values.data() for vector of bools
-            std::copy(raw_data_ptr, raw_data_ptr + length, values.begin());
-            //std::memcpy(values.data(), raw_data_ptr, length * sizeof(float));
-            op.reset(new ROperator_Constant<bool>("bool",values, shape, input_name, output_name));
+            //values are int32 in ONNX
+            std::vector<int8_t> values(length);
+            if (t.int32_data_size() == int(length)) {
+               for (size_t i = 0; i < length; i++) {
+                  auto val = t.int32_data(i);
+                  if (val < 0 || val > 1)
+                     throw std::runtime_error("TMVA::SOFIE ONNX Parser Constant has invalid boolean value " + std::to_string(val));
+                  values[i] = static_cast<int8_t>(val);
+               }
+            } else
+               throw std::runtime_error("TMVA::SOFIE ONNX Parser COnstant : invalid tensor data values");
+
+            op.reset(new ROperator_Constant<int8_t>("bool",values, shape, input_name, output_name));
             break;
          }
          default:
diff --git a/tmva/sofie_parsers/src/ParseNonZero.cxx b/tmva/sofie_parsers/src/ParseNonZero.cxx
diff --git a/tmva/sofie_parsers/src/RModelParser_ONNX.cxx b/tmva/sofie_parsers/src/RModelParser_ONNX.cxx

-Original file line number
+Diff line change
@@ @@ -0,0 +1,13 @@ @@
++
 +onnx-example:Z
++
 +dataoutput"NonZero	TestGraphZ
 +data
++
++
++
 +b
 +output
++
++
 +B