Neel-Shah-29 · Neel-Shah-29 · Aug 10, 2023
diff --git a/tmva/sofie/CMakeLists.txt b/tmva/sofie/CMakeLists.txt
@@ -42,6 +42,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTTMVASofie
    TMVA/ROperator_LayerNormalization.hxx
    TMVA/ROperator_Expand.hxx
    TMVA/ROperator_Erf.hxx
+   TMVA/ROperator_Range.hxx
    TMVA/SOFIE_common.hxx
    TMVA/SOFIEHelpers.hxx
   SOURCES

diff --git a/tmva/sofie/inc/TMVA/OperatorList.hxx b/tmva/sofie/inc/TMVA/OperatorList.hxx
@@ -29,3 +29,5 @@
 #include "TMVA/ROperator_Gather.hxx"
 #include "TMVA/ROperator_Swish.hxx"
 #include "TMVA/ROperator_Erf.hxx"
+#include "TMVA/ROperator_Range.hxx"
+
diff --git a/tmva/sofie/inc/TMVA/RModel.hxx b/tmva/sofie/inc/TMVA/RModel.hxx
@@ -41,22 +41,22 @@ private:
    std::unordered_map<std::string, TensorInfo> fReadyInputTensorInfos;
    std::unordered_map<std::string, InitializedTensor> fInitializedTensors;
    std::unordered_map<std::string, TensorInfo> fIntermediateTensorInfos;
-   std::vector<std::string> fOutputTensorNames;
-   std::vector<std::string> fInputTensorNames;  //input tensor names using ONNX order
+   std::unordered_map<std::string, DynamicTensorInfo> fDynamicTensorInfos;
 
-   std::vector<std::unique_ptr<ROperator>> fOperators;
+   std::unordered_set<std::string> fNeededBlasRoutines;
+   const std::unordered_set<std::string> fAllowedStdLib = {"vector", "algorithm", "cmath"};
+   std::unordered_set<std::string> fNeededStdLib = {"vector"};
+   std::unordered_set<std::string> fCustomOpHeaders;
 
    std::string fName="UnnamedModel";
    std::string fFileName; //file name of original model file for identification
    std::string fParseTime; //UTC date and time string at parsing
-
-
    std::string fGC; //generated code
-   std::unordered_set<std::string> fNeededBlasRoutines;
 
-   const std::unordered_set<std::string> fAllowedStdLib = {"vector", "algorithm", "cmath"};
-   std::unordered_set<std::string> fNeededStdLib = {"vector"};
-   std::unordered_set<std::string> fCustomOpHeaders;
+   std::vector<std::unique_ptr<ROperator>> fOperators;
+   std::vector<std::string> fOutputTensorNames;
+   std::vector<std::string> fInputTensorNames;  //input tensor names using ONNX order
+
    bool fUseWeightFile = true;
    bool fUseSession = true;
 
@@ -75,6 +75,7 @@ public:
    RModel(std::string name, std::string parsedtime);
 
    const std::vector<size_t>& GetTensorShape(std::string name);
+   const std::vector<Dim>& GetDynamicTensorShape(std::string name);
    const ETensorType& GetTensorType(std::string name);
 
    bool CheckIfTensorAlreadyExist(std::string tensor_name);
@@ -84,7 +85,9 @@ public:
    void AddInitializedTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape, std::shared_ptr<void> data);
    // Check if a tensor is initialized
    bool IsInitializedTensor(const std::string& name) const;
+   bool IsDynamicTensor(const std::string& name) const;
    void AddIntermediateTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape);
