ROperator_Gemm.hxx

Go to the documentation of this file.

#ifndef TMVA_SOFIE_ROPERATOR_GEMM
#define TMVA_SOFIE_ROPERATOR_GEMM
 
 
#include "TMVA/SOFIE_common.hxx"
#include "TMVA/ROperator.hxx"
#include "TMVA/RModel.hxx"
 
#include <sstream>
#include <algorithm>
#include <iterator>
#include <iomanip>
#include <limits>
#include <cassert>
 
namespace TMVA{
namespace Experimental{
namespace SOFIE{
 
 
   template <typename T>
   class ROperator_Gemm final : public ROperator
   {
 
   private:
      bool fIsDynamic = false;
      bool fBroadcastBias = false;
      bool fCheckBiasShapeAtRuntime = false; // flag to identify the need to do a run time check of bias shape compatibility in case of dynamic shapes and uni-directional broadcasting
      bool fBiasBroadcastAssumed = false;    // Initialize assumed a broadcast: the integer shape of Y was unknown
 
      float fAttrAlpha = 1.0;
      float fAttrBeta = 1.0;
      int_t fAttrTransA = 0;
      int_t fAttrTransB = 0;
 
      std::string fNA;
      std::string fNB;
      std::string fNC = "";
      std::string fNY;
      std::string fType;
      EActivationType fActivation;
      std::vector<Dim> fShapeA;
      std::vector<Dim> fShapeB;
      std::vector<size_t> fShapeC;
      std::vector<Dim> fDimShapeC;
      std::vector<Dim> fShapeY;
      RModel * fModel = nullptr;
 
   public:
 
      ROperator_Gemm(){}
      ROperator_Gemm(float alpha, float beta, int_t transA, int_t transB, std::string nameA, std::string nameB, std::string nameY, EActivationType activation=EActivationType::UNDEFINED):
         fAttrAlpha(alpha), fAttrBeta(beta), fAttrTransA(transA), fAttrTransB(transB), fNA(UTILITY::Clean_name(nameA)),
         fNB(UTILITY::Clean_name(nameB)), fNY(UTILITY::Clean_name(nameY))
      {
         fActivation = activation;
         fType = "float";
         static_assert(std::is_same_v<T, float>,
                  "TMVA::SOFIE - Unsupported type parsing a Gemm operator");
         fInputTensorNames = { fNA, fNB };
         fOutputTensorNames = { fNY };
      }
 
      ROperator_Gemm(float alpha, float beta, int_t transA, int_t transB, std::string nameA, std::string nameB, std::string nameC, std::string nameY, EActivationType activation=EActivationType::UNDEFINED):
         fAttrAlpha(alpha), fAttrBeta(beta), fAttrTransA(transA), fAttrTransB(transB), fNA(UTILITY::Clean_name(nameA)),
         fNB(UTILITY::Clean_name(nameB)), fNC(UTILITY::Clean_name(nameC)), fNY(UTILITY::Clean_name(nameY)), fActivation(activation)
      {
         fActivation = activation;
         fType = "float";
 
         fInputTensorNames = {fNA, fNB, fNC};
         fOutputTensorNames = { fNY };
      }
 
