root/TMVA/ROperator_ScatterND.hxx

0001 #ifndef TMVA_SOFIE_ROPERATOR_ScatterND
0002 #define TMVA_SOFIE_ROPERATOR_ScatterND
0003
0004 #include "TMVA/SOFIE_common.hxx"
0005 #include "TMVA/ROperator.hxx"
0006 #include "TMVA/RModel.hxx"
0007
0008 #include <sstream>
0009 #include <stdexcept>
0010 #include <string>
0011
0012 namespace TMVA{
0013 namespace Experimental{
0014 namespace SOFIE{
0015
0016 class ROperator_ScatterND final : public ROperator
0017 {
0018 private:
0019
0020
0021    std::string fNX;
0022    std::string fNI;
0023    std::string fNU;
0024    std::string fNY;
0025    std::string fReduction;
0026
0027    std::vector<Dim> fShapeX;
0028    std::vector<Dim> fShapeI;
0029    std::vector<Dim> fShapeY;
0030
0031
0032    std::vector<int64_t> fIndices;  // indices vector in case they are known at initialization
0033
0034    std::string fType;
0035
0036
0037 public:
0038    ROperator_ScatterND(){}
0039    ROperator_ScatterND(const std::string & nameX, const std::string & nameI, const std::string & nameU, const std::string & nameY,
0040                         std::string reduction):
0041       fNX(UTILITY::Clean_name(nameX)), fNI(UTILITY::Clean_name(nameI)), fNU(UTILITY::Clean_name(nameU)),
0042       fNY(UTILITY::Clean_name(nameY)), fReduction(reduction)
0043    {
0044       fInputTensorNames = { fNX, fNI, fNU };
0045       fOutputTensorNames = { fNY };
0046    }
0047
0048    void Initialize(RModel& model) override {
0049
0050        // input must be a graph input, or already initialized intermediate tensor
0051       if (!model.CheckIfTensorAlreadyExist(fNX)){
0052          throw std::runtime_error(std::string("TMVA SOFIE ScatterND Op Input Tensor ") + fNX + "is not found in model");
0053       }
0054       if (!model.CheckIfTensorAlreadyExist(fNI)) {
0055          throw std::runtime_error(std::string("TMVA SOFIE ScatterND Op Input Tensor ") + fNI + "is not found in model");
0056       }
0057       if (!model.CheckIfTensorAlreadyExist(fNU)) {
0058          throw std::runtime_error(std::string("TMVA SOFIE ScatterND Op Input Tensor ") + fNU + "is not found in model");
0059       }
0060       //tbd check for constant tensors
0061
0062       fShapeX = model.GetDimTensorShape(fNX);
0063       fShapeI = model.GetDimTensorShape(fNI);
0064       auto shapeU = model.GetDimTensorShape(fNU);
0065
0066       //  Validate inputs if fShapeI last is not dynamic
0067
0068       //if (!model.IsDynamicTensor(fNI)) {
0069       const size_t r = fShapeX.size();       // rank of data
0070       const size_t q = fShapeI.size();    // rank of indices
0071       if (!(fShapeI.back().isParam) ) {
0072          const size_t k = fShapeI.back().dim;             // index depth
0073
0074          if (k > r)
0075             throw std::invalid_argument(
0076                "ScatterND: last dim of indices (" + std::to_string(k) +
0077                ") must be <= rank of data (" + std::to_string(r) + ")");
0078
0079          // Expected updates rank = q - 1 + r - k
0080          int64_t expected_updates_rank = q - 1 + r - k;
0081          if ((int64_t) shapeU.size() != expected_updates_rank)
0082             throw std::invalid_argument("ScatterND: updates rank mismatch");
0083       } else {
0084          //  Assumption is that last dimension of index shape is known (is not dynamic)
0085          throw std::runtime_error("TMVA SOFIE ScatterND : Index_shape(-1) is not known. This case is not supported");
0086       }
0087
0088       // output shape is equal to input shape
0089       fShapeY = fShapeX;
0090
0091       model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShapeY);
0092       if (model.Verbose()) {
0093          std::cout << "ScatterElements: input: " << ConvertDimShapeToString(fShapeX)
0094                                                 << " indices " << ConvertDimShapeToString(fShapeI)
0095                                                 << " update " <<  ConvertDimShapeToString(shapeU);
0096          std::cout << "\t----> " << ConvertDimShapeToString(fShapeY) << std::endl;
0097       }
0098    }
0099
0100    std::string Generate(std::string opName) override {
0101       if (fIsOutputConstant) {
0102          // no code to generate here for constant output. Tensor output is defined in Session constructor
0103          return "//---------------------------------------\n";
0104       }
0105       opName = "op_" + opName;
0106       std::stringstream out;
0107       out << "//--------- ScatterND " << opName << " --> " << ConvertDimShapeToString(fShapeY) << "\n";
0108
0109       size_t r = fShapeX.size();
0110
0111       // Strides
0112       auto stridesX = UTILITY::ComputeStrideFromShape(fShapeX);
0113       auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY);
0114       auto stridesI = UTILITY::ComputeStrideFromShape(fShapeI);
0115
0116       // case input_index_shape == rank of input
0117       size_t k = fShapeI.back().dim;
0118
0119       // Total number of index tuples = product of indices dims except last
0120       std::vector<Dim> shapeIndFirst(fShapeI.begin(), fShapeI.begin()+ fShapeI.size()-1);
0121       auto num_index_tuples = ConvertDimShapeToLength(shapeIndFirst);
0122
0123       //slice size (is product of input from k to r)
0124       std::vector<Dim> shapeSlice(fShapeX.begin()+k, fShapeX.end());
0125       auto slice_size = ConvertDimShapeToLength(shapeSlice);
0126
0127       auto data_length = ConvertDimShapeToLength(fShapeX);
0128
0129       //step1: input->output
0130       out << SP << "// Step 1: copy input data to output\n";
0131       out << SP << "std::copy(tensor_" << fNX << ", tensor_" << fNX << " + " << data_length << ", tensor_" << fNY << ");\n";
0132
0133       // Step 2: Emit strides as a static constexpr array
0134       out << SP << "// Step 2: data strides (row-major)\n";
0135       //to do: use static constexpr for defined strides
0136       out << SP << "size_t " << opName << "_data_strides[" << r << "] = {";
0137       for (size_t i = 0; i < r; ++i)
0138          out << stridesX[i] << (i + 1 < r ? ", " : "");
0139       out << "};\n\n";
0140
0141       // Step 3: Scatter loop
0142       out << SP << "// Step 3: scatter updates into output\n";
0143       out << SP << "for (int64_t idx = 0; idx < " << num_index_tuples << "; idx++) {\n";
0144
0145       // Resolve flat data offset from k-dimensional index tuple
0146       out << SP << SP << "int64_t data_offset = 0;\n";
0147       for (size_t dim = 0; dim < k; ++dim) {
0148          out << SP << SP << "{\n";
0149          out << SP << SP << SP << "int64_t coord = tensor_" << fNI
0150              << "[idx * " << k << " + " << dim << "];\n";
0151          // Support negative indices
0152          out << SP << SP << SP << "if (coord < 0) coord += " << fShapeX[dim] << ";\n";
0153          out << SP << SP << SP << "data_offset += coord * "
0154                << opName << "_data_strides[" << dim << "];\n";
0155          out << SP << SP << "}\n";
0156       }
0157
0158       // Apply updates with reduction
0159       out << SP << SP << "for (int64_t s = 0; s < " << slice_size << "; s++) {\n";
0160       out << SP << SP << SP << "auto upd = tensor_" << fNU
0161          << "[idx * " << slice_size << " + s];\n";
0162
0163       if (fReduction.empty() || fReduction == "none") {
0164          out << SP << SP << SP << "tensor_" << fNY << "[data_offset + s] = upd;\n";
0165       } else if (fReduction == "add") {
0166          out << SP << SP << SP << "tensor_" << fNY<< "[data_offset + s] += upd;\n";
0167       } else if (fReduction == "mul") {
0168          out << SP << SP << SP << "tensor_" << fNY << "[data_offset + s] *= upd;\n";
0169       } else if (fReduction == "min") {
0170          out << SP << SP << SP << "tensor_" << fNY<< "[data_offset + s] = "
0171                << "std::min(tensor_" << fNY << "[data_offset + s], upd);\n";
0172       } else if (fReduction == "max") {
0173          out << SP << SP << SP << "tensor_" << fNY << "[data_offset + s] = "
0174             << "std::max(tensor_" << fNY << "[data_offset + s], upd);\n";
0175       } else {
0176          throw std::runtime_error(
0177             "TMVA SOFIE ScatterND: unsupported reduction '" + fReduction + "'");
0178       }
0179
0180       out << SP << SP << "}\n";  // end slice loop
0181       out << SP << "}\n";        // end index tuple loop
0182
0183       return out.str();
0184    }
0185
0186 };
0187
0188 }//SOFIE
0189 }//Experimental
0190 }//TMVA
0191
0192
0193 #endif //TMVA_SOFIE_ROPERATOR_RELU