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
cf7d14bc9c2a1a51de5c02a657d00a1943fbef55

Branches

Tags

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

Clone

HTTPS

Download ZIP

onnxruntime-extensions/includes

includes/custom_op_lite.h

1115lines · modecode

RawDownload
Latest commit unavailable.
1// Copyright (c) Microsoft Corporation. All rights reserved.
2// Licensed under the MIT License.
3
4#pragma once
5
6#include <optional>
7#include <numeric>
8#include "tensor_api.h"
9#include "onnxruntime_cpp_api_legacy.hpp"
10
11namespace Ort {
12namespace Custom {
13
14class OrtKernelArg {
15public:
16 OrtKernelArg(const OrtW::CustomOpApi& api,
17 OrtKernelContext& ctx,
18 size_t indice,
19 bool is_input) : api_(api), ctx_(ctx), indice_(indice) {
20 if (is_input) {
21 const OrtValue* const_value = api.KernelContext_GetInput(&ctx, indice);
22 const OrtMemoryInfo* mem_info = {};
23 api.ThrowOnError(api.GetOrtApi().GetTensorMemoryInfo(const_value, &mem_info));
24 if (mem_info) {
25 api.ThrowOnError(api.GetOrtApi().MemoryInfoGetName(mem_info, &mem_type_));
26 }
27 }
28 }
29
30 bool IsCpuTensor() const {
31 return strcmp("Cpu", mem_type_) == 0;
32 }
33
34protected:
35 const OrtW::CustomOpApi& api_;
36 OrtKernelContext& ctx_;
37 size_t indice_;
38 const char* mem_type_ = "Cpu";
39};
40
41class OrtKernelContextStorage : public ITensorStorage {
42public:
43 OrtKernelContextStorage(const OrtW::CustomOpApi& api,
44 OrtKernelContext& ctx,
45 size_t indice,
46 bool is_input) : api_(api), ctx_(ctx), indice_(indice) {
47 if (is_input){
48 auto input_count = api.KernelContext_GetInputCount(&ctx);
49 if (indice >= input_count) {
50 ORTX_CXX_API_THROW("invalid indice", ORT_RUNTIME_EXCEPTION);
51 }
52 const_value_ = api.KernelContext_GetInput(&ctx, indice);
53 auto* info = api.GetTensorTypeAndShape(const_value_);
54 shape_ = api.GetTensorShape(info);
55 api.ReleaseTensorTypeAndShapeInfo(info);
56 }
57 }
58
59 const std::vector<int64_t>& Shape() const override {
60 if (!IsInitialized())
61 ORTX_CXX_API_THROW("Tensor not initialized", ORT_RUNTIME_EXCEPTION);
62 return *shape_;
63 }
64
65 virtual bool IsInitialized() const override {
66 return shape_.has_value();
67 }
68
69 const void* DataRaw() const override {
70 return api_.GetTensorRawData(const_value_);
71 }
72
73 void* Initialize(const std::vector<int64_t>& shape, size_t element_size) override {
74 if (!const_value_) {
75 const_value_ = api_.KernelContext_GetOutput(&ctx_, indice_, shape.data(), shape.size());
76 }
77 return api_.GetTensorMutableRawData(const_cast<OrtValue*>(const_value_));
78 }
79
80private:
81 const OrtW::CustomOpApi& api_;
82 OrtKernelContext& ctx_;
83 size_t indice_;
84 const OrtValue* const_value_{}; // for input
85 std::optional<std::vector<int64_t>> shape_;
86};
87
88template <typename T>
89class OrtTensor : public OrtKernelArg, public Tensor<T> {
90public:
91 OrtTensor(const OrtW::CustomOpApi& api,
92 OrtKernelContext& ctx,
93 size_t indice,
94 bool is_input) : OrtKernelArg(api, ctx, indice, is_input),
95 Tensor<T>(std::make_unique<OrtKernelContextStorage>(api, ctx, indice, is_input)){
96 }
97};
98
99class OrtStringTensorStorage : public IStringTensorStorage<std::string>{
100public:
101 using strings = std::vector<std::string>;
102 OrtStringTensorStorage(const OrtW::CustomOpApi& api,
103 OrtKernelContext& ctx,
104 size_t indice,
105 bool is_input) : api_(api), ctx_(ctx), indice_(indice){
106 if (is_input) {
107 auto input_count = api_.KernelContext_GetInputCount(&ctx_);
108 if (indice >= input_count) {
109 ORTX_CXX_API_THROW("invalid indice", ORT_RUNTIME_EXCEPTION);
110 }
111
112 auto* const_value = api_.KernelContext_GetInput(&ctx_, indice);
113 auto* info = api_.GetTensorTypeAndShape(const_value);
114 shape_ = api_.GetTensorShape(info);
115 api_.ReleaseTensorTypeAndShapeInfo(info);
116
117 size_t num_chars;
118 OrtW::ThrowOnError(api_.GetOrtApi(), api_.GetOrtApi().GetStringTensorDataLength(const_value, &num_chars));
119 std::vector<char> chars(num_chars + 1, '\0');
120 assert((*shape_).size() == 1);
121 auto num_strings = (*shape_)[0];
122 std::vector<size_t> offsets((*shape_)[0]);
123 OrtW::ThrowOnError(api_.GetOrtApi(), api_.GetOrtApi().GetStringTensorContent(const_value,
124 (void*)chars.data(),
125 num_chars,
126 offsets.data(),
127 offsets.size()));
128 auto upper_bound = static_cast<int64_t>(num_strings) - 1;
129 input_strings_.resize(num_strings);
130 for (int64_t i = upper_bound; i >= 0; --i) {
131 if (i < upper_bound) {
132 chars[offsets[i + 1]] = '\0';
133 }
134 input_strings_[i] = chars.data() + offsets[i];
135 }
136 }
137 }
138
139 const std::vector<int64_t>& Shape() const override {
140 if (!IsInitialized())
141 ORTX_CXX_API_THROW("Tensor not initialized", ORT_RUNTIME_EXCEPTION);
142 return *shape_;
143 }
144
145 virtual const void* DataRaw() const override {
146 if (input_strings_.size() != 1) {
147 ORTX_CXX_API_THROW("DataRaw() only applies to string scalar", ORT_RUNTIME_EXCEPTION);
148 }
149 return reinterpret_cast<const void*>(input_strings_[0].c_str());
150 }
151
152 virtual bool IsInitialized() const override {
153 return shape_.has_value();
154 }
155
156 virtual void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) override {
157 std::vector<const char*> raw;
158 for (const auto& s : ss) {
159 raw.push_back(s.data());
160 }
161 auto* output = api_.KernelContext_GetOutput(&ctx_, indice_, dims.data(), dims.size());
162 OrtW::ThrowOnError(api_.GetOrtApi(), api_.GetOrtApi().FillStringTensor(output, raw.data(), raw.size()));
163 }
164
165 virtual void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) override {
166 auto* output = api_.KernelContext_GetOutput(&ctx_, indice_, dims.data(), dims.size());
167 OrtW::ThrowOnError(api_.GetOrtApi(), api_.GetOrtApi().FillStringTensor(output, ss.data(), ss.size()));
168 }
169
170 const strings& Data() const override {
171 return input_strings_;
172 }
173
174private:
175 const OrtW::CustomOpApi& api_;
176 OrtKernelContext& ctx_;
177 size_t indice_;
178 std::vector<std::string> input_strings_;
179 std::optional<std::vector<int64_t>> shape_;
180};
181
182
183class OrtStringViewTensorStorage : public IStringTensorStorage<std::string_view>{
184public:
185 using strings = std::vector<std::string_view>;
186 OrtStringViewTensorStorage(const OrtW::CustomOpApi& api,
187 OrtKernelContext& ctx,
188 size_t indice,
189 bool is_input) : api_(api), ctx_(ctx), indice_(indice){
190 if (is_input) {
191 auto input_count = api_.KernelContext_GetInputCount(&ctx_);
192 if (indice >= input_count) {
193 ORTX_CXX_API_THROW("invalid indice", ORT_RUNTIME_EXCEPTION);
194 }
195 auto* const_value = api_.KernelContext_GetInput(&ctx_, indice);
196 auto* info = api_.GetTensorTypeAndShape(const_value);
197 shape_ = api_.GetTensorShape(info);
198 api_.ReleaseTensorTypeAndShapeInfo(info);
199
200 size_t num_chars;
201 OrtW::ThrowOnError(api_.GetOrtApi(), api_.GetOrtApi().GetStringTensorDataLength(const_value, &num_chars));
202 chars_.resize(num_chars + 1, '\0');
203
204 auto num_strings = static_cast<size_t>((*shape_)[0]);
205 if (num_strings) {
206 std::vector<size_t> offsets(num_strings);
207 OrtW::ThrowOnError(api_.GetOrtApi(), api_.GetOrtApi().GetStringTensorContent(const_value,
208 (void*)chars_.data(),
209 num_chars,
210 offsets.data(),
211 offsets.size()));
212 offsets.push_back(num_chars);
213 for (size_t i = 0; i < num_strings; ++i) {
214 input_string_views_.emplace_back(chars_.data() + offsets[i], offsets[i + 1] - offsets[i]);
215 }
216 }
217 }
218 }
219
220 const std::vector<int64_t>& Shape() const override {
221 if (!IsInitialized())
222 ORTX_CXX_API_THROW("Tensor not initialized", ORT_RUNTIME_EXCEPTION);
223 return *shape_;
224 }
225
226 virtual const void* DataRaw() const override {
227 if (input_string_views_.size() != 1) {
228 ORTX_CXX_API_THROW("DataRaw() only applies to string scalar", ORT_RUNTIME_EXCEPTION);
229 }
230 return reinterpret_cast<const void*>(input_string_views_[0].data());
231 }
232
233 virtual bool IsInitialized() const override {
234 return shape_.has_value();
235 }
236
237 virtual void SetStringOutput(const strings& ss, const std::vector<int64_t>& dims) override {
238 ORTX_CXX_API_THROW("Set output for string view tensor is not supported", ORT_RUNTIME_EXCEPTION);
239 }
240
241 virtual void SetStringOutput(const std::vector<const char*>& ss, const std::vector<int64_t>& dims) override {
242 ORTX_CXX_API_THROW("Set output for string view tensor is not supported", ORT_RUNTIME_EXCEPTION);
243 }
244
245 const strings& Data() const override {
246 return input_string_views_;
247 }
248
249private:
250 const OrtW::CustomOpApi& api_;
251 OrtKernelContext& ctx_;
252 size_t indice_;
253 std::vector<char> chars_; // for input
254 std::vector<std::string_view> input_string_views_; // for input
255 std::optional<std::vector<int64_t>> shape_;
256};
257
258// to make the metaprogramming magic happy.
259template <>
260class OrtTensor<std::string> : public OrtKernelArg,
261 public Tensor<std::string>{
262public:
263 OrtTensor(const OrtW::CustomOpApi& api,
264 OrtKernelContext& ctx,
265 size_t indice,
266 bool is_input) : OrtKernelArg(api, ctx, indice, is_input),
267 Tensor<std::string>(std::make_unique<OrtStringTensorStorage>(api, ctx, indice, is_input)) {}
268};
269
270template <>
271class OrtTensor<std::string_view> : public OrtKernelArg,
272 public Tensor<std::string_view>{
273public:
274 OrtTensor(const OrtW::CustomOpApi& api,
275 OrtKernelContext& ctx,
276 size_t indice,
277 bool is_input) : OrtKernelArg(api, ctx, indice, is_input),
278 Tensor<std::string_view>(std::make_unique<OrtStringViewTensorStorage>(api, ctx, indice, is_input)) {}
279};
280
281using TensorPtr = std::unique_ptr<Custom::Arg>;
282using TensorPtrs = std::vector<TensorPtr>;
283
284// Represent variadic input or output
285struct Variadic : public OrtKernelArg, public Arg {
286 Variadic(const OrtW::CustomOpApi& api,
287 OrtKernelContext& ctx,
288 size_t indice,
289 bool is_input) : OrtKernelArg(api, ctx, indice, is_input) {
290#if ORT_API_VERSION < 14
291 ORTX_CXX_API_THROW("Variadic input or output only supported after onnxruntime 1.14", ORT_RUNTIME_EXCEPTION);
292#endif
293 if (is_input) {
294 auto input_count = api.KernelContext_GetInputCount(&ctx_);
295 for (size_t ith_input = 0; ith_input < input_count; ++ith_input) {
296 auto* const_value = api_.KernelContext_GetInput(&ctx_, ith_input);
297 auto* info = api_.GetTensorTypeAndShape(const_value);
298 auto type = api_.GetTensorElementType(info);
299 api_.ReleaseTensorTypeAndShapeInfo(info);
300 TensorPtr tensor;
301 switch (type) {
302 case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL:
303 tensor = std::make_unique<Custom::OrtTensor<bool>>(api, ctx, ith_input, true);
304 break;
305 case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
306 tensor = std::make_unique<Custom::OrtTensor<float>>(api, ctx, ith_input, true);
307 break;
308 case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE:
309 tensor = std::make_unique<Custom::OrtTensor<double>>(api, ctx, ith_input, true);
310 break;
311 case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8:
312 tensor = std::make_unique<Custom::OrtTensor<uint8_t>>(api, ctx, ith_input, true);
313 break;
314 case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8:
315 tensor = std::make_unique<Custom::OrtTensor<int8_t>>(api, ctx, ith_input, true);
316 break;
317 case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16:
318 tensor = std::make_unique<Custom::OrtTensor<uint16_t>>(api, ctx, ith_input, true);
319 break;
320 case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16:
321 tensor = std::make_unique<Custom::OrtTensor<int16_t>>(api, ctx, ith_input, true);
322 break;
323 case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32:
324 tensor = std::make_unique<Custom::OrtTensor<uint32_t>>(api, ctx, ith_input, true);
325 break;
326 case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32:
327 tensor = std::make_unique<Custom::OrtTensor<int32_t>>(api, ctx, ith_input, true);
328 break;
329 case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64:
330 tensor = std::make_unique<Custom::OrtTensor<uint64_t>>(api, ctx, ith_input, true);
331 break;
332 case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64:
333 tensor = std::make_unique<Custom::OrtTensor<int64_t>>(api, ctx, ith_input, true);
334 break;
335 case ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING:
336 tensor = std::make_unique<Custom::OrtTensor<std::string>>(api, ctx, ith_input, true);
337 break;
338 default:
339 ORTX_CXX_API_THROW("unknow input type", ORT_RUNTIME_EXCEPTION);
340 break;
341 }
342 tensors_.emplace_back(tensor.release());
343 } // for
344 } else {
345 // a Variadic used for output is populated by the Compute so leave tensors_ empty here
346 }
347 }
348 template <typename T>
349 T* AllocateOutput(size_t ith_output, const std::vector<int64_t>& shape) {
350 auto tensor = std::make_unique<OrtTensor<T>>(api_, ctx_, ith_output, false);
351 auto raw_output = tensor.get()->Allocate(shape);
352 tensors_.emplace_back(tensor.release());
353 return raw_output;
354 }
355 Tensor<std::string>& AllocateStringTensor(size_t ith_output) {
356 auto tensor = std::make_unique<OrtTensor<std::string>>(api_, ctx_, ith_output, false);
357 Tensor<std::string>& output = *tensor;
358 tensors_.emplace_back(tensor.release());
359 return output;
360 }
361 const void* DataRaw() const {
362 ORTX_CXX_API_THROW("DataRaw() cannot be applied to Variadic", ORT_RUNTIME_EXCEPTION);
363 return nullptr;
364 }
365 size_t SizeInBytes() const {
366 ORTX_CXX_API_THROW("SizeInBytes() cannot be applied to Variadic", ORT_RUNTIME_EXCEPTION);
367 return 0;
368 }
369 size_t Size() const {
370 return tensors_.size();
371 }
372 const TensorPtr& operator[](size_t indice) const {
373 return tensors_.at(indice);
374 }
375
376 private:
377 TensorPtrs tensors_;
378};
379
380class OrtGraphKernelContext : public KernelContext {
381public:
382 OrtGraphKernelContext(const OrtApi& api, const OrtKernelContext& ctx) : api_(api) {
383 OrtMemoryInfo* info;
384 OrtW::ThrowOnError(api, api.CreateCpuMemoryInfo(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault, &info));
385 OrtW::ThrowOnError(api, api.KernelContext_GetAllocator(&ctx, info, &allocator_));
386 api.ReleaseMemoryInfo(info);
387 }
388
389 virtual ~OrtGraphKernelContext(){
390 if (allocator_){
391 api_.ReleaseAllocator(allocator_);
392 }
393 }
394
395 void* AllocScratchBuffer(size_t size) override{
396 return allocator_->Alloc(allocator_, size);
397 }
398
399 void FreeScratchBuffer(void* p) override {
400 if (p){
401 allocator_->Free(allocator_, p);
402 }
403 }
404
405private:
406 const OrtApi& api_;
407 OrtAllocator* allocator_;
408};
409
410#ifdef USE_CUDA
411
412enum CudaResource {
413 cuda_handle_t = 10000,
414 cudnn_handle_t,
415 cublas_handle_t,
416 deferred_cpu_allocator_t,
417 // below are cuda ep options
418 device_id_t,
419};
420
421struct CudaContext {
422 static const int cuda_resource_ver = 1;
423 void Init(const OrtW::CustomOpApi& api, const OrtKernelContext& ctx) {
424 const auto& ort_api = api.GetOrtApi();
425 ort_api.KernelContext_GetResource(&ctx, cuda_resource_ver, CudaResource::cuda_handle_t, &cuda_stream);
426 if (!cuda_stream) {
427 ORTX_CXX_API_THROW("Failed to fetch cuda stream from context", ORT_RUNTIME_EXCEPTION);
428 }
429 ort_api.KernelContext_GetResource(&ctx, cuda_resource_ver, CudaResource::cublas_handle_t, &cublas);
430 if (!cublas) {
431 ORTX_CXX_API_THROW("Failed to fetch cublas handle from context", ORT_RUNTIME_EXCEPTION);
432 }
433 void* resource = nullptr;
434 OrtStatusPtr result = ort_api.KernelContext_GetResource(&ctx, cuda_resource_ver, CudaResource::device_id_t, &resource);
435 if (result) {
436 ORTX_CXX_API_THROW("Failed to fetch device id from context", ORT_RUNTIME_EXCEPTION);
437 }
438 memcpy(&device_id, &resource, sizeof(int));
439 }
440 void* cuda_stream = {};
441 void* cublas = {};
442 int device_id = 0;
443};
444
445
446class OrtGraphCudaKernelContext : public CUDAKernelContext {
447public:
448 static const int cuda_resource_ver = 1;
449
450 OrtGraphCudaKernelContext(const OrtApi& api, const OrtKernelContext& ctx) : api_(api) {
451 api.KernelContext_GetResource(&ctx, cuda_resource_ver, CudaResource::cuda_handle_t, &cuda_stream_);
452 if (!cuda_stream_) {
453 ORTX_CXX_API_THROW("Failed to fetch cuda stream from context", ORT_RUNTIME_EXCEPTION);
454 }
455 api.KernelContext_GetResource(&ctx, cuda_resource_ver, CudaResource::cublas_handle_t, &cublas_);
456 if (!cublas_) {
457 ORTX_CXX_API_THROW("Failed to fetch cublas handle from context", ORT_RUNTIME_EXCEPTION);
458 }
459 void* resource = nullptr;
460 OrtStatusPtr result = api.KernelContext_GetResource(&ctx, cuda_resource_ver, CudaResource::device_id_t, &resource);
461 if (result) {
462 ORTX_CXX_API_THROW("Failed to fetch device id from context", ORT_RUNTIME_EXCEPTION);
463 }
464 memcpy(&device_id_, &resource, sizeof(int));
465
466 OrtMemoryInfo* info;
467 OrtW::ThrowOnError(api, api.CreateCpuMemoryInfo(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault, &info));
468 OrtW::ThrowOnError(api, api.KernelContext_GetAllocator(&ctx, info, &cpu_allocator_));
469 api.ReleaseMemoryInfo(info);
470
471 OrtMemoryInfo* cuda_mem_info;
472 OrtW::ThrowOnError(api, api.CreateMemoryInfo("GPU", OrtAllocatorType::OrtArenaAllocator, device_id_, OrtMemType::OrtMemTypeDefault, &cuda_mem_info));
473 OrtW::ThrowOnError(api, api.KernelContext_GetAllocator(&ctx, cuda_mem_info, &cuda_allocator_));
474 api.ReleaseMemoryInfo(cuda_mem_info);
475
476 }
477
478 virtual ~OrtGraphCudaKernelContext(){
479 if (cpu_allocator_){
480 api_.ReleaseAllocator(cpu_allocator_);
481 }
482 if (cuda_allocator_){
483 api_.ReleaseAllocator(cuda_allocator_);
484 }
485 }
486
487 void* AllocScratchBuffer(size_t size) override{
488 return cpu_allocator_->Alloc(cpu_allocator_, size);
489 }
490
491 void FreeScratchBuffer(void* p) override {
492 if (p){
493 cpu_allocator_->Free(cpu_allocator_, p);
494 }
495 }
496
497 void* AllocCudaScratchBuffer(size_t size) override {
498 return cuda_allocator_->Alloc(cuda_allocator_, size);
499 }
500
501 void FreeCudaScratchBuffer(void* p) override {
502 if (p){
503 cuda_allocator_->Free(cuda_allocator_, p);
504 }
505 }
506
507 void* GetCudaStream() const override {
508 return cuda_stream_;
509 }
510
511 void* GetCublasHandle() const override {
512 return cublas_;
513 }
514
515 int GetCudaDeviceId() const override {
516 return device_id_;
517 }
518
519private:
520 const OrtApi& api_;
521 OrtAllocator* cpu_allocator_;
522 OrtAllocator* cuda_allocator_;
523 void* cuda_stream_ = {};
524 void* cublas_ = {};
525 int device_id_ = 0;
526};
527
528#endif
529
530// using mf16_t = uint16_t;
531
532struct OrtLiteCustomOp : public OrtCustomOp {
533 // CreateTuple
534 template <size_t ith_input, size_t ith_output, typename... Ts>
535 static typename std::enable_if<sizeof...(Ts) == 0, std::tuple<>>::type
536 CreateTuple(const OrtW::CustomOpApi*, OrtKernelContext*, std::vector<TensorPtr>&, size_t, size_t, const std::string&) {
537 return std::make_tuple();
538 }
539
540 template <size_t ith_input, size_t ith_output, typename T, typename... Ts>
541 static typename std::enable_if<std::is_same<T, OrtKernelContext*>::value, std::tuple<T, Ts...>>::type
542 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) {
543 std::tuple<T> current = std::tuple<OrtKernelContext*>{context};
544 auto next = CreateTuple<ith_input, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep);
545 return std::tuple_cat(current, next);
546 }
547
548#ifdef USE_CUDA
549 template <size_t ith_input, size_t ith_output, typename T, typename... Ts>
550 static typename std::enable_if<std::is_same<T, const CudaContext&>::value, std::tuple<T, Ts...>>::type
551 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) {
552 thread_local CudaContext cuda_context;
553 cuda_context.Init(*api, *context);
554 std::tuple<T> current = std::tuple<const CudaContext&>{cuda_context};
555 auto next = CreateTuple<ith_input, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep);
556 return std::tuple_cat(current, next);
557 }
558#endif
559
560 template <size_t ith_input, size_t ith_output, typename T, typename... Ts>
561 static typename std::enable_if<std::is_same<T, KernelContext&>::value, std::tuple<T, Ts...>>::type
562 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) {
563 tensors.push_back(std::make_unique<OrtGraphCudaKernelContext>(api->GetOrtApi(), *context));
564 std::tuple<T> current = std::tuple<T>{reinterpret_cast<T>(*tensors.back().get())};
565 auto next = CreateTuple<ith_input, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep);
566 return std::tuple_cat(current, next);
567 }
568
569 template <size_t ith_input, size_t ith_output, typename T, typename... Ts>
570 static typename std::enable_if<std::is_same<T, CUDAKernelContext&>::value, std::tuple<T, Ts...>>::type
571 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) {
572 tensors.push_back(std::make_unique<OrtGraphKernelContext>(api->GetOrtApi(), *context));
573 std::tuple<T> current = std::tuple<T>{reinterpret_cast<T>(*tensors.back().get())};
574 auto next = CreateTuple<ith_input, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep);
575 return std::tuple_cat(current, next);
576 }
577
578#if ORT_API_VERSION >= 14
579 template <size_t ith_input, size_t ith_output, typename T, typename... Ts>
580 static typename std::enable_if<std::is_same<T, const Variadic*>::value, std::tuple<T, Ts...>>::type
581 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) {
582 tensors.push_back(std::make_unique<Variadic>(*api, *context, ith_input, true));
583 std::tuple<T> current = std::tuple<T>{reinterpret_cast<T>(tensors.back().get())};
584 auto next = CreateTuple<ith_input + 1, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep);
585 return std::tuple_cat(current, next);
586 }
587
588 template <size_t ith_input, size_t ith_output, typename T, typename... Ts>
589 static typename std::enable_if<std::is_same<T, const Variadic&>::value, std::tuple<T, Ts...>>::type
590 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) {
591 tensors.push_back(std::make_unique<Variadic>(*api, *context, ith_input, true));
592 std::tuple<T> current = std::tuple<T>{reinterpret_cast<T>(*tensors.back().get())};
593 auto next = CreateTuple<ith_input + 1, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep);
594 return std::tuple_cat(current, next);
595 }
596
597 template <size_t ith_input, size_t ith_output, typename T, typename... Ts>
598 static typename std::enable_if<std::is_same<T, Variadic*>::value, std::tuple<T, Ts...>>::type
599 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) {
600 tensors.push_back(std::make_unique<Variadic>(*api, *context, ith_output, false));
601 std::tuple<T> current = std::tuple<T>{reinterpret_cast<T>(tensors.back().get())};
602 auto next = CreateTuple<ith_input, ith_output + 1, Ts...>(api, context, tensors, num_input, num_output, ep);
603 return std::tuple_cat(current, next);
604 }
605
606 template <size_t ith_input, size_t ith_output, typename T, typename... Ts>
607 static typename std::enable_if<std::is_same<T, Variadic&>::value, std::tuple<T, Ts...>>::type
608 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) {
609 tensors.push_back(std::make_unique<Variadic>(*api, *context, ith_output, false));
610 std::tuple<T> current = std::tuple<T>{reinterpret_cast<T>(*tensors.back().get())};
611 auto next = CreateTuple<ith_input, ith_output + 1, Ts...>(api, context, tensors, num_input, num_output, ep);
612 return std::tuple_cat(current, next);
613 }
614#endif
615
616#define CREATE_TUPLE_INPUT(data_type) \
617 template <size_t ith_input, size_t ith_output, typename T, typename... Ts> \
618 static typename std::enable_if<std::is_same<T, const Custom::Tensor<data_type>*>::value, std::tuple<T, Ts...>>::type \
619 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) { \
620 tensors.push_back(std::make_unique<Custom::OrtTensor<data_type>>(*api, *context, ith_input, true)); \
621 std::tuple<T> current = std::tuple<T>{reinterpret_cast<T>(tensors.back().get())}; \
622 auto next = CreateTuple<ith_input + 1, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep); \
623 return std::tuple_cat(current, next); \
624 } \
625 template <size_t ith_input, size_t ith_output, typename T, typename... Ts> \
626 static typename std::enable_if<std::is_same<T, const Custom::Tensor<data_type>&>::value, std::tuple<T, Ts...>>::type \
627 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) { \
628 tensors.push_back(std::make_unique<Custom::OrtTensor<data_type>>(*api, *context, ith_input, true)); \
629 std::tuple<T> current = std::tuple<T>{reinterpret_cast<T>(*tensors.back().get())}; \
630 auto next = CreateTuple<ith_input + 1, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep); \
631 return std::tuple_cat(current, next); \
632 } \
633 template <size_t ith_input, size_t ith_output, typename T, typename... Ts> \
634 static typename std::enable_if<std::is_same<T, std::optional<const Custom::Tensor<data_type>*>>::value, std::tuple<T, Ts...>>::type \
635 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) { \
636 if (ith_input < num_input) { \
637 tensors.push_back(std::make_unique<Custom::OrtTensor<data_type>>(*api, *context, ith_input, true)); \
638 std::tuple<T> current = std::tuple<T>{reinterpret_cast<Custom::Tensor<data_type>*>(tensors.back().get())}; \
639 auto next = CreateTuple<ith_input + 1, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep); \
640 return std::tuple_cat(current, next); \
641 } else { \
642 std::tuple<T> current = std::tuple<T>{}; \
643 auto next = CreateTuple<ith_input + 1, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep); \
644 return std::tuple_cat(current, next); \
645 } \
646 } \
647 template <size_t ith_input, size_t ith_output, typename T, typename... Ts> \
648 static typename std::enable_if<std::is_same<T, const Custom::Span<data_type>*>::value, std::tuple<T, Ts...>>::type \
649 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) { \
650 tensors.push_back(std::make_unique<Custom::OrtTensor<data_type>>(*api, *context, ith_input, true)); \
651 if (!reinterpret_cast<Custom::OrtTensor<data_type>*>(tensors.back().get())->IsCpuTensor()) { \
652 ORTX_CXX_API_THROW("span input could only be applied to CPU tensor", ORT_FAIL); \
653 } \
654 std::tuple<T> current = std::tuple<T>{&reinterpret_cast<Custom::Tensor<data_type>*>(tensors.back().get())->AsSpan()}; \
655 auto next = CreateTuple<ith_input + 1, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep); \
656 return std::tuple_cat(current, next); \
657 } \
658 template <size_t ith_input, size_t ith_output, typename T, typename... Ts> \
659 static typename std::enable_if<std::is_same<T, const Custom::Span<data_type>&>::value, std::tuple<T, Ts...>>::type \
660 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) { \
661 tensors.push_back(std::make_unique<Custom::OrtTensor<data_type>>(*api, *context, ith_input, true)); \
662 if (!reinterpret_cast<Custom::OrtTensor<data_type>*>(tensors.back().get())->IsCpuTensor()) { \
663 ORTX_CXX_API_THROW("span input could only be applied to CPU tensor", ORT_FAIL); \
664 } \
665 std::tuple<T> current = std::tuple<T>{reinterpret_cast<Custom::Tensor<data_type>*>(tensors.back().get())->AsSpan()}; \
666 auto next = CreateTuple<ith_input + 1, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep); \
667 return std::tuple_cat(current, next); \
668 } \
669 template <size_t ith_input, size_t ith_output, typename T, typename... Ts> \
670 static typename std::enable_if<std::is_same<T, std::optional<const Custom::Span<data_type>*>>::value, std::tuple<T, Ts...>>::type \
671 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) { \
672 if (ith_input < num_input) { \
673 tensors.push_back(std::make_unique<Custom::OrtTensor<data_type>>(*api, *context, ith_input, true)); \
674 if (!reinterpret_cast<Custom::OrtTensor<data_type>*>(tensors.back().get())->IsCpuTensor()) { \
675 ORTX_CXX_API_THROW("span input could only be applied to CPU tensor", ORT_FAIL); \
676 } \
677 std::tuple<T> current = std::tuple<T>{&reinterpret_cast<Custom::Tensor<data_type>*>(tensors.back().get())->AsSpan()}; \
678 auto next = CreateTuple<ith_input + 1, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep); \
679 return std::tuple_cat(current, next); \
680 } else { \
681 std::tuple<T> current = std::tuple<T>{}; \
682 auto next = CreateTuple<ith_input + 1, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep); \
683 return std::tuple_cat(current, next); \
684 } \
685 } \
686 template <size_t ith_input, size_t ith_output, typename T, typename... Ts> \
687 static typename std::enable_if<std::is_same<T, data_type>::value, std::tuple<T, Ts...>>::type \
688 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) { \
689 tensors.push_back(std::make_unique<Custom::OrtTensor<data_type>>(*api, *context, ith_input, true)); \
690 if (!reinterpret_cast<Custom::OrtTensor<data_type>*>(tensors.back().get())->IsCpuTensor()) { \
691 ORTX_CXX_API_THROW("scalar input could only be applied to CPU tensor", ORT_FAIL); \
692 } \
693 std::tuple<T> current = std::tuple<T>{reinterpret_cast<Custom::Tensor<data_type>*>(tensors.back().get())->AsScalar()}; \
694 auto next = CreateTuple<ith_input + 1, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep); \
695 return std::tuple_cat(current, next); \
696 } \
697 template <size_t ith_input, size_t ith_output, typename T, typename... Ts> \
698 static typename std::enable_if<std::is_same<T, std::optional<data_type>>::value, std::tuple<T, Ts...>>::type \
699 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) { \
700 if (ith_input < num_input) { \
701 tensors.push_back(std::make_unique<Custom::OrtTensor<data_type>>(*api, *context, ith_input, true)); \
702 if (!reinterpret_cast<Custom::OrtTensor<data_type>*>(tensors.back().get())->IsCpuTensor()) { \
703 ORTX_CXX_API_THROW("scalar input could only be applied to CPU tensor", ORT_FAIL); \
704 } \
705 std::tuple<T> current = std::tuple<T>{reinterpret_cast<Custom::Tensor<data_type>*>(tensors.back().get())->AsScalar()}; \
706 auto next = CreateTuple<ith_input + 1, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep); \
707 return std::tuple_cat(current, next); \
708 } else { \
709 std::tuple<T> current = std::tuple<T>{}; \
710 auto next = CreateTuple<ith_input + 1, ith_output, Ts...>(api, context, tensors, num_input, num_output, ep); \
711 return std::tuple_cat(current, next); \
712 } \
713 }
714#define CREATE_TUPLE_OUTPUT(data_type) \
715 template <size_t ith_input, size_t ith_output, typename T, typename... Ts> \
716 static typename std::enable_if<std::is_same<T, Custom::Tensor<data_type>*>::value, std::tuple<T, Ts...>>::type \
717 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) { \
718 tensors.push_back(std::make_unique<Custom::OrtTensor<data_type>>(*api, *context, ith_output, false)); \
719 std::tuple<T> current = std::tuple<T>{reinterpret_cast<T>(tensors.back().get())}; \
720 auto next = CreateTuple<ith_input, ith_output + 1, Ts...>(api, context, tensors, num_input, num_output, ep); \
721 return std::tuple_cat(current, next); \
722 } \
723 template <size_t ith_input, size_t ith_output, typename T, typename... Ts> \
724 static typename std::enable_if<std::is_same<T, Custom::Tensor<data_type>&>::value, std::tuple<T, Ts...>>::type \
725 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) { \
726 tensors.push_back(std::make_unique<Custom::OrtTensor<data_type>>(*api, *context, ith_output, false)); \
727 std::tuple<T> current = std::tuple<T>{reinterpret_cast<T>(*tensors.back().get())}; \
728 auto next = CreateTuple<ith_input, ith_output + 1, Ts...>(api, context, tensors, num_input, num_output, ep); \
729 return std::tuple_cat(current, next); \
730 } \
731 template <size_t ith_input, size_t ith_output, typename T, typename... Ts> \
732 static typename std::enable_if<std::is_same<T, std::optional<Custom::Tensor<data_type>*>>::value, std::tuple<T, Ts...>>::type \
733 CreateTuple(const OrtW::CustomOpApi* api, OrtKernelContext* context, std::vector<TensorPtr>& tensors, size_t num_input, size_t num_output, const std::string& ep) { \
734 if (ith_output < num_output) { \
735 tensors.push_back(std::make_unique<Custom::OrtTensor<data_type>>(*api, *context, ith_output, false)); \
736 std::tuple<T> current = std::tuple<T>{reinterpret_cast<Custom::Tensor<data_type>*>(tensors.back().get())}; \
737 auto next = CreateTuple<ith_input, ith_output + 1, Ts...>(api, context, tensors, num_input, num_output, ep); \
738 return std::tuple_cat(current, next); \
739 } else { \
740 std::tuple<T> current = std::tuple<T>{}; \
741 auto next = CreateTuple<ith_input, ith_output + 1, Ts...>(api, context, tensors, num_input, num_output, ep); \
742 return std::tuple_cat(current, next); \
743 } \
744 }
745#define CREATE_TUPLE(data_type) \
746 CREATE_TUPLE_INPUT(data_type) \
747 CREATE_TUPLE_OUTPUT(data_type)
748
749 CREATE_TUPLE(bool)
750#if ORT_API_VERSION >= 16
751 CREATE_TUPLE(MFloat16)
752 CREATE_TUPLE(BFloat16)
753#endif
754 CREATE_TUPLE(float)
755 CREATE_TUPLE(double)
756 CREATE_TUPLE(int8_t)
757 CREATE_TUPLE(int16_t)
758 CREATE_TUPLE(int32_t)
759 CREATE_TUPLE(int64_t)
760 CREATE_TUPLE(uint8_t)
761 CREATE_TUPLE(uint16_t)
762 CREATE_TUPLE(uint32_t)
763 CREATE_TUPLE(uint64_t)
764 CREATE_TUPLE(std::string)
765 CREATE_TUPLE_INPUT(std::string_view)
766
767 // ParseArgs ...
768 template <typename... Ts>
769 static typename std::enable_if<0 == sizeof...(Ts)>::type
770 ParseArgs(std::vector<ONNXTensorElementDataType>&, std::vector<ONNXTensorElementDataType>&) {
771 }
772
773 template <typename T, typename... Ts>
774 static typename std::enable_if<0 <= sizeof...(Ts) && std::is_same<T, OrtKernelContext*>::value>::type
775 ParseArgs(std::vector<ONNXTensorElementDataType>& input_types, std::vector<ONNXTensorElementDataType>& output_types) {
776 ParseArgs<Ts...>(input_types, output_types);
777 }
778
779#ifdef USE_CUDA
780 template <typename T, typename... Ts>
781 static typename std::enable_if<0 <= sizeof...(Ts) && std::is_same<T, const CudaContext&>::value>::type
782 ParseArgs(std::vector<ONNXTensorElementDataType>& input_types, std::vector<ONNXTensorElementDataType>& output_types) {
783 ParseArgs<Ts...>(input_types, output_types);
784 }
785#endif
786
787 template <typename T, typename... Ts>
788 static typename std::enable_if<0 <= sizeof...(Ts) && std::is_same<T, KernelContext&>::value>::type
789 ParseArgs(std::vector<ONNXTensorElementDataType>& input_types, std::vector<ONNXTensorElementDataType>& output_types) {
790 ParseArgs<Ts...>(input_types, output_types);
791 }
792
793 template <typename T, typename... Ts>
794 static typename std::enable_if<0 <= sizeof...(Ts) && std::is_same<T, CUDAKernelContext&>::value>::type
795 ParseArgs(std::vector<ONNXTensorElementDataType>& input_types, std::vector<ONNXTensorElementDataType>& output_types) {
796 ParseArgs<Ts...>(input_types, output_types);
797 }
798
799#if ORT_API_VERSION >= 14
800 template <typename T, typename... Ts>
801 static typename std::enable_if<0 <= sizeof...(Ts) && std::is_same<T, const Variadic&>::value>::type
802 ParseArgs(std::vector<ONNXTensorElementDataType>& input_types, std::vector<ONNXTensorElementDataType>& output_types) {
803 if (!input_types.empty()) {
804 ORTX_CXX_API_THROW("for op has variadic input, only one input is allowed", ORT_RUNTIME_EXCEPTION);
805 }
806 input_types.push_back(ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED);
807 ParseArgs<Ts...>(input_types, output_types);
808 }
809
810 template <typename T, typename... Ts>
811 static typename std::enable_if<0 <= sizeof...(Ts) && std::is_same<T, const Variadic*>::value>::type
812 ParseArgs(std::vector<ONNXTensorElementDataType>& input_types, std::vector<ONNXTensorElementDataType>& output_types) {
813 if (!input_types.empty()) {
814 ORTX_CXX_API_THROW("for op has variadic input, only one input is allowed", ORT_RUNTIME_EXCEPTION);
815 }
816 input_types.push_back(ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED);
817 ParseArgs<Ts...>(input_types, output_types);
818 }
819
820 template <typename T, typename... Ts>
821 static typename std::enable_if<0 <= sizeof...(Ts) && std::is_same<T, Variadic&>::value>::type
822 ParseArgs(std::vector<ONNXTensorElementDataType>& input_types, std::vector<ONNXTensorElementDataType>& output_types) {
823 if (!output_types.empty()) {
824 ORTX_CXX_API_THROW("for op has variadic output, only one output is allowed", ORT_RUNTIME_EXCEPTION);
825 }
826 output_types.push_back(ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED);
827 ParseArgs<Ts...>(input_types, output_types);
828 }
829
830 template <typename T, typename... Ts>
831 static typename std::enable_if<0 <= sizeof...(Ts) && std::is_same<T, Variadic*>::value>::type
832 ParseArgs(std::vector<ONNXTensorElementDataType>& input_types, std::vector<ONNXTensorElementDataType>& output_types) {
833 if (!output_types.empty()) {
834 ORTX_CXX_API_THROW("for op has variadic output, only one output is allowed", ORT_RUNTIME_EXCEPTION);
835 }
836 output_types.push_back(ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED);
837 ParseArgs<Ts...>(input_types, output_types);
838 }
839#endif
840
841#define PARSE_INPUT_BASE(pack_type, onnx_type) \
842 template <typename T, typename... Ts> \
843 static typename std::enable_if<0 <= sizeof...(Ts) && std::is_same<T, pack_type>::value>::type \
844 ParseArgs(std::vector<ONNXTensorElementDataType>& input_types, std::vector<ONNXTensorElementDataType>& output_types) { \
845 input_types.push_back(onnx_type); \
846 ParseArgs<Ts...>(input_types, output_types); \
847 } \
848 template <typename T, typename... Ts> \
849 static typename std::enable_if<0 <= sizeof...(Ts) && std::is_same<T, std::optional<pack_type>>::value>::type \
850 ParseArgs(std::vector<ONNXTensorElementDataType>& input_types, std::vector<ONNXTensorElementDataType>& output_types) { \
851 input_types.push_back(onnx_type); \
852 ParseArgs<Ts...>(input_types, output_types); \
853 }
854
855#define PARSE_INPUT(data_type, onnx_type) \
856 PARSE_INPUT_BASE(const Custom::Tensor<data_type>*, onnx_type) \
857 PARSE_INPUT_BASE(const Custom::Tensor<data_type>&, onnx_type) \
858 PARSE_INPUT_BASE(const Custom::Span<data_type>*, onnx_type) \
859 PARSE_INPUT_BASE(const Custom::Span<data_type>&, onnx_type) \
860 PARSE_INPUT_BASE(data_type, onnx_type)
861
862#define PARSE_OUTPUT(data_type, onnx_type) \
863 template <typename T, typename... Ts> \
864 static typename std::enable_if<0 <= sizeof...(Ts) && std::is_same<T, Custom::Tensor<data_type>*>::value>::type \
865 ParseArgs(std::vector<ONNXTensorElementDataType>& input_types, std::vector<ONNXTensorElementDataType>& output_types) { \
866 output_types.push_back(onnx_type); \
867 ParseArgs<Ts...>(input_types, output_types); \
868 } \
869 template <typename T, typename... Ts> \
870 static typename std::enable_if<0 <= sizeof...(Ts) && std::is_same<T, Custom::Tensor<data_type>&>::value>::type \
871 ParseArgs(std::vector<ONNXTensorElementDataType>& input_types, std::vector<ONNXTensorElementDataType>& output_types) { \
872 output_types.push_back(onnx_type); \
873 ParseArgs<Ts...>(input_types, output_types); \
874 } \
875 template <typename T, typename... Ts> \
876 static typename std::enable_if<0 <= sizeof...(Ts) && std::is_same<T, std::optional<Custom::Tensor<data_type>*>>::value>::type \
877 ParseArgs(std::vector<ONNXTensorElementDataType>& input_types, std::vector<ONNXTensorElementDataType>& output_types) { \
878 output_types.push_back(onnx_type); \
879 ParseArgs<Ts...>(input_types, output_types); \
880 }
881
882#define PARSE_ARGS(data_type, onnx_type) \
883 PARSE_INPUT(data_type, onnx_type) \
884 PARSE_OUTPUT(data_type, onnx_type)
885
886 PARSE_ARGS(bool, ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL)
887#if ORT_API_VERSION >= 16
888 PARSE_ARGS(MFloat16, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16)
889 PARSE_ARGS(BFloat16, ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16)
890#endif
891 PARSE_ARGS(float, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT)
892 PARSE_ARGS(double, ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE)
893 PARSE_ARGS(int8_t, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8)
894 PARSE_ARGS(int16_t, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16)
895 PARSE_ARGS(int32_t, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32)
896 PARSE_ARGS(int64_t, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64)
897 PARSE_ARGS(uint8_t, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8)
898 PARSE_ARGS(uint16_t, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16)
899 PARSE_ARGS(uint32_t, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32)
900 PARSE_ARGS(uint64_t, ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64)
901 PARSE_ARGS(std::string, ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING)
902 PARSE_ARGS(std::string_view, ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING) // todo - remove string_view output
903
904 OrtLiteCustomOp(const char* op_name,
905 const char* execution_provider) : op_name_(op_name),
906 execution_provider_(execution_provider) {
907 // Zero out OrtCustomOp so that any added func pointers are nullptr for forwards compatibility
908 memset(&this->version, 0, sizeof(OrtCustomOp));
909
910 int act_ver = GetActiveOrtAPIVersion();
911 OrtCustomOp::version = act_ver < ORT_API_VERSION ? act_ver : ORT_API_VERSION;
912
913 OrtCustomOp::GetName = [](const OrtCustomOp* op) { return static_cast<const OrtLiteCustomOp*>(op)->op_name_.c_str(); };
914 OrtCustomOp::GetExecutionProviderType = [](const OrtCustomOp* op) { return ((OrtLiteCustomOp*)op)->execution_provider_.c_str(); };
915
916 OrtCustomOp::GetInputTypeCount = [](const OrtCustomOp* op) {
917 auto self = reinterpret_cast<const OrtLiteCustomOp*>(op);
918 return self->input_types_.size();
919 };
920
921 OrtCustomOp::GetInputType = [](const OrtCustomOp* op, size_t indice) {
922 auto self = reinterpret_cast<const OrtLiteCustomOp*>(op);
923 return self->input_types_[indice];
924 };
925
926 OrtCustomOp::GetOutputTypeCount = [](const OrtCustomOp* op) {
927 auto self = reinterpret_cast<const OrtLiteCustomOp*>(op);
928 return self->output_types_.size();
929 };
930
931 OrtCustomOp::GetOutputType = [](const OrtCustomOp* op, size_t indice) {
932 auto self = reinterpret_cast<const OrtLiteCustomOp*>(op);
933 return self->output_types_[indice];
934 };
935
936#if ORT_API_VERSION >= 14
937 OrtCustomOp::GetInputCharacteristic = [](const OrtCustomOp* op, size_t) {
938 auto self = reinterpret_cast<const OrtLiteCustomOp*>(op);
939 return (self->input_types_.empty() || self->input_types_[0] != ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED) ? INPUT_OUTPUT_OPTIONAL : INPUT_OUTPUT_VARIADIC;
940 };
941
942 OrtCustomOp::GetOutputCharacteristic = [](const OrtCustomOp* op, size_t) {
943 auto self = reinterpret_cast<const OrtLiteCustomOp*>(op);
944 return (self->output_types_.empty() || self->output_types_[0] != ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED) ? INPUT_OUTPUT_OPTIONAL : INPUT_OUTPUT_VARIADIC;
945 };
946
947 OrtCustomOp::GetVariadicInputMinArity = [](const OrtCustomOp*) {
948 return 1;
949 };
950
951 OrtCustomOp::GetVariadicInputHomogeneity = [](const OrtCustomOp*) {
952 return 0;
953 };
954
955 OrtCustomOp::GetVariadicOutputMinArity = [](const OrtCustomOp*) {
956 return 1;
957 };
958
959 OrtCustomOp::GetVariadicOutputHomogeneity = [](const OrtCustomOp*) {
960 return 0;
961 };
962
963 OrtCustomOp::GetInputMemoryType = [](const OrtCustomOp*, size_t) {
964 return OrtMemTypeDefault;
965 };
966#else
967 OrtCustomOp::GetInputCharacteristic = [](const OrtCustomOp*, size_t) {
968 return INPUT_OUTPUT_OPTIONAL;
969 };
970
971 OrtCustomOp::GetOutputCharacteristic = [](const OrtCustomOp* op, size_t) {
972 return INPUT_OUTPUT_OPTIONAL;
973 };
974#endif
975 }
976
977 const std::string op_name_;
978 const std::string execution_provider_;
979
980 std::vector<ONNXTensorElementDataType> input_types_;
981 std::vector<ONNXTensorElementDataType> output_types_;
982};
983
984template <typename... Args>
985struct OrtLiteCustomFunc : public OrtLiteCustomOp {
986 using ComputeFn = void (*)(Args...);
987 using MyType = OrtLiteCustomFunc<Args...>;
988
989 struct Kernel {
990 ComputeFn compute_fn_{};
991 std::string ep_{};
992 std::unique_ptr<OrtW::CustomOpApi> api_;
993 };
994
995 OrtLiteCustomFunc(const char* op_name,
996 const char* execution_provider,
997 ComputeFn compute_fn) : OrtLiteCustomOp(op_name, execution_provider),
998 compute_fn_(compute_fn) {
999 ParseArgs<Args...>(input_types_, output_types_);
1000
1001 OrtCustomOp::KernelCompute = [](void* op_kernel, OrtKernelContext* context) {
1002 auto kernel = reinterpret_cast<Kernel*>(op_kernel);
1003 std::vector<TensorPtr> tensors;
1004 auto t = CreateTuple<0, 0, Args...>(kernel->api_.get(),
1005 context,
1006 tensors,
1007 kernel->api_->KernelContext_GetInputCount(context),
1008 kernel->api_->KernelContext_GetOutputCount(context),
1009 kernel->ep_);
1010 std::apply([kernel](Args const&... t_args) { kernel->compute_fn_(t_args...); }, t);
1011 };
1012
1013 OrtCustomOp::CreateKernel = [](const OrtCustomOp* this_, const OrtApi* ort_api, const OrtKernelInfo* info) {
1014 auto kernel = std::make_unique<Kernel>();
1015 auto self = static_cast<const OrtLiteCustomFunc*>(this_);
1016 kernel->compute_fn_ = self->compute_fn_;
1017 kernel->ep_ = self->execution_provider_;
1018 kernel->api_ = std::make_unique<OrtW::CustomOpApi>(*ort_api);
1019 return reinterpret_cast<void*>(kernel.release());
1020 };
1021
1022 OrtCustomOp::KernelDestroy = [](void* op_kernel) {
1023 delete reinterpret_cast<Kernel*>(op_kernel);
1024 };
1025 }
1026
1027 ComputeFn compute_fn_;
1028};
1029
1030class OrtAttributeReader {
1031public:
1032 OrtAttributeReader(const OrtApi& api, const OrtKernelInfo& info) : base_kernel_(api, info) {
1033 }
1034
1035 template <class T>
1036 T TryToGetAttributeWithDefault(const char* name, const T& default_value) const noexcept {
1037 return base_kernel_.TryToGetAttributeWithDefault(name, default_value);
1038 }
1039
1040private:
1041 BaseKernel base_kernel_;
1042};
1043
1044template <typename CustomOp>
1045struct OrtLiteCustomStruct : public OrtLiteCustomOp {
1046 template <typename... Args>
1047 using CustomComputeFn = void (CustomOp::*)(Args...) const;
1048 using MyType = OrtLiteCustomStruct<CustomOp>;
1049
1050 struct Kernel {
1051 std::unique_ptr<CustomOp> custom_op_;
1052 std::string ep_{};
1053 std::unique_ptr<OrtW::CustomOpApi> api_;
1054 };
1055
1056 OrtLiteCustomStruct(const char* op_name,
1057 const char* execution_provider) : OrtLiteCustomOp(op_name,
1058 execution_provider) {
1059 init(&CustomOp::Compute);
1060 }
1061
1062 template <typename... Args>
1063 void init(CustomComputeFn<Args...>) {
1064 ParseArgs<Args...>(input_types_, output_types_);
1065
1066 OrtCustomOp::KernelCompute = [](void* op_kernel, OrtKernelContext* context) {
1067 auto kernel = reinterpret_cast<Kernel*>(op_kernel);
1068 std::vector<TensorPtr> tensors;
1069 auto t = CreateTuple<0, 0, Args...>(kernel->api_.get(),
1070 context,
1071 tensors,
1072 kernel->api_->KernelContext_GetInputCount(context),
1073 kernel->api_->KernelContext_GetOutputCount(context),
1074 kernel->ep_);
1075 std::apply([kernel](Args const&... t_args) { kernel->custom_op_->Compute(t_args...); }, t);
1076 };
1077
1078 OrtCustomOp::CreateKernel = [](const OrtCustomOp* this_, const OrtApi* ort_api, const OrtKernelInfo* info) {
1079 auto kernel = std::make_unique<Kernel>();
1080
1081 if constexpr (std::is_constructible<CustomOp, const OrtApi&, const OrtKernelInfo&>::value){
1082 kernel->custom_op_ = std::make_unique<CustomOp>(*ort_api, *info);
1083 }
1084 else {
1085 kernel->custom_op_ = std::make_unique<CustomOp>(OrtAttributeReader(*ort_api, *info));
1086 }
1087 auto self = static_cast<const MyType*>(this_);
1088 kernel->ep_ = self->execution_provider_;
1089 kernel->api_ = std::make_unique<OrtW::CustomOpApi>(*ort_api);
1090 return reinterpret_cast<void*>(kernel.release());
1091 };
1092
1093 OrtCustomOp::KernelDestroy = [](void* op_kernel) {
1094 delete reinterpret_cast<Kernel*>(op_kernel);
1095 };
1096 }
1097};
1098
1099template <typename... Args>
1100OrtLiteCustomOp* CreateLiteCustomOp(const char* op_name,
1101 const char* execution_provider,
1102 void (*custom_compute_fn)(Args...)) {
1103 using LiteOp = OrtLiteCustomFunc<Args...>;
1104 return std::make_unique<LiteOp>(op_name, execution_provider, custom_compute_fn).release();
1105}
1106
1107template <typename CustomOp>
1108OrtLiteCustomOp* CreateLiteCustomOp(const char* op_name,
1109 const char* execution_provider) {
1110 using LiteOp = OrtLiteCustomStruct<CustomOp>;
1111 return std::make_unique<LiteOp>(op_name, execution_provider).release();
1112}
1113
1114} // namespace Custom
1115} // namespace Ort
1116