[tmva][sofie] Add in RModel managing of Shape tensors types

Include in RMOdel a new type of tensor, shape tensors , where their values (a shape) is stored in RModel as vector of Dims (so shape can be parametrised) and a shape of rank 1 or 0 ( in case of scalars).   Shapes which are fully known can instead be stored as constant tensors

This allows to treat and manage the chain of operator treting shapes when the shapes are express as parameters (are dynamic).

We then include then these functions in RModel:

- AddShapeTensor given name and shape values (vector of Dim) and also a boolean flag specifying if tensor is a scalar. In this case the shape of teh shape tensor is a scalar (tensor with rank 0).
- GetShapeTensorValues given name and retuirning vector of Dims
- IsShapeTensor gievn the name

We also include in RTensor a convenient template function GetTensorData, working for both constant tensors and ShapeTensors.

In this commit we apply the needed fixes to use Shape Tensor in order to parge the ATLAS gnn model

Author: lmoneta

math/mathcore/inc/Fit/FitUtil.h

@@ -238,7 +238,7 @@ namespace FitUtil {
         // Figure out the size of the SIMD vectors.
         constexpr static int vecSize = sizeof(ROOT::Double_v) / sizeof(double);
         double xBuffer[vecSize];
-        std::vector<ROOT::Double_v> xx(fDim);
+        ROOT::Double_v xx[fDim];
         for (unsigned int i = 0; i < fDim; ++i) {
            // The Load() function reads multiple values from the pointed-to
            // memory into xx. This is why we have to copy the input values from
@@ -250,7 +250,7 @@ namespace FitUtil {
            }
            vecCore::Load<ROOT::Double_v>(xx[i], xBuffer);
         }
-        auto res = (*f)(xx.data(), p);
+        auto res = (*f)(xx, p);
         return vecCore::Get<ROOT::Double_v>(res, 0);
      }
 

tmva/sofie/inc/TMVA/RModel.hxx

@@ -25,11 +25,14 @@ private:
    std::unordered_map<std::string, InitializedTensor> fInitializedTensors;
    std::unordered_map<std::string, TensorInfo> fIntermediateTensorInfos;
    std::unordered_map<std::string, DynamicTensorInfo> fDynamicTensorInfos;
+   std::unordered_map<std::string, std::pair<std::vector<Dim>, bool>> fShapeTensors; // constant tensors describing a shape
    std::unordered_map<std::string, std::string> fShapeParams; // parameters defining the dynamic shape (e.g. batch size), store also its default value
    std::vector<std::string> fDimShapeNames; // parameter names used to define the shapes
    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)
@@ -52,11 +55,14 @@ public:
 
    int Verbose() const { return fVerbose;}
 
-   const std::vector<size_t> & GetTensorShape(const std::string & name) const;
+   std::vector<size_t> GetTensorShape(const std::string & name) const;
    std::vector<Dim> GetDimTensorShape(const std::string & name) const;
    std::vector<Dim> GetDynamicTensorShape(const std::string & name) const ;
 
-   const ETensorType &GetTensorType(std::string name) const;
+   // get the values for the tensor representing a shape
+   const std::vector<Dim> & GetShapeTensorValues(const std::string & tensor_name) const;
+
+   ETensorType GetTensorType(std::string name) const;
 
 
    bool CheckIfTensorAlreadyExist(std::string tensor_name);
@@ -73,6 +79,7 @@ public:
    void AddConstantTensor(std::string tensor_name, ETensorType type, std::vector<std::size_t> shape,
                              std::shared_ptr<void> data);
 
+
    template<class T>
    void AddConstantTensor(const std::string & name, const std::vector<size_t> & shape, const T * data) {
       size_t length = ConvertShapeToLength(shape);
@@ -99,6 +106,9 @@ public:
       AddInitializedTensor(tensor_name,  GetTemplatedType(T()), shape, data);
    }
 
