Logo ROOT  
Reference Guide
 
Loading...
Searching...
No Matches
ROperator_TopK.hxx
Go to the documentation of this file.
1#ifndef TMVA_SOFIE_ROPERATOR_TOPK
2#define TMVA_SOFIE_ROPERATOR_TOPK
3
5#include "TMVA/ROperator.hxx"
6#include "TMVA/RModel.hxx"
7
8#include <sstream>
9
10namespace TMVA {
11namespace Experimental {
12namespace SOFIE {
13
14template <typename T>
16
17private:
21
22 size_t fK;
23 std::string fNK;
24 std::string fNX;
25 std::string fNVal;
26 std::string fNInd;
27 std::vector<size_t> fShapeX;
28 std::vector<size_t> fShapeY;
29 std::string fType;
30
31public:
33 ROperator_TopK(int attr_axis, int attr_largest, int attr_sorted, std::string nameK, std::string nameX, std::string nameVal, std::string nameInd)
37 fNK(UTILITY::Clean_name(nameK)),
38 fNX(UTILITY::Clean_name(nameX)),
39 fNVal(UTILITY::Clean_name(nameVal)),
40 fNInd(UTILITY::Clean_name(nameInd)){
43 }
44
45 std::vector<ETensorType> TypeInference(std::vector<ETensorType> input) override {
47 return {ret, ret};
48 }
49
50 std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {
51 if (input.size() != 2) {
52 throw std::runtime_error("TMVA SOFIE TopK Op Shape Inference needs exactly 2 input tensors");
53 }
54
55 auto shape = input[0]; // Shape format: [ m x n x o x p ... ]
56
57 // set the dimension at the specified axis to k (fAttrAxis is checked before that is in the correct range
58 shape[fAttrAxis] = fK; // Modified shape: [ m x n x k x p ... ]
59 return {shape, shape};
60 }
61
62
63 void Initialize(RModel& model) override {
64 if (model.CheckIfTensorAlreadyExist(fNX) == false) {
65 // input must be a graph input, or already initialized intermediate tensor
66 throw std::runtime_error("TMVA SOFIE TopK Op Input Tensor is not found in model");
67 }
68 if (model.CheckIfTensorAlreadyExist(fNK) == false) {
69 // input must be a graph input, or already initialized intermediate tensor
70 throw std::runtime_error("TMVA SOFIE TopK Op Input Tensor i.e. K is not found in model");
71 }
72
73 fShapeX = model.GetTensorShape(fNX);
74 auto fShapeK = model.GetTensorShape(fNK);
75 auto kptr = static_cast<int64_t *>(model.GetInitializedTensorData(fNK).get());
76 fK = *kptr;
79 if(static_cast<size_t>(fAttrAxis) >= fShapeX.size()){
80 throw
81 std::runtime_error("TMVA::SOFIE ONNX TopK op axis = "+ std::to_string(fAttrAxis) +" value exeeds size of tensor " +fNX+" of size "+fShapeX.size()+" .");
82 }
83 // fK cannot be larger that axis dimension
84 fK = std::min(fK, fShapeX[fAttrAxis]);
85 // if(fK>fShapeX[fAttrAxis]){
86 // throw
87 // std::runtime_error("TMVA::SOFIE ONNX TopK op k = "+ std::to_string(fK) +" value exeeds value of tensor " +fNX+" of size "+fShapeX.size()+" at axis= "+std::to_string(fAttrAxis)+".");
88 // }
89 // fShapeX = model.GetTensorShape(fNX); // [ m x n x o x p ... ]
90 // if(k[0]>=fShapeX.size()){
91 // throw
92 // std::runtime_error("TMVA::SOFIE ONNX TopK op k = "+ std::to_string(k[0]) +"value exeeds size of tensor " +fNX+" of size "+fShapeX.size()+" .");
93 // }
94 // fShapeY.push_back(2);
95 // for (auto i : fShapeX)
96 // fShapeY.push_back(i); // [ 2 x m x n x o x p ... ]
97 // size_t axis = fAttrAxis < 0 ? fShapeX.size() + fAttrAxis : fAttrAxis;
98 // fShapeY[axis] = k[0]; // [ 2 x m x n x K x p ... ]
100
101 // for(int i=0;i<fShapeX.size();i++)
102 // std::cout<<fShapeX[i]<<" ";
103 // std::cout<<"\ny size -> "<<fShapeY.size()<<std::endl;
104
105
107 // output indices should be an int64 tensor
110 }
111
112 std::string Generate(std::string OpName) override {
113 OpName = "op_" + OpName;
114 if (fShapeX.empty()) {
115 throw std::runtime_error("TMVA SOFIE Operator TopK called to Generate without being initialized first");
116 }
117 std::stringstream out;
118 size_t size = fShapeX.size();
119 size_t axis = fAttrAxis < 0 ? size + fAttrAxis : fAttrAxis;
120 out << "\n" << SP << "//------ TopK\n";
121
125 // we perform loop on dimension before sorted axis and after sorted axis
126 size_t n_before = (axis>0) ? length/strideX[axis-1] : 1;
127 size_t n_after = strideX[axis];
128 size_t n_elements = fShapeX[axis]; // number of elements to be sorted
129
130 // }
131 out << SP << "{\n"; // to define a separate scope for the operator code
132 out << SP << "std::vector<std::pair<float,int64_t>> elements(" << n_elements << ");\n";
133 // loop on elements before
134 if (n_before > 1) {
135 out << SP << "for (size_t i = 0; i < " << n_before << "; i++) {\n";
136 out << SP << SP << "size_t xoffset = i*" << strideX[axis-1] << ";\n";
137 out << SP << SP << "size_t yoffset = i*" << strideY[axis-1] << ";\n";
138 out << SP;
139 } else {
140 out << SP << "size_t xoffset = 0;\n";
141 out << SP << "size_t yoffset = 0;\n";
142 }
143 if (n_after > 1)
144 out << SP << "for (size_t j = 0; j < " << n_after << "; j++) {\n";
145 else
146 out << SP << "const size_t j = 0;\n";
147
148 // copy elements to be sorted in vector of pair
149 out << SP << SP << "for (size_t l = 0; l < " << n_elements << "; l++) {\n";
150 out << SP << SP << SP << "elements[l] = std::make_pair(tensor_" << fNX << "[xoffset + " << strideX[axis] << "*l + j], l);\n";
151 out << SP << SP << "}\n";
152
153 if (fAttrSorted) {
154 if (fAttrLargest) {
155 out<<SP<<SP << "std::partial_sort(elements.begin(),elements.begin()+" << fK << ",elements.end()," <<
156 "[](std::pair<float,int64_t>a,std::pair<float,int64_t>b){return (a.first!=b.first) ? (a.first>b.first) : a.second < b.second;});\n";
157
158 } else
159 out<<SP<<SP << "std::partial_sort(elements.begin(),elements.begin()+" << fK << ",elements.end()," <<
160 "[](std::pair<float,int64_t>a,std::pair<float,int64_t>b){return (a.first!=b.first) ? (a.first<b.first) : a.second < b.second;});\n";
161 } else
162 // in this case we don;t need to return sorted elements, so we keep same order as before
163 out<<SP<<SP << "std::partial_sort(elements.begin(),elements.begin()+" << fK << ",elements.end());\n";
164
165 // copy the selected elements in the output
166 out << SP << SP << "for (size_t l = 0; l < " << fK << "; l++) {\n";
167 out << SP << SP << SP << "tensor_" << fNVal << "[yoffset + " << strideY[axis] << "*l + j] = elements[l].first;\n";
168 out << SP << SP << SP << "tensor_" << fNInd << "[yoffset + " << strideY[axis] << "*l + j] = elements[l].second;\n";
169 out << SP << SP << "}\n";
170 if (n_after > 1) out << SP << SP << "}\n";
171 if (n_before> 1) out << SP << "}\n";
172 out << SP << "}\n"; // end operator scope
173 return out.str();
174 }
175};
176
177} // nameSPace SOFIE
178} // nameSPace Experimental
179} // nameSPace TMVA
180
181#endif // TMVA_SOFIE_ROPERATOR_TOPK
size_t size(const MatrixT &matrix)
retrieve the size of a square matrix
ROOT::Detail::TRangeCast< T, true > TRangeDynCast
TRangeDynCast is an adapter class that allows the typed iteration through a TCollection.
Option_t Option_t TPoint TPoint const char GetTextMagnitude GetFillStyle GetLineColor GetLineWidth GetMarkerStyle GetTextAlign GetTextColor GetTextSize void input
Option_t Option_t TPoint TPoint const char GetTextMagnitude GetFillStyle GetLineColor GetLineWidth GetMarkerStyle GetTextAlign GetTextColor GetTextSize void char Point_t Rectangle_t WindowAttributes_t Float_t Float_t Float_t Int_t Int_t UInt_t UInt_t Rectangle_t Int_t Int_t Window_t TString Int_t GCValues_t GetPrimarySelectionOwner GetDisplay GetScreen GetColormap GetNativeEvent const char const char dpyName wid window const char font_name cursor keysym reg const char only_if_exist regb h Point_t winding char text const char depth char const char Int_t count const char ColorStruct_t color const char Pixmap_t Pixmap_t PictureAttributes_t attr const char char ret_data h unsigned char height h length
void AddIntermediateTensor(std::string tensor_name, ETensorType type, std::vector< Dim > dim_shape)
Definition RModel.cxx:200
bool CheckIfTensorAlreadyExist(std::string tensor_name)
Definition RModel.cxx:95
const ETensorType & GetTensorType(std::string name) const
Definition RModel.cxx:67
const std::vector< size_t > & GetTensorShape(std::string name) const
Definition RModel.cxx:29
std::shared_ptr< void > GetInitializedTensorData(std::string tensor_name)
Definition RModel.cxx:261
void SetNotWritableInitializedTensor(const std::string &tensor_name)
Definition RModel.cxx:270
std::string Generate(std::string OpName) override
std::vector< std::vector< size_t > > ShapeInference(std::vector< std::vector< size_t > > input) override
std::vector< ETensorType > TypeInference(std::vector< ETensorType > input) override
ROperator_TopK(int attr_axis, int attr_largest, int attr_sorted, std::string nameK, std::string nameX, std::string nameVal, std::string nameInd)
void Initialize(RModel &model) override
std::vector< std::string_view > fInputTensorNames
Definition ROperator.hxx:46
const std::string SP
space used to correctly indent the generated C++ code
Definition ROperator.hxx:42
std::vector< std::string_view > fOutputTensorNames
Definition ROperator.hxx:47
std::vector< size_t > ComputeStrideFromShape(const std::vector< size_t > &shape)
compute stride of a tensor given its shape (assume layout is row-major)
std::string ConvertTypeToString(ETensorType type)
std::size_t ConvertShapeToLength(std::vector< size_t > shape)
create variable transformations