microsoft/onnxruntime-extensions
Publicmirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable
shared/ortcustomops.cc
111lines · modecode
| 1 | // Copyright (c) Microsoft Corporation. All rights reserved. |
| 2 | // Licensed under the MIT License. |
| 3 | |
| 4 | #include <set> |
| 5 | |
| 6 | #include "onnxruntime_extensions.h" |
| 7 | #include "ocos.h" |
| 8 | |
| 9 | |
| 10 | class ExternalCustomOps { |
| 11 | public: |
| 12 | ExternalCustomOps() { |
| 13 | } |
| 14 | |
| 15 | static ExternalCustomOps& instance() { |
| 16 | static ExternalCustomOps g_instance; |
| 17 | return g_instance; |
| 18 | } |
| 19 | |
| 20 | void Add(const OrtCustomOp* c_op) { |
| 21 | op_array_.push_back(c_op); |
| 22 | } |
| 23 | |
| 24 | const OrtCustomOp* GetNextOp(size_t& idx) { |
| 25 | if (idx >= op_array_.size()) { |
| 26 | return nullptr; |
| 27 | } |
| 28 | |
| 29 | return op_array_[idx++]; |
| 30 | } |
| 31 | |
| 32 | ExternalCustomOps(ExternalCustomOps const&) = delete; |
| 33 | void operator=(ExternalCustomOps const&) = delete; |
| 34 | |
| 35 | private: |
| 36 | std::vector<const OrtCustomOp*> op_array_; |
| 37 | }; |
| 38 | |
| 39 | extern "C" bool ORT_API_CALL AddExternalCustomOp(const OrtCustomOp* c_op) { |
| 40 | ExternalCustomOps::instance().Add(c_op); |
| 41 | return true; |
| 42 | } |
| 43 | |
| 44 | extern "C" OrtStatus* ORT_API_CALL RegisterCustomOps(OrtSessionOptions* options, const OrtApiBase* api) { |
| 45 | OrtCustomOpDomain* domain = nullptr; |
| 46 | const OrtApi* ortApi = api->GetApi(ORT_API_VERSION); |
| 47 | std::set<std::string> pyop_nameset; |
| 48 | OrtStatus* status = nullptr; |
| 49 | |
| 50 | #if defined(PYTHON_OP_SUPPORT) |
| 51 | if (status = RegisterPythonDomainAndOps(options, ortApi)){ |
| 52 | return status; |
| 53 | } |
| 54 | #endif // PYTHON_OP_SUPPORT |
| 55 | |
| 56 | if (status = ortApi->CreateCustomOpDomain(c_OpDomain, &domain)) { |
| 57 | return status; |
| 58 | } |
| 59 | |
| 60 | #if defined(PYTHON_OP_SUPPORT) |
| 61 | size_t count = 0; |
| 62 | const OrtCustomOp* c_ops = FetchPyCustomOps(count); |
| 63 | while (c_ops != nullptr) { |
| 64 | if (status = ortApi->CustomOpDomain_Add(domain, c_ops)) { |
| 65 | return status; |
| 66 | } else { |
| 67 | pyop_nameset.emplace(c_ops->GetName(c_ops)); |
| 68 | } |
| 69 | ++count; |
| 70 | c_ops = FetchPyCustomOps(count); |
| 71 | } |
| 72 | #endif |
| 73 | |
| 74 | static std::vector<FxLoadCustomOpFactory> c_factories = { |
| 75 | LoadCustomOpClasses<CustomOpClassBegin> |
| 76 | #if defined(ENABLE_TF_STRING) |
| 77 | , LoadCustomOpClasses_Text |
| 78 | #endif // ENABLE_TF_STRING |
| 79 | #if defined(ENABLE_MATH) |
| 80 | , LoadCustomOpClasses_Math |
| 81 | #endif |
| 82 | #if defined(ENABLE_TOKENIZER) |
| 83 | , LoadCustomOpClasses_Tokenizer |
| 84 | #endif |
| 85 | }; |
| 86 | |
| 87 | for (auto fx : c_factories) { |
| 88 | auto ops = fx(); |
| 89 | while (*ops != nullptr) { |
| 90 | if (pyop_nameset.find((*ops)->GetName(*ops)) == pyop_nameset.end()) { |
| 91 | if (status = ortApi->CustomOpDomain_Add(domain, *ops)) { |
| 92 | return status; |
| 93 | } |
| 94 | } |
| 95 | ++ops; |
| 96 | } |
| 97 | } |
| 98 | |
| 99 | size_t idx = 0; |
| 100 | const OrtCustomOp* e_ops = ExternalCustomOps::instance().GetNextOp(idx); |
| 101 | while (e_ops != nullptr) { |
| 102 | if (pyop_nameset.find(e_ops->GetName(e_ops)) == pyop_nameset.end()) { |
| 103 | if (status = ortApi->CustomOpDomain_Add(domain, e_ops)) { |
| 104 | return status; |
| 105 | } |
| 106 | e_ops = ExternalCustomOps::instance().GetNextOp(idx); |
| 107 | } |
| 108 | } |
| 109 | |
| 110 | return ortApi->AddCustomOpDomain(options, domain); |
| 111 | } |
| 112 | |