GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/onnxruntime-extensions

Public

mirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
b341d66ebb794e190613545b11bfa2be0e9c4440

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

onnxruntime-extensions/includes

includes/tensor_api.h

494lines · modecode

RawDownload
Latest commit unavailable.
1#include <optional>
2#include <numeric>
3#include <type_traits>
4#include "onnxruntime_customop.hpp"
5#include "onnxruntime_f16.h"
6
7namespace Ort {
8namespace Custom {
9
10// this is for the ORT custom op template magic
11class Arg {
12};
13
14template <typename T>
15struct Span {
16 const T* data_ = {};
17 size_t size_ = {};
18 void Assign(const T* data, size_t size) {
19 data_ = data;
20 size_ = size;
21 }
22 size_t size() const { return size_; }
23 T operator[](size_t indice) const {
24 return data_[indice];
25 }
26 const T* data() const { return data_; }
27};
28
29
30#if ORT_API_VERSION >= 16
31
32template <>
33struct Span<MFloat16> {
34 const MFloat16* data_ = {};
35 size_t size_ = {};
36 void Assign(const MFloat16* data, size_t size) {
37 data_ = data;
38 size_ = size;
39 }
40 size_t size() const { return size_; }
41 MFloat16 operator[](size_t indice) const {
42 return data_[indice];
43 }
44 const MFloat16* data() const { return data_; }
45};
46
47template <>
48struct Span<BFloat16> {
49 const BFloat16* data_ = {};
50 size_t size_ = {};
51 void Assign(const BFloat16* data, size_t size) {
52 data_ = data;
53 size_ = size;
54 }
55 size_t size() const { return size_; }
56 BFloat16 operator[](size_t indice) const {
57 return data_[indice];
58 }
59 const BFloat16* data() const { return data_; }
60};
61
62#endif
63
64class ITensorStorage{
65public:
66 virtual const std::vector<int64_t>& Shape() const = 0;
67 virtual const void* DataRaw() const = 0;
68 virtual bool IsInitialized() const = 0;
69 virtual void* Initialize(const std::vector<int64_t>& shape, size_t element_size) = 0;
70};
71
72class IAllocator {
73public:
74 virtual void* Alloc(size_t size) = 0;
75 virtual void Free(void* p) = 0;
76};
77
78// TODO: remove this
79
80class TestAllocator : public IAllocator {
81public:
82 void* Alloc(size_t size) override {
83 return malloc(size);
84 }
85
86 void Free(void* p) override {
87 if (p){
88 free(p);
89 }
90 }
91};
92
93class OrtEagerTensorStorage : public ITensorStorage {
94public:
95 OrtEagerTensorStorage(const std::vector<int64_t>& shape,
96 void* buffer) : buffer_(buffer), shape_(shape){
97
98 }
99
100 OrtEagerTensorStorage(IAllocator* allocator) : allocator_(allocator){
101 }
102
103 virtual ~OrtEagerTensorStorage(){
104 if (allocator_ && buffer_)
105 allocator_->Free(buffer_);
106 }
107
108 const std::vector<int64_t>& Shape() const override {
109 if (!IsInitialized())
110 ORTX_CXX_API_THROW("Tensor not initialized", ORT_RUNTIME_EXCEPTION);
111 return *shape_;
112 }
113
114 virtual bool IsInitialized() const override {
115 return shape_.has_value();
116 }
117
118 const void* DataRaw() const override {
119 return buffer_;
120 }
121
122 void* Initialize(const std::vector<int64_t>& shape, size_t element_size) override {
123 if (IsInitialized())
124 return buffer_;
125 assert(allocator_);
126 shape_ = shape;
127 int64_t n_elem = std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>());
128 auto buffer_size = n_elem * element_size;
129 buffer_ = allocator_->Alloc(buffer_size);
130 return buffer_;
131 }
132
133private:
134 void* buffer_ {};
135 std::optional<std::vector<int64_t>> shape_;
136 // caller need to make sure the allocator is alive
137 IAllocator* allocator_;
138};
139
140template <typename T>
141class Tensor : public Arg {
142 public:
143 using TT = typename std::remove_reference<T>::type;
144 Tensor(std::unique_ptr<ITensorStorage> tensor_storage) : storage_(std::move(tensor_storage)){
145 }
146
147 Tensor(const std::vector<int64_t>& shape, void* buffer) : Tensor(std::make_unique<OrtEagerTensorStorage>(shape, buffer)) {}
148
149 Tensor(IAllocator* allocator) : storage_(std::make_unique<OrtEagerTensorStorage>(allocator)){}
150
151 virtual ~Tensor() = default;
152
153 operator bool() const {
154 return storage_->IsInitialized();
155 }
156
157 const std::vector<int64_t>& Shape() const {
158 return storage_->Shape();
159 }
160
161 int64_t NumberOfElement() const {
162 auto& shape = storage_->Shape();
163 return std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>());
164 }
165
166 std::string Shape2Str() const {
167 if (storage_->IsInitialized()) {
168 std::string shape_str;
169 auto& shape = storage_->Shape();
170 for (const auto& dim : shape) {
171 shape_str.append(std::to_string(dim));
172 shape_str.append(", ");
173 }
174 return shape_str;
175 } else {
176 return "empty";
177 }
178 }
179
180 const TT* Data() const {
181#if ORT_API_VERSION >= 16
182 if constexpr (std::is_same<TT, MFloat16>::value || std::is_same<TT, BFloat16>::value)
183 return reinterpret_cast<const TT*>(storage_->DataRaw());
184 else
185#endif
186 return static_cast<const TT*>(storage_->DataRaw());
187 }
188
189 const void* DataRaw() const {
190 return storage_->DataRaw();
191 }
192
193 size_t SizeInBytes() const {
194 return NumberOfElement() * sizeof(TT);
195 }
196
197 TT* Allocate(const std::vector<int64_t>& shape) {
198 // it should be OK to allocate multiple times
199 void* buffer = storage_->Initialize(shape, sizeof(TT));
200#if ORT_API_VERSION >= 16
201 if constexpr (std::is_same<TT, MFloat16>::value || std::is_same<TT, BFloat16>::value)
202 return reinterpret_cast<TT*>(buffer);
203 else
204#endif
205 return static_cast<TT*>(buffer);
206 }
207
208 const Span<T>& AsSpan() {
209#if ORT_API_VERSION >= 16
210 if constexpr (std::is_same<TT, MFloat16>::value || std::is_same<TT, BFloat16>::value) {
211 ORTX_CXX_API_THROW("AsSpan for MFloat16 / BFloat16 not implemented", ORT_RUNTIME_EXCEPTION);
212 }
213 else{
214#endif
215 auto& shape = storage_->Shape();
216 if (shape.size() != 1) {
217 ORTX_CXX_API_THROW("to get a span, shape must be 1-D, actual shape: " + Shape2Str(), ORT_RUNTIME_EXCEPTION);
218 }
219 span_.Assign(Data(), shape[0]);
220 return span_;
221#if ORT_API_VERSION >= 16
222 }
223#endif
224 }
225
226 const T& AsScalar() {
227#if ORT_API_VERSION >= 16
228 if constexpr (std::is_same<TT, MFloat16>::value || std::is_same<TT, BFloat16>::value) {
229 ORTX_CXX_API_THROW("AsScalar for MFloat16 / BFloat16 not implemented", ORT_RUNTIME_EXCEPTION);
230 }
231 else{
232#endif
233 auto& shape = storage_->Shape();
234 if ((shape.size() == 1 && shape[0] != 1) || shape.size() > 1) {
235 ORTX_CXX_API_THROW("to get a scalar, shape must be {1}, actual shape: " + Shape2Str(), ORT_RUNTIME_EXCEPTION);
236 }
237 return *Data();
238#if ORT_API_VERSION >= 16
239 }
240#endif
241 }
242
243 private:
244 std::unique_ptr<ITensorStorage> storage_;
245 Span<T> span_;
246};
247
248template<typename T>
249class IStringTensorStorage{
250public:
251 using strings = std::vector<T>;
252 virtual const std::vector<int64_t>& Shape() const = 0;
253 virtual const void* DataRaw() const = 0;
254 virtual const strings& Data() const = 0;
255 virtual bool IsInitialized() const = 0;
256 virtual void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) = 0;
257 virtual void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) = 0;
258};
259
260template<typename T>
261class EagerStringTensorStorage : public IStringTensorStorage<T>{
262public:
263 using strings = std::vector<T>;
264 EagerStringTensorStorage(const strings& ss) : input_strings_(ss), shape_(std::vector<int64_t>{ss.size()}){}
265
266 EagerStringTensorStorage() {}
267
268 const std::vector<int64_t>& Shape() const override {
269 if (!IsInitialized())
270 ORTX_CXX_API_THROW("Tensor not initialized", ORT_RUNTIME_EXCEPTION);
271 return *shape_;
272 }
273
274 virtual const void* DataRaw() const override {
275 if (input_strings_.size() != 1) {
276 ORTX_CXX_API_THROW("DataRaw() only applies to string scalar", ORT_RUNTIME_EXCEPTION);
277 }
278 if constexpr (std::is_same<std::string_view, T>::value)
279 return reinterpret_cast<const void*>(input_strings_[0].data());
280 else
281 return reinterpret_cast<const void*>(input_strings_[0].c_str());
282 }
283
284 virtual bool IsInitialized() const override {
285 return shape_.has_value();
286 }
287
288 virtual void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) override {
289 if constexpr (std::is_same<std::string_view, T>::value)
290 ORTX_CXX_API_THROW("Set output for string view tensor is not supported", ORT_RUNTIME_EXCEPTION);
291 input_strings_.assign(ss.begin(), ss.end());
292 shape_ = dims;
293 }
294
295 const strings& Data() const override {
296 return input_strings_;
297 }
298
299 virtual void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) override {
300 if constexpr (std::is_same<std::string_view, T>::value)
301 ORTX_CXX_API_THROW("Set output for string view tensor is not supported", ORT_RUNTIME_EXCEPTION);
302
303 for (const char* s : ss){
304 input_strings_.push_back(s);
305 }
306 shape_ = dims;
307 }
308
309private:
310 std::vector<T> input_strings_;
311 std::optional<std::vector<int64_t>> shape_;
312};
313
314template <>
315class Tensor<std::string> : public Arg {
316 public:
317 using strings = std::vector<std::string>;
318
319 Tensor(std::unique_ptr<IStringTensorStorage<std::string>> storage) : storage_(std::move(storage)) {}
320
321 Tensor(const strings& ss) : storage_(std::make_unique<EagerStringTensorStorage<std::string>>(ss)) {}
322
323 Tensor() : storage_(std::make_unique<EagerStringTensorStorage<std::string>>()) {}
324
325 const strings& Data() const {
326 return storage_->Data();
327 }
328
329 const std::vector<int64_t>& Shape() const {
330 return storage_->Shape();
331 }
332
333 int64_t NumberOfElement() const {
334 auto& shape = storage_->Shape();
335 return std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>());
336 }
337
338 std::string Shape2Str() const {
339 if (storage_->IsInitialized()) {
340 std::string shape_str;
341 auto& shape = storage_->Shape();
342 for (const auto& dim : shape) {
343 shape_str.append(std::to_string(dim));
344 shape_str.append(", ");
345 }
346 return shape_str;
347 } else {
348 return "empty";
349 }
350 }
351
352 const void* DataRaw() const {
353 return storage_->DataRaw();
354 }
355
356 size_t SizeInBytes() const {
357 auto& ss = storage_->Data();
358 if (ss.size() != 1) {
359 ORTX_CXX_API_THROW("SizeInBytes() only applies to string scalar", ORT_RUNTIME_EXCEPTION);
360 }
361 return ss[0].size();
362 }
363
364 void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) {
365 storage_->SetStringOutput(ss, dims);
366 }
367 void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) {
368 storage_->SetStringOutput(ss, dims);
369 }
370 const Span<std::string>& AsSpan() {
371 ORTX_CXX_API_THROW("span for TensorT of string not implemented", ORT_RUNTIME_EXCEPTION);
372 }
373 const std::string& AsScalar() {
374 auto& ss = storage_->Data();
375 if (ss.size() != 1) {
376 ORTX_CXX_API_THROW("to get a scalar, shape must be {1}, actual shape: " + Shape2Str(), ORT_RUNTIME_EXCEPTION);
377 }
378 return ss[0];
379 }
380
381 private:
382 std::unique_ptr<IStringTensorStorage<std::string>> storage_;
383};
384
385
386template <>
387class Tensor<std::string_view> : public Arg {
388 public:
389 using strings = std::vector<std::string_view>;
390
391 Tensor(std::unique_ptr<IStringTensorStorage<std::string_view>> storage) : storage_(std::move(storage)) {}
392
393 Tensor(const strings& ss) : storage_(std::make_unique<EagerStringTensorStorage<std::string_view>>(ss)) {}
394
395 const strings& Data() const {
396 return storage_->Data();
397 }
398
399 const std::vector<int64_t>& Shape() const {
400 return storage_->Shape();
401 }
402
403 int64_t NumberOfElement() const {
404 auto& shape = storage_->Shape();
405 return std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>());
406 }
407
408 std::string Shape2Str() const {
409 if (storage_->IsInitialized()) {
410 std::string shape_str;
411 auto& shape = storage_->Shape();
412 for (const auto& dim : shape) {
413 shape_str.append(std::to_string(dim));
414 shape_str.append(", ");
415 }
416 return shape_str;
417 } else {
418 return "empty";
419 }
420 }
421
422 const void* DataRaw() const {
423 return storage_->DataRaw();
424 }
425
426 size_t SizeInBytes() const {
427 auto& ss = storage_->Data();
428 if (ss.size() != 1) {
429 ORTX_CXX_API_THROW("SizeInBytes() only applies to string scalar", ORT_RUNTIME_EXCEPTION);
430 }
431 return ss[0].size();
432 }
433
434 void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) {
435 storage_->SetStringOutput(ss, dims);
436 }
437 void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) {
438 storage_->SetStringOutput(ss, dims);
439 }
440 const Span<std::string_view>& AsSpan() {
441 ORTX_CXX_API_THROW("span for TensorT of string not implemented", ORT_RUNTIME_EXCEPTION);
442 }
443 const std::string_view& AsScalar() {
444 auto& ss = storage_->Data();
445 if (ss.size() != 1) {
446 ORTX_CXX_API_THROW("to get a scalar, shape must be {1}, actual shape: " + Shape2Str(), ORT_RUNTIME_EXCEPTION);
447 }
448 return ss[0];
449 }
450
451 private:
452 std::unique_ptr<IStringTensorStorage<std::string_view>> storage_;
453};
454
455
456template<typename ...Args>
457class NamedArgumentDict{
458public:
459 using ValueTuple = std::tuple<Args...>;
460
461 NamedArgumentDict(const std::vector<const char*>& keys, const std::tuple<Args...>& args) : names_(keys), entries_(args) {
462 }
463
464 template<typename T>
465 T TryToGetAttributeWithDefault(const char* name, const T& default_value) const {
466 return TryToGetAttributeWithDefaultInternal<0>(name, default_value);
467 }
468
469private:
470 template<size_t I, typename T>
471 typename std::enable_if<I == sizeof...(Args), T>::type
472 TryToGetAttributeWithDefaultInternal(const char* name, const T& default_value) const {
473 return default_value;
474 }
475
476 template<size_t I, typename T>
477 typename std::enable_if<I < sizeof...(Args), T>::type
478 TryToGetAttributeWithDefaultInternal(const char* name, const T& default_value) const {
479 if (names_[I] == name){
480 if constexpr (std::is_same<std::tuple_element_t<I, ValueTuple>, T>::value)
481 return std::get<I>(entries_);
482 else
483 throw std::runtime_error("name matched but type is not");
484 }
485 return TryToGetAttributeWithDefaultInternal<I+1>(name, default_value);
486 }
487
488 std::vector<const char*> names_;
489 std::tuple<Args...> entries_;
490
491};
492
493}
494}
495