microsoft/onnxruntime-extensions

Public

mirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable

CodeCommitsIssuesPull requestsActionsInsightsSecurity
leca/gqa

Branches

Tags

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

Clone

HTTPS

Download ZIP

includes/onnxruntime_no_customop.h

126lines · modecode

1// Copyright (c) Microsoft Corporation. All rights reserved.
2// Licensed under the MIT License.
3
4// This file defines API which depends on ONNXRuntime, but not including Custom Op and related facilities
5// Custom Op and related classes, functions and macros are in onnxruntime_customop.hpp
6#pragma once
7#include "exceptions.h"
8
9// namespace of ORT ABI Wrapper
10namespace OrtW {
11
12class API {
13 // To use ONNX C ABI in a way like OrtW::API::CreateStatus.
14 public:
15 static API& instance(const OrtApi* ort_api = nullptr) noexcept {
16 static API self(ort_api);
17 return self;
18 }
19
20 static OrtStatusPtr CreateStatus(OrtErrorCode code, _In_ const char* msg) noexcept {
21 return instance()->CreateStatus(code, msg);
22 }
23
24 static void ReleaseStatus(OrtStatusPtr ptr) noexcept {
25 instance()->ReleaseStatus(ptr);
26 }
27
28 template <typename T>
29 static OrtStatusPtr KernelInfoGetAttribute(const OrtKernelInfo& info, const char* name, T& value) noexcept;
30
31 static void ThrowOnError(OrtStatusPtr ptr) {
32 OrtW::ThrowOnError(instance().api_, ptr);
33 }
34
35 // Caller is responsible for releasing OrtMemoryInfo object
36 static OrtStatusPtr CreateOrtMemoryInfo(const char* name, enum OrtAllocatorType type, int id, enum OrtMemType mem_type, OrtMemoryInfo** out) noexcept {
37 return instance()->CreateMemoryInfo(name, type, id, mem_type, out);
38 }
39
40 static void ReleaseMemoryInfo(OrtMemoryInfo* mem_info) {
41 return instance()->ReleaseMemoryInfo(mem_info);
42 }
43
44#if ORT_API_VERSION >= 18
45 static OrtStatusPtr KernelContextGetScratchBuffer(const OrtKernelContext* context, const OrtMemoryInfo* mem_info, size_t count_or_bytes, void** out) {
46 return instance()->KernelContext_GetScratchBuffer(context, mem_info, count_or_bytes, out);
47 }
48#endif
49 private:
50 const OrtApi* operator->() const {
51 return &api_;
52 }
53
54 API(const OrtApi* api) : api_(*api) {
55 if (api == nullptr) {
56 ORTX_CXX_API_THROW("ort-extensions internal error: ORT-APIs used before RegisterCustomOps", ORT_RUNTIME_EXCEPTION);
57 }
58 }
59
60 const OrtApi& api_;
61};
62
63template <>
64inline OrtStatusPtr API::KernelInfoGetAttribute<int64_t>(const OrtKernelInfo& info, const char* name, int64_t& value) noexcept {
65 return instance()->KernelInfoGetAttribute_int64(&info, name, &value);
66}
67
68template <>
69inline OrtStatusPtr API::KernelInfoGetAttribute<float>(const OrtKernelInfo& info, const char* name, float& value) noexcept {
70 return instance()->KernelInfoGetAttribute_float(&info, name, &value);
71}
72
73template <>
74inline OrtStatusPtr API::KernelInfoGetAttribute<std::string>(const OrtKernelInfo& info, const char* name, std::string& value) noexcept {
75 size_t size = 0;
76 std::string out;
77 // Feed nullptr for the data buffer to query the true size of the string attribute
78 OrtStatus* status = instance()->KernelInfoGetAttribute_string(&info, name, nullptr, &size);
79 if (status == nullptr) {
80 out.resize(size);
81 status = instance()->KernelInfoGetAttribute_string(&info, name, &out[0], &size);
82 out.resize(size - 1); // remove the terminating character '\0'
83 }
84
85 if (status == nullptr) {
86 value = std::move(out);
87 }
88
89 return status;
90}
91
92template <class T>
93inline OrtStatusPtr GetOpAttribute(const OrtKernelInfo& info, const char* name, T& value) noexcept {
94 if (auto status = API::KernelInfoGetAttribute(info, name, value); status) {
95 // Ideally, we should know which kind of error code can be ignored, but it is not available now.
96 // Just ignore all of them.
97 API::ReleaseStatus(status);
98 }
99
100 return nullptr;
101}
102
103template <class T>
104inline T GetOpAttributeOrDefault(const OrtKernelInfo& info, const char* name, const T& default_value) noexcept {
105 T ret;
106 if (API::KernelInfoGetAttribute(info, name, ret)) {
107 ret = default_value;
108 }
109 return ret;
110}
111
112inline OrtStatusPtr CreateStatus(const char* msg, OrtErrorCode code) {
113 return API::CreateStatus(code, msg);
114}
115
116inline OrtStatusPtr CreateStatus(const std::string& msg, OrtErrorCode code) {
117 return API::CreateStatus(code, msg.c_str());
118}
119
120inline void ReleaseStatus(OrtStatusPtr& status) {
121 API::ReleaseStatus(status);
122 status = nullptr;
123}
124
125} // namespace OrtW
126
127