[tmva][sofie] Add support for Dynamic tensors in more operators

Include now also BatchNormalization, COnstantOfShape, Gather , Reshape and Sigmoid

Fix the broadcasting of dynamic tensors in BasicBinary

Author: lmoneta

tmva/sofie/inc/TMVA/ROperator_BasicBinary.hxx

@@ -55,6 +55,7 @@ template<typename T, EBasicBinaryOperator Op>
 class ROperator_BasicBinary final : public ROperator{
 private:
 
+   int fBroadcastFlag = 0;
    std::string fNA;
    std::string fNB;
    std::string fNBroadcastedA;
@@ -114,12 +115,14 @@ public:
       // case of known shapes
       if (!fShapeA.empty() && !fShapeB.empty()) {
          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
-            bool broadcastA = ret.first > 1;
-            bool broadcastB = ret.first == 1 || ret.first == 3;
+            bool broadcastA = ret.first & 2;
+            bool broadcastB = ret.first & 1;
             // Broadcast A to Y
             if (broadcastA) {
                fNBroadcastedA = "Broadcasted" + fNA + "to" + fNY;
@@ -191,17 +194,28 @@ public:
       else {
          // case A or B have dynamic shapes. We need to broadcast if shape are not same
          auto ret = UTILITY::MultidirectionalBroadcastShape(fDimShapeA, fDimShapeB);
+         fBroadcastFlag = ret.first;
          fDimShapeY = ret.second;
-         if (ret.first > 1) {
+         std::cout << BinaryOperatorTrait<T, Op>::Name() << " : checking for Dim shapes " << fBroadcastFlag << "  " << ConvertShapeToString(fDimShapeY) << std::endl;
+         if (ret.first & 2) {
             // case we broadcast A
             fNBroadcastedA = "Broadcasted" + fNA + "to" + fNY;
             model.AddIntermediateTensor(fNBroadcastedA, model.GetTensorType(fNA), fDimShapeY);
          }
-         if (ret.first == 1 || ret.first == 3)  {
+         if (ret.first & 1)  {
             // case we broadcast B
             fNBroadcastedB = "Broadcasted" + fNB + "to" + fNY;
             model.AddIntermediateTensor(fNBroadcastedB, model.GetTensorType(fNB), fDimShapeY);
          }
+         // case of all parametric shapes and we know only at run time
+         // we don't add in this case an intermediate tensor for broadcasting
+         // if (ret.first == 4) {
+         //    for (auto & d : fDimShapeY) {
+         //       if (d.isParam && d.param.find("broadcast") != std::string::npos) {
+         //          d.param += fNY;
+         //       }
+         //    }
+         // }
          // add output tensor
          model.AddIntermediateTensor(fNY, model.GetTensorType(fNA), fDimShapeY);
       }
@@ -212,11 +226,11 @@ public:
       return out.str();
    }
 
-   std::string Generate(std::string OpName) override {
+   std::string Generate(std::string opName) override {
 
       if (fIsOutputConstant) return "";
 
-      OpName = "op_" + OpName;
+      opName = "op_" + opName;
 
       if (fDimShapeY.empty()) {
          throw std::runtime_error("TMVA SOFIE Binary Op called to Generate without being initialized first");
@@ -225,21 +239,55 @@ public:
       out << SP << "\n//------ " << BinaryOperatorTrait<T,Op>::Name() << "\n";
       auto length = ConvertDimShapeToLength(fDimShapeY);
       std::string typeName = TensorType<T>::Name();
+      // we need to check if we can broadcast (case flag has bit 4 set)
+      if (fBroadcastFlag & 4) {
+         // need to check if shapes are the same
+         auto lengthA = ConvertDimShapeToLength(fDimShapeA);
+         auto lengthB = ConvertDimShapeToLength(fDimShapeB);
+         out << SP << "if (" << lengthA << "!=" << lengthB << ") {\n";
+         // check if A->B or B->A
+         //bool broadcastable = true;
+         for (size_t i = 0; i < fDimShapeY.size(); i++) {
+            if (fBroadcastFlag & 5 && fDimShapeY[i] == fDimShapeA[i] && fDimShapeA[i].dim > 1 && fDimShapeB[i].isParam) {
+               // B->A B[i] needs to be 1
+               out << SP << SP << "if (" << fDimShapeB[i] << "!= 1)\n";
+               out << SP << SP << SP << "throw std::runtime_error(\"SOFIE - Cannot broadcast B->A in operator "
+                                        << opName << "\");\n";
+            }
+            if (fBroadcastFlag & 6 && fDimShapeY[i] == fDimShapeB[i] && fDimShapeB[i].dim > 1 && fDimShapeA[i].isParam) {
+               //A-> B A[i] needs to be 1
+               out << SP << SP << "if (" << fDimShapeA[i] << "!= 1)\n";
+               out << SP << SP << SP << "throw std::runtime_error(\"SOFIE - Cannot broadcast A->B in operator "
+                                        << opName << "\");\n";
+            }
+            else if (fDimShapeA[i].isParam && fDimShapeB[i].isParam) {
+               // both shapes are parametric and we broadcast to maximum
+               // we allocate here output vector
+               out << SP << SP << "if (" << fDimShapeA[i] << " != " << fDimShapeB[i] << " && ("
+                   << fDimShapeA[i] << " != 1 || " << fDimShapeB[i] << " != 1))\n";
+               out << SP << SP << "throw std::runtime_error(\"SOFIE - Cannot broadcast shapes in operator "
+                                        << opName << "\");\n";
+            }
+         }
+      } else {
+         out << SP << "{\n";
+      }
       // Broadcast A if it's uninitialized
       // use broadcasting function where we pass an already allocated tensor to minimize memory allocations
       if (!fNBroadcastedA.empty()) {
-         out << SP << "// Broadcasting uninitialized tensor " << fNA << "\n";
-         out << SP  << "TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<" << typeName << ">(tensor_" << fNA << ", "
+         out << SP << SP << "// Broadcasting uninitialized tensor " << fNA << "\n";
+         out << SP << SP  << "TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<" << typeName << ">(tensor_" << fNA << ", "
                   << ConvertDimShapeToString(fDimShapeA) << ", " << ConvertDimShapeToString(fDimShapeY)
                   << ", fTensor_" << fNBroadcastedA << ");\n";
       }
       // Broadcast B if it's uninitialized
       if (!fNBroadcastedB.empty()) {
-         out << SP << "// Broadcasting uninitialized tensor " << fNB << "\n";
-         out << SP << "TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<" << typeName << ">(tensor_" << fNB << ", "
+         out << SP << SP << "// Broadcasting uninitialized tensor " << fNB << "\n";
+         out << SP << SP << "TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<" << typeName << ">(tensor_" << fNB << ", "
                   << ConvertDimShapeToString(fDimShapeB) << ", " << ConvertDimShapeToString(fDimShapeY)
                   << ", fTensor_" << fNBroadcastedB << ");\n";
       }
+      out << SP << "}\n"; // end if on broadcasting
       const std::string& nameA = fNBroadcastedA.empty()? fNA : fNBroadcastedA;
       const std::string& nameB = fNBroadcastedB.empty()? fNB : fNBroadcastedB;
       out << SP << "for (size_t id = 0; id < " << length << " ; id++){\n";

tmva/sofie/inc/TMVA/ROperator_BatchNormalization.hxx

@@ -32,12 +32,12 @@ private:
    std::string fNY;
    EActivationType fActivation;
 
-   std::vector<size_t> fShapeX;
+   std::vector<Dim> fShapeX;
    std::vector<size_t> fShapeScale;
    std::vector<size_t> fShapeB;
    std::vector<size_t> fShapeMean;
    std::vector<size_t> fShapeVar;
-   std::vector<size_t> fShapeY;
+   std::vector<Dim> fShapeY;
 
    std::string fType;
 
@@ -109,7 +109,7 @@ public:
             std::runtime_error("TMVA SOFIE BatchNormalization op Input Tensor " + fNVar + " fnv is not found in model");
       }
 
-      fShapeX = model.GetTensorShape(fNX);
+      fShapeX = model.GetDimTensorShape(fNX);
 
       if (fShapeX.size() <  2 || fShapeX.size() > 4) {
          throw
@@ -123,16 +123,17 @@ public:
       fShapeY = fShapeX;
       model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeY);
 
-      if (fShapeB.size() == 1) {
+      if (fShapeB.size() == 1 && !model.IsDynamicTensor(fNX)) {
+         auto shapeX = model.GetTensorShape(fNX);
          // Broadcast scale, bias, input_mean and input_var to shape_X
          auto original_B = model.GetInitializedTensorData(fNB);
          auto original_S = model.GetInitializedTensorData(fNScale);
          auto original_M = model.GetInitializedTensorData(fNMean);
          auto original_V = model.GetInitializedTensorData(fNVar);
-         size_t batchSize = fShapeX[0];
-         size_t channels = fShapeX[1];
-         size_t height = (fShapeX.size() > 2) ? fShapeX[2] : 1;
-         size_t width = (fShapeX.size() > 3) ? fShapeX[3] : 1;
+         size_t batchSize = shapeX[0];
+         size_t channels = shapeX[1];
+         size_t height = (shapeX.size() > 2) ? shapeX[2] : 1;
+         size_t width = (shapeX.size() > 3) ? shapeX[3] : 1;
          size_t n = batchSize * channels * height * width;
          if (fType == "float") {
             float *original_bias = static_cast<float *>(original_B.get());
@@ -181,6 +182,8 @@ public:
             fShapeMean = model.GetTensorShape(fNMean);
             fShapeVar = model.GetTensorShape(fNVar);
          }
+      } else {
+         // we need to broadcast at run time
       }
    }
 
@@ -192,15 +195,15 @@ public:
 
       std::stringstream out;
       //// Batch Norm op
-      size_t batchSize = fShapeX[0];
-      size_t channels = fShapeX[1];
-      size_t height = (fShapeX.size() > 2) ? fShapeX[2] : 1;
-      size_t width = (fShapeX.size() > 3) ? fShapeX[3] : 1;
-      size_t n = batchSize * channels * height * width;
+      std::string batchSize = fShapeX[0].GetVal();
+      std::string channels = fShapeX[1].GetVal();
+      std::string height = (fShapeX.size() > 2) ? fShapeX[2].GetVal() : "1";
+      std::string width = (fShapeX.size() > 3) ? fShapeX[3].GetVal() : "1";
+      auto n = ConvertDimShapeToLength(fShapeX);
 
       //// copy X into Y
       out << "\n\n//---- BatchNorm\n";
-      out << SP << "constexpr int " << OpName << "_N =" << batchSize * channels * height * width << ";\n";
+      out << SP << "constexpr int " << OpName << "_N =" << n << ";\n";
       out << SP << "constexpr int "<<OpName<< "_incx = 1;\n";
       out << SP << "constexpr int "<<OpName<< "_incy = 1;\n";
       out << SP << "BLAS::scopy_(&" << OpName << "_N, " << "tensor_" << fNX << ", &" << OpName << "_incx," << "tensor_" << fNY << ", &" << OpName << "_incy);\n\n";
@@ -222,7 +225,7 @@ public:
          << "tensor_" << fNY << ", &" << OpName << "_incy);\n\n";
 
       if(fActivation == EActivationType::RELU){
-         out << SP << "for (int id = 0; id < " << ConvertShapeToLength(fShapeY) << " ; id++){\n";
+         out << SP << "for (int id = 0; id < " << n << " ; id++){\n";
          out << SP << SP << "tensor_" << fNY << "[id] = ((tensor_" << fNY << "[id] > 0 )? tensor_" << fNY << "[id] : 0);\n";
          out << SP << "}\n";
       }

tmva/sofie/inc/TMVA/ROperator_Constant.hxx

@@ -20,6 +20,8 @@ private:
    std::string fNX;
    std::string fNY;
    std::vector<size_t> fShape;
+   std::vector<Dim> fDimShape;
+   std::vector<Dim> fDimOutputShape;
    std::vector<T> fValues;
    std::string fAttrType;
    bool fIsConstantOfShape = false;
@@ -58,28 +60,47 @@ public:
          }
          // get output shape from input values:
          // can work only if input is a constant or initialized tensor (or dynamic one)
-         auto dptr = model.GetInitializedTensorData(fNX);
-         auto input_tensor = static_cast<int64_t *>(dptr.get());
-         auto input_shape = model.GetTensorShape(fNX);
-         if (input_shape.size() > 1 )
-            throw std::runtime_error("TMVA SOFIE ConstantOfShape Op Input Tensor has invalid shape");
-         if (input_tensor != nullptr && !input_shape.empty()) {
-            fShape = std::vector<size_t> (input_shape[0]);
-            for (size_t i = 0; i < fShape.size(); i++)
-               fShape[i] = input_tensor[i];
-         } else
-            fShape = {1};  // scalar case
-
-         length = ConvertShapeToLength(fShape);
-         if (fValues.size() != 1)
-            throw std::runtime_error("TMVA SOFIE ConstantOfShape Op value Tensor has invalid size " + std::to_string(fValues.size()));
-
-         T value = fValues[0];
-         fValues = std::vector<T>(length, value);
+         if (model.IsInitializedTensor(fNX) || model.IsConstantTensor(fNX)) {
+            fIsOutputConstant = true;
+            auto dptr = model.GetInitializedTensorData(fNX);
+            auto input_tensor = static_cast<int64_t *>(dptr.get());
+            auto input_shape = model.GetTensorShape(fNX);
+            if (input_shape.size() > 1 )
+               throw std::runtime_error("TMVA SOFIE ConstantOfShape Op Input Tensor has invalid shape");
+            if (input_tensor != nullptr && !input_shape.empty()) {
+               fShape = std::vector<size_t> (input_shape[0]);
+               for (size_t i = 0; i < fShape.size(); i++)
+                  fShape[i] = input_tensor[i];
+            } else
+               fShape = {1};  // scalar case
+
+            length = ConvertShapeToLength(fShape);
+            if (fValues.size() != 1)
+               throw std::runtime_error("TMVA SOFIE ConstantOfShape Op value Tensor has invalid size " + std::to_string(fValues.size()));
+
+            T value = fValues[0];
+            fValues = std::vector<T>(length, value);
+         }
+         else {
+            // case of non constant tensors- we need to do at run time
+            fDimShape = model.GetDimTensorShape(fNX);
+            if (fDimShape.size() > 1 )
+               throw std::runtime_error("TMVA SOFIE ConstantOfShape Op Input Tensor has invalid shape");
+            if (!fDimShape[0].isParam) {
+               fDimOutputShape.resize(fDimShape[0].dim);
+               for (size_t i = 0; i < fDimShape[0].dim; i++) {
+                  fDimOutputShape[i] = Dim{ std::string("s_") + fNY + "_" + std::to_string(i)};
+               }
+            }
+            else {
+               throw std::runtime_error("TMVA SOFIE ConstantOfShape Op Input Tensor has not defied shape");
+            }
+         }
 
       } else {
          // case of constant operator
          // in case of standard constant the shape is provided as input
+         fIsOutputConstant = true;
          length = ConvertShapeToLength(fShape);
          if (length != fValues.size())
             throw std::runtime_error("TMVA SOFIE Constant Op has invalid shape : " + ConvertShapeToString(fShape) +
@@ -90,18 +111,41 @@ public:
       // but keep its initialization in the generated code. The values might also be needed in initializing the
       // following operators using as input Constant or ConstantOfShape
        // resize fValues to shape length
-      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;
+      if (fIsOutputConstant) {
+         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;
+         }
+      } else {
+         model.AddIntermediateTensor(fNY, ConvertStringToType(TensorType<T>::Name()), fDimOutputShape);
       }
    }
 
-   std::string Generate(std::string /* OpName */) override {
+   std::string Generate(std::string opName) override {
       // no code to generate here. Tensor are defined in Session constructor
-      return "//---------------------------------------\n";
+      if (fIsOutputConstant) {
+         if (fNX.empty())
+            return "// ---- Constant (no-op) \n";
+         else
+            return "// ---- ConstantOfShape (no-op) \n";
+      }
+      // Only ConstantOfShape might require generation code
+      // generate constant tensor according to input
+      std::stringstream out;
+      out << "\n//--------- ConstantOfShape " << opName << "\n";
+       // set shape values
+      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
+      out << SP << "if (" << length << " > fTensor_" << fNY << ".size())\n";
+      out << SP << SP << "fTensor_" << fNY << ".resize(" << length  << ");\n";
+      out << SP << "std::fill(fTensor_" << fNY << ".begin(), fTensor_" << fNY << ".end(), " << fValues[0] << ");\n";
+      return out.str();
    }
 };
 

tmva/sofie/inc/TMVA/ROperator_Gather.hxx

@@ -138,6 +138,8 @@ public:
       // Indices of shape q
       size_t q = fShapeIndices.size();
       // Strides
+      std::cout << "shapes of Gather " << ConvertShapeToString(fShapeX) << "  " <<
+            ConvertShapeToString(fShapeY) << "  " << ConvertShapeToString(fShapeIndices) << std::endl;
       auto stridesX = UTILITY::ComputeStrideFromShape(fShapeX);
       auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY);
       auto stridesIndices = UTILITY::ComputeStrideFromShape(fShapeIndices);