+   void AddShapeTensor(const std::string & name, const std::vector<Dim> & shapeValues, bool scalar = false);
+
+
    // add and initialize subgraph to the model
    void InitializeSubGraph(std::shared_ptr<RModel>  graph);
 
@@ -115,6 +125,8 @@ public:
    bool IsDimInputTensor(const std::string &name) const;
    // check if tensor is a fully specified input tensor
    bool IsReadyInputTensor(const std::string &name) const;
+   /// check if a tensor is a shape tensor
+   bool IsShapeTensor(const std::string & name) const;
 
    // Add intermediate tensor
    void AddIntermediateTensor(std::string tensor_name, ETensorType type, std::vector<Dim> dim_shape);
@@ -131,6 +143,9 @@ public:
                                 std::shared_ptr<void> data);
    std::shared_ptr<void> GetInitializedTensorData(std::string tensor_name);
 
+   template<class T>
+   std::vector<T> GetTensorData(const std::string & name);
+
    void Initialize(int batchSize = -1, bool verbose = false);
    void Initialize(const std::map<std::string,size_t> & inputParams, bool verbose = false);
 
@@ -204,6 +219,23 @@ public:
    ClassDefNV(RModel, 3);
 };
 
+// need to implement here templated member functions and its specialization
+
+
+template<class T>
+inline std::vector<T> RModel::GetTensorData(const std::string & name) {
+   if (!IsInitializedTensor(name)) return std::vector<T>{};
+   T * data = static_cast<T*>(GetInitializedTensorData(name).get());
+   size_t size = ConvertShapeToLength(GetTensorShape(name));
+   return std::vector<T>(data, data+size);
+}
+
+template<>
+inline std::vector<Dim> RModel::GetTensorData<Dim>(const std::string & name) {
+   if (!IsShapeTensor(name)) return std::vector<Dim>{};
+   return GetShapeTensorValues(name);
+}
+
 } // namespace SOFIE
 } // namespace Experimental
 } // namespace TMVA

tmva/sofie/inc/TMVA/ROperator.hxx

@@ -42,7 +42,8 @@ protected:
    const std::string SP = "   ";    ///< space used to correctly indent the generated C++ code
    bool fUseSession = false;        ///< flag to identify if using the session class
    bool fIsOutputConstant = false;  ///< flag to identify if operator has a constant output (no need to generate code)
-   
+   bool fIsOutputParamShape = false;     ///< flag to identify of the output represents a parametric shape (can be knwon at compile time)
+
    mutable std::vector<std::string_view> fInputTensorNames;
    mutable std::vector<std::string_view> fOutputTensorNames;
 
@@ -54,7 +55,7 @@ public:
    std::span<const std::string_view> GetOpOutputTensors() const {
       return fOutputTensorNames;
    }
-   
+
 };
 
 

tmva/sofie/inc/TMVA/ROperator_BasicBinary.hxx

@@ -113,11 +113,10 @@ public:
          fShapeB = model.GetTensorShape(fNB);
          fDimShapeB = ConvertShapeToDim(fShapeB);
       }
-      std::cout <<  BinaryOperatorTrait<T, Op>::Name() << "  ";
-      if (dynamicInputs & 1)
-         std::cout <<  fNA << " is dynamic " << ConvertShapeToString(fDimShapeA) << "  ";
-      if (dynamicInputs & 2)
-         std::cout <<  fNB << " is dynamic " << ConvertShapeToString(fDimShapeB) << "  ";
+      if (dynamicInputs & 1 && model.Verbose() )
+         std::cout <<  BinaryOperatorTrait<T, Op>::Name() << " : input " << fNA << " is dynamic " << ConvertShapeToString(fDimShapeA) << "  ";
+      if (dynamicInputs & 2 && model.Verbose())
+         std::cout <<  BinaryOperatorTrait<T, Op>::Name() << " : input " << fNB << " is dynamic " << ConvertShapeToString(fDimShapeB) << "  ";
       std::cout << std::endl;
       // check if need to broadcast at initialization time if shapes are known and different
       // (we could broadcast the tensor tensor to maximum values of dynamic shapes - to be done)