      template <typename U>
      std::vector<U> DoShapeInference(const std::vector<std::vector<U>> & input){
         if (input.size() > 3) throw std::runtime_error("TMVA SOFIE Gemm Op Shape Inference only need 2 or 3 input tensor");
         // accept tensor with input dimensions > 2
         // example: A = (d1,d2,...,N1,N2)  B = (d1,d2,...,N2,N3)    --> Y = (d1,d2,..,N1,N3)
         for (auto& i: input){
            if (i.size() < 2){
               throw std::runtime_error("TMVA SOFIE Gemm Op Shape Inference only accept input tensor with >=2 dimensions");
            }
         }
 
         // when there are 3 inputs shape of Y is the one of C
         if (input.size() == 3){
            //shape of C is shape of Y
            return input[2];
         }
         // ioffset cannot be less than 2
         int ioffset = input[0].size()-2;  // in case of tensors with dim > 2
 
         std::vector<U> s_a(input[0].begin() + ioffset, input[0].begin() + ioffset + 2);
         std::vector<U> s_b(input[1].begin() + ioffset, input[1].begin() + ioffset + 2);
         // reverse in case of transpose
         if (fAttrTransA){
            std::reverse(s_a.begin(), s_a.end());
         }
         if (fAttrTransB){
            std::reverse(s_b.begin(), s_b.end());
         }
         std::vector<U> s_y;
         s_y.reserve(input[0].size());
         if (input[0].size() > 2 && input[1].size() == input[0].size()) {
            // in case of dim > 2 first dimensions are equal to the input ones not
            // equal to 1 (e.g. (1,2,3) * (2,3,4) -> (2,2,4))
            // here could probably use the Broadcasting function  UTILITY::MultidirectionalBroadcastShape
            for (size_t i = 0; i < input[0].size()-2; i++) {
               Dim valueA = input[0][i];
               Dim valueB = input[1][i];
               if (valueA.GetVal() != valueB.GetVal()) {
                  if (valueB.GetVal() == "1")
                     s_y.push_back(input[0][i]);
                  else if (valueA.GetVal() == "1")
                     s_y.push_back(input[1][i]);
                  else if (!valueA.isParam && !valueB.isParam)
                     throw std::runtime_error("TMVA SOFIE Gemm Op - invalid input shapes " + valueA.GetVal() + " and "
                        + valueB.GetVal());
                  else if (valueA.isParam && valueB.isParam){
                      // check which parameter is first in RModel list
                     auto & dimNames = fModel->GetDimShapeNames();
                     auto p1 = std::find(dimNames.begin(), dimNames.end(), valueA.param);
                     auto p2 = std::find(dimNames.begin(), dimNames.end(), valueB.param);
                     if (p1 < p2) s_y.push_back(input[0][i]);
                     else  s_y.push_back(input[1][i]);
                  }
                  else if (!valueA.isParam)
                     s_y.push_back(input[0][i]);
                  else if (!valueB.isParam)
                     s_y.push_back(input[1][i]);
                  else
                     throw std::runtime_error("TMVA SOFIE Gemm Op - invalid input shapes " + valueA.GetVal() + " and "
                        + valueB.GetVal());
               }
               else
                  s_y.push_back(input[0][i]);
            }
         }
 
         s_y.push_back(s_a[0]);
         s_y.push_back(s_b[1]);
         return s_y;
      }
 
      std::vector<Dim> DynamicShapeInference(const std::vector<std::vector<Dim>> & input){
         return DoShapeInference<Dim>(input);
      }
 
 
 
