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
75f0dbffa990fa07bb0826d0b8faef9fd44f461f

Branches

Tags

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

Clone

HTTPS

Download ZIP

onnxruntime-extensions/include/custom_op

include/custom_op/tensor_api.h

604lines · modecode

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