ROperator_TopK.hxx

Go to the documentation of this file.

#ifndef TMVA_SOFIE_ROPERATOR_TOPK
#define TMVA_SOFIE_ROPERATOR_TOPK
 
#include "TMVA/SOFIE_common.hxx"
#include "TMVA/ROperator.hxx"
#include "TMVA/RModel.hxx"
 
#include <sstream>
 
namespace TMVA {
namespace Experimental {
namespace SOFIE {
 
template <typename T>
class ROperator_TopK final : public ROperator {
 
private:
   int fAttrAxis;
   int fAttrLargest;
   int fAttrSorted;
 
   Dim fK;
   std::string fNK;
   std::string fNX;
   std::string fNVal;
   std::string fNInd;
   std::vector<Dim> fShapeX;
   std::vector<Dim> fShapeY;
   std::string fType;
 
public:
   ROperator_TopK() {}
   ROperator_TopK(int attr_axis, int attr_largest, int attr_sorted, std::string nameK, std::string nameX, std::string nameVal, std::string nameInd)
      : fAttrAxis(attr_axis),
        fAttrLargest(attr_largest),
        fAttrSorted(attr_sorted),
        fNK(UTILITY::Clean_name(nameK)),
        fNX(UTILITY::Clean_name(nameX)),
        fNVal(UTILITY::Clean_name(nameVal)),
        fNInd(UTILITY::Clean_name(nameInd)){
            fInputTensorNames = { fNX, fNK };
            fOutputTensorNames = { fNVal, fNInd };
        }
 
   std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
      ETensorType ret = input[0];
      return {ret, ret};
   }
 
   void Initialize(RModel& model) override {
      if (model.CheckIfTensorAlreadyExist(fNX) == false) {
         // input must be a graph input, or already initialized intermediate tensor
         throw std::runtime_error("TMVA SOFIE TopK Op Input Tensor is not found in model");
      }
      if (model.CheckIfTensorAlreadyExist(fNK) == false) {
         // input must be a graph input, or already initialized intermediate tensor
         throw std::runtime_error("TMVA SOFIE TopK Op Input Tensor i.e. K is not found in model");
      }
 
      fShapeX = model.GetDimTensorShape(fNX);
      // K can either be an initialized tensor or a shape tensor, in which case its value is
      // known only symbolically (e.g. it depends on one of the input dimensions)
      Dim kdim;
      if (model.IsShapeTensor(fNK)) {
         auto &kvalues = model.GetShapeTensorValues(fNK);
         if (kvalues.size() != 1)
            throw std::runtime_error("TMVA SOFIE TopK Op input tensor K = " + fNK + " must be a single value");
         kdim = kvalues[0];
      } else if (model.IsInitializedTensor(fNK)) {
         auto kptr = static_cast<int64_t *>(model.GetInitializedTensorData(fNK).get());
         kdim = Dim{static_cast<size_t>(*kptr)};
         model.SetNotWritableInitializedTensor(fNK);
      } else {
         throw std::runtime_error("TMVA SOFIE TopK Op input tensor K = " + fNK +
                                  " must be known at initialization time");
      }
      fAttrAxis = fAttrAxis < 0 ? fShapeX.size() + fAttrAxis : fAttrAxis;
      if(static_cast<size_t>(fAttrAxis) >=  fShapeX.size()){
         throw
            std::runtime_error("TMVA::SOFIE ONNX TopK op axis = "+ std::to_string(fAttrAxis) +" value exeeds size of tensor " +fNX+" of size "+fShapeX.size()+" .");
      }
      // fK cannot be larger that axis dimension
      if (kdim.isParam || fShapeX[fAttrAxis].isParam)
         fK = Dim{std::string("std::min(size_t(" + kdim.GetVal() + "), size_t(" + fShapeX[fAttrAxis].GetVal() + "))"),
                  static_cast<size_t>(-1)};
      else
         fK = Dim{std::min(kdim.dim, fShapeX[fAttrAxis].dim)};
 
      // output shape is equal to input shape apart for value in fAttrAxis
      fShapeY = fShapeX;
      fShapeY[fAttrAxis] = Dim{fK};
 
      model.AddIntermediateTensor(fNVal, model.GetTensorType(fNX), fShapeY);
 
      // output indices should be an int64 tensor
      model.AddIntermediateTensor(fNInd, ETensorType::INT64, fShapeY);
      fType = ConvertTypeToString(model.GetTensorType(fNX));
      model.AddNeededStdLib("algorithm");
      model.AddNeededStdLib("cstdint");
      model.AddNeededStdLib("cstring");
 
      if (model.Verbose()) {
         std::cout << "TopK " << fNX << "  " << ConvertDimShapeToString(fShapeX)
                      << "---> " << fNVal << " " <<  ConvertDimShapeToString(fShapeY) << std::endl;
      }
   }
 
