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/pyfunc.cc

378lines · modecode

RawDownload
Latest commit unavailable.
1// Copyright (c) Microsoft Corporation. All rights reserved.
2// Licensed under the MIT License.
3
4#include <vector>
5#include <fstream>
6#include <mutex>
7
8#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
9#define PY_ARRAY_UNIQUE_SYMBOL ocos_python_ARRAY_API
10#include <numpy/arrayobject.h>
11
12#include <pybind11/iostream.h>
13#include <pybind11/pybind11.h>
14#include <pybind11/stl.h>
15#include <pybind11/functional.h>
16#include <pybind11/numpy.h>
17#include <thread>
18#include "../utils.h"
19#include "pykernel.h"
20
21namespace py = pybind11;
22
23static int to_numpy(ONNXTensorElementDataType dt) {
24 switch (dt) {
25 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
26 return NPY_FLOAT;
27 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8:
28 return NPY_UINT8;
29 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8:
30 return NPY_INT8;
31 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16:
32 return NPY_UINT16;
33 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16:
34 return NPY_INT16;
35 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32:
36 return NPY_INT32;
37 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64:
38 return NPY_INT64;
39 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL:
40 return NPY_BOOL;
41 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16:
42 return NPY_FLOAT16;
43 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE:
44 return NPY_DOUBLE;
45 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32:
46 return NPY_UINT32;
47 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64:
48 return NPY_UINT64;
49 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64:
50 return NPY_COMPLEX64;
51 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128:
52 return NPY_COMPLEX128;
53 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING:
54 return NPY_OBJECT;
55 default:
56 throw std::runtime_error("No corresponding Numpy data type/Tensor data Type.");
57 }
58}
59
60static size_t element_size(ONNXTensorElementDataType dt) {
61 switch (dt) {
62 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
63 return sizeof(float);
64 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8:
65 return sizeof(uint8_t);
66 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8:
67 return sizeof(int8_t);
68 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16:
69 return sizeof(uint16_t);
70 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16:
71 return sizeof(int16_t);
72 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32:
73 return sizeof(int32_t);
74 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64:
75 return sizeof(int64_t);
76 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL:
77 return sizeof(bool);
78 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16:
79 return sizeof(uint16_t);
80 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE:
81 return sizeof(double);
82 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32:
83 return sizeof(uint32_t);
84 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64:
85 return sizeof(uint64_t);
86 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64:
87 return sizeof(_C_float_complex);
88 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128:
89 return sizeof(_C_double_complex);
90 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING:
91 return sizeof(std::string*);
92 default:
93 throw std::runtime_error("No corresponding Numpy data type/Tensor data Type.");
94 }
95}
96
97static ONNXTensorElementDataType from_numpy(int dt) {
98 switch (dt) {
99 case NPY_FLOAT:
100 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;
101 case NPY_UINT8:
102 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8;
103 case NPY_INT8:
104 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8;
105 case NPY_UINT16:
106 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16;
107 case NPY_INT16:
108 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16;
109 case NPY_INT32:
110 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32;
111 case NPY_INT64:
112 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64;
113 case NPY_BOOL:
114 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL;
115 case NPY_FLOAT16:
116 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16;
117 case NPY_DOUBLE:
118 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE;
119 case NPY_UINT32:
120 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32;
121 case NPY_UINT64:
122 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64;
123 case NPY_COMPLEX64:
124 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64;
125 case NPY_COMPLEX128:
126 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128;
127 case NPY_OBJECT:
128 case NPY_STRING:
129 return ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING;
130 default:
131 throw std::runtime_error("No corresponding ONNX data type/Tensor data Type.");
132 }
133}
134
135struct PyCustomOpDefImpl : public PyCustomOpDef {
136 typedef std::vector<int64_t> shape_t;
137 static int64_t calc_size_from_shape(const shape_t& sp) {
138 size_t c = 1;
139 for (auto it = sp.begin(); it != sp.end(); ++it) {
140 c *= *it;
141 }
142 return c;
143 }
144
145 static py::object BuildPyObjFromTensor(const void* p, const shape_t& shape, ONNXTensorElementDataType dtype) {
146 std::vector<npy_intp> npy_dims;
147 for (auto n : shape) {
148 npy_dims.push_back(n);
149 }
150 const int numpy_type = to_numpy(dtype);
151 py::object obj = py::reinterpret_steal<py::object>(PyArray_SimpleNew(
152 static_cast<int>(shape.size()), npy_dims.data(), numpy_type));
153 void* out_ptr = static_cast<void*>(
154 PyArray_DATA(reinterpret_cast<PyArrayObject*>(obj.ptr())));
155
156 if (dtype == ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING) {
157 py::object* outObj = static_cast<py::object*>(out_ptr);
158 auto size = calc_size_from_shape(shape);
159 const std::string* src = (const std::string*)p;
160 for (int i = 0; i < size; i++, src++) {
161 outObj[i] = py::cast(*src);
162 }
163 } else {
164 size_t size_type = element_size(dtype);
165 memcpy(out_ptr, p, size_type * calc_size_from_shape(shape));
166 }
167 return obj;
168 }
169
170 static py::object InvokePyFunction(uint64_t id, const py::object& feed) {
171 return (*op_invoker)(id, feed);
172 }
173
174 using callback_t = std::function<py::object(uint64_t id, const py::object&)>;
175 static std::auto_ptr<callback_t> op_invoker;
176};
177
178std::auto_ptr<PyCustomOpDefImpl::callback_t> PyCustomOpDefImpl::op_invoker;
179// static py::function g_pyfunc_caller;
180// static std::mutex op_mutex;
181// static std::condition_variable op_cv;
182// static bool is_ready = false;
183
184typedef struct {
185 const OrtValue* input_X;
186 ONNXTensorElementDataType dtype;
187 std::vector<int64_t> dimensions;
188} InputInformation;
189
190void PyCustomOpKernel::Compute(OrtKernelContext* context) {
191 // std::unique_lock<std::mutex> lck(op_mutex);
192 // is_ready = true;
193 // op_cv.notify_all();
194 // std::this_thread::sleep_for(std::chrono::milliseconds(5000));
195 size_t n_inputs = ort_.KernelContext_GetInputCount(context);
196 size_t n_outputs = ort_.KernelContext_GetOutputCount(context);
197
198 // Setup inputs
199 std::vector<InputInformation> inputs;
200 inputs.reserve(n_inputs);
201 for (size_t index = 0; index < n_inputs; ++index) {
202 const OrtValue* input_X = ort_.KernelContext_GetInput(context, index);
203 std::vector<int64_t> i_dimensions;
204 OrtTensorTypeAndShapeInfo* i_info = ort_.GetTensorTypeAndShape(input_X);
205 i_dimensions = ort_.GetTensorShape(i_info);
206 ONNXTensorElementDataType i_dtype = ort_.GetTensorElementType(i_info);
207 ort_.ReleaseTensorTypeAndShapeInfo(i_info);
208 inputs.push_back(InputInformation{input_X, i_dtype, i_dimensions});
209 }
210
211 /* Acquire GIL before calling Python code, due to it was released in sess.run */
212 py::gil_scoped_acquire acquire;
213
214 // TODO: Direct-Buffer-Access doesn't work for some reason.
215 // OrtTensorTypeAndShapeInfo* output_info = ort_.GetTensorTypeAndShape(output);
216 // int64_t size = ort_.GetTensorShapeElementCount(output_info);
217 // ort_.ReleaseTensorTypeAndShapeInfo(output_info);
218 // py::buffer_info buf(
219 // const_cast<void *>(X), /* Pointer to buffer */
220 // sizeof(float), /* Size of one scalar */
221 // py::format_descriptor<float>::format(), /* Python struct-style format descriptor */
222 // 2, /* Number of dimensions */
223 // {2, 3}, /* Buffer dimensions */
224 // {sizeof(float) * dimensions.data()[1], /* Strides (in bytes) for each index */
225 // sizeof(float)});
226
227 {
228 py::list pyinputs;
229 for (auto it = inputs.begin(); it != inputs.end(); ++it) {
230 py::object input0 = PyCustomOpDefImpl::BuildPyObjFromTensor(
231 (const void*)ort_.GetTensorData<float>(it->input_X), it->dimensions, it->dtype);
232 pyinputs.append(input0);
233 }
234
235 // Call python function id, shape, flat coefficient.
236 py::tuple fetch = PyCustomOpDefImpl::InvokePyFunction(obj_id_, pyinputs);
237 int64_t rid = fetch[0].cast<int64_t>();
238 assert(rid == obj_id_);
239
240 // Setup output.
241 for (size_t no = 0; no < n_outputs; ++no) {
242 auto dims = fetch[1 + no * 2].cast<std::vector<int64_t>>();
243 OrtValue* output = ort_.KernelContext_GetOutput(context, no, dims.data(), dims.size());
244 OrtTensorTypeAndShapeInfo* o_info = ort_.GetTensorTypeAndShape(output);
245 ONNXTensorElementDataType o_dtype = ort_.GetTensorElementType(o_info);
246 const void* Y = (const void*)ort_.GetTensorData<float>(output);
247 ort_.ReleaseTensorTypeAndShapeInfo(o_info);
248 void* out = (void*)ort_.GetTensorMutableData<float>(output);
249
250 if (o_dtype == ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING) {
251 auto retval = fetch[2 + no * 2].cast<std::vector<std::string>>();
252 std::string* type_outPtr = (std::string*)out;
253 std::string* end = type_outPtr + retval.size();
254 const std::string* source = (const std::string*)retval.data();
255 for (; type_outPtr != end; ++type_outPtr, ++source) {
256 *type_outPtr = *source;
257 }
258 } else {
259 py::array retval = fetch[2 + no * 2].cast<py::array>();
260 if (element_size(o_dtype) != retval.itemsize()) {
261 switch (o_dtype) {
262 case ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
263 retval = fetch[2 + no * 2].cast<py::array_t<float>>();
264 break;
265 default:
266 throw std::runtime_error(MakeString(
267 "Type mismatch between declared output element size (",
268 element_size(o_dtype), ") and python element size (",
269 retval.itemsize(), ")"));
270 }
271 }
272 size_t size = element_size(o_dtype);
273 memcpy(out, retval.data(), size * retval.size());
274 }
275 }
276
277 py::gil_scoped_release release;
278
279 // TODO: the return value from the python callback function doesn't work in pybind11&numpy.
280 // py::gil_scoped_acquire acquire;
281 // int64_t rid = fetch[0].cast<int64_t>();
282 // assert(rid == obj_id_);
283 // size_t ntp = fetch.size() - 1;
284 // py::object ret_val = fetch[ntp];
285 // auto p2 = PyArray_FROM_O(ret_val.ptr());
286 // PyArrayObject* darray = reinterpret_cast<PyArrayObject*>(p2);
287 // std::vector<int64_t> dims;
288 // const int npy_type = PyArray_TYPE(darray);
289 // {
290 // int ndim = PyArray_NDIM(darray);
291 // const npy_intp* npy_dims = PyArray_DIMS(darray);
292 // dims.resize(ndim);
293 // std::copy(npy_dims, npy_dims+ndim, dims.begin());
294 // }
295
296 // auto element_type = PyCustomOpDefImpl::from_numpy(npy_type);
297 // OrtValue* output = ort_.KernelContext_GetOutput(context, 0, dims.data(), dims.size());
298 // float* out = ort_.GetTensorMutableData<float>(output);
299 // const void* pyOut = PyData(darray, dd);
300
301 // memcpy(out, pyOut, PyArray_NBYTES((PyArrayObject*)NULL));
302 }
303}
304
305std::vector<PyCustomOpFactory>& PyCustomOpDef_python_operator_list() {
306 static std::vector<PyCustomOpFactory> lst_custom_opdef;
307 return lst_custom_opdef;
308}
309
310void PyCustomOpDef::AddOp(const PyCustomOpDef* cod) {
311 // No need to protect against concurrent access, GIL is doing that.
312 PyCustomOpDef_python_operator_list().push_back(PyCustomOpFactory(cod));
313}
314
315const PyCustomOpFactory* PyCustomOpDef_FetchPyCustomOps(size_t count) {
316 // The result must stay alive
317 std::vector<PyCustomOpFactory>& copy = PyCustomOpDef_python_operator_list();
318 if (count < copy.size())
319 return &(copy[count]);
320 return nullptr;
321}
322
323const OrtCustomOp* FetchPyCustomOps(size_t& count) {
324 auto ptr = PyCustomOpDef_FetchPyCustomOps(count);
325 if (ptr == nullptr)
326 return nullptr;
327 return ptr;
328}
329
330// static std::ofstream logger;
331static int init_numpy() {
332 import_array();
333 // logger.open("./ggtest.log.txt", std::ofstream::out | std::ofstream::app);
334 // logger << "first line." << std::endl;
335 return 0;
336}
337
338void AddGlobalMethods(pybind11::module& m) {
339 m.def("add_custom_op", [](const PyCustomOpDef& cod) { PyCustomOpDef::AddOp(&cod); });
340}
341
342void AddObjectMethods(pybind11::module& m) {
343 pybind11::class_<PyCustomOpDef>(m, "PyCustomOpDef")
344 .def(pybind11::init<>())
345 .def_readwrite("op_type", &PyCustomOpDef::op_type)
346 .def_readwrite("obj_id", &PyCustomOpDef::obj_id)
347 .def_readwrite("input_types", &PyCustomOpDef::input_types)
348 .def_readwrite("output_types", &PyCustomOpDef::output_types)
349 .def_static("install_hooker", [](py::object obj) {
350 std::auto_ptr<PyCustomOpDefImpl::callback_t> s_obj(new PyCustomOpDefImpl::callback_t(obj));
351 PyCustomOpDefImpl::op_invoker = s_obj; })
352 .def_readonly_static("undefined", &PyCustomOpDef::undefined)
353 .def_readonly_static("dt_float", &PyCustomOpDef::dt_float)
354 .def_readonly_static("dt_uint8", &PyCustomOpDef::dt_uint8)
355 .def_readonly_static("dt_int8", &PyCustomOpDef::dt_int8)
356 .def_readonly_static("dt_uint16", &PyCustomOpDef::dt_uint16)
357 .def_readonly_static("dt_int16", &PyCustomOpDef::dt_int16)
358 .def_readonly_static("dt_int32", &PyCustomOpDef::dt_int32)
359 .def_readonly_static("dt_int64", &PyCustomOpDef::dt_int64)
360 .def_readonly_static("dt_string", &PyCustomOpDef::dt_string)
361 .def_readonly_static("dt_bool", &PyCustomOpDef::dt_bool)
362 .def_readonly_static("dt_float16", &PyCustomOpDef::dt_float16)
363 .def_readonly_static("dt_double", &PyCustomOpDef::dt_double)
364 .def_readonly_static("dt_uint32", &PyCustomOpDef::dt_uint32)
365 .def_readonly_static("dt_uint64", &PyCustomOpDef::dt_uint64)
366 .def_readonly_static("dt_complex64", &PyCustomOpDef::dt_complex64)
367 .def_readonly_static("dt_complex128", &PyCustomOpDef::dt_complex128)
368 .def_readonly_static("dt_bfloat16 =", &PyCustomOpDef::dt_bfloat16);
369}
370
371PYBIND11_MODULE(_ortcustomops, m) {
372 m.doc() = "pybind11 stateful interface to ORT Custom Ops library";
373 //TODO: RegisterExceptions(m);
374
375 init_numpy();
376 AddGlobalMethods(m);
377 AddObjectMethods(m);
378}
379