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
aad7b671cbfc0b7af5c3b35283bdb6c6742b61c7

Branches

Tags

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

Clone

HTTPS

Download ZIP

onnxruntime-extensions/include

include/ocos.h

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