@@ -126,7 +125,6 @@ public:
          auto ret = UTILITY::MultidirectionalBroadcastShape(fShapeA, fShapeB);
          fBroadcastFlag = ret.first;
          fShapeY = ret.second;
-         std::cout << BinaryOperatorTrait<T, Op>::Name() << " : checking for defined shapes " << fBroadcastFlag << "  " << ConvertShapeToString(fShapeY) << std::endl;
          bool broadcast =  ret.first > 0;
          if (broadcast) {
             // Y is the common shape of A and B
@@ -212,7 +210,6 @@ public:
          auto ret = UTILITY::MultidirectionalBroadcastShape(fDimShapeA, fDimShapeB);
          fBroadcastFlag = ret.first;
          fDimShapeY = ret.second;
-         std::cout << BinaryOperatorTrait<T, Op>::Name() << " : checking for Dim shapes " << fBroadcastFlag << "  " << ConvertShapeToString(fDimShapeY) << std::endl;
          // case of all parametric shapes and MultiDirectionalBroadcastShape  return the max of the 2
          // need to do before we declare the output tensor shape and the broadcasted ones
          if (ret.first & 4) {

tmva/sofie/inc/TMVA/ROperator_Comparision.hxx

@@ -157,22 +157,95 @@ public:
          fShapeY = fShapeX1;
       }
       // case of constant tensors
-      if (model.IsInitializedTensor(fNX1) && model.IsInitializedTensor(fNX2) ) {
+      T * data1 = nullptr;
+      T * data2 = nullptr;
+      std::vector<Dim> shapeData1;
+      std::vector<Dim> shapeData2;
+      size_t length = ConvertShapeToLength(fShapeY);
+      bool *  outData = new bool[length];
+      if (model.IsInitializedTensor(fNX1)) {
+         data1 = static_cast<T *>(model.GetInitializedTensorData(fNX1).get());
+      } else if (model.IsShapeTensor(fNX1)) {
+         shapeData1 = model.GetShapeTensorValues(fNX1);
+      }
+      if (model.IsInitializedTensor(fNX2)) {
+         data2 = static_cast<T *>(model.GetInitializedTensorData(fNX2).get());
+      } else if (model.IsShapeTensor(fNX2)) {
+         shapeData2 = model.GetShapeTensorValues(fNX2);
+      }
+      if (data1 && data2) {
          fIsOutputConstant = true;
-         auto data1 = static_cast<T *>(model.GetInitializedTensorData(fNX1).get());
-         auto data2 = static_cast<T *>(model.GetInitializedTensorData(fNX2).get());
-         size_t length = ConvertShapeToLength(fShapeY);
-         bool * outData = new bool[length];
          for (size_t i = 0; i < length; i++)
             outData[i] = ComparisionTrait<T,Op>::Result(data1[i], data2[i]);
          model.AddConstantTensor(fNY, fShapeY, outData);
          if (model.Verbose())
             std::cout <<  ComparisionTrait<T,Op>::Name() << " op ---> " << fNY << "  " << ConvertShapeToString(fShapeY) << " : "
                << ConvertValuesToString(length,outData) << std::endl;
-         delete [] outData;
-      } else {
+      } else if ((data1 || !shapeData1.empty()) && (data2 || !shapeData2.empty())) {
+         fIsOutputConstant = true;
+         if (data1 && !data2) {
+            // data 1 is constant and data2 is shape
+            for (size_t i = 0; i < length; i++) {
+               if (shapeData2[i].isParam) {
+                  if (shapeData2[i].dim == size_t(-1) || data1[i] > 0) {
+                     fIsOutputConstant = false;
+                     break;
+                  } else {
+                     // assume a comparison is done with .dim = 0
+                     shapeData2[i].dim = 0;
+                  }
+               }
+               outData[i] = ComparisionTrait<T,Op>::Result(data1[i], static_cast<T>(shapeData2[i].dim));
+            }
+         } else if (!data1 && data2) {
+            // data 1 is shape and dat2 is constant
+            for (size_t i = 0; i < length; i++) {
+               if (shapeData1[i].isParam) {
+                  if (shapeData1[i].dim == size_t(-1) || data2[i] > 0) {
+                     fIsOutputConstant = false;
+                     break;
+                  } else {
+                     // assume a comparison is done with .dim = 0
+                     shapeData1[i].dim = 0;
+                  }
+               }
+               outData[i] = ComparisionTrait<T,Op>::Result(static_cast<T>(shapeData1[i].dim), data2[i]);
+            }
+         } else if (!shapeData1.empty() && !shapeData2.empty() ) {
+            // both data1 and data2 are shape tensors
+            for (size_t i = 0; i < length; i++) {
+               if (!shapeData1[i].isParam && !shapeData2[i].isParam) {
+                  outData[i] = ComparisionTrait<T,Op>::Result(shapeData1[i].dim, shapeData2[i].dim);
+               }
+               else if (shapeData1[i].isParam && shapeData2[i].isParam) {
+                  if (shapeData1[i].param == shapeData2[i].param)
+                     outData[i] = ComparisionTrait<int,Op>::Result(1,1); // comparison of two equal value
+                  else {
+                     fIsOutputConstant = false;
+                     break;
+                  }
+               }
+               else {
+                  fIsOutputConstant = false;
+                  break;
+               }
+            }
+         }
+         if (fIsOutputConstant) {
+            model.AddConstantTensor(fNY, fShapeY, outData);
+            if (model.Verbose())
+               std::cout <<  ComparisionTrait<T,Op>::Name() << " op ---> " << fNY << "  " << ConvertShapeToString(fShapeY) << " : "
+                  << ConvertValuesToString(length,outData) << " (constant) " << std::endl;
+
+         }
+      }
+      delete [] outData;
+      if (!fIsOutputConstant) {
          model.AddIntermediateTensor(fNY, ETensorType::BOOL , fShapeY);
+         if (model.Verbose())
+               std::cout <<  ComparisionTrait<T,Op>::Name() << " op ---> " << fNY << "  " << ConvertShapeToString(fShapeY) << std::endl;
       }
+
       // check if this is not output operators to add a specific line for definining the tensor_xxx variable
       const auto & outputTensorNames = model.GetOutputTensorNames();
       fIsModelOutput = false;

tmva/sofie/inc/TMVA/ROperator_Concat.hxx

@@ -193,12 +193,18 @@
                std::cout << "Output of concat operator has shape " << ConvertDimShapeToString(fOutputShape) << std::endl;
 
             // check if concat has constant inputs , axis 0(concat contigous memory and type is integer)
+            bool isOutputShape = false;
             if (model.GetTensorType(fInputs[0]) == ETensorType::INT64 && fAxis == 0) {
                fIsOutputConstant = true;
+               isOutputShape = true;
+
                for ( auto & input : fInputs) {
                   if (!model.IsInitializedTensor(input)) {
                      fIsOutputConstant = false;
-                     break;
+                     if (!model.IsShapeTensor(input)) {
+                        isOutputShape = false;
+                        break;
+                     }
                   }
                }
                if (fIsOutputConstant) {
@@ -217,8 +223,35 @@
                   model.AddConstantTensor<int64_t>(fOutput, outputShape, outputData.data());
                   if (model.Verbose()) {
                      std::cout << "output of Concat is a constant tensor " << ConvertShapeToString(outputShape) << " : "
-                     << ConvertValuesToString(outputData) << std::endl;
+                     << ConvertValuesToString(outputData) << " (constant)" << std::endl;
+                  }
+               } else if (isOutputShape) {
+                  auto outputShape = ConvertShapeToInt(fOutputShape);  // conversion must be possible
+                  std::vector<Dim> outputData(ConvertShapeToLength(outputShape));
+                  size_t offset = 0;
+                  for ( auto & input : fInputs) {
+                     std::vector<Dim> inputData;
+                     auto inputShape = model.GetTensorShape(input); // shape is not dynamic
+                     size_t inputLength = ConvertShapeToLength(inputShape); // shape can be a scalar
+                     if (model.IsShapeTensor(input))
+                        inputData = model.GetShapeTensorValues(input);
+                     else if (model.IsConstantTensor(input)) {
+                        inputData.resize(inputLength);
+                        auto intData = static_cast<int64_t*>(model.GetInitializedTensorData(input).get());
+                        for (size_t i = 0; i < inputData.size(); i++)
+                           inputData[i] = Dim{ static_cast<size_t>(intData[i])};
+                     }
+                     std::cout << "concatanating input data " << inputLength << "  " << inputData[0] << std::endl;
+                     std::copy(inputData.begin(), inputData.end(), outputData.begin() + offset );
+                     offset += inputLength;
+                  }
+                  // add output tensor
+                  model.AddShapeTensor(fOutput,outputData, false); // cannot be a  scalar
+                  if (model.Verbose()) {
+                     std::cout << "output of Concat is a shape tensor " << ConvertShapeToString(outputShape) << " : "
+                     << ConvertShapeToString(outputData) << " (shape)" <<  std::endl;
                   }
+                  fIsOutputConstant = true;
                }
             }
             if (!fIsOutputConstant) {

tmva/sofie/inc/TMVA/ROperator_Constant.hxx

@@ -25,6 +25,7 @@ private:
    std::vector<T> fValues;
    std::string fAttrType;
    bool fIsConstantOfShape = false;
+   bool fIsUndefinedInputShape = false;
 
 public:
    ROperator_Constant(){}
@@ -52,6 +53,7 @@ public:
    void Initialize(RModel& model) override {
        //input must be a graph input, or already initialized intermediate tensor
       size_t length = 1;
+      /// ConstantOfShape-------------
       if (!fNX.empty()) {
          // case of ConstantOfShape (since no inputs in case of Constant operator)
          fIsConstantOfShape  = true;
@@ -81,8 +83,13 @@ public:
             T value = fValues[0];
             fValues = std::vector<T>(length, value);
          }
-         else {
-            // case of non constant tensors- we need to do at run time
+         else if (model.IsShapeTensor(fNX)) {
+            // case tensor values representing output shapes are  known
+            fDimOutputShape = model.GetShapeTensorValues(fNX);
+         } else {
+            // case of not known shape tensors- we need to do at run time
+            // not sure if we ever encounter this case
+            fIsUndefinedInputShape = true;
             fDimShape = model.GetDimTensorShape(fNX);
             if (fDimShape.size() > 1 )
                throw std::runtime_error("TMVA SOFIE ConstantOfShape Op Input Tensor has invalid shape");
@@ -138,9 +145,11 @@ public:
       // generate constant tensor according to input
 
       out << "\n//--------- ConstantOfShape " << opName << " --> " << ConvertShapeToString(fDimOutputShape) << "\n";
-       // set shape values
-      for (size_t i = 0; i < fDimOutputShape.size(); i++) {
-         out << SP << "size_t " << fDimOutputShape[i].param << " = " << "tensor_" << fNX << "[" << i << "];\n";
+       // set shape values if needed
+      if (fIsUndefinedInputShape) {
+         for (size_t i = 0; i < fDimOutputShape.size(); i++) {
+            out << SP << "size_t " << fDimOutputShape[i].param << " = " << "tensor_" << fNX << "[" << i << "];\n";
+         }
       }
       auto length = ConvertDimShapeToLength(fDimOutputShape);
       // vector is already allocated- fill with values