+   void AddDynamicTensor(std::string tensor_name, ETensorType type, std::vector<Dim> shape);
    void AddBlasRoutines(std::vector<std::string> routines) {
       for (auto &routine : routines) {
          fNeededBlasRoutines.insert(routine);

diff --git a/tmva/sofie/inc/TMVA/ROperator_Range.hxx b/tmva/sofie/inc/TMVA/ROperator_Range.hxx
@@ -0,0 +1,91 @@
+#ifndef TMVA_SOFIE_ROPERATOR_RANGE
+#define TMVA_SOFIE_ROPERATOR_RANGE
+
+#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_Range final : public ROperator
+{
+private:
+
+   std::string fNStart;
+   std::string fNLimit;
+   std::string fNDelta;
+   std::string fNOutput;
+   std::vector<Dim> fShape;
+   std::string fType;
+
+public:
+   ROperator_Range(){}
+
+   ROperator_Range(std::string start, std::string limit, std::string delta, std::string nameOutput):
+      fNStart(start), fNLimit(limit), fNDelta(delta),
+      fNOutput(UTILITY::Clean_name(nameOutput)) {
+      if (std::is_same<T, float>::value) {
+         fType = "float";
+      } else if (std::is_same<T, int64_t>::value) {
+         fType = "int64_t";
+      } else {
+         throw
+            std::runtime_error("TMVA::SOFIE - Unsupported type by Range operator");
+      }
+   }
+
+   std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
+      return input;
+   }
+
+   std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
+      auto ret = input; //suggest copy to compiler
+      return ret;
+   }
+
+   void Initialize(RModel& model) override {
+       //input must be a graph input, or already initialized intermediate tensor
+      if (!model.CheckIfTensorAlreadyExist(fNStart)) {
+         throw
+            std::runtime_error("TMVA SOFIE Range Op Input Tensor " + fNStart + "is not found in model");
+      }
+      if (!model.CheckIfTensorAlreadyExist(fNLimit)) {
+         throw
+            std::runtime_error("TMVA SOFIE Range Op Input Tensor " + fNLimit + "is not found in model");
+      }
+      if (!model.CheckIfTensorAlreadyExist(fNDelta)) {
+         throw
+            std::runtime_error("TMVA SOFIE Range Op Input Tensor " + fNDelta + "is not found in model");
+      }
+      ETensorType type = ConvertStringToType(fType);
+      fShape = {Dim{true, 0, "range_size"}};
+      model.AddDynamicTensor(fNOutput, type, fShape);
+   }
+
+   std::string Generate(std::string OpName) override {
+      OpName = "op_" + OpName;
+      if (fShape.empty()) {
+         throw std::runtime_error("TMVA SOFIE Range operator called to Generate without being initialized first");
+      }
+      std::stringstream out;
+      out << "\n//------ Range\n";
+      std::string sizeName = fShape[0].param;
+      out << SP << "size_t " << sizeName << " = static_cast<size_t>(std::max(std::ceil((static_cast<float>(*tensor_" << fNLimit << ") - static_cast<float>(*tensor_" << fNStart << ")) / static_cast<float>(*tensor_" << fNDelta << ")), 0.0f));\n";
+      out << SP << "tensor_" << fNOutput << ".resize(" << sizeName << ");\n";
+      out << SP << "for (size_t i = 0; i < " << sizeName << "; i++) {\n";
+      out << SP << SP << "tensor_" << fNOutput << "[i] = *tensor_" << fNStart << " + i * (*tensor_" << fNDelta << ");\n";
+      out << SP << "}\n";
+      return out.str();
+   }
+};
+
+}//SOFIE
+}//Experimental
+}//TMVA
+
+#endif //TMVA_SOFIE_ROPERATOR_RANGE
diff --git a/tmva/sofie/inc/TMVA/SOFIE_common.hxx b/tmva/sofie/inc/TMVA/SOFIE_common.hxx
@@ -50,10 +50,17 @@ struct TensorInfo{
    std::vector<size_t> shape;
 };
 
+struct DynamicTensorInfo{
+   ETensorType type;
+   std::vector<Dim> shape;
+};
+
 std::size_t ConvertShapeToLength(std::vector<size_t> shape);
 
 std::string ConvertShapeToString(std::vector<size_t> shape);
 