      void Initialize(RModel& model) override {
         //TODO: propagate A or B as specified by ONNX standard
         fModel = &model;
 
         if ((model.CheckIfTensorAlreadyExist(fNA) == false) || (model.CheckIfTensorAlreadyExist(fNB) == false) ){   //input must be a graph input, or already initialized intermediate tensor
            throw std::runtime_error("TMVA SOFIE Gemm Op Input Tensor " + fNA + " or " + fNB + " is not found in model");
         }
         if (fNC != ""){
            if (model.CheckIfTensorAlreadyExist(fNC) == false){   //input must be a graph input, or already initialized intermediate tensor
               throw std::runtime_error("TMVA SOFIE Gemm Op Input Tensor " + fNC + " is not found in model");
            }
         }
         if (model.IsDynamicTensor(fNA) || model.IsDimInputTensor(fNA) ) {
            fShapeA = model.GetDynamicTensorShape(fNA);
            fIsDynamic = true;
         } else {
            auto shapeA_int = model.GetTensorShape(fNA);
            fShapeA = ConvertShapeToDim(shapeA_int);
         }
         // case A is of dim1 we prepend a 1 but we need to remove later
         bool prependOne = false;
         if (fShapeA.size() == 1) {
            fShapeA.insert(fShapeA.begin(), Dim(1));
            prependOne = true;
         }
 
         if (model.IsDynamicTensor(fNB) || model.IsDimInputTensor(fNB)) {
            fShapeB = model.GetDynamicTensorShape(fNB);
            fIsDynamic = true;
         }
         else {
            auto shapeB_int = model.GetTensorShape(fNB);
            fShapeB = ConvertShapeToDim(shapeB_int);
         }
         // case B is dim1 we append a 1 but we need to remove later
         bool appendOne = false;
         if (fShapeB.size() == 1) {
            fShapeB.insert(fShapeB.end(), Dim(1));
            appendOne = true;
         }
         // assume if not shape is 2 that extra values are 1.
         // implement also MatMul case where we stack matrices (see numpy.matmul)
         if (fShapeA.size() != fShapeB.size()) {
            // if different dimensions we prepend 1 values
            if (fShapeA.size() < fShapeB.size()) {
               fShapeA.insert(fShapeA.begin(), fShapeB.size()-fShapeA.size(), Dim(1));
            } else if (fShapeB.size() < fShapeA.size()) {
               fShapeB.insert(fShapeB.begin(), fShapeA.size()-fShapeB.size(), Dim(1));
            }
         }
 
         fShapeY = DynamicShapeInference({fShapeA, fShapeB});
         std::vector<size_t> shapeY = ConvertShapeToInt(fShapeY);
 
         // bias is normally not dynamic (not support it for time being)
         if (fNC != ""){
            if (model.IsDynamicTensor(fNC))
               fDimShapeC = model.GetDynamicTensorShape(fNC);
            else {
               fShapeC = model.GetTensorShape(fNC);
               fDimShapeC = ConvertShapeToDim(fShapeC);
            }
            // for dynamic outputs broadcasting is always needed
            bool broadcast_needed = false;
            if (fIsDynamic && shapeY.empty()) {
               broadcast_needed = true;
               fBiasBroadcastAssumed = true;
            } else
               // consider broadcasting also if they have different length
               broadcast_needed = (fShapeC != shapeY);
 
 
            if (broadcast_needed) {
               fBroadcastBias = true;
               // check if broadcasting is compatible and note that prepend 1 to shapeC
               auto r = UTILITY::MultidirectionalBroadcastShape(fShapeY, fDimShapeC);
               // return flag must not have bit equal to 2 since this is a unidirectional broadcast of C->Y
               //
               if ((r.first & 2) == 2) {
                  throw std::runtime_error("TMVA SOFIE Gemm Op - bias tensor of shape " + ConvertDimShapeToString(fDimShapeC) + " cannot be uni-directional broadcasted to " + ConvertDimShapeToString(fShapeY));
               } else if (r.first  == 4) {
                  // we need to do a run time check of bias shape if it is compatible
                  fCheckBiasShapeAtRuntime = true;
               }
               fShapeC = ConvertShapeToInt(fDimShapeC);
            }
         }
 
         // remove appended or prepended value of 1 in Y
         if (prependOne) {
            if (fIsDynamic)
               fShapeY.erase(fShapeY.begin());
            else
               shapeY.erase(shapeY.begin());
         }
         if (appendOne) {
            if (fIsDynamic)
               fShapeY.erase(fShapeY.end()-1);
            else
               shapeY.erase(shapeY.end()-1);
         }
 
         // Constant-fold Gemm/MatMul when A, B (and C) are all initializers, following the
         // ROperator_BasicBinary pattern (compute now, skip Generate() entirely). Only full
         // constant folding is handled; propagating just A or B (see the TODO above) would
         // need a different mechanism than fIsOutputConstant's all-or-nothing fold.
         bool canFold = !fIsDynamic
            && model.IsInitializedTensor(fNA)
            && model.IsInitializedTensor(fNB)
            && (fNC.empty() || model.IsInitializedTensor(fNC))
            && fShapeA.size() <= 2       // exclude stacked/batched MatMul
            && !fBroadcastBias           // exclude bias requiring run-time broadcast
            && !fCheckBiasShapeAtRuntime;
 
         if (canFold) {
            auto shapeA_i = ConvertShapeToInt(fShapeA);
            auto shapeB_i = ConvertShapeToInt(fShapeB);
            size_t dimA = shapeA_i.size();
            size_t dimB = shapeB_i.size();
            size_t m = fAttrTransA ? shapeA_i[dimA - 1] : shapeA_i[dimA - 2];
            size_t k = fAttrTransA ? shapeA_i[dimA - 2] : shapeA_i[dimA - 1];
            size_t n = fAttrTransB ? shapeB_i[dimB - 2] : shapeB_i[dimB - 1];
 
            auto dataA = static_cast<T *>(model.GetInitializedTensorData(fNA).get());
            auto dataB = static_cast<T *>(model.GetInitializedTensorData(fNB).get());
 
            // plain host-side 2D matrix multiply: Y = alpha * op(A) * op(B)
            std::vector<T> dataY(m * n, T(0));
            for (size_t i = 0; i < m; i++) {
               for (size_t j = 0; j < n; j++) {
                  T sum{};
                  for (size_t p = 0; p < k; p++) {
                     T aVal = fAttrTransA ? dataA[p * m + i] : dataA[i * k + p];
                     T bVal = fAttrTransB ? dataB[j * k + p] : dataB[p * n + j];
                     sum += aVal * bVal;
                  }
                  dataY[i * n + j] = static_cast<T>(fAttrAlpha) * sum;
               }
            }
            // Y += beta * C (fBroadcastBias is false here, so C already matches Y's length)
            if (!fNC.empty()) {
               auto dataC = static_cast<T *>(model.GetInitializedTensorData(fNC).get());
               for (size_t idx = 0; idx < dataY.size(); idx++)
                  dataY[idx] += static_cast<T>(fAttrBeta) * dataC[idx];
            }
            // fuse ReLU now since Generate() will be skipped entirely for a constant output
            if (fActivation == EActivationType::RELU) {
               for (auto &v : dataY)
                  v = std::max(v, T(0));
            }
 
            model.AddConstantTensor<T>(fNY, shapeY, dataY.data());
            // flag the operand tensors to not be written in the generated code or weight file
            model.SetNotWritableInitializedTensor(fNA);
            model.SetNotWritableInitializedTensor(fNB);
            if (!fNC.empty())
               model.SetNotWritableInitializedTensor(fNC);
            fIsOutputConstant = true;
 
            if (model.Verbose()) {
               std::cout << "Gemm (or MatMul) " << fNA << " , " << fNB;
               if (!fNC.empty())
                  std::cout << " , " << fNC;
               std::cout << " ---> " << fNY << " (constant) " << ConvertShapeToString(shapeY) << std::endl;
            }
            return;
         }
 
         if (!fIsDynamic)
            model.AddIntermediateTensor(fNY, model.GetTensorType(fNA), shapeY);
         else
            model.AddDynamicTensor(fNY, model.GetTensorType(fNA), fShapeY);
 
         if (model.Verbose()){
            std::cout << "Gemm (or MatMul) " << " ---> " << fNY << " shape ";
            if (fIsDynamic)
               std::cout << ConvertDimShapeToString(fShapeY) << std::endl;
            else
               std::cout << ConvertShapeToString(shapeY) << std::endl;
         }
 
         model.AddNeededStdLib("algorithm");
 
         // register the inference helper functions used by the generated code
         if (fType == "float")
            model.AddNeededHelperFunction("Gemm_Call");
         // bias handling emits Copy / Fill, fused activation emits Relu
         if (fNC != "") {
            model.AddNeededHelperFunction("Copy");
            model.AddNeededHelperFunction("Fill");
         }
         if (fActivation == EActivationType::RELU)
            model.AddNeededHelperFunction("Relu");
      }
 
