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
2ef88b0bda3208844b2bee2ea682bc76a4085c63

Branches

Tags

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

Clone

HTTPS

Download ZIP

onnxruntime-extensions/ocos/pyfunc

ocos/pyfunc/pykernel.h

87lines · 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 <vector>
7#include <map>
8#include "ocos.h"
9
10struct PyCustomOpDef {
11 std::string op_type;
12 uint64_t obj_id;
13 std::vector<int> input_types;
14 std::vector<int> output_types;
15
16 static void AddOp(const PyCustomOpDef* cod);
17
18 static const int undefined = ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
19 static const int dt_float = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; // maps to c type float
20 static const int dt_uint8 = ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8; // maps to c type uint8_t
21 static const int dt_int8 = ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8; // maps to c type int8_t
22 static const int dt_uint16 = ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16; // maps to c type uint16_t
23 static const int dt_int16 = ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16; // maps to c type int16_t
24 static const int dt_int32 = ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32; // maps to c type int32_t
25 static const int dt_int64 = ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64; // maps to c type int64_t
26 static const int dt_string = ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING; // maps to c++ type std::string
27 static const int dt_bool = ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL;
28 static const int dt_float16 = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16;
29 static const int dt_double = ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE; // maps to c type double
30 static const int dt_uint32 = ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32; // maps to c type uint32_t
31 static const int dt_uint64 = ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64; // maps to c type uint64_t
32 static const int dt_complex64 = ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64; // complex with float32 real and imaginary components
33 static const int dt_complex128 = ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128; // complex with float64 real and imaginary components
34 static const int dt_bfloat16 = ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16; // Non-IEEE floating-point format based on IEEE754 single-precision
35};
36
37struct PyCustomOpKernel {
38 PyCustomOpKernel(OrtApi api, uint64_t id)
39 : api_(api),
40 ort_(api_),
41 obj_id_(id) {
42 }
43
44 void Compute(OrtKernelContext* context);
45
46 private:
47 OrtApi api_; // keep a copy of the struct, whose ref is used in the ort_
48 Ort::CustomOpApi ort_;
49 uint64_t obj_id_;
50};
51
52struct PyCustomOpFactory : Ort::CustomOpBase<PyCustomOpFactory, PyCustomOpKernel> {
53 PyCustomOpFactory(const PyCustomOpDef* opdef) {
54 if (opdef == nullptr)
55 throw std::runtime_error("Python definition is empty.");
56 opdef_ = opdef;
57 }
58
59 void* CreateKernel(OrtApi api, const OrtKernelInfo* /* info */) {
60 return new PyCustomOpKernel(api, opdef_->obj_id);
61 };
62
63 const char* GetName() const {
64 return opdef_->op_type.c_str();
65 };
66
67 size_t GetInputTypeCount() const {
68 return opdef_->input_types.size();
69 };
70
71 ONNXTensorElementDataType GetInputType(size_t idx) const {
72 return static_cast<ONNXTensorElementDataType>(opdef_->input_types[idx]);
73 };
74
75 size_t GetOutputTypeCount() const {
76 return opdef_->output_types.size();
77 };
78
79 ONNXTensorElementDataType GetOutputType(size_t idx) const {
80 return static_cast<ONNXTensorElementDataType>(opdef_->output_types[idx]);
81 }
82
83 const PyCustomOpDef* opdef_;
84};
85
86std::vector<PyCustomOpFactory>& PyCustomOpDef_python_operator_list();
87const PyCustomOpFactory* PyCustomOpDef_FetchPyCustomOps(size_t count);
88