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
ff3c4cd104aa7fffd045b4a3ea27462f18ad0c98

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

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