      std::string Generate(std::string opName) override {
         if (fIsOutputConstant)
            return ""; // no op for constant tensors
 
         opName = "op_" + opName;
 
         // if (fShapeA.empty() || fShapeB.empty() || fShapeY.empty() || (fNC != "" && fShapeC.empty())) {
         //    throw std::runtime_error("TMVA SOFIE Gemm Op called to Generate without being initialized first");
         // }
         std::stringstream out;
         out << "\n//--------- Gemm " << opName << " " << ConvertDimShapeToString(fShapeA) << " * " << ConvertDimShapeToString(fShapeB)
             << " -> " << ConvertDimShapeToString(fShapeY) << "\n";
         // need to consider case A and B have dim > 2 (for MatMul)
         int64_t dimA = fShapeA.size();
         int64_t dimB = fShapeB.size();
         int64_t dimY = fShapeY.size();
         int64_t dimC = fDimShapeC.size();
         if (dimA != dimB || dimA != dimY || (fBroadcastBias && dimC != dimY)) {
             std::cout << " shape A " << ConvertDimShapeToString(fShapeA)
                       << " shape B " << ConvertDimShapeToString(fShapeB)
                       << " shape C " << ConvertDimShapeToString(fDimShapeC)
                       << " shape Y " << ConvertDimShapeToString(fShapeY) << std::endl;
             throw std::runtime_error("TMVA SOFIE Gemm(MatMul) has invalid shape for inputs or output");
         }
         auto m = (fAttrTransA ? fShapeA[dimA-1].GetVal() : fShapeA[dimA-2].GetVal());
         auto n = (fAttrTransB ? fShapeB[dimB-2].GetVal() : fShapeB[dimB-1].GetVal());
         auto k = (fAttrTransA ? fShapeA[dimA-2].GetVal() : fShapeA[dimA-1].GetVal());
         // size of A: if (transposeA) is m*k else k*m
         // size of B  n*k
         std::vector<Dim> sY = {fShapeY[dimY-2], fShapeY[dimY-1]};
         // extra dimensions in case of stacked MatMul
         std::vector<Dim> sExtraY;
         for (int64_t i = 0; i < dimY-2; i++) {
            sExtraY.push_back(fShapeY[i]);
         }
         auto lengthGemm = ConvertDimShapeToLength(sY); // size of the Gemm operation
         auto lengthExtra_Y = ConvertDimShapeToLength(sExtraY); // extra length in case input tensors are of dim>2 (MatMul)
         std::string lengthExtra_C;
         std::vector<Dim> sExtraC;
         std::vector<Dim> sC;
         bool haveExtraC = false;
         if (dimC > 2) {
            sC = {fDimShapeC[dimC-2], fDimShapeC[dimC-1]};
            for (int64_t i = 0; i < dimC-2; i++) {
               sExtraC.push_back(fDimShapeC[i]);
            }
            lengthExtra_C = ConvertDimShapeToLength(sExtraC);
            if (lengthExtra_C != "1") haveExtraC = true;
         } else if (dimC > 0) {
            for (int64_t i = 0; i < dimC; i++) {
               sC.push_back(fDimShapeC[i]);
            }
         }
 
         // case bias is present
         if (!fNC.empty()){
             // when the 2 last dims of bias and Y are not compatible we need to perform a run time broadcast
             if (sC != sY)
                fBroadcastBias = true;
             else if (fBiasBroadcastAssumed && sExtraC == sExtraY)
                // C has exactly the shape of Y, nothing to broadcast. Only revisit the
                // assumption Initialize had to make while the shape of Y was still unknown:
                // a bias it did compare and found to need broadcasting keeps it.
                fBroadcastBias = false;
             if (!fBroadcastBias) {
                // add a check in case broadcasting was not needed or done outside of session
                // C should have smaller dimension of Y
                if (!fIsDynamic) {
                   if ((std::stoi(lengthGemm) != std::stoi(ConvertDimShapeToLength(sC))) ||
                       (haveExtraC && std::stoi(lengthExtra_Y) != std::stoi(lengthExtra_C)))
                      throw std::runtime_error("TMVA SOFIE Gemm Op " + opName + " Bias tensor " + fNC +
                                               " has not correct size " + ConvertShapeToString(fShapeC) +
                                               " output length " + lengthGemm);
                } else {
                   // add a dynamic check (C should not be a dynamic tensor)
                   out << SP << "assert(" << lengthGemm << " == " << ConvertDimShapeToLength(sC) << ");\n";
                   if (haveExtraC)
                      out << SP << "assert(" << lengthExtra_Y << " == " << lengthExtra_C << ");\n";
                }
             }
         } else {
            fBroadcastBias = false;
            //in this case fAttrBeta needs to be equal to zero otherwise second time we run we will use
            // the previous result
            if (fAttrBeta != 0) {
               // some model don't have bias but Beta is not zero - force it to zero
               fAttrBeta = 0;
               std::cout << "WARNING: TMVA SOFIE Gemm Op " + opName + " Bias tensor is not present but beta value in Gemm is not zero - force it to zero\n";
            }
         }
 
         // include MatMul case where we stack the Gemm operations
         // exclude case where we have only 1's in the additional dims
         bool doStackMul = dimY > 2 && ( fIsDynamic  || std::stoi(lengthExtra_Y) > 1);
         // compute input offset for stack multiplications
         std::string lengthExtra_A;
         std::string lengthExtra_B;
         std::string increment_A;
         std::string increment_B;
 
         if (doStackMul) {
            std::vector<Dim> sA(fShapeA.begin(), fShapeA.begin()+dimA-2);
            std::vector<Dim> sB(fShapeB.begin(), fShapeB.begin()+dimB-2);
            std::vector<Dim> mA = {fShapeA[dimA-2], fShapeA[dimA-1]};
            std::vector<Dim> mB = {fShapeB[dimB-2], fShapeB[dimB-1]};
            lengthExtra_A = ConvertDimShapeToLength(sA);
            lengthExtra_B = ConvertDimShapeToLength(sB);
            // if A ( b, m, k) and B (b, k, n) these are the strides of A and B ( m*k for A and n*k for B )
            increment_A = ConvertDimShapeToLength(mA);
            increment_B = ConvertDimShapeToLength(mB);
         }
         bool extraA = (doStackMul && lengthExtra_A != "1");
         bool extraB = (doStackMul && lengthExtra_B != "1");
         bool extraC = (doStackMul && haveExtraC && !fBroadcastBias);
         // run time check for bias broadcasting
         std::string biasShapeType = opName + "_biasShapeType";
         if (fBroadcastBias && fCheckBiasShapeAtRuntime) {
            // create a flag according to bias shape:
            // = 1 for (1,Y2)
            // = 2 for (Y1,1)
            // = 3 for a scalar
            out << SP << "int " << biasShapeType << " = 0;\n";
            // case vector of columns
            if (sC[0].GetVal() != "1" && sC[1].GetVal() != sY[1].GetVal())
               out << SP << "if (" << sC[0] << " == 1 && " << sC[1] << " == " << sY[1] << ")\n";
            else if (sC[0].GetVal() == "1")
               out << SP << "if (" << sC[1] << " == " << sY[1] << ")\n";
            else if (sC[1].GetVal() == sY[1].GetVal())
               out << SP << "if (" << sC[0] << " == 1)\n";
 
            out << SP << SP << biasShapeType << " = 1;\n";
 
            // case vector of rows
            if (sC[1].GetVal() != "1" && sC[0].GetVal() != sY[0].GetVal())
               out << SP << "else if (" << sC[1] << " == 1 && " << sC[0] << " == " << sY[0] << ")\n";
            else if (sC[1].GetVal() == "1")
                out << SP << "else if (" << sC[0] << " == " << sY[0] << ")\n";
            else if (sC[0].GetVal() == sY[0].GetVal())
               out << SP << "else if (" << sC[1] << " == 1)\n";
 
            out << SP << SP << biasShapeType << " = 2;\n";
 
            // case scalar
            if (sC[0].GetVal() != "1" && sC[1].GetVal() != "1")
               out << SP << "else if (" << sC[0] << " == 1 && " << sC[1] << " == 1 )\n";
            else if (sC[0].GetVal() == "1")
               out << SP << "else if (" << sC[1] << " == 1)\n";
            else if (sC[1].GetVal() == "1")
               out << SP << "else if (" << sC[0] << " == 1)\n";
            out << SP << SP << biasShapeType << " = 3;\n";
            out << SP << "else\n";
            out << SP << SP << "throw std::runtime_error(\"TMVA SOFIE Gemm Op - bias tensor "
                                 << ConvertDimShapeToString(fDimShapeC) << " cannot be broadcasted to "
                                 << ConvertDimShapeToString(fShapeY) << "\");\n";
         }
         auto SP2 = SP;
         if (doStackMul) {
            out << SP << "size_t " << opName << "_y_offset = 0;\n"; // needed if we stack the gemm operations
            if (extraA)
               out << SP << "size_t " << opName << "_A_offset = 0;\n";
            if (extraB)
               out << SP << "size_t " << opName << "_B_offset = 0;\n";
            if (extraC)
               out << SP << "size_t " << opName << "_C_offset = 0;\n";
            out << SP << "for (size_t i = 0; i < " << lengthExtra_Y << "; i++){\n";
            SP2 += SP;
         }
         // do the bias broadcasting at run time by
         // initializing output Y vector with bias values
         if (fBroadcastBias) {
 
            fAttrBeta = 1.;
 
            // loop on first output dimension
            out << SP2 << "for (size_t j = 0; j < " << sY[0] << "; j++) { \n";
            out << SP2 << SP << "size_t y_index = ";
            if (doStackMul) // add offset in case of stack multiplications (not sure if bias is present in these cases)
               out <<  opName << "_y_offset + ";
            if (sY[1].GetVal() != "1")
               out << sY[1] << " * j;\n";
            else
               out << "j;\n";
 
            std::string prefix = SP2 + SP;
            std::string target = "tensor_" + fNY;
            if (sC.size() != 2) {
               throw std::runtime_error("TMVA SOFIE Gemm Op - invalid rank for bias tensor " + ConvertDimShapeToString(fDimShapeC) + ConvertDimShapeToString(sC));
            } if (sC[0].GetVal() == "1" && sC[1].GetVal() == sY[1].GetVal()) {
               out << prefix << "Copy(" << target << " + y_index, tensor_" << fNC << ", " << sY[1] << ");\n";
            } else if (sC[1].GetVal() == "1" && sC[0].GetVal() == sY[0].GetVal()) {
               out << prefix << "Fill(" << target << " + y_index, tensor_" << fNC << "[j], " << sY[1] << ");\n";
            } else if (sC[0].GetVal() == "1" && sC[1].GetVal() == "1") {
               // scalar case
               out << prefix << "Fill(" << target << " + y_index, tensor_" << fNC << "[0], " << sY[1] << ");\n";
            } else if (fCheckBiasShapeAtRuntime) {
               // in the generic dynamic case we check at run time that bias is compatible
               // we check that bias[0] = 1 or equal to SY[0] and that bias[1] = 1 or equal to SY[1]
               // tbd: this run-time check coul;d be moved outside the loop for better run time efficiency
               out << SP2 << SP << "if (" << biasShapeType << " == 1)\n";   // case vector of columns
               out << SP << prefix << "Copy(" << target << " + y_index, tensor_" << fNC << ", " << sY[1] << ");\n";
               out << SP2 << SP << "else if (" << biasShapeType << " == 2)\n";  // case vector of rows
               out << SP << prefix << "Fill(" << target << " + y_index, tensor_" << fNC << "[j], " << sY[1] << ");\n";
               out << SP2 << SP << "else \n";  // scalar case
               out << SP << prefix << "Fill(" << target << " + y_index, tensor_" << fNC << "[0], " << sY[1] << ");\n";
            } else {
               throw std::runtime_error("TMVA SOFIE Gemm Op - invalid shape for bias tensor " + ConvertDimShapeToString(fDimShapeC));
            }
 
            out << SP2 << "}\n";
         }
 
         if (fType == "float"){
 
            out << SP2 << "Gemm_Call(" << "tensor_" << fNY;
             if (doStackMul) out << " + " << opName << "_y_offset";
            out <<   ", "
             << (fAttrTransB ? "true, " : "false, ")
             << (fAttrTransA ? "true, " : "false, ")
             << n << ", " << m << ", " << k << ", ";
            out << std::setprecision(std::numeric_limits<float>::max_digits10) << fAttrAlpha << ", tensor_" << fNB;
            if (extraB) out << " + " << opName << "_B_offset";
            out << ", tensor_" << fNA;
            if (extraA) out << " + " << opName << "_A_offset";
            out << ", " << std::setprecision(std::numeric_limits<float>::max_digits10) << fAttrBeta << ",";
            // in the case of bias and no broadcasting needed - I need to add bias as an extra tensor in Gemm call
            if (!fNC.empty() && !fBroadcastBias) {
               out << "tensor_" << fNC;
               if (extraC) {
                  out << " + " << opName << "_C_offset";
               }
            } else {
               out << "nullptr";
            }
            out << ");\n";
 
         }
 
         if (doStackMul) {
            out << SP << SP <<  opName << "_y_offset += " << lengthGemm << ";\n";
            if (lengthExtra_A != "1")
               out << SP << SP << opName << "_A_offset += " << increment_A << ";\n";
            if (lengthExtra_B != "1")
               out << SP << SP << opName << "_B_offset += " << increment_B << ";\n";
            if (extraC)
               // increment_C is lengthGEmm
               out << SP << SP << opName << "_C_offset += " << lengthGemm << ";\n";
            out << SP << "}\n"; // end of loop on the stacked multiplication
         }
 
         // fuse with Relu
         if(fActivation == EActivationType::RELU){
               out << SP << "//--- applying RELU to output\n";
               std::string tnsr = "tensor_" + fNY;
               std::string reluSize = ConvertDimShapeToLength(fShapeY);
               out << SP << "Relu(" << tnsr << ", " << tnsr << ", " << reluSize << ");\n";
         }
 
         return out.str();
      }
 
      std::vector<std::string> GetBlasRoutines() override { return {"Gemm", "Gemv"}; }
 
   };
 
 
}//SOFIE
}//Experimental
}//TMVA
 
 
#endif //TMVA_SOFIE_ROPERATOR_GEMM