GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/onnxruntime-extensions

Public

mirrored fromhttps://github.com/microsoft/onnxruntime-extensionsAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
edgchen1/update_xcode_fix

Branches

Tags

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

Clone

HTTPS

Download ZIP

onnxruntime-extensions/includes

includes/ocos.h

173lines · 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#define _SILENCE_CXX17_CODECVT_HEADER_DEPRECATION_WARNING
7#include <string>
8#include <algorithm>
9#include <functional>
10#include <iterator>
11#include <vector>
12
13#include "onnxruntime_customop.hpp"
14
15// A helper API to support test kernels.
16// Must be invoked before RegisterCustomOps.
17extern "C" bool ORT_API_CALL AddExternalCustomOp(const OrtCustomOp* c_op);
18
19constexpr const char* c_OpDomain = "ai.onnx.contrib";
20constexpr const char* c_ComMsExtOpDomain = "com.microsoft.extensions";
21
22struct BaseKernel {
23 BaseKernel(const OrtApi& api, const OrtKernelInfo& info) noexcept : api_(api), info_(info), ort_(api_) {
24 }
25
26 template <class T>
27 bool TryToGetAttribute(const char* name, T& value) const noexcept;
28
29 template <class T>
30 T TryToGetAttributeWithDefault(const char* name, const T& default_value) const noexcept {
31 T result = default_value;
32 TryToGetAttribute(name, result);
33 return result;
34 }
35
36 void SetOutput(OrtKernelContext* ctx, size_t output_idx, const std::vector<int64_t>& dim,
37 const std::vector<int64_t>& data);
38
39 protected:
40 OrtErrorCode GetErrorCodeAndRelease(OrtStatusPtr status) const noexcept;
41 const OrtApi& api_;
42 OrtW::CustomOpApi ort_;
43 const OrtKernelInfo& info_;
44};
45
46struct OrtTensorDimensions : std::vector<int64_t> {
47 OrtTensorDimensions() = default;
48 OrtTensorDimensions(const OrtW::CustomOpApi& ort, const OrtValue* value) {
49 OrtTensorTypeAndShapeInfo* info = ort.GetTensorTypeAndShape(value);
50 std::vector<int64_t>::operator=(ort.GetTensorShape(info));
51 ort.ReleaseTensorTypeAndShapeInfo(info);
52 }
53
54 int64_t Size() const {
55 int64_t s = 1;
56 for (auto it = begin(); it != end(); ++it)
57 s *= *it;
58 return s;
59 }
60
61 bool IsScalar() const {
62 return empty();
63 }
64
65 bool IsVector() const {
66 return size() == 1;
67 }
68};
69
70template <typename... Args>
71class CuopContainer {
72 public:
73 CuopContainer() : op_instances_({[]() { return std::make_shared<Args>(); }()...}) {
74 ocos_list_.reserve(op_instances_.size());
75 std::transform(op_instances_.begin(), op_instances_.end(), std::back_inserter(ocos_list_),
76 [](const std::shared_ptr<OrtCustomOp>& custom_op) { return custom_op.get(); });
77 }
78
79 const std::vector<const OrtCustomOp*>& GetCustomOps() const {
80 return ocos_list_;
81 }
82
83 private:
84 std::vector<const OrtCustomOp*> ocos_list_;
85 std::vector<std::shared_ptr<OrtCustomOp>> op_instances_; // use shared_ptr to capture type specific deleter
86};
87
88#define CustomCpuFunc(name, f) []() { return std::shared_ptr<ortc::OrtLiteCustomOp>(ortc::CreateLiteCustomOp(name, "CPUExecutionProvider", f)); }
89#define CustomCpuStruct(name, s) []() { return std::shared_ptr<ortc::OrtLiteCustomOp>(ortc::CreateLiteCustomOp<s>(name, "CPUExecutionProvider")); }
90
91template <typename F>
92void AppendCustomOp(std::vector<std::shared_ptr<OrtCustomOp>>& ops,
93 F arg) {
94 ops.emplace_back(std::move(arg()));
95}
96
97template <typename T, typename... Args>
98void AppendCustomOp(std::vector<std::shared_ptr<OrtCustomOp>>& ops,
99 T arg, Args... args) {
100 AppendCustomOp(ops, arg);
101 AppendCustomOp(ops, args...);
102}
103
104class OrtOpLoader {
105 public:
106 template <typename... Args>
107 OrtOpLoader(Args... args) {
108 LoadOps(args...);
109 for (auto& ptr : op_instances_) {
110 if (ptr)
111 ocos_list_.push_back(ptr.get());
112 }
113 }
114
115 const std::vector<const OrtCustomOp*>& GetCustomOps() const {
116 return ocos_list_;
117 }
118
119 private:
120 template <typename T>
121 void LoadOps(T fn) {
122 AppendCustomOp(op_instances_, fn);
123 }
124
125 template <typename T, typename... Args>
126 void LoadOps(T fn, Args... args) {
127 AppendCustomOp(op_instances_, fn);
128 AppendCustomOp(op_instances_, args...);
129 }
130
131 std::vector<const OrtCustomOp*> ocos_list_;
132 std::vector<std::shared_ptr<OrtCustomOp>> op_instances_;
133};
134
135struct CustomOpClassBegin {
136};
137
138using FxLoadCustomOpFactory = std::function<const std::vector<const OrtCustomOp*>&()>;
139
140template <typename _Begin_place_holder, typename... Args>
141const std::vector<const OrtCustomOp*>& LoadCustomOpClasses() {
142 static CuopContainer<Args...> ctr; // Let C++ runtime take cares of the MP initializing.
143 return ctr.GetCustomOps();
144}
145
146#if defined(PYTHON_OP_SUPPORT)
147const OrtCustomOp* FetchPyCustomOps(size_t& count);
148OrtStatusPtr RegisterPythonDomainAndOps(OrtSessionOptions*, const OrtApi*);
149#endif
150
151#ifdef ENABLE_MATH
152extern FxLoadCustomOpFactory LoadCustomOpClasses_Math;
153#endif // ENABLE_MATH
154
155#ifdef ENABLE_TOKENIZER
156extern FxLoadCustomOpFactory LoadCustomOpClasses_Tokenizer;
157#endif // ENABLE_TOKENIZER
158
159#ifdef ENABLE_TF_STRING
160extern FxLoadCustomOpFactory LoadCustomOpClasses_Text;
161#endif // ENABLE_TF_STRING
162
163#ifdef ENABLE_CV2
164extern FxLoadCustomOpFactory LoadCustomOpClasses_CV2;
165#endif // ENABLE_OPENCV
166
167#ifdef ENABLE_VISION
168extern FxLoadCustomOpFactory LoadCustomOpClasses_Vision;
169#endif
170
171#ifdef ENABLE_DR_LIBS
172extern FxLoadCustomOpFactory LoadCustomOpClasses_Audio;
173#endif
174