   std::string Generate(std::string OpName) override {
      OpName = "op_" + OpName;
      if (fShapeX.empty()) {
         throw std::runtime_error("TMVA SOFIE Operator TopK called to Generate without being initialized first");
      }
      std::stringstream out;
      size_t size = fShapeX.size();
      size_t axis = fAttrAxis < 0 ? size + fAttrAxis : fAttrAxis;
      out << "\n" << SP << "//------ TopK\n";
 
      auto length=ConvertDimShapeToLength(fShapeX);
      auto strideX = UTILITY::ComputeStrideFromShape(fShapeX);
      auto strideY = UTILITY::ComputeStrideFromShape(fShapeY);
      // we perform loop on dimension before sorted axis and after sorted axis
      std::vector<Dim> shape_before(fShapeX.begin(), fShapeX.begin() + axis);   // input shape before axis
      std::string n_before = (axis>0) ? ConvertDimShapeToLength(shape_before) : "1";
      std::string n_after = strideX[axis].GetVal();
      std::string n_elements = fShapeX[axis].GetVal(); // number of elements to be sorted
 
      // }
      out << SP << "{\n"; // to define a separate scope for the operator code
 
      // Ties are broken by the element index, so no two entries ever compare equivalent:
      // the ordering is total and the selected set is therefore unique. That is what makes
      // the (unstable) std::nth_element below safe - it cannot pick a different set from a
      // full sort.
      //
      // For float that ordering can be expressed as a single unsigned integer. Flipping the
      // sign bit on positives and every bit on negatives maps a (non-NaN) float onto a
      // uint32 whose unsigned order matches the float order; putting the element index in
      // the low 32 bits then reproduces "ties by smaller index" exactly. A comparison
      // becomes one 64-bit instruction instead of a two-field comparator call, and an
      // element is 8 bytes instead of 16, which halves what nth_element has to move.
      // Wider types cannot pack a value and an index into 64 bits, so they keep the pairs.
      bool packed = (fType == "float");
      // the index has to fit in the low 32 bits
      if (packed && !fShapeX[fAttrAxis].isParam && fShapeX[fAttrAxis].dim > 0xFFFFFFFFULL)
         packed = false;
 
      std::string pairType = "std::pair<" + fType + ",int64_t>";
      if (packed) {
         out << SP << "std::vector<uint64_t> elements(" << n_elements << ");\n";
         if (fShapeX[fAttrAxis].isParam) {
            out << SP << "if (static_cast<unsigned long long>(" << n_elements << ") > 0xFFFFFFFFULL)\n";
            out << SP << SP << "throw std::runtime_error(\"TMVA SOFIE TopK - reduced axis is longer "
                << "than the 2^32 limit of the packed index\");\n";
         }
      } else {
         out << SP << "std::vector<" << pairType << "> elements(" << n_elements << ");\n";
         // taking the pairs by const reference avoids copying them on every comparison
         out << SP << "auto " << OpName << "_cmp = [](const " << pairType << " &a, const " << pairType << " &b) {\n";
         out << SP << SP << "return (a.first != b.first) ? (a.first " << (fAttrLargest ? ">" : "<")
             << " b.first) : a.second < b.second;\n";
         out << SP << "};\n";
      }
      // loop on elements before
      if (n_before != "1") {
         out << SP << "for (size_t i = 0; i < " << n_before << "; i++) {\n";
         out << SP << SP << "size_t xoffset = i*" << strideX[axis-1] << ";\n";
         out << SP << SP << "size_t yoffset = i*" << strideY[axis-1] << ";\n";
         out << SP;
      } else {
         out << SP << "size_t xoffset = 0;\n";
         out << SP << "size_t yoffset = 0;\n";
      }
      if (n_after !=  "1")
         out << SP << "for (size_t j = 0; j < " << n_after << "; j++) {\n";
      else
         out << SP << "const size_t j = 0;\n";
 
      // copy the elements to be sorted into the working buffer
      out << SP << SP << "for (size_t l = 0; l < " << n_elements << "; l++) {\n";
      if (packed) {
         out << SP << SP << SP << "uint32_t b_ = 0;\n";
         out << SP << SP << SP << "std::memcpy(&b_, &tensor_" << fNX << "[xoffset + " << strideX[axis]
             << "*l + j], sizeof(b_));\n";
         out << SP << SP << SP << "b_ ^= (b_ & 0x80000000u) ? 0xFFFFFFFFu : 0x80000000u;\n";
         if (fAttrLargest)
            out << SP << SP << SP << "b_ = ~b_;\n"; // reverse the value order, keep index ascending
         out << SP << SP << SP << "elements[l] = (static_cast<uint64_t>(b_) << 32) | static_cast<uint32_t>(l);\n";
      } else {
         out << SP << SP << SP << "elements[l] = std::make_pair(tensor_" << fNX << "[xoffset + " << strideX[axis]
             << "*l + j], l);\n";
      }
      out << SP << SP << "}\n";
 
      // Move the K selected elements to the front in linear time, then order just those.
      // std::partial_sort would be O(n log K) with heap operations over the whole range.
      std::string cmp = packed ? "" : (", " + OpName + "_cmp");
      out << SP << SP << "std::nth_element(elements.begin(), elements.begin() + (" << fK << "), elements.end()" << cmp
          << ");\n";
      // The ONNX spec leaves the order unspecified when sorted=0, but we sort anyway: it is
      // only O(K log K) and it keeps the generated code reproducible across standard libraries.
      out << SP << SP << "std::sort(elements.begin(), elements.begin() + (" << fK << ")" << cmp << ");\n";
 
      // copy the selected elements in the output
      out << SP << SP << "for (size_t l = 0; l < " << fK << "; l++) {\n";
      if (packed) {
         out << SP << SP << SP << "uint32_t b_ = static_cast<uint32_t>(elements[l] >> 32);\n";
         if (fAttrLargest)
            out << SP << SP << SP << "b_ = ~b_;\n";
         out << SP << SP << SP << "b_ ^= (b_ & 0x80000000u) ? 0x80000000u : 0xFFFFFFFFu;\n";
         out << SP << SP << SP << fType << " v_;\n";
         out << SP << SP << SP << "std::memcpy(&v_, &b_, sizeof(v_));\n";
         out << SP << SP << SP << "tensor_" << fNVal << "[yoffset + " << strideY[axis] << "*l + j] = v_;\n";
         out << SP << SP << SP << "tensor_" << fNInd << "[yoffset + " << strideY[axis]
             << "*l + j] = static_cast<int64_t>(static_cast<uint32_t>(elements[l]));\n";
      } else {
         out << SP << SP << SP << "tensor_" << fNVal << "[yoffset + " << strideY[axis]
             << "*l + j] = elements[l].first;\n";
         out << SP << SP << SP << "tensor_" << fNInd << "[yoffset + " << strideY[axis]
             << "*l + j] = elements[l].second;\n";
      }
      out << SP << SP << "}\n";
      if (n_after != "1") out << SP << SP << "}\n";
      if (n_before != "1") out << SP << "}\n";
      out << SP << "}\n"; // end operator scope
      return out.str();
   }
};
 
} // namespace SOFIE
} // namespace Experimental
} // namespace TMVA
 
#endif // TMVA_SOFIE_ROPERATOR_TOPK