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
cd5ea11aaa691097dc26c5eeba1985aec58c2afa

Branches

Tags

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

Clone

HTTPS

Download ZIP

onnxruntime-extensions/shared/lib

shared/lib/ortcustomops.cc

174lines · modecode

RawDownload
Latest commit unavailable.
1// Copyright (c) Microsoft Corporation. All rights reserved.
2// Licensed under the MIT License.
3
4#include <mutex>
5#include <set>
6
7#include "onnxruntime_extensions.h"
8#include "ocos.h"
9
10using namespace OrtW;
11
12struct OrtCustomOpDomainDeleter {
13 explicit OrtCustomOpDomainDeleter(const OrtApi* ort_api) {
14 ort_api_ = ort_api;
15 }
16 void operator()(OrtCustomOpDomain* domain) const {
17 ort_api_->ReleaseCustomOpDomain(domain);
18 }
19
20 const OrtApi* ort_api_;
21};
22
23using OrtCustomOpDomainUniquePtr = std::unique_ptr<OrtCustomOpDomain, OrtCustomOpDomainDeleter>;
24static std::vector<OrtCustomOpDomainUniquePtr> ort_custom_op_domain_container;
25static std::mutex ort_custom_op_domain_mutex;
26
27static void AddOrtCustomOpDomainToContainer(OrtCustomOpDomain* domain, const OrtApi* ort_api) {
28 std::lock_guard<std::mutex> lock(ort_custom_op_domain_mutex);
29 auto ptr = std::unique_ptr<OrtCustomOpDomain, OrtCustomOpDomainDeleter>(domain, OrtCustomOpDomainDeleter(ort_api));
30 ort_custom_op_domain_container.push_back(std::move(ptr));
31}
32
33class ExternalCustomOps {
34 public:
35 ExternalCustomOps() {
36 }
37
38 static ExternalCustomOps& instance() {
39 static ExternalCustomOps g_instance;
40 return g_instance;
41 }
42
43 void Add(const OrtCustomOp* c_op) {
44 op_array_.push_back(c_op);
45 }
46
47 const OrtCustomOp* GetNextOp(size_t& idx) {
48 if (idx >= op_array_.size()) {
49 return nullptr;
50 }
51
52 return op_array_[idx++];
53 }
54
55 ExternalCustomOps(ExternalCustomOps const&) = delete;
56 void operator=(ExternalCustomOps const&) = delete;
57
58 private:
59 std::vector<const OrtCustomOp*> op_array_;
60};
61
62extern "C" bool ORT_API_CALL AddExternalCustomOp(const OrtCustomOp* c_op) {
63 ExternalCustomOps::instance().Add(c_op);
64 return true;
65}
66
67extern "C" ORTX_EXPORT OrtStatus* ORT_API_CALL RegisterCustomOps(OrtSessionOptions* options, const OrtApiBase* api) {
68 OrtCustomOpDomain* domain = nullptr;
69 const OrtApi* ortApi = api->GetApi(ORT_API_VERSION);
70 std::set<std::string> pyop_nameset;
71 OrtStatus* status = nullptr;
72
73#if defined(PYTHON_OP_SUPPORT)
74 if (status = RegisterPythonDomainAndOps(options, ortApi); status) {
75 return status;
76 }
77#endif // PYTHON_OP_SUPPORT
78
79 if (status = ortApi->CreateCustomOpDomain(c_OpDomain, &domain); status) {
80 return status;
81 }
82
83 AddOrtCustomOpDomainToContainer(domain, ortApi);
84
85#if defined(PYTHON_OP_SUPPORT)
86 size_t count = 0;
87 const OrtCustomOp* c_ops = FetchPyCustomOps(count);
88 while (c_ops != nullptr) {
89 if (status = ortApi->CustomOpDomain_Add(domain, c_ops); status) {
90 return status;
91 } else {
92 pyop_nameset.emplace(c_ops->GetName(c_ops));
93 }
94 ++count;
95 c_ops = FetchPyCustomOps(count);
96 }
97#endif
98
99 static std::vector<FxLoadCustomOpFactory> c_factories = {
100 LoadCustomOpClasses<CustomOpClassBegin>
101#if defined(ENABLE_TF_STRING)
102 ,
103 LoadCustomOpClasses_Text
104#endif // ENABLE_TF_STRING
105#if defined(ENABLE_MATH)
106 ,
107 LoadCustomOpClasses_Math
108#endif
109#if defined(ENABLE_TOKENIZER)
110 ,
111 LoadCustomOpClasses_Tokenizer
112#endif
113#if defined(ENABLE_CV2)
114 ,
115 LoadCustomOpClasses_CV2
116#endif
117 };
118
119 for (const auto& fx : c_factories) {
120 const auto& ops = fx();
121 for (const OrtCustomOp* op : ops) {
122 if (pyop_nameset.find(op->GetName(op)) == pyop_nameset.end()) {
123 if (status = ortApi->CustomOpDomain_Add(domain, op); status) {
124 return status;
125 }
126 }
127 }
128 }
129
130 size_t idx = 0;
131 const OrtCustomOp* e_ops = ExternalCustomOps::instance().GetNextOp(idx);
132 while (e_ops != nullptr) {
133 if (pyop_nameset.find(e_ops->GetName(e_ops)) == pyop_nameset.end()) {
134 if (status = ortApi->CustomOpDomain_Add(domain, e_ops); status) {
135 return status;
136 }
137 e_ops = ExternalCustomOps::instance().GetNextOp(idx);
138 }
139 }
140
141 if (status = ortApi->AddCustomOpDomain(options, domain); status) {
142 return status;
143 }
144
145 // Create domain for ops using the new domain name.
146 if (status = ortApi->CreateCustomOpDomain(c_ComMsExtOpDomain, &domain); status) {
147 return status;
148 }
149
150 AddOrtCustomOpDomainToContainer(domain, ortApi);
151
152 static std::vector<FxLoadCustomOpFactory> new_domain_factories = {
153 LoadCustomOpClasses<CustomOpClassBegin>
154#if defined(ENABLE_VISION)
155 ,
156 LoadCustomOpClasses_Vision
157#endif
158#if defined(ENABLE_TOKENIZER)
159 ,
160 LoadCustomOpClasses_Tokenizer
161#endif
162 };
163
164 for (const auto& fx : new_domain_factories) {
165 const auto& ops = fx();
166 for (const OrtCustomOp* op : ops) {
167 if (status = ortApi->CustomOpDomain_Add(domain, op); status) {
168 return status;
169 }
170 }
171 }
172
173 return ortApi->AddCustomOpDomain(options, domain);
174}
175