root/TMVA/ROperator_NonZero.hxx

0001 #ifndef TMVA_SOFIE_ROPERATOR_NONZERO
0002 #define TMVA_SOFIE_ROPERATOR_NONZERO
0003
0004 #include "TMVA/SOFIE_common.hxx"
0005 #include "TMVA/ROperator.hxx"
0006 #include "TMVA/RModel.hxx"
0007
0008 #include <sstream>
0009
0010 namespace TMVA{
0011 namespace Experimental{
0012 namespace SOFIE{
0013
0014 template<class T>
0015 class ROperator_NonZero final : public ROperator
0016 {
0017
0018 private:
0019
0020    std::string fNX;
0021    std::string fNY;
0022    std::vector<Dim> fShapeX;
0023    std::vector<Dim> fShapeY;
0024
0025 public:
0026    ROperator_NonZero(){}
0027    ROperator_NonZero(std::string nameX, std::string nameY):
0028       fNX(UTILITY::Clean_name(nameX)), fNY(UTILITY::Clean_name(nameY)){
0029          fInputTensorNames = { fNX };
0030          fOutputTensorNames = { fNY };
0031       }
0032
0033
0034
0035    void Initialize(RModel& model) override {
0036       if (model.CheckIfTensorAlreadyExist(fNX) == false){   //input must be a graph input, or already initialized intermediate tensor
0037          throw std::runtime_error("TMVA SOFIE NonZero Op Input Tensor " + fNX + " is not found in model");
0038       }
0039
0040
0041       // case input is constant
0042       if (model.IsConstantTensor(fNX)) {
0043          // compute output directly
0044          T * data = static_cast<T*>(model.GetInitializedTensorData(fNX).get());
0045          // shape is fully known
0046          auto shapeX = model.GetTensorShape(fNX);
0047          std::vector<size_t> shapeY(2);
0048          shapeY[0] = shapeX.size();
0049          auto length = ConvertShapeToLength(shapeX);
0050          auto strides = UTILITY::ComputeStrideFromShape(shapeX);
0051          std::vector<std::vector<int64_t>> nonzero_indices;
0052          for (size_t i = 0; i < length; i++) {
0053             if (data[i] != 0) {
0054                // get indices
0055                size_t flat_index = i;
0056                std::vector<int64_t> indices(shapeX.size());
0057                for (size_t j = 0; j < shapeX.size(); ++j) {
0058                   indices[j] = flat_index / strides[j];
0059                   flat_index %= strides[j];
0060                }
0061                nonzero_indices.emplace_back(indices);
0062             }
0063          }
0064          shapeY[1] = nonzero_indices.size();
0065          std::vector<int64_t> dataY(shapeY[0]* shapeY[1]);
0066          size_t k = 0;
0067          for (size_t i = 0; i < shapeY[0]; i++) {
0068             for (size_t j = 0; j < shapeY[1]; j++) {
0069                dataY[k] = nonzero_indices[j][i];
0070                k++;
0071             }
0072          }
0073          if (dataY.empty()) {
0074             // no zero elements found
0075             dataY.resize(1);
0076             shapeY.clear();  // use an empty shape
0077          }
0078
0079          model.AddConstantTensor(fNY, shapeY, dataY);
0080          if (model.Verbose()) {
0081             std::cout << "NonZero : " << fNX << " -> " << fNY << " " << ConvertShapeToString(shapeY)
0082                      << " : " << ConvertValuesToString(dataY) << std::endl;
0083          }
0084          fIsOutputConstant = true;
0085
0086       } else {
0087
0088          fShapeX = model.GetDimTensorShape(fNX);
0089
0090          // output shape(-1) depends on number of elements of non zero values
0091          // first dim is rank of input
0092          fShapeY.resize(2);
0093          fShapeY[0] = fShapeX.size();
0094
0095          // identify as -1 since we will declare maximum as size of input
0096          fShapeY[1] = Dim{std::string("v_NonZero_") + fNX, static_cast<size_t>(-1)};
0097
0098          model.AddIntermediateTensor(fNY, ETensorType::INT64, fShapeY);
0099          if (model.Verbose()) {
0100             std::cout << "NonZero : " << fNX << " -> " << fNY << " " << ConvertDimShapeToString(fShapeY) << std::endl;
0101          }
0102       }
0103    }
0104
0105    std::string GenerateSessionMembersCode(std::string /*opName*/) override {
0106       if (fIsOutputConstant) return "";
0107       // define output value used as max non zero with max size = input shape * N
0108       auto inputLength = ConvertDimShapeToLength(fShapeX);
0109       std::stringstream out;
0110       out << SP << "size_t fV_NonZero_" << fNX << " = " << inputLength << ";\n";
0111       return out.str();
0112    }
0113
0114    std::string Generate(std::string opName) override {
0115       if (fIsOutputConstant) {
0116          return "";
0117       }
0118       opName = "op_" + opName;
0119       if (fShapeX.empty()) {
0120          throw std::runtime_error("TMVA SOFIE Operator NonZero called to Generate without being initialized first");
0121       }
0122       std::stringstream out;
0123       auto intShapeX = ConvertShapeToInt(fShapeX);
0124       size_t inputLength = 0;
0125       std::string s_inputLength = ConvertDimShapeToLength(fShapeX);
0126       if (!intShapeX.empty())
0127          inputLength = ConvertShapeToLength(intShapeX);
0128
0129       size_t dims = fShapeX.size();
0130       out << "\n//------ NonZero\n";
0131
0132       std::string vnonzero = "v_NonZero_" + fNX;
0133
0134       // loop on input indices
0135       out << SP << "size_t offset_" << opName << " = 0;\n";
0136       out << SP << "size_t " << vnonzero << " = 0;\n";
0137       for (size_t j = 0; j < dims; j++) {
0138          std::string index = "i_" + std::to_string(j);
0139          for (size_t k = 0; k <= j; k++) out << SP;
0140          out << "for (size_t " << index << " = 0; " << index << " < " << fShapeX[j] << "; " << index << "++) {\n";
0141       }
0142       for (size_t k = 0; k <= dims; k++) out << SP;
0143       out << "if (tensor_" << fNX << "[offset_" << opName << "++]) {\n";
0144       for (size_t j = 0; j < dims; j++) {
0145          for (size_t k = 0; k <= dims+1; k++) out << SP;
0146          out << "tensor_" << fNY << "[";
0147          if (j > 0) {
0148             if (inputLength > 0) {
0149                out << inputLength * j;
0150             } else {
0151                out << s_inputLength;
0152                if (j > 1) out << " * " << j;
0153             }
0154             out << " + ";
0155          }
0156          out << vnonzero << "] = i_" << j << ";\n";
0157       }
0158       for (size_t k = 0; k <= dims+1; k++) out << SP;
0159       out << vnonzero << "++;\n";
0160       for (size_t k = 0; k <= dims; k++) out << SP;
0161       out << "}\n";
0162       //end loops
0163       for (size_t j = dims; j > 0; j--) {
0164          for (size_t k = 0; k <j; k++) out << SP;
0165          out << "}\n";
0166       }
0167       // now we need to rearrange the vector if nonzero is less than length of input
0168       out << SP << "if (" << vnonzero << " < " << s_inputLength << "){\n";
0169       for (size_t j = 1; j < dims; j++) {
0170          out << SP << SP << "std::copy(tensor_" << fNY;
0171          if (j>0) out << " + " << s_inputLength;
0172          if (j>1) out << " * " << j;
0173          out << ", tensor_" << fNY;
0174          if (j>0) out << " + " << s_inputLength;
0175          if (j>1) out << " * " << j;
0176          out << " + " << vnonzero << ", tensor_" <<  fNY;
0177          if (j>0) out << " + " << vnonzero;
0178          if (j>1) out << "* " << j;
0179          out << ");\n";
0180       }
0181       out << SP << "}\n";
0182
0183       return out.str();
0184    }
0185
0186 };
0187
0188 }//SOFIE
0189 }//Experimental
0190 }//TMVA
0191
0192
0193 #endif //TMVA_SOFIE_ROPERATOR_NonZero