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
a0c26255112b7a56e48abc62046477e0c7ea6152

Branches

Tags

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

Clone

HTTPS

Download ZIP

onnxruntime-extensions/includes

includes/ocos.h

134lines · 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 <algorithm>
7#include <cassert>
8#include <functional>
9#include <iterator>
10#include <string>
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";
21template <typename... Args>
22class CuopContainer {
23 public:
24 CuopContainer() : op_instances_({[]() { return std::make_shared<Args>(); }()...}) {
25 ocos_list_.reserve(op_instances_.size());
26 std::transform(op_instances_.begin(), op_instances_.end(), std::back_inserter(ocos_list_),
27 [](const std::shared_ptr<OrtCustomOp>& custom_op) { return custom_op.get(); });
28 }
29
30 const std::vector<const OrtCustomOp*>& GetCustomOps() const {
31 return ocos_list_;
32 }
33
34 private:
35 std::vector<const OrtCustomOp*> ocos_list_;
36 std::vector<std::shared_ptr<OrtCustomOp>> op_instances_; // use shared_ptr to capture type specific deleter
37};
38
39#define CustomCpuFunc(name, f) []() { return std::shared_ptr<ortc::OrtLiteCustomOp>(ortc::CreateLiteCustomOp(name, "CPUExecutionProvider", f)); }
40#define CustomCpuStruct(name, s) []() { return std::shared_ptr<ortc::OrtLiteCustomOp>(ortc::CreateLiteCustomOp<s>(name, "CPUExecutionProvider")); }
41#define CustomAzureStruct(name, s) []() { return std::shared_ptr<ortc::OrtLiteCustomOp>(ortc::CreateLiteCustomOp<s>(name, "AzureExecutionProvider")); }
42
43#define CustomCpuFuncV2(name, f) []() { return std::shared_ptr<ortc::OrtLiteCustomOp>(ortc::CreateLiteCustomOpV2(name, "CPUExecutionProvider", f)); }
44#define CustomCpuStructV2(name, s) []() { return std::shared_ptr<ortc::OrtLiteCustomOp>(ortc::CreateLiteCustomOpV2<s>(name, "CPUExecutionProvider")); }
45
46
47template <typename F>
48void AppendCustomOp(std::vector<std::shared_ptr<OrtCustomOp>>& ops,
49 F arg) {
50 ops.emplace_back(std::move(arg()));
51}
52
53template <typename T, typename... Args>
54void AppendCustomOp(std::vector<std::shared_ptr<OrtCustomOp>>& ops,
55 T arg, Args... args) {
56 AppendCustomOp(ops, arg);
57 AppendCustomOp(ops, args...);
58}
59
60class OrtOpLoader {
61 public:
62 template <typename... Args>
63 OrtOpLoader(Args... args) {
64 LoadOps(args...);
65 for (auto& ptr : op_instances_) {
66 if (ptr)
67 ocos_list_.push_back(ptr.get());
68 }
69 }
70
71 const std::vector<const OrtCustomOp*>& GetCustomOps() const {
72 return ocos_list_;
73 }
74
75 private:
76 template <typename T>
77 void LoadOps(T fn) {
78 AppendCustomOp(op_instances_, fn);
79 }
80
81 template <typename T, typename... Args>
82 void LoadOps(T fn, Args... args) {
83 AppendCustomOp(op_instances_, fn);
84 AppendCustomOp(op_instances_, args...);
85 }
86
87 std::vector<const OrtCustomOp*> ocos_list_;
88 std::vector<std::shared_ptr<OrtCustomOp>> op_instances_;
89};
90
91struct CustomOpClassNull {
92};
93
94template <typename _Begin_place_holder = CustomOpClassNull, typename... Args>
95const std::vector<const OrtCustomOp*>& LoadCustomOpClasses() {
96 static CuopContainer<Args...> ctr; // Let C++ runtime take cares of the MP initializing.
97 return ctr.GetCustomOps();
98}
99
100using CustomOpArray = const std::vector<const OrtCustomOp*>;
101using FxLoadCustomOpFactory = std::function<CustomOpArray&()>;
102
103#if defined(PYTHON_OP_SUPPORT)
104const OrtCustomOp* FetchPyCustomOps(size_t& count);
105OrtStatusPtr RegisterPythonDomainAndOps(OrtSessionOptions*, const OrtApi*);
106#endif
107
108#ifdef ENABLE_MATH
109extern FxLoadCustomOpFactory LoadCustomOpClasses_Math;
110#endif // ENABLE_MATH
111
112#ifdef ENABLE_TOKENIZER
113extern FxLoadCustomOpFactory LoadCustomOpClasses_Tokenizer;
114#endif // ENABLE_TOKENIZER
115
116#ifdef ENABLE_TF_STRING
117extern FxLoadCustomOpFactory LoadCustomOpClasses_Text;
118#endif // ENABLE_TF_STRING
119
120#ifdef ENABLE_CV2
121extern FxLoadCustomOpFactory LoadCustomOpClasses_CV2;
122#endif // ENABLE_OPENCV
123
124#ifdef ENABLE_VISION
125extern FxLoadCustomOpFactory LoadCustomOpClasses_Vision;
126#endif
127
128#ifdef ENABLE_DR_LIBS
129extern FxLoadCustomOpFactory LoadCustomOpClasses_Audio;
130#endif
131
132#if ENABLE_AZURE
133extern FxLoadCustomOpFactory LoadCustomOpClasses_Azure;
134#endif
135