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;
 
      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 };
      }
 
      std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
         ETensorType out = input[0];
         return {out};
      }
 
      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<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
         std::vector<std::vector<size_t>> ret;
         ret.push_back(DoShapeInference<size_t>(input));
         return ret;
      }
      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;
            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 be equal to 1 since this is a unidirectional broadcast of C->Y
               if (r.first > 1) {
                  throw std::runtime_error("TMVA SOFIE Gemm Op - bias tensor of shape " + ConvertDimShapeToString(fDimShapeC) + " cannot be uni-directional broadcasted to " + ConvertDimShapeToString(fShapeY));
               }
               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);
         }
 
         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");
      }
 
      std::string Generate(std::string opName) override {
         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;
            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) {
               throw std::runtime_error("TMVA SOFIE Gemm Op " + opName + " Bias tensor is not present but beta value in Gemm is not zero");
            }
         }
 
         // 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);
         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.;
            out << SP2 << "for (size_t j = 0; j < " << sY[0] << "; j++) { \n";
            out << SP2 << SP << "size_t y_index = ";
            if (doStackMul) // add offset in caseof 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 + "TMVA::Experimental::SOFIE::";
            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 {
               throw std::runtime_error("TMVA SOFIE Gemm Op - invalid shape for bias tensor " + ConvertDimShapeToString(fDimShapeC));
            }
 
            out << SP2 << "}\n";
         }
 
         if (fType == "float"){
 
            out << SP2 << "TMVA::Experimental::SOFIE::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 << "TMVA::Experimental::SOFIE::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