From b8c8fe91f55343925a5580b4cea204a544f1744d Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Tue, 15 Dec 2020 17:19:17 +0800 Subject: [PATCH] ort's to_dlpack. --- onnxruntime/python/dl_convertor.cc | 124 +++++++++++++++ onnxruntime/python/dl_convertor.h | 20 +++ onnxruntime/python/dlpack.h | 141 ++++++++++++++++++ .../onnxruntime_inference_collection.py | 3 + .../python/onnxruntime_pybind_state.cc | 18 +++ .../orttraining/python/training/ortmodule.py | 53 ++++--- 6 files changed, 332 insertions(+), 27 deletions(-) create mode 100644 onnxruntime/python/dl_convertor.cc create mode 100644 onnxruntime/python/dl_convertor.h create mode 100644 onnxruntime/python/dlpack.h diff --git a/onnxruntime/python/dl_convertor.cc b/onnxruntime/python/dl_convertor.cc new file mode 100644 index 0000000000..e30d84b8e4 --- /dev/null +++ b/onnxruntime/python/dl_convertor.cc @@ -0,0 +1,124 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "python/dl_convertor.h" + +namespace onnxruntime { +namespace python { + +DLDataType get_dlpack_data_type(const OrtValue& ml_value) { + ORT_ENFORCE(ml_value.IsTensor(), "Only OrtValues that are Tensors are currently supported"); + DLDataType dtype; + dtype.lanes = 1; + const Tensor& tensor = ml_value.Get(); + switch (tensor.GetElementType()) { + case ONNX_NAMESPACE::TensorProto_DataType_DOUBLE: + dtype.code = DLDataTypeCode::kDLFloat; + dtype.bits = sizeof(double); + break; + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: + dtype.code = DLDataTypeCode::kDLFloat; + dtype.bits = sizeof(float); + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT8: + dtype.code = DLDataTypeCode::kDLInt; + dtype.bits = sizeof(int8_t); + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT16: + dtype.code = DLDataTypeCode::kDLInt; + dtype.bits = sizeof(int16_t); + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT32: + dtype.code = DLDataTypeCode::kDLInt; + dtype.bits = sizeof(int); + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + dtype.code = DLDataTypeCode::kDLInt; + dtype.bits = sizeof(int64_t); + break; + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: + dtype.code = DLDataTypeCode::kDLFloat; + dtype.bits = sizeof(MLFloat16); + break; + // Currently bool is same as uint8 on both code and bits. + // PyTorch's to_dlpack also does this, but in from_dlpack, + // a torch.uint8 tensor is generated. This limitation from + // PyTorch means we cannot create a torch.bool tensor + // from a DLPack tensor. + case ONNX_NAMESPACE::TensorProto_DataType_BOOL: + dtype.code = DLDataTypeCode::kDLUInt; + dtype.bits = sizeof(bool); + break; + case ONNX_NAMESPACE::TensorProto_DataType_UINT8: + dtype.code = DLDataTypeCode::kDLUInt; + dtype.bits = sizeof(uint8_t); + break; + case ONNX_NAMESPACE::TensorProto_DataType_UINT16: + dtype.code = DLDataTypeCode::kDLUInt; + dtype.bits = sizeof(uint16_t); + break; + case ONNX_NAMESPACE::TensorProto_DataType_UINT32: + dtype.code = DLDataTypeCode::kDLUInt; + dtype.bits = sizeof(uint32_t); + break; + case ONNX_NAMESPACE::TensorProto_DataType_UINT64: + dtype.code = DLDataTypeCode::kDLUInt; + dtype.bits = sizeof(uint64_t); + break; + default: + ORT_THROW("Unexpected data type of ", tensor.GetElementType()); + } + + dtype.bits *= 8; // bits. + return dtype; +} + +DLContext get_dlpack_context(const OrtValue& ml_value, const int64_t& device_id) { + ORT_ENFORCE(ml_value.IsTensor(), "Only OrtValues that are Tensors are currently supported"); + DLContext ctx; + ctx.device_id = device_id; + const Tensor& tensor = ml_value.Get(); + const auto& location = tensor.Location(); + switch (location.device.Type()) { + case OrtDevice::CPU: + ctx.device_type = DLDeviceType::kDLCPU; + break; + case OrtDevice::GPU: + ctx.device_type = DLDeviceType::kDLGPU; + break; + default: + ORT_THROW("Cannot pack tensors on this device."); + } + + return ctx; +} + +struct OrtDLManagedTensor { + OrtValue handle; + DLManagedTensor tensor; +}; + +void deleter(DLManagedTensor* arg) { delete static_cast(arg->manager_ctx); } + +// This function returns a shared_ptr to memory managed DLpack tensor +// constructed out of OrtValue. +DLManagedTensor* ort_value_to_dlpack(const OrtValue& ml_value) { + ORT_ENFORCE(ml_value.IsTensor(), "Only OrtValues that are Tensors are currently supported"); + OrtDLManagedTensor* ort_dlmanaged_tensor(new OrtDLManagedTensor); + const Tensor& tensor = ml_value.Get(); + ort_dlmanaged_tensor->handle = ml_value; + ort_dlmanaged_tensor->tensor.manager_ctx = ort_dlmanaged_tensor; + ort_dlmanaged_tensor->tensor.deleter = &deleter; + ort_dlmanaged_tensor->tensor.dl_tensor.data = const_cast(tensor.DataRaw()); + ort_dlmanaged_tensor->tensor.dl_tensor.ctx = get_dlpack_context(ml_value, tensor.Location().device.Id()); + ort_dlmanaged_tensor->tensor.dl_tensor.ndim = tensor.Shape().NumDimensions(); + ort_dlmanaged_tensor->tensor.dl_tensor.dtype = get_dlpack_data_type(ml_value); + ort_dlmanaged_tensor->tensor.dl_tensor.shape = + tensor.Shape().NumDimensions() > 0 ? const_cast(&tensor.Shape()[0]) : nullptr; + ort_dlmanaged_tensor->tensor.dl_tensor.strides = nullptr; + ort_dlmanaged_tensor->tensor.dl_tensor.byte_offset = 0; + return &(ort_dlmanaged_tensor->tensor); +} + +} // namespace python +} // namespace onnxruntime diff --git a/onnxruntime/python/dl_convertor.h b/onnxruntime/python/dl_convertor.h new file mode 100644 index 0000000000..51289ec8e1 --- /dev/null +++ b/onnxruntime/python/dl_convertor.h @@ -0,0 +1,20 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/framework/ml_value.h" +#include "python/dlpack.h" + +// this convertor will: +// 1) take a OrtValue object and wrap it in the DLPack tensor + +namespace onnxruntime { +namespace python { + +DLManagedTensor* ort_value_to_dlpack(const OrtValue& ml_value); +DLDataType get_dlpack_data_type(const OrtValue& ml_value); +DLContext get_dlpack_context(const OrtValue& ml_value, const int64_t& device_id); + +} // namespace python +} // namespace onnxruntime diff --git a/onnxruntime/python/dlpack.h b/onnxruntime/python/dlpack.h new file mode 100644 index 0000000000..e3d55ea21d --- /dev/null +++ b/onnxruntime/python/dlpack.h @@ -0,0 +1,141 @@ +/*! + * Copyright (c) 2017 by Contributors + * \file dlpack.h + * \brief The common header of DLPack. + */ +#ifndef DLPACK_DLPACK_H_ +#define DLPACK_DLPACK_H_ + +#ifdef __cplusplus +#define DLPACK_EXTERN_C extern "C" +#else +#define DLPACK_EXTERN_C +#endif + +/*! \brief The current version of dlpack */ +#define DLPACK_VERSION 010 + +/*! \brief DLPACK_DLL prefix for windows */ +#ifdef _WIN32 +#ifdef DLPACK_EXPORTS +#define DLPACK_DLL __declspec(dllexport) +#else +#define DLPACK_DLL __declspec(dllimport) +#endif +#else +#define DLPACK_DLL +#endif + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif +/*! + * \brief The device type in DLContext. + */ +typedef enum { + kDLCPU = 1, + kDLGPU = 2, + // kDLCPUPinned = kDLCPU | kDLGPU + kDLCPUPinned = 3, + kDLOpenCL = 4, + kDLMetal = 8, + kDLVPI = 9, + kDLROCM = 10, +} DLDeviceType; + +/*! + * \brief A Device context for Tensor and operator. + */ +typedef struct { + /*! \brief The device type used in the device. */ + DLDeviceType device_type; + /*! \brief The device index */ + int device_id; +} DLContext; + +/*! + * \brief The type code options DLDataType. + */ +typedef enum { + kDLInt = 0U, + kDLUInt = 1U, + kDLFloat = 2U, +} DLDataTypeCode; + +/*! + * \brief The data type the tensor can hold. + * + * Examples + * - float: type_code = 2, bits = 32, lanes=1 + * - float4(vectorized 4 float): type_code = 2, bits = 32, lanes=4 + * - int8: type_code = 0, bits = 8, lanes=1 + */ +typedef struct { + /*! + * \brief Type code of base types. + * We keep it uint8_t instead of DLDataTypeCode for minimal memory + * footprint, but the value should be one of DLDataTypeCode enum values. + * */ + uint8_t code; + /*! + * \brief Number of bits, common choices are 8, 16, 32. + */ + uint8_t bits; + /*! \brief Number of lanes in the type, used for vector types. */ + uint16_t lanes; +} DLDataType; + +/*! + * \brief Plain C Tensor object, does not manage memory. + */ +typedef struct { + /*! + * \brief The opaque data pointer points to the allocated data. + * This will be CUDA device pointer or cl_mem handle in OpenCL. + * This pointer is always aligns to 256 bytes as in CUDA. + */ + void* data; + /*! \brief The device context of the tensor */ + DLContext ctx; + /*! \brief Number of dimensions */ + int ndim; + /*! \brief The data type of the pointer*/ + DLDataType dtype; + /*! \brief The shape of the tensor */ + int64_t* shape; + /*! + * \brief strides of the tensor, + * can be NULL, indicating tensor is compact. + */ + int64_t* strides; + /*! \brief The offset in bytes to the beginning pointer to data */ + uint64_t byte_offset; +} DLTensor; + +/*! + * \brief C Tensor object, manage memory of DLTensor. This data structure is + * intended to facilitate the borrowing of DLTensor by another framework. It is + * not meant to transfer the tensor. When the borrowing framework doesn't need + * the tensor, it should call the deleter to notify the host that the resource + * is no longer needed. + */ +typedef struct DLManagedTensor { + /*! \brief DLTensor which is being memory managed */ + DLTensor dl_tensor; + /*! \brief the context of the original host framework of DLManagedTensor in + * which DLManagedTensor is used in the framework. It can also be NULL. + */ + void * manager_ctx; + /*! \brief Destructor signature void (*)(void*) - this should be called + * to destruct manager_ctx which holds the DLManagedTensor. It can be NULL + * if there is no way for the caller to provide a reasonable destructor. + */ + void (*deleter)(struct DLManagedTensor * self); +} DLManagedTensor; +#ifdef __cplusplus +} // DLPACK_EXTERN_C +#endif +#endif // DLPACK_DLPACK_H_ diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index 4a36123df9..42cd824f00 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -441,3 +441,6 @@ class OrtValue: Valid only for OrtValues holding Tensors. Throws for OrtValues holding non-Tensors. ''' return self._ortvalue.numpy() + + def to_dlpack(self): + return self._ortvalue.to_dlpack() diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index efa1e75460..f8b2d2150a 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -22,6 +22,7 @@ #include "core/session/IOBinding.h" #include "core/session/abi_session_options_impl.h" #include "core/platform/env.h" +#include "python/dl_convertor.h" struct OrtStatus { OrtErrorCode code; @@ -1086,6 +1087,18 @@ void addOpSchemaSubmodule(py::module& m) { #endif //onnxruntime_PYBIND_EXPORT_OPSCHEMA +void dlpack_capsule_destructor(PyObject* data) { + DLManagedTensor* dlmanged_tensor = (DLManagedTensor*)PyCapsule_GetPointer(data, "dltensor"); + if (dlmanged_tensor) { + // the dlmanged_tensor has not been consumed, call deleter ourselves. + dlmanged_tensor->deleter(const_cast(dlmanged_tensor)); + } else { + // the dlmanged_tensor has been consumed, + // PyCapsule_GetPointer has set an error indicator. + PyErr_Clear(); + } +} + void addObjectMethods(py::module& m, Environment& env) { py::enum_(m, "GraphOptimizationLevel") .value("ORT_DISABLE_ALL", GraphOptimizationLevel::ORT_DISABLE_ALL) @@ -1292,6 +1305,11 @@ void addObjectMethods(py::module& m, Environment& env) { GetPyObjFromTensor(ml_value->Get(), obj, nullptr, nullptr); #endif return obj; + }) + .def("to_dlpack", [](OrtValue* ml_value) -> py::object { + DLManagedTensor* dlmanaged_tensor = ort_value_to_dlpack(*ml_value); + return py::reinterpret_steal( + PyCapsule_New(dlmanaged_tensor, "dltensor", dlpack_capsule_destructor)); }); py::class_ session_io_binding(m, "SessionIOBinding"); diff --git a/orttraining/orttraining/python/training/ortmodule.py b/orttraining/orttraining/python/training/ortmodule.py index a0a59ce42f..f7c97eed3e 100644 --- a/orttraining/orttraining/python/training/ortmodule.py +++ b/orttraining/orttraining/python/training/ortmodule.py @@ -11,6 +11,7 @@ from inspect import signature # Needed to re-implement PyTorch's cpu,cuda,to methods from torch import Tensor, device, dtype +from torch.utils.dlpack import from_dlpack from typing import Union, Tuple, Any, Callable, Iterator, Set, Optional, overload, TypeVar, Mapping, Dict from onnxruntime.capi import _pybind_state as C @@ -23,7 +24,7 @@ __TEMP_ENABLE_METHOD_TIMING__ = False # Needed to re-implement PyTorch's cpu,cuda,to methods T = TypeVar('T', bound='Module') -def _create_iobinding(io_binding, inputs, model, output_buffers, device): +def _create_iobinding(io_binding, inputs, model, device): '''Creates IO binding for a `model` inputs and output''' def _get_device_index(device): if type(device) == str: @@ -39,13 +40,7 @@ def _create_iobinding(io_binding, inputs, model, output_buffers, device): inputs[idx].data_ptr()) for value_info in model.graph.output: - name = value_info.name - output_tensor = output_buffers[name] - io_binding.bind_output(name, output_tensor.device.type, - _get_device_index(device), - _utils.dtype_torch_to_numpy(output_tensor.dtype), - list(output_tensor.size()), - output_tensor.data_ptr()) + io_binding.bind_output(value_info.name, device if device.find(':') == -1 else device[:device.find(':')]) def _onnx_value_info_to_buffer_tensor(value_info, device): '''Create a torch zeroed tensor with the same shape and type of `value_info`''' @@ -54,6 +49,18 @@ def _onnx_value_info_to_buffer_tensor(value_info, device): dtype = _utils.dtype_onnx_to_torch(value_info.type.tensor_type.elem_type) return torch.zeros(shape, device=device, dtype=dtype) +# TODO: PyTorch's to_dlpack() uses same config for both torch.bool and torch.uint8, +# and convert the config to torch.uint8 tensor duing from_dlpack(). So a boolean tensor +# from forward graph outputs will be converted to torch.uint8 tensor. When this tensor +# is feeded to backward graph as input, it will cause data type mismatch issue during +# inference session running. We cannot change the from_dlpack() in PyTorch side, so we +# have to handle this specially, which will introduce a cast here and there is data copied. +# Always cast from torch.uint8 to torch.bool is not logically right, we need to check the +# real data type of the inputs in the backeard graph, and perform the cast only necessary. +def _ort_output_to_torch_tensor(ort_output): + tensor = from_dlpack(ort_output.to_dlpack()) + return tensor.to(torch.bool) if tensor.dtype == torch.uint8 else tensor + class ORTModule(torch.nn.Module): def __init__(self, module): @@ -76,13 +83,11 @@ class ORTModule(torch.nn.Module): self._onnx_forward = None self._forward_session = None self._forward_io_binding = None - self._forward_output_buffers = {} # Backward pass self._onnx_backward = None self._backward_session = None self._backward_io_binding = None - self._backward_output_buffers = {} # Log level self._loglevel = getattr(logging, 'WARNING') @@ -157,8 +162,8 @@ class ORTModule(torch.nn.Module): ''' if not self._onnx_forward or self._require_export: self._require_export = False - self._onnx_training = ORTModule._get_forward_graph(self._original_module, *inputs, **kwargs) + # TODO: PyTorch exporter bug: changes the initializer order initializer_names = [p[0] for p in self._original_module.named_parameters()] @@ -186,13 +191,7 @@ class ORTModule(torch.nn.Module): # IO binding # TODO: we should try to reuse the output buffers as some of the output tensors are same sizes, expecially the backward graph outputs. self._forward_io_binding = self._forward_session.io_binding() - self._forward_output_buffers = {} - for output in self._onnx_forward.graph.output: - self._forward_output_buffers[output.name] = _onnx_value_info_to_buffer_tensor(output, str(self._device)) self._backward_io_binding = self._backward_session.io_binding() - self._backward_output_buffers = {} - for output in self._onnx_backward.graph.output: - self._backward_output_buffers[output.name] = _onnx_value_info_to_buffer_tensor(output, str(self._device)) if self._save_onnx: onnx.save(self._onnx_forward, self._save_onnx_prefix + '_forward.onnx') @@ -217,11 +216,11 @@ class ORTModule(torch.nn.Module): # Use IO binding _create_iobinding(self._forward_io_binding, inputs, self._onnx_forward, - self._forward_output_buffers, str(self._device)) # Run self._forward_session.run_with_iobinding(self._forward_io_binding) + forward_outputs = self._forward_io_binding.get_outputs() # Stash tensors needed by backward forward_input_dict = self._convert_forward_input_list_to_dict(*inputs) @@ -229,13 +228,14 @@ class ORTModule(torch.nn.Module): for name in self._onnx_graphs_info.backward_user_input_names) ctx_initializers = tuple(forward_input_dict[name] \ for name in self._onnx_graphs_info.backward_intializer_names_as_input) - ctx_intermediates = tuple(self._forward_output_buffers[name] \ - for name in self._onnx_graphs_info.intermediate_tensor_names) + ctx_intermediates = tuple(_ort_output_to_torch_tensor(forward_output) \ + for forward_output in forward_outputs[len(self._onnx_graphs_info.user_output_names):]) ctx.save_for_backward(*[*ctx_inputs, *ctx_initializers, *ctx_intermediates]) # Return model output - outputs = tuple(self._forward_output_buffers[name] for name in self._onnx_graphs_info.user_output_names) - return outputs[0] if len(outputs) == 1 else outputs + user_outputs = tuple(_ort_output_to_torch_tensor(forward_output) \ + for forward_output in forward_outputs[:len(self._onnx_graphs_info.user_output_names)]) + return user_outputs[0] if len(user_outputs) == 1 else user_outputs @staticmethod def backward(ctx, *grad_output): @@ -253,16 +253,16 @@ class ORTModule(torch.nn.Module): backward_grad_output = tuple(grad_output_dict[name] for name in self._onnx_graphs_info.backward_output_grad_names) _create_iobinding(self._backward_io_binding, [*ctx.saved_tensors, *backward_grad_output], self._onnx_backward, - self._backward_output_buffers, str(self._device)) # Run self._backward_session.run_with_iobinding(self._backward_io_binding) + backward_outputs = self._backward_io_binding.get_outputs() # Return input and initializer gradients results = [torch.tensor([1])] * len(self._onnx_graphs_info.user_input_names) - results += [self._backward_output_buffers[name] \ - for name in self._onnx_graphs_info.initializer_grad_names_to_train] + results += [_ort_output_to_torch_tensor(backward_output) \ + for backward_output in backward_outputs[:len(self._onnx_graphs_info.initializer_grad_names_to_train)]] return tuple(results) proc_inputs = [data for data in inputs if data is not None] @@ -365,11 +365,10 @@ class ORTModule(torch.nn.Module): # Ignore optional *inputs explicitly specified as None sig = signature(module.forward) all_input_names = sig.parameters.keys() - # input_names = [name for idx, name in enumerate(all_input_names) if inputs[idx] is not None] input_names = [] dynamic_axes = {} for input_idx, name in enumerate(all_input_names): - if inputs[input_idx] is not None: + if input_idx < len(inputs) and inputs[input_idx] is not None: input_names.append(name) dynamic_axes[name] = {} for dim_idx in range(len(inputs[input_idx].shape)):