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
rel-0.5

Branches

Tags

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

Clone

HTTPS

Download ZIP

onnxruntime-extensions/pyop

pyop/pykernel.h

94lines · modeblame

RawDownload
Latest commit unavailable.
28b2ab10Wenbing Li5 years ago1// Copyright (c) Microsoft Corporation. All rights reserved.
2// Licensed under the MIT License.
3
4#pragma once
3ef90b46Xavier Dupré5 years ago5
28b2ab10Wenbing Li5 years ago6#include <vector>
7#include <map>
8#include "ocos.h"
9
10struct PyCustomOpDef {
11std::string op_type;
5320af1eWenbing Li3 years ago12uint64_t obj_id = 0;
28b2ab10Wenbing Li5 years ago13std::vector<int> input_types;
14std::vector<int> output_types;
a0769463Xavier Dupré5 years ago15std::vector<std::string> attrs;
28b2ab10Wenbing Li5 years ago16
e36205eeXavier Dupré5 years ago17static void AddOp(const PyCustomOpDef* cod);
28b2ab10Wenbing Li5 years ago18
5a3c6295Wenbing Li5 years ago19// no initializer here to avoid gcc whole-archive
20static const int undefined;
21static const int dt_float;
22static const int dt_uint8;
23static const int dt_int8;
24static const int dt_uint16;
25static const int dt_int16;
26static const int dt_int32;
27static const int dt_int64;
28static const int dt_string;
29static const int dt_bool;
30static const int dt_float16;
31static const int dt_double;
32static const int dt_uint32;
33static const int dt_uint64;
34static const int dt_complex64;
35static const int dt_complex128;
36static const int dt_bfloat16;
28b2ab10Wenbing Li5 years ago37};
38
39struct PyCustomOpKernel {
6c3b496eWenbing Li4 years ago40PyCustomOpKernel(OrtApi api, const OrtKernelInfo* info, uint64_t id, const std::vector<std::string>& attrs);
28b2ab10Wenbing Li5 years ago41void Compute(OrtKernelContext* context);
42
43private:
44OrtApi api_; // keep a copy of the struct, whose ref is used in the ort_
45Ort::CustomOpApi ort_;
46uint64_t obj_id_;
a0769463Xavier Dupré5 years ago47std::map<std::string, std::string> attrs_values_;
28b2ab10Wenbing Li5 years ago48};
49
50struct PyCustomOpFactory : Ort::CustomOpBase<PyCustomOpFactory, PyCustomOpKernel> {
6c3b496eWenbing Li4 years ago51PyCustomOpFactory() {
52// STL vector needs it.
53}
54
55PyCustomOpFactory(const PyCustomOpDef* opdef, const std::string& domain, const std::string& op) {
e36205eeXavier Dupré5 years ago56if (opdef == nullptr)
57throw std::runtime_error("Python definition is empty.");
28b2ab10Wenbing Li5 years ago58opdef_ = opdef;
6c3b496eWenbing Li4 years ago59op_domain_ = domain;
60op_type_ = op;
28b2ab10Wenbing Li5 years ago61}
62
a0769463Xavier Dupré5 years ago63void* CreateKernel(OrtApi api, const OrtKernelInfo* info) const {
64return new PyCustomOpKernel(api, info, opdef_->obj_id, opdef_->attrs);
28b2ab10Wenbing Li5 years ago65};
66
67const char* GetName() const {
6c3b496eWenbing Li4 years ago68return op_type_.c_str();
28b2ab10Wenbing Li5 years ago69};
70
71size_t GetInputTypeCount() const {
e36205eeXavier Dupré5 years ago72return opdef_->input_types.size();
28b2ab10Wenbing Li5 years ago73};
74
75ONNXTensorElementDataType GetInputType(size_t idx) const {
e36205eeXavier Dupré5 years ago76return static_cast<ONNXTensorElementDataType>(opdef_->input_types[idx]);
28b2ab10Wenbing Li5 years ago77};
78
a0769463Xavier Dupré5 years ago79const std::vector<std::string>& GetAttributesNames() const {
80return opdef_->attrs;
81}
82
28b2ab10Wenbing Li5 years ago83size_t GetOutputTypeCount() const {
e36205eeXavier Dupré5 years ago84return opdef_->output_types.size();
28b2ab10Wenbing Li5 years ago85};
86
87ONNXTensorElementDataType GetOutputType(size_t idx) const {
e36205eeXavier Dupré5 years ago88return static_cast<ONNXTensorElementDataType>(opdef_->output_types[idx]);
89}
28b2ab10Wenbing Li5 years ago90
5320af1eWenbing Li3 years ago91const PyCustomOpDef* opdef_ = nullptr;
6c3b496eWenbing Li4 years ago92std::string op_type_;
93std::string op_domain_;
28b2ab10Wenbing Li5 years ago94};