+std::string ConvertDynamicShapeToLength(std::vector<Dim> shape);
+
 struct InitializedTensor{
    ETensorType fType;
    std::vector<std::size_t> fShape;

diff --git a/tmva/sofie/src/RModel.cxx b/tmva/sofie/src/RModel.cxx
@@ -2,6 +2,7 @@
 #include <algorithm>
 #include <cctype>
 #include <memory>
+#include <string>
 
 #include "TMVA/RModel.hxx"
 #include "TMVA/SOFIE_common.hxx"
@@ -55,7 +56,7 @@ namespace SOFIE{
       fName = fFileName.substr(0, fFileName.rfind("."));
    }
 
-   const std::vector<size_t>& RModel::GetTensorShape(std::string name){
+   const std::vector<size_t>& RModel::GetTensorShape(std::string name) {
       auto f = fReadyInputTensorInfos.find(name);
       if (f != fReadyInputTensorInfos.end()){
          return f->second.shape;
@@ -76,6 +77,15 @@ namespace SOFIE{
       throw std::runtime_error("TMVA SOFIE tensor [" + name + "] for which the shape is requested is not found");
    }
 
+   const std::vector<Dim>& RModel::GetDynamicTensorShape(std::string name){
+      auto f = fDynamicTensorInfos.find(name);
+      if (f != fDynamicTensorInfos.end()){
+         return f->second.shape;
+      }
+
+      throw std::runtime_error("TMVA SOFIE tensor [" + name + "] for which the shape is requested is not found");
+   }
+
    const ETensorType& RModel::GetTensorType(std::string name){
       auto f = fReadyInputTensorInfos.find(name);
       if (f != fReadyInputTensorInfos.end()){
@@ -93,6 +103,10 @@ namespace SOFIE{
       if (f4 != fIntermediateTensorInfos.end()){
          return f4->second.type;
       }
+      auto f5 = fDynamicTensorInfos.find(name);
+      if (f5 != fDynamicTensorInfos.end()){
+         return f5->second.type;
+      }
 
       throw std::runtime_error("TMVA SOFIE tensor [" + name + "] for which the type is requested is not found");
    }
@@ -101,6 +115,7 @@ namespace SOFIE{
       if (fReadyInputTensorInfos.find(tensor_name) != fReadyInputTensorInfos.end())  return true;
       if (fInitializedTensors.find(tensor_name) != fInitializedTensors.end()) return true;
       if (fIntermediateTensorInfos.find(tensor_name) != fIntermediateTensorInfos.end()) return true;
+      if (fDynamicTensorInfos.find(tensor_name) != fDynamicTensorInfos.end()) return true;
       return false;
    }
 
@@ -156,6 +171,11 @@ namespace SOFIE{
       return fInitializedTensors.find(name) != fInitializedTensors.end();
    }
 
+   bool RModel::IsDynamicTensor(const std::string& tensorName) const {
+      std::string name = UTILITY::Clean_name(tensorName);
+      return fDynamicTensorInfos.find(name) != fDynamicTensorInfos.end();
+   }
+
    void RModel::AddIntermediateTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape){
       tensor_name = UTILITY::Clean_name(tensor_name);
       if (CheckIfTensorAlreadyExist(tensor_name)){
@@ -165,6 +185,15 @@ namespace SOFIE{
       fIntermediateTensorInfos[tensor_name] = new_tensor;
    }
 
+   void RModel::AddDynamicTensor(std::string tensor_name, ETensorType type, std::vector<Dim> shape){
+      tensor_name = UTILITY::Clean_name(tensor_name);
+      if (CheckIfTensorAlreadyExist(tensor_name)){
+         throw std::runtime_error("TMVA-SOFIE: intermediate tensor with name " + tensor_name + " already exists \n");
+      }
+      DynamicTensorInfo new_tensor {type, shape};
+      fDynamicTensorInfos[tensor_name] = new_tensor;
+   }
+
    void RModel::AddOutputTensorNameList(std::vector<std::string> outputtensornames){
       for(auto& it : outputtensornames){
          fOutputTensorNames.push_back(UTILITY::Clean_name(it));
@@ -230,9 +259,9 @@ namespace SOFIE{
       }
 
 
-      for (auto& i : fOperators){
+      for (auto& op : fOperators) {
          //std::cout << "initialize operator  " << typeid(*i).name() << std::endl;
-         i->Initialize(*this);
+         op->Initialize(*this);
       }
    }
 
@@ -301,7 +330,7 @@ namespace SOFIE{
       if (fUseSession) {
          fGC += "struct Session {\n";
       }
-      for (auto& i: fInitializedTensors){
+      for (auto& i: fInitializedTensors) {
          if (i.second.fType == ETensorType::FLOAT){
             size_t length = 1;
             for (auto & dim: i.second.fShape){
@@ -325,6 +354,7 @@ namespace SOFIE{
 
          }
       }
+
       for (auto&i: fIntermediateTensorInfos){
          size_t length = ConvertShapeToLength(i.second.shape);
          if (i.second.type == ETensorType::FLOAT){
@@ -340,6 +370,17 @@ namespace SOFIE{
             fGC += "int64_t * tensor_" + i.first + " = fTensor_" + i.first  + ".data();\n";
          }
       }
+
+      for (auto&i: fDynamicTensorInfos) {
+         if (i.second.type == ETensorType::FLOAT) {
+            fGC += "std::vector<float> tensor_" + i.first  + ";\n";
+         } else if (i.second.type == ETensorType::DOUBLE) {
+            fGC += "std::vector<double> tensor_" + i.first  + ";\n";
+         } else if (i.second.type == ETensorType::INT64) {
+            fGC += "std::vector<int64_t> tensor_" + i.first  + ";\n";
+         }
+      }
+
       if (fUseSession) {
          // add here specific operator code that needs to define session data members
          fGC += "\n";
@@ -377,21 +418,32 @@ namespace SOFIE{
       std::string outputType;
       if (outputSize == 1) {
          auto f = fIntermediateTensorInfos.find(fOutputTensorNames[0]);
-         if (f == fIntermediateTensorInfos.end()){
-            throw std::runtime_error("TMVA-SOFIE: output tensor " + fOutputTensorNames[0] + " not found when trying to get its info");
-         }else{
+         if (f != fIntermediateTensorInfos.end()) {
             outputType = ConvertTypeToString(f->second.type);
             fGC += "std::vector<" + outputType + "> ";
+         } else {
+            auto f2 = fDynamicTensorInfos.find(fOutputTensorNames[0]);
+            if (f2 != fDynamicTensorInfos.end()) {
+               outputType = ConvertTypeToString(f2->second.type);
+               fGC += "std::vector<" + outputType + "> ";
+            } else {
+               throw std::runtime_error("TMVA-SOFIE: output tensor " + fOutputTensorNames[0] + " not found when trying to get its info");
+            }
          }
       } else {
          std::vector<ETensorType> outputTensorsTypes(outputSize);
          for (size_t i = 0; i < outputSize; i++) {
             auto f = fIntermediateTensorInfos.find(fOutputTensorNames[i]);
-            if (f == fIntermediateTensorInfos.end()) {
-               throw std::runtime_error("TMVA-SOFIE: output tensor " + fOutputTensorNames[i]
-                  + " not found when trying to get its info");
-            } else {
+            if (f != fIntermediateTensorInfos.end()) {
                outputTensorsTypes[i] = f->second.type;
+            } else {
+               auto f2 = fDynamicTensorInfos.find(fOutputTensorNames[i]);
+               if (f2 != fDynamicTensorInfos.end()) {
+                  outputTensorsTypes[i] = f2->second.type;
+               } else {
+                  throw std::runtime_error("TMVA-SOFIE: output tensor " + fOutputTensorNames[i]
+                     + " not found when trying to get its info");
+               }
             }
          }
          // assume all output types are the same
@@ -437,10 +489,14 @@ namespace SOFIE{
          fGC+= (fOperators[id]->Generate(std::to_string(id)));
       }
       if (outputSize == 1) {
-         size_t outputLength = ConvertShapeToLength(GetTensorShape(fOutputTensorNames[0]));
-
-         fGC += SP + "std::vector<" + outputType + "> ret (tensor_" + fOutputTensorNames[0] + ", tensor_" + fOutputTensorNames[0] + " + " +
-               std::to_string(outputLength) + ");\n";
+         std::string tensorName = fOutputTensorNames[0];
+         if (IsDynamicTensor(tensorName)) {
+            fGC += SP + "std::vector<" + outputType + "> ret (tensor_" + tensorName + ");\n";
+         } else {
+            size_t outputLength = ConvertShapeToLength(GetTensorShape(fOutputTensorNames[0]));
+            fGC += SP + "std::vector<" + outputType + "> ret (tensor_" + fOutputTensorNames[0] + ", tensor_" + fOutputTensorNames[0] + " + " +
+                  std::to_string(outputLength) + ");\n";
+         }
       } else {
          for (size_t i = 0; i < outputSize; i++) {
             if (!fOutputTensorNames[i].empty()) {

diff --git a/tmva/sofie/src/SOFIE_common.cxx b/tmva/sofie/src/SOFIE_common.cxx
@@ -59,7 +59,7 @@ ETensorType ConvertStringToType(std::string type){
    if(type == "float32" || type == "float" || type == "Float"){
      return ETensorType::FLOAT;
    }
-   else if(type == "int64"){
+   else if(type == "int64" || type == "int64_t"){
      return ETensorType::INT64;
    }
    else if (type == "double" || type == "float64"){
@@ -81,6 +81,24 @@ std::string ConvertShapeToString(std::vector<size_t> shape) {
    return out.str();
 }
 
+std::string ConvertDynamicShapeToLength(std::vector<Dim> shape) {
+   std::string length;
+   if (shape[0].isParam) {
+      length += shape[0].param;
+   } else {
+      length += std::to_string(shape[0].dim);
+   }
+   for (size_t i = 1; i < shape.size(); i++) {
+      length += " * ";
+      if (shape[i].isParam) {
+         length += shape[i].param;
+      } else {
+         length += std::to_string(shape[i].dim);
+      }
+   }
+   return length;
+}
+
 namespace{
 template<typename T>
 static inline void copy_vector_data(int_t no_of_copies, int_t input_size, T* input, T* target){  //only visible within this translation unit