ort's to_dlpack.

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<Tensor>();
+  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<Tensor>();
+  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<OrtDLManagedTensor*>(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<Tensor>();
+  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<void*>(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<int64_t*>(&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

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

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 <stdint.h>
+#include <stddef.h>
+
+#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_

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()

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<DLManagedTensor*>(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_<GraphOptimizationLevel>(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<Tensor>(), 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<py::object>(
+            PyCapsule_New(dlmanaged_tensor, "dltensor", dlpack_capsule_destructor));
       });
 
   py::class_<SessionIOBinding> session_io_binding(m, "SessionIOBinding");

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)):