From 9f3f77532a77206bf823141d111748e7056c1ac3 Mon Sep 17 00:00:00 2001 From: justinchuby <11205048+justinchuby@users.noreply.github.com> Date: Thu, 5 Nov 2026 03:43:20 +0000 Subject: [PATCH 01/11] Test PyTorch DLPack dtype interchange in inference CI Cover all byte-aligned ORT dtypes in both directions, document and test packed dtype incompatibilities, and retain zero-copy ownership without a redundant OrtValue copy. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/dlpack/dlpack_converter.cc | 7 +- onnxruntime/core/dlpack/dlpack_converter.h | 6 +- .../onnxruntime_inference_collection.py | 13 +- .../python/onnxruntime_pybind_ortvalue.cc | 9 +- .../python/onnxruntime_pybind_state_common.cc | 5 +- .../python/onnxruntime_pybind_state_common.h | 7 +- .../python/onnxruntime_test_python_dlpack.py | 209 ++++++++++++++++++ tools/ci_build/build.py | 9 + 8 files changed, 241 insertions(+), 24 deletions(-) create mode 100644 onnxruntime/test/python/onnxruntime_test_python_dlpack.py diff --git a/onnxruntime/core/dlpack/dlpack_converter.cc b/onnxruntime/core/dlpack/dlpack_converter.cc index 0858be4842328..8784c0fe567ed 100644 --- a/onnxruntime/core/dlpack/dlpack_converter.cc +++ b/onnxruntime/core/dlpack/dlpack_converter.cc @@ -176,12 +176,11 @@ DLManagedTensor* OrtValueToDlpack(OrtValue& ort_value) { ort_dlmanaged_tensor->tensor.dl_tensor.device = GetDlpackDevice(ort_value, tensor.Location().device.Id()); ort_dlmanaged_tensor->tensor.dl_tensor.ndim = static_cast(tensor.Shape().NumDimensions()); ort_dlmanaged_tensor->tensor.dl_tensor.dtype = GetDlpackDataType(ort_value); - ort_dlmanaged_tensor->tensor.dl_tensor.shape = - tensor.Shape().NumDimensions() > 0 ? &const_cast(tensor.Shape())[0] : nullptr; + const auto dims = tensor.Shape().GetDims(); + ort_dlmanaged_tensor->tensor.dl_tensor.shape = dims.empty() ? nullptr : const_cast(dims.data()); auto& strides = ort_dlmanaged_tensor->strides; - strides.reserve(tensor.Shape().NumDimensions()); + strides.reserve(dims.size()); // Empty tensors have no addressable elements, so zero strides avoid overflowing dimension products. - const auto dims = tensor.Shape().GetDims(); const bool is_empty = std::find(dims.begin(), dims.end(), 0) != dims.end(); SafeInt stride = 1; for (size_t i = tensor.Shape().NumDimensions(); i > 0; --i) { diff --git a/onnxruntime/core/dlpack/dlpack_converter.h b/onnxruntime/core/dlpack/dlpack_converter.h index 647e38180349a..3469df0f084cf 100644 --- a/onnxruntime/core/dlpack/dlpack_converter.h +++ b/onnxruntime/core/dlpack/dlpack_converter.h @@ -11,10 +11,8 @@ namespace onnxruntime { namespace dlpack { -// This convertor will take an OrtValue and wrap it as a DLPack tensor -// This may create a new ownership to the underlying tensor in OrtValue, -// so we do pass-by-value here. We don't use pass-by-reference because -// it implies no new ownership. +// Wrap an OrtValue as a DLPack tensor, retaining a copy of the OrtValue +// until the consumer calls the DLPack deleter. DLManagedTensor* OrtValueToDlpack(OrtValue& ort_value); // Imports a contiguous tensor, preserving its DLPack dtype (including kDLBool). diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index d260822355244..f31b9d1244c7c 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -1529,9 +1529,10 @@ def __dlpack__(self, *, stream=None): ``from_dlpack(ortvalue)`` (e.g. ``torch.from_dlpack``, ``jax.dlpack.from_dlpack``, ``numpy.from_dlpack``). - The OrtValue must hold a contiguous tensor. No data is copied; - the consumer shares memory with this OrtValue, which must remain - alive while the capsule is in use. + The OrtValue must hold a contiguous tensor. No data is copied. + The capsule and its consumer retain the tensor allocation, so this + OrtValue can be released. For externally owned buffers (for example, + a NumPy-backed CPU OrtValue), the external owner must remain alive. WebGPU OrtValues cannot be exported through DLPack because their data is stored in opaque WebGPU buffers rather than CUDA-addressable memory. @@ -1571,6 +1572,12 @@ def from_dlpack(cls, data, /) -> OrtValue: including float8, packed float4, int4/uint4, and int2/uint2. String, complex, and float6 tensors are not supported. + Interchange also depends on the other framework's representation. + PyTorch 2.10 exchanges all byte-aligned ORT numeric and boolean types, + including bfloat16 and the five float8 types. Its float4 uses two lanes (unsupported here), + and its int2/int4 shell dtypes export as int8/uint8, not packed types. + PyTorch cannot consume ORT's packed float4, int4/uint4, or int2/uint2. + No data is copied; the new OrtValue shares memory with the source. :param data: A tensor object supporting the DLPack protocol, or a raw diff --git a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc index 0d751fbf6b99d..e557f15eaa877 100644 --- a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc @@ -508,16 +508,15 @@ void addOrtValueMethods(pybind11::module& m) { .def("to_dlpack", [](OrtValue* ort_value) -> py::object { ORT_ENFORCE(!IsWebGpuBuffer(*ort_value), "DLPack export is not supported for WebGPU OrtValues."); return py::reinterpret_steal(ToDlpack(*ort_value)); }, - "Returns a DLPack representing the tensor. This method does not copy the pointer shape, " - "instead, it copies the pointer value. The OrtValue must be persist until the dlpack structure " - "is consumed.") + "Returns a DLPack capsule sharing the tensor buffer without copying. " + "The capsule retains the OrtValue; externally owned buffers must remain alive.") .def_static("from_dlpack", [](py::object data) { return FromDlpack(data.ptr()); }, py::arg("data"), "Converts a tensor from a external library into an OrtValue by means of the __dlpack__ protocol.") .def("__dlpack__", [](OrtValue* ort_value, py::object /* stream */) -> py::object { ORT_ENFORCE(!IsWebGpuBuffer(*ort_value), "DLPack export is not supported for WebGPU OrtValues."); return py::reinterpret_steal(ToDlpack(*ort_value)); }, py::arg("stream") = py::none(), "Returns a DLPack representing the tensor (part of __dlpack__ protocol). " - "This method does not copy the pointer shape, instead, it copies the pointer value. " - "The OrtValue must persist until the dlpack structure is consumed.") + "The capsule retains the OrtValue and shares its buffer without copying; " + "externally owned buffers must remain alive.") .def("__dlpack_device__", [](const OrtValue* ort_value) -> py::tuple { ORT_ENFORCE(ort_value->IsTensor(), "Only tensor type OrtValues are supported"); ORT_ENFORCE(!IsWebGpuBuffer(*ort_value), "DLPack export is not supported for WebGPU OrtValues."); diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.cc b/onnxruntime/python/onnxruntime_pybind_state_common.cc index 44da3cf46c815..9e35f3d8fd838 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.cc +++ b/onnxruntime/python/onnxruntime_pybind_state_common.cc @@ -56,10 +56,7 @@ void DlpackCapsuleDestructor(PyObject* data) { } } -// Allocate a new Capsule object, which takes the ownership of OrtValue. -// Caller is responsible for releasing. -// This function calls OrtValueToDlpack(...). -PyObject* ToDlpack(OrtValue ort_value) { +PyObject* ToDlpack(OrtValue& ort_value) { DLManagedTensor* dlmanaged_tensor = dlpack::OrtValueToDlpack(ort_value); return PyCapsule_New(dlmanaged_tensor, "dltensor", DlpackCapsuleDestructor); } diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.h b/onnxruntime/python/onnxruntime_pybind_state_common.h index 5e5ee306d6652..d99c605a38505 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.h +++ b/onnxruntime/python/onnxruntime_pybind_state_common.h @@ -541,10 +541,9 @@ bool CheckIfTensor(const std::vector& def_list, #if defined(ENABLE_DLPACK) -// Allocate a new Capsule object, which takes the ownership of OrtValue. -// Caller is responsible for releasing. -// This function calls OrtValueToDlpack(...). -PyObject* ToDlpack(OrtValue ort_value); +// Allocate a capsule retaining an OrtValue reference via OrtValueToDlpack. +// Caller is responsible for releasing the capsule. +PyObject* ToDlpack(OrtValue& ort_value); // Consume a Capsule object and claims the ownership of its underlying tensor to // create a OrtValue. This function calls DlpackToOrtValue(...) to do the conversion. diff --git a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py new file mode 100644 index 0000000000000..530c18230a261 --- /dev/null +++ b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py @@ -0,0 +1,209 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Inference-only DLPack interchange tests, required by the transformers-enabled CI legs. + +PyTorch is mandatory here: missing dependencies or DLPack support must fail, not skip. +The dtype contract below targets the PyTorch 2.10 version installed by build.py. +""" + +import ctypes +import gc +import math +import unittest +import weakref + +import torch +from onnx import TensorProto + +import onnxruntime as ort + +FLOAT8_DTYPES = ( + (torch.float8_e4m3fn, TensorProto.FLOAT8E4M3FN), + (torch.float8_e4m3fnuz, TensorProto.FLOAT8E4M3FNUZ), + (torch.float8_e5m2, TensorProto.FLOAT8E5M2), + (torch.float8_e5m2fnuz, TensorProto.FLOAT8E5M2FNUZ), + (torch.float8_e8m0fnu, TensorProto.FLOAT8E8M0), +) +INTERCHANGE_DTYPES = ( + (torch.float32, TensorProto.FLOAT), + (torch.float64, TensorProto.DOUBLE), + (torch.float16, TensorProto.FLOAT16), + (torch.bfloat16, TensorProto.BFLOAT16), + (torch.bool, TensorProto.BOOL), + (torch.int8, TensorProto.INT8), + (torch.int16, TensorProto.INT16), + (torch.int32, TensorProto.INT32), + (torch.int64, TensorProto.INT64), + (torch.uint8, TensorProto.UINT8), + (torch.uint16, TensorProto.UINT16), + (torch.uint32, TensorProto.UINT32), + (torch.uint64, TensorProto.UINT64), + *FLOAT8_DTYPES, +) +PACKED_DTYPES = ( + (torch.int4, TensorProto.INT4, TensorProto.INT8, 4, "Unsupported kInt bits 4"), + (torch.uint4, TensorProto.UINT4, TensorProto.UINT8, 4, "Unsupported kUInt bits 4"), + (torch.int2, TensorProto.INT2, TensorProto.INT8, 2, "Unsupported kInt bits 2"), + (torch.uint2, TensorProto.UINT2, TensorProto.UINT8, 2, "Unsupported kUInt bits 2"), + (torch.float4_e2m1fn_x2, TensorProto.FLOAT4E2M1, None, 4, "Unsupported kDLFloat4_e2m1fn lanes 1"), +) +SHAPES = ((), (2, 3), (2, 1, 3), (0,), (2, 0, 3)) + + +def tensor_bytes(tensor): + # Reshape first because PyTorch cannot reinterpret a scalar as a smaller dtype. + return tensor.reshape(-1).view(torch.uint8) + + +def make_tensor(dtype, shape): + count = math.prod(shape) + if dtype in (entry[0] for entry in FLOAT8_DTYPES): + return torch.arange(count, dtype=torch.int64).to(torch.uint8).view(dtype).reshape(shape) + if dtype == torch.bool: + values = [False, True] + elif dtype.is_floating_point: + values = [-2.5, -0.0, 0.0, 1.5, float("inf"), float("nan")] + else: + limits = torch.iinfo(dtype) + values = [limits.min, limits.max, 0, 1, 2, 3] + return torch.tensor((values * ((count + len(values) - 1) // len(values)))[:count], dtype=dtype).reshape(shape) + + +class TestDlpackPyTorch(unittest.TestCase): + def assert_tensor(self, actual, expected): + self.assertEqual(actual.dtype, expected.dtype) + self.assertEqual(actual.shape, expected.shape) + self.assertEqual(actual.device, expected.device) + if actual.numel(): + # Exact bits cover NaNs, signed zero and shell dtypes without arithmetic. + # Empty tensors may have zero strides, which PyTorch cannot reinterpret. + torch.testing.assert_close(tensor_bytes(actual), tensor_bytes(expected), rtol=0, atol=0) + + def test_pytorch_to_ort(self): + for dtype, onnx_type in INTERCHANGE_DTYPES: + for shape in SHAPES: + for capsule in (False, True): + with self.subTest(dtype=dtype, shape=shape, capsule=capsule): + source = make_tensor(dtype, shape) + expected = source.clone() + imported = ort.OrtValue.from_dlpack(torch.utils.dlpack.to_dlpack(source) if capsule else source) + self.assertEqual(imported.element_type(), onnx_type) + self.assertEqual(imported.shape(), list(shape)) + self.assertEqual(imported.__dlpack_device__(), (1, 0)) + if source.numel(): + self.assertEqual(imported.data_ptr(), source.data_ptr()) + self.assertEqual( + ctypes.string_at(imported.data_ptr(), source.numel() * source.element_size()), + tensor_bytes(expected).numpy().tobytes(), + ) + del source + gc.collect() + result = torch.from_dlpack(imported) + del imported + gc.collect() + self.assert_tensor(result, expected) + + def test_ort_to_pytorch(self): + for dtype, onnx_type in INTERCHANGE_DTYPES: + for shape in SHAPES: + for capsule in (False, True): + with self.subTest(dtype=dtype, shape=shape, capsule=capsule): + expected = make_tensor(dtype, shape) + # ORT owns this allocation independently of the expected PyTorch tensor. + source = ort.OrtValue.ortvalue_from_shape_and_type(list(shape), onnx_type) + if expected.numel(): + ctypes.memmove( + source.data_ptr(), expected.data_ptr(), expected.numel() * expected.element_size() + ) + result = torch.from_dlpack(source._ortvalue.to_dlpack() if capsule else source) + self.assert_tensor(result, expected) + if expected.numel(): + self.assertEqual(result.data_ptr(), source.data_ptr()) + tensor_bytes(result)[0] ^= 1 + tensor_bytes(expected)[0] ^= 1 + self.assertEqual( + ctypes.string_at(source.data_ptr(), 1), tensor_bytes(expected)[:1].numpy().tobytes() + ) + source_ref = weakref.ref(source) + del source + gc.collect() + self.assertIsNone(source_ref()) + self.assert_tensor(result, expected) + + def test_float8_all_bit_patterns(self): + for dtype, onnx_type in FLOAT8_DTYPES: + with self.subTest(dtype=dtype): + source = torch.arange(256, dtype=torch.int64).to(torch.uint8).view(dtype).reshape(16, 16) + imported = ort.OrtValue.from_dlpack(source) + self.assertEqual(imported.element_type(), onnx_type) + self.assert_tensor(torch.from_dlpack(imported), source) + owned = ort.OrtValue.ortvalue_from_shape_and_type([16, 16], onnx_type) + ctypes.memmove(owned.data_ptr(), source.data_ptr(), 256) + self.assert_tensor(torch.from_dlpack(owned), source) + + def test_capsule_keeps_ort_allocation_alive(self): + for dtype, onnx_type in INTERCHANGE_DTYPES: + with self.subTest(dtype=dtype): + expected = make_tensor(dtype, (2, 3)) + source = ort.OrtValue.ortvalue_from_shape_and_type([2, 3], onnx_type) + pointer = source.data_ptr() + ctypes.memmove(pointer, expected.data_ptr(), expected.numel() * expected.element_size()) + capsule = source._ortvalue.to_dlpack() + source_ref = weakref.ref(source) + del source + gc.collect() + self.assertIsNone(source_ref()) + result = torch.from_dlpack(capsule) + self.assertEqual(result.data_ptr(), pointer) + self.assert_tensor(result, expected) + + def test_packed_dtype_boundaries(self): + for dtype, onnx_type, storage_type, bits, error in PACKED_DTYPES: + with self.subTest(dtype=dtype): + # ORT's shape counts logical sub-byte elements, including an odd tail. + source = ort.OrtValue.ortvalue_from_shape_and_type([5], onnx_type) + num_bytes = (5 * bits + 7) // 8 + ctypes.memset(source.data_ptr(), 0x35, num_bytes) + imported = ort.OrtValue.from_dlpack(source) + self.assertEqual(imported.element_type(), onnx_type) + self.assertEqual(imported.shape(), [5]) + self.assertEqual(imported.data_ptr(), source.data_ptr()) + with self.assertRaisesRegex(BufferError, error): + torch.from_dlpack(source) + del source + gc.collect() + self.assertEqual(ctypes.string_at(imported.data_ptr(), num_bytes), b"\x35" * num_bytes) + + pytorch_source = torch.arange(5, dtype=torch.uint8).view(dtype) + if storage_type is None: + # PyTorch float4 has two lanes per element; ORT requires one. + with self.assertRaisesRegex(RuntimeError, "ORT does not support lanes != 1"): + ort.OrtValue.from_dlpack(pytorch_source) + else: + # PyTorch int2/int4 shell dtypes export as full bytes, not packed elements. + imported = ort.OrtValue.from_dlpack(pytorch_source) + self.assertEqual(imported.element_type(), storage_type) + self.assertEqual(imported.shape(), [5]) + self.assertEqual(imported.data_ptr(), pytorch_source.data_ptr()) + self.assertEqual(ctypes.string_at(imported.data_ptr(), 5), bytes(range(5))) + + def test_contiguous_slice_and_singleton_strides(self): + sources = ( + torch.arange(20, dtype=torch.float32)[5:11].reshape(2, 3), + torch.empty_strided((2, 1, 3), (3, 99, 1), dtype=torch.float32).fill_(2), + ) + for source in sources: + with self.subTest(shape=source.shape, strides=source.stride()): + self.assertTrue(source.is_contiguous()) + imported = ort.OrtValue.from_dlpack(source) + self.assertEqual(imported.data_ptr(), source.data_ptr()) + self.assert_tensor(torch.from_dlpack(imported), source) + + def test_noncontiguous_input_is_rejected(self): + source = torch.arange(6, dtype=torch.float32).reshape(2, 3).t() + with self.assertRaisesRegex(RuntimeError, "ORT only supports contiguous tensor"): + ort.OrtValue.from_dlpack(source) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index d7bcd6ce578bf..91747ef97c998 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -1960,6 +1960,15 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs): run_subprocess( [sys.executable, "onnxruntime_test_python.py"], cwd=cwd, dll_path=dll_path, python_path=python_path ) + if args.enable_transformers_tool_test and not args.disable_contrib_ops: + # These inference-only tests require the pinned PyTorch installed above. + # Keep them mandatory on these legs: missing DLPack/dtypes must fail, not skip. + run_subprocess( + [sys.executable, "onnxruntime_test_python_dlpack.py"], + cwd=cwd, + dll_path=dll_path, + python_path=python_path, + ) run_subprocess( [sys.executable, "onnxruntime_test_python_fpa_intb_tune.py"], cwd=cwd, From 650f904fed9e1e1573353dd781b770a307a4285c Mon Sep 17 00:00:00 2001 From: justinchuby <11205048+justinchuby@users.noreply.github.com> Date: Thu, 5 Nov 2026 04:22:36 +0000 Subject: [PATCH 02/11] Use object-based DLPack interchange for OrtValue Remove the legacy Python to_dlpack method and accept protocol objects consistently in native and public from_dlpack. Migrate callers, preserve capsule compatibility and training autograd semantics, and cover native protocol surfaces and producer errors. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../onnxruntime_inference_collection.py | 8 +-- .../python/onnxruntime_pybind_ortvalue.cc | 15 ++--- .../python/onnxruntime_pybind_state_common.cc | 15 ++++- .../python/onnxruntime_pybind_state_common.h | 5 +- .../test/python/onnxruntime_test_python.py | 13 ++--- .../python/onnxruntime_test_python_dlpack.py | 56 ++++++++++++++++++- .../python/training/ortmodule/_utils.py | 3 +- .../test/python/orttraining_test_ortvalue.py | 24 +++++--- 8 files changed, 96 insertions(+), 43 deletions(-) diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index f31b9d1244c7c..1636c8c8d99ee 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -1585,13 +1585,7 @@ def from_dlpack(cls, data, /) -> OrtValue: :return: An OrtValue wrapping the tensor data. :raises TypeError: If the object does not provide a valid DLPack capsule. """ - # If the input supports the __dlpack__ protocol, call it to get the capsule. - if hasattr(data, "__dlpack__"): - capsule = data.__dlpack__() - else: - capsule = data - - return cls(C.OrtValue.from_dlpack(capsule)) + return cls(C.OrtValue.from_dlpack(data)) def update_inplace(self, data) -> None: """ diff --git a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc index e557f15eaa877..05bb8961ded9a 100644 --- a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc @@ -505,12 +505,7 @@ void addOrtValueMethods(pybind11::module& m) { #endif }) #if defined(ENABLE_DLPACK) - .def("to_dlpack", [](OrtValue* ort_value) -> py::object { - ORT_ENFORCE(!IsWebGpuBuffer(*ort_value), "DLPack export is not supported for WebGPU OrtValues."); - return py::reinterpret_steal(ToDlpack(*ort_value)); }, - "Returns a DLPack capsule sharing the tensor buffer without copying. " - "The capsule retains the OrtValue; externally owned buffers must remain alive.") - .def_static("from_dlpack", [](py::object data) { return FromDlpack(data.ptr()); }, py::arg("data"), "Converts a tensor from a external library into an OrtValue by means of the __dlpack__ protocol.") + .def_static("from_dlpack", &FromDlpack, py::arg("data"), "Imports a tensor from a DLPack protocol object or legacy capsule without copying.") .def("__dlpack__", [](OrtValue* ort_value, py::object /* stream */) -> py::object { ORT_ENFORCE(!IsWebGpuBuffer(*ort_value), "DLPack export is not supported for WebGPU OrtValues."); return py::reinterpret_steal(ToDlpack(*ort_value)); }, py::arg("stream") = py::none(), @@ -542,7 +537,7 @@ void addOrtValueMethods(pybind11::module& m) { v->push_back(ortvalue); }) #if defined(ENABLE_DLPACK) - .def("push_back", [](std::vector* v, py::object dlpack_tensor) { v->push_back(FromDlpack(dlpack_tensor.ptr())); }, "Add a new OrtValue after being ownership was transferred from the DLPack structure.", py::arg("dlpack_tensor")) + .def("push_back", [](std::vector* v, py::object data) { v->push_back(FromDlpack(std::move(data))); }, "Add an OrtValue imported from a DLPack protocol object or legacy capsule.", py::arg("dlpack_tensor")) .def("push_back_batch", [](std::vector* v, std::vector& torch_tensors, std::vector& data_ptrs, std::vector& element_types, const std::vector>& shapes, const std::vector& devices) { for (size_t i = 0; i < torch_tensors.size(); ++i) { py::object& element_type = element_types.at(i); @@ -647,12 +642,12 @@ consuming the DLPack structure or return a list of capsule if this function is N :param to_tensor: this function takes a capsule holding a pointer onto a DLPack structure and returns a new tensor which becomes the new owner of the data. This function takes one python object and - returns a new python object. It fits the same signature as `torch.utils.from_dlpack`, + returns a new python object. It fits the same signature as `torch.from_dlpack`, if None, the method returns a capsule for every new DLPack structure. :return: a list containing the new tensors or a the new capsules if *to_tensor* is None -This method is used to replace `tuple(torch._C._from_dlpack(ov.to_dlpack()) for ov in ort_values)` -by a faster instruction `tuple(ort_values.to_dlpack(torch._C._from_dlpack))`. This loop +This method is used to replace `tuple(torch.from_dlpack(ov) for ov in ort_values)` +by a faster instruction `tuple(ort_values.to_dlpacks(torch.from_dlpack))`. This loop is difficult to parallelize as it goes through the GIL many times. It creates many tensors acquiring ownership of existing OrtValue. This method saves one object creation and an C++ allocation diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.cc b/onnxruntime/python/onnxruntime_pybind_state_common.cc index 9e35f3d8fd838..287c1dedb8a3a 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.cc +++ b/onnxruntime/python/onnxruntime_pybind_state_common.cc @@ -61,9 +61,18 @@ PyObject* ToDlpack(OrtValue& ort_value) { return PyCapsule_New(dlmanaged_tensor, "dltensor", DlpackCapsuleDestructor); } -// Consume a Capsule object and claims the ownership of its underlying tensor to -// create a OrtValue. This function calls DlpackToOrtValue(...) to do the conversion. -OrtValue FromDlpack(PyObject* dlpack_tensor) { +OrtValue FromDlpack(py::object data) { + if (!PyCapsule_CheckExact(data.ptr())) { + auto method = py::reinterpret_steal(PyObject_GetAttrString(data.ptr(), "__dlpack__")); + if (method) { + data = method(); + } else if (PyErr_ExceptionMatches(PyExc_AttributeError)) { + PyErr_Clear(); + } else { + throw py::error_already_set(); + } + } + PyObject* dlpack_tensor = data.ptr(); if (!PyCapsule_IsValid(dlpack_tensor, "dltensor")) { throw py::type_error("from_dlpack expected a valid DLPack capsule named 'dltensor'"); } diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.h b/onnxruntime/python/onnxruntime_pybind_state_common.h index d99c605a38505..9574ed33d8738 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.h +++ b/onnxruntime/python/onnxruntime_pybind_state_common.h @@ -545,9 +545,8 @@ bool CheckIfTensor(const std::vector& def_list, // Caller is responsible for releasing the capsule. PyObject* ToDlpack(OrtValue& ort_value); -// Consume a Capsule object and claims the ownership of its underlying tensor to -// create a OrtValue. This function calls DlpackToOrtValue(...) to do the conversion. -OrtValue FromDlpack(PyObject* dlpack_tensor); +// Import a DLPack protocol object or legacy capsule, transferring tensor ownership. +OrtValue FromDlpack(pybind11::object data); // Destructor for Capsule object holding a DLPack structure. void DlpackCapsuleDestructor(PyObject* data); diff --git a/onnxruntime/test/python/onnxruntime_test_python.py b/onnxruntime/test/python/onnxruntime_test_python.py index 0672b801b44e1..c4068224efc6d 100644 --- a/onnxruntime/test/python/onnxruntime_test_python.py +++ b/onnxruntime/test/python/onnxruntime_test_python.py @@ -1477,8 +1477,7 @@ def test_ort_value_dlpack(self): self.assertEqual(numpy_arr_input.shape, tuple(ortvalue.shape())) ptr = ortvalue._ortvalue.data_ptr() - dlp = ortvalue._ortvalue.to_dlpack() - ortvalue2 = C.OrtValue.from_dlpack(dlp) + ortvalue2 = C.OrtValue.from_dlpack(ortvalue) self.assertEqual(ptr, ortvalue2.data_ptr()) new_array = ortvalue2.numpy() np.testing.assert_equal(numpy_arr_input, new_array) @@ -1506,13 +1505,11 @@ def test_ort_value_dlpack_zero_size(self): with self.subTest(shape=shape): arr = np.zeros(shape, dtype=np.float32) # Test via numpy __dlpack__ protocol - dlp = arr.__dlpack__() - ortvalue = C.OrtValue.from_dlpack(dlp) + ortvalue = C.OrtValue.from_dlpack(arr) self.assertEqual(list(shape), list(ortvalue.shape())) # Test round-trip: OrtValue -> dlpack -> OrtValue ort_input = onnxrt.OrtValue.ortvalue_from_numpy(arr) - dlp2 = ort_input._ortvalue.to_dlpack() - ortvalue2 = C.OrtValue.from_dlpack(dlp2) + ortvalue2 = C.OrtValue.from_dlpack(ort_input._ortvalue) self.assertEqual(list(shape), list(ortvalue2.shape())) def test_ort_value_array_protocol(self): @@ -1692,9 +1689,11 @@ def __dlpack__(self): with self.assertRaisesRegex(TypeError, error): onnxrt.OrtValue.from_dlpack(InvalidDlpackProvider()) + with self.assertRaisesRegex(TypeError, error): + C.OrtValue.from_dlpack(InvalidDlpackProvider()) source = onnxrt.OrtValue.ortvalue_from_numpy(np.array([1.0], dtype=np.float32)) - capsule = source._ortvalue.to_dlpack() + capsule = source.__dlpack__() C.OrtValue.from_dlpack(capsule) with self.assertRaisesRegex(TypeError, error): C.OrtValue.from_dlpack(capsule) diff --git a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py index 530c18230a261..5cd375f3c0740 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py +++ b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py @@ -16,6 +16,7 @@ from onnx import TensorProto import onnxruntime as ort +from onnxruntime.capi import _pybind_state as C FLOAT8_DTYPES = ( (torch.float8_e4m3fn, TensorProto.FLOAT8E4M3FN), @@ -86,7 +87,7 @@ def test_pytorch_to_ort(self): with self.subTest(dtype=dtype, shape=shape, capsule=capsule): source = make_tensor(dtype, shape) expected = source.clone() - imported = ort.OrtValue.from_dlpack(torch.utils.dlpack.to_dlpack(source) if capsule else source) + imported = ort.OrtValue.from_dlpack(source.__dlpack__() if capsule else source) self.assertEqual(imported.element_type(), onnx_type) self.assertEqual(imported.shape(), list(shape)) self.assertEqual(imported.__dlpack_device__(), (1, 0)) @@ -115,7 +116,7 @@ def test_ort_to_pytorch(self): ctypes.memmove( source.data_ptr(), expected.data_ptr(), expected.numel() * expected.element_size() ) - result = torch.from_dlpack(source._ortvalue.to_dlpack() if capsule else source) + result = torch.from_dlpack(source.__dlpack__() if capsule else source) self.assert_tensor(result, expected) if expected.numel(): self.assertEqual(result.data_ptr(), source.data_ptr()) @@ -148,7 +149,7 @@ def test_capsule_keeps_ort_allocation_alive(self): source = ort.OrtValue.ortvalue_from_shape_and_type([2, 3], onnx_type) pointer = source.data_ptr() ctypes.memmove(pointer, expected.data_ptr(), expected.numel() * expected.element_size()) - capsule = source._ortvalue.to_dlpack() + capsule = source.__dlpack__() source_ref = weakref.ref(source) del source gc.collect() @@ -157,6 +158,55 @@ def test_capsule_keeps_ort_allocation_alive(self): self.assertEqual(result.data_ptr(), pointer) self.assert_tensor(result, expected) + def test_ortvalue_dlpack_surface(self): + for value_type in (ort.OrtValue, C.OrtValue): + with self.subTest(value_type=value_type): + self.assertFalse(hasattr(value_type, "to_dlpack")) + for method in ("from_dlpack", "__dlpack__", "__dlpack_device__"): + self.assertTrue(callable(getattr(value_type, method))) + + def test_native_ortvalue_protocol(self): + for dtype, onnx_type in INTERCHANGE_DTYPES: + with self.subTest(dtype=dtype): + source = make_tensor(dtype, (2, 3)) + native = C.OrtValue.from_dlpack(source) + self.assertEqual(native.element_type(), onnx_type) + self.assertEqual(native.data_ptr(), source.data_ptr()) + self.assertEqual(native.__dlpack_device__(), (1, 0)) + self.assert_tensor(torch.from_dlpack(native), source) + public = ort.OrtValue.from_dlpack(native) + self.assertEqual(public.data_ptr(), source.data_ptr()) + imported = C.OrtValue.from_dlpack(public) + self.assertEqual(imported.data_ptr(), source.data_ptr()) + self.assert_tensor(torch.from_dlpack(imported), source) + values = C.OrtValueVector() + values.push_back(source) + self.assertEqual(values[0].data_ptr(), source.data_ptr()) + self.assert_tensor(torch.from_dlpack(values[0]), source) + values.push_back(public) + for consumer in (torch.from_dlpack, C.OrtValue.from_dlpack): + for result in values.to_dlpacks(consumer): + self.assertEqual(result.data_ptr(), source.data_ptr()) + self.assert_tensor(torch.from_dlpack(result), source) + + def test_from_dlpack_propagates_producer_errors(self): + class InvalidProducer: + def __dlpack__(self): + raise ValueError("DLPack producer failed") + + class InvalidProtocolAttribute: + @property + def __dlpack__(self): + raise ValueError("DLPack lookup failed") + + for value_type in (ort.OrtValue, C.OrtValue): + for producer, error in ( + (InvalidProducer(), "DLPack producer failed"), + (InvalidProtocolAttribute(), "DLPack lookup failed"), + ): + with self.subTest(value_type=value_type, producer=producer), self.assertRaisesRegex(ValueError, error): + value_type.from_dlpack(producer) + def test_packed_dtype_boundaries(self): for dtype, onnx_type, storage_type, bits, error in PACKED_DTYPES: with self.subTest(dtype=dtype): diff --git a/orttraining/orttraining/python/training/ortmodule/_utils.py b/orttraining/orttraining/python/training/ortmodule/_utils.py index 419856bc88cc7..3aeada12de7aa 100644 --- a/orttraining/orttraining/python/training/ortmodule/_utils.py +++ b/orttraining/orttraining/python/training/ortmodule/_utils.py @@ -50,7 +50,8 @@ def _ortvalue_from_torch_tensor(torch_tensor: torch.Tensor) -> C.OrtValue: """Converts a torch tensor to an OrtValue.""" if torch_tensor.device.type == "ort": return C.aten_ort_tensor_to_ort_value(torch_tensor) - return C.OrtValue.from_dlpack(to_dlpack(torch_tensor)) + # The protocol rejects requires_grad tensors; ORTModule manages autograd separately. + return C.OrtValue.from_dlpack(torch_tensor.detach() if torch_tensor.requires_grad else torch_tensor) def _ortvalues_to_torch_tensor( diff --git a/orttraining/orttraining/test/python/orttraining_test_ortvalue.py b/orttraining/orttraining/test/python/orttraining_test_ortvalue.py index 8caab593ae43b..655db137b4109 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortvalue.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortvalue.py @@ -32,13 +32,12 @@ def C_ort_from_dlpack(dlpack): return from_dlpack(dlpack) -def _torch_tensor_from_dl_pack(dlpack, device): - return from_dlpack(dlpack) if device.type != "ort" else C_ort_from_dlpack(dlpack) +def _torch_tensor_from_dl_pack(tensor, device): + return from_dlpack(tensor) if device.type != "ort" else C_ort_from_dlpack(tensor.__dlpack__()) def _ortvalue_to_torch_tensor(ortvalue, device): - dlpack_tensor = ortvalue.to_dlpack() - return _torch_tensor_from_dl_pack(dlpack_tensor, device) + return _torch_tensor_from_dl_pack(ortvalue, device) def _ortvalues_to_torch_tensor(vect, device): @@ -52,8 +51,7 @@ def testOrtValueDlPack_float32(self): self.assertEqual(numpy_arr_input.shape, tuple(ortvalue.shape())) ptr = ortvalue._ortvalue.data_ptr() - dlp = ortvalue._ortvalue.to_dlpack() - ortvalue2 = C_OrtValue.from_dlpack(dlp) + ortvalue2 = C_OrtValue.from_dlpack(ortvalue) self.assertEqual(ptr, ortvalue2.data_ptr()) new_array = ortvalue2.numpy() assert_almost_equal(numpy_arr_input, new_array) @@ -73,8 +71,7 @@ def testOrtValueDlPack_bool(self): self.assertEqual(numpy_arr_input.shape, tuple(ortvalue.shape())) ptr = ortvalue._ortvalue.data_ptr() - dlp = ortvalue._ortvalue.to_dlpack() - ortvalue2 = C_OrtValue.from_dlpack(dlp) + ortvalue2 = C_OrtValue.from_dlpack(ortvalue) self.assertEqual(ptr, ortvalue2.data_ptr()) new_array = ortvalue2.numpy() assert_almost_equal(numpy_arr_input, new_array) @@ -182,11 +179,20 @@ def test_torch_dlpack_bool_and_uint8_preserve_dtype(self): original = torch.tensor([1, 0, 1, 0], dtype=dtype) ortvalue = _utils._ortvalue_from_torch_tensor(original) self.assertEqual(ortvalue.data_ptr(), original.data_ptr()) - result = from_dlpack(ortvalue.to_dlpack()) + result = from_dlpack(ortvalue) self.assertEqual(result.dtype, dtype) self.assertEqual(result.data_ptr(), original.data_ptr()) torch.testing.assert_close(original, result) + def test_torch_dlpack_preserves_autograd_source(self): + original = torch.tensor([1.0, 2.0], requires_grad=True) + ortvalue = _utils._ortvalue_from_torch_tensor(original) + result = from_dlpack(ortvalue) + self.assertEqual(result.data_ptr(), original.data_ptr()) + self.assertTrue(original.requires_grad) + self.assertFalse(result.requires_grad) + torch.testing.assert_close(result, original) + def testOrtValueVectorDlPackNone_cpu(self): self.OrtValueVectorDlPackOrtValue(None, None, "cpu") From becf59dc566ed95402252e1c5c951ce7e8152e3f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 4 Nov 2026 22:32:53 -0800 Subject: [PATCH 03/11] Harden OrtValue DLPack ownership and CUDA handoffs Retain NumPy storage in native values, fence CUDA imports and exports, and make batched capsules independently owned and exception-safe. Remove unused dlpack_at and cover ownership, callback failures, vector mutation, and CUDA stream handoffs. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../providers/cuda/cuda_provider_factory.cc | 14 + .../providers/cuda/cuda_provider_factory.h | 2 + .../onnxruntime_inference_collection.py | 26 +- .../python/onnxruntime_pybind_ortvalue.cc | 117 +++----- .../python/onnxruntime_pybind_state_common.cc | 67 ++++- .../python/onnxruntime_pybind_state_common.h | 2 +- .../cuda/test_cases/cuda_test_provider.cc | 5 + .../test/python/onnxruntime_test_python.py | 2 +- .../python/onnxruntime_test_python_dlpack.py | 267 ++++++++++++++++++ 9 files changed, 411 insertions(+), 91 deletions(-) diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index e95d1ecf5b8d4..2be9b3d2920d5 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -228,6 +228,20 @@ struct ProviderInfo_CUDA_Impl final : ProviderInfo_CUDA { return num_devices; } + Status SynchronizeDevice(int device_id) override { + int original_device; + CUDA_RETURN_IF_ERROR(cudaGetDevice(&original_device)); + if (original_device != device_id) { + CUDA_RETURN_IF_ERROR(cudaSetDevice(device_id)); + } + const auto sync_error = cudaDeviceSynchronize(); + // Restore the caller's device even when synchronization reports an error. + const auto restore_error = original_device != device_id ? cudaSetDevice(original_device) : cudaSuccess; + CUDA_RETURN_IF_ERROR(sync_error); + CUDA_RETURN_IF_ERROR(restore_error); + return Status::OK(); + } + void CUDAExecutionProviderInfo__FromProviderOptions(const ProviderOptions& options, CUDAExecutionProviderInfo& info) override { info = CUDAExecutionProviderInfo::FromProviderOptions(options); } diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.h b/onnxruntime/core/providers/cuda/cuda_provider_factory.h index 1a4b19cb100d3..ca6581a8c62df 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.h +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.h @@ -61,6 +61,8 @@ struct ProviderInfo_CUDA { ORT_NOT_IMPLEMENTED(__FUNCTION__, " is only implements in test code path."); } + virtual Status SynchronizeDevice(int device_id) = 0; + protected: ~ProviderInfo_CUDA() = default; // Can only be destroyed through a subclass instance }; diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index 1636c8c8d99ee..c2d538e346d8d 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -1308,8 +1308,10 @@ def ortvalue_from_numpy( cls, numpy_obj: np.ndarray, /, device_type="cpu", device_id=0, vendor_id: int | OrtDeviceVendorId = -1 ) -> OrtValue: """ - Factory method to construct an OrtValue (which holds a Tensor) from a given Numpy object - A copy of the data in the Numpy object is held by the OrtValue only if the device is NOT cpu + Construct an OrtValue (which holds a Tensor) from a NumPy array. + Contiguous CPU arrays are shared without copying; noncontiguous arrays + are copied to contiguous storage. Other devices copy the data to the + target device. The native OrtValue retains any NumPy storage it uses. :param numpy_obj: The Numpy object to construct the OrtValue from :param device_type: e.g. cpu, cuda, cann, cpu by default @@ -1331,10 +1333,10 @@ def ortvalue_from_numpy( def ortvalue_from_numpy_with_onnx_type(cls, data: np.ndarray, /, onnx_element_type: int) -> OrtValue: """ This method creates an instance of OrtValue on top of the numpy array. - No data copy is made and the lifespan of the resulting OrtValue should never - exceed the lifespan of bytes object. The API attempts to reinterpret - the data type which is expected to be the same size. This is useful + No data copy is made; the native OrtValue retains the NumPy array. + The API reinterprets the data type, which must have the same element size. This is useful when we want to use an ONNX data type that is not supported by numpy. + The array must be C-contiguous; use ``numpy.ascontiguousarray`` otherwise. :param data: numpy.ndarray. :param onnx_element_type: a valid onnx TensorProto::DataType enum value @@ -1531,13 +1533,17 @@ def __dlpack__(self, *, stream=None): The OrtValue must hold a contiguous tensor. No data is copied. The capsule and its consumer retain the tensor allocation, so this - OrtValue can be released. For externally owned buffers (for example, - a NumPy-backed CPU OrtValue), the external owner must remain alive. + OrtValue can be released. NumPy-backed OrtValues also retain the NumPy + allocation through the consumer's lifetime. WebGPU OrtValues cannot be exported through DLPack because their data is stored in opaque WebGPU buffers rather than CUDA-addressable memory. - :param stream: Optional stream on which the tensor data is accessible. - Currently unused; included for protocol compliance. + :param stream: CUDA consumer stream: ``None`` or ``1`` for the legacy + default, ``2`` for the per-thread default, or a positive stream handle. + Since OrtValues do not track their producer stream, CUDA exports + synchronize the entire tensor device before returning. ``-1`` skips + synchronization; the caller must ensure the data is ready. CPU exports + require ``None``. :return: A PyCapsule holding a DLManagedTensor. """ return self._ortvalue.__dlpack__(stream=stream) @@ -1579,6 +1585,8 @@ def from_dlpack(cls, data, /) -> OrtValue: PyTorch cannot consume ORT's packed float4, int4/uint4, or int2/uint2. No data is copied; the new OrtValue shares memory with the source. + CUDA imports synchronize the tensor's device before returning, so ORT + can safely use the storage on its own streams. :param data: A tensor object supporting the DLPack protocol, or a raw DLPack PyCapsule. diff --git a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc index 05bb8961ded9a..378bff1b1eabc 100644 --- a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc @@ -24,6 +24,21 @@ bool IsWebGpuBuffer(const OrtValue& ort_value) { ort_value.Get().Location().name == WEBGPU_BUFFER; } +void RetainNumpyOwner(OrtValue& ort_value, const py::object& array) { + OrtValue tensor_owner = ort_value; + auto* tensor = tensor_owner.GetMutable(); + auto type = tensor_owner.Type(); + auto array_owner = tensor->OwnsBuffer() ? py::object{} : array; + // Retain the owner in the native value so copies, vectors and DLPack consumers + // remain safe independently of the Python wrapper, including destruction without the GIL. + ort_value.Init(tensor, type, + [tensor_owner = std::move(tensor_owner), array_owner = std::move(array_owner)](void*) mutable { + py::gil_scoped_acquire acquire; + tensor_owner = OrtValue{}; + array_owner = py::object{}; + }); +} + std::unique_ptr OrtValueFromShapeAndType(const std::vector& shape, MLDataType element_type, const OrtDevice& device) { @@ -107,8 +122,7 @@ std::unique_ptr OrtValueFromShapeAndTypeWithMemoryInfo(const std::vect void addOrtValueMethods(pybind11::module& m) { py::class_ ortvalue_binding(m, "OrtValue"); ortvalue_binding - // Factory method to create an OrtValue (Tensor) from the given Numpy object - // The Tensor allocates and manages its own memory (on the specified device) and copies data from the Numpy data buffer + // Share contiguous CPU NumPy storage; copy other inputs to the requested device. .def_static("ortvalue_from_numpy", [](const py::object& array_on_cpu, const OrtDevice& device) { if (!IsNumericNumpyArray(array_on_cpu)) { throw std::runtime_error("Creation of OrtValues is currently only supported from non-string numpy arrays"); @@ -122,6 +136,7 @@ void addOrtValueMethods(pybind11::module& m) { // Likewise, there is no need to specify the name (as the name was previously used to lookup the def list) CreateGenericMLValue(nullptr, GetAllocator(), "", array_on_cpu, ml_value.get(), true); + RetainNumpyOwner(*ml_value, array_on_cpu); } else if (device.Type() == OrtDevice::GPU) { #if USE_DML if (device.Vendor() == OrtDevice::VendorIds::MICROSOFT) { @@ -280,6 +295,11 @@ void addOrtValueMethods(pybind11::module& m) { // Create an ortvalue value on top of the numpy array, but interpret the data // as a different type with the same element size. .def_static("ortvalue_from_numpy_with_onnx_type", [](py::array& data, int32_t onnx_element_type) -> std::unique_ptr { + if (!(data.flags() & py::array::c_style)) { + throw py::value_error( + "ortvalue_from_numpy_with_onnx_type requires a C-contiguous NumPy array. " + "Use numpy.ascontiguousarray(data) before importing."); + } if (!ONNX_NAMESPACE::TensorProto_DataType_IsValid(onnx_element_type)) { ORT_THROW("Not a valid ONNX Tensor data type: ", onnx_element_type); } @@ -296,6 +316,7 @@ void addOrtValueMethods(pybind11::module& m) { auto ort_value = std::make_unique(); Tensor::InitOrtValue(element_type, GetShape(data), const_cast(data.data()), cpu_allocator->Info(), *ort_value); + RetainNumpyOwner(*ort_value, data); return ort_value; }) // Factory method to create an OrtValue from the given shape and numpy element type on the specified device. @@ -506,12 +527,12 @@ void addOrtValueMethods(pybind11::module& m) { }) #if defined(ENABLE_DLPACK) .def_static("from_dlpack", &FromDlpack, py::arg("data"), "Imports a tensor from a DLPack protocol object or legacy capsule without copying.") - .def("__dlpack__", [](OrtValue* ort_value, py::object /* stream */) -> py::object { + .def("__dlpack__", [](OrtValue* ort_value, py::object stream) -> py::object { ORT_ENFORCE(!IsWebGpuBuffer(*ort_value), "DLPack export is not supported for WebGPU OrtValues."); - return py::reinterpret_steal(ToDlpack(*ort_value)); }, py::arg("stream") = py::none(), + return py::reinterpret_steal(ToDlpack(*ort_value, std::move(stream))); }, py::arg("stream") = py::none(), "Returns a DLPack representing the tensor (part of __dlpack__ protocol). " - "The capsule retains the OrtValue and shares its buffer without copying; " - "externally owned buffers must remain alive.") + "The capsule retains the tensor and its NumPy owner without copying. " + "CUDA exports synchronize the tensor's device unless stream=-1 requests no synchronization.") .def("__dlpack_device__", [](const OrtValue* ort_value) -> py::tuple { ORT_ENFORCE(ort_value->IsTensor(), "Only tensor type OrtValues are supported"); ORT_ENFORCE(!IsWebGpuBuffer(*ort_value), "DLPack export is not supported for WebGPU OrtValues."); @@ -574,11 +595,6 @@ void addOrtValueMethods(pybind11::module& m) { .def("__len__", [](const std::vector& v) { return v.size(); }) .def("__iter__", [](const std::vector& v) { return py::make_iterator(v.cbegin(), v.cend()); }, py::keep_alive<0, 1>()) .def("__getitem__", [](const std::vector& v, const size_t idx) { return v.at(idx); }, py::keep_alive<0, 1>()) -#if defined(ENABLE_DLPACK) - .def("dlpack_at", [](std::vector* v, const size_t idx) { - ORT_ENFORCE(!IsWebGpuBuffer(v->at(idx)), "DLPack export is not supported for WebGPU OrtValues."); - return py::reinterpret_steal(ToDlpack(v->at(idx))); }) -#endif .def("element_type_at", [](std::vector* v, const size_t idx) -> int32_t { return GetTensorProtoType(v->at(idx)); }, "Returns an integer equal to the ONNX proto type of the tensor at position i. " "This integer is one type defined by ONNX TensorProto_DataType " @@ -587,71 +603,32 @@ void addOrtValueMethods(pybind11::module& m) { py::arg("idx")) #if defined(ENABLE_DLPACK) .def("to_dlpacks", [](const std::vector& v, py::object to_tensor) -> py::list { - if (v.size() == 0) - return py::list(); - - py::list list_dlpacks; - PyObject* obj; - - py::gil_scoped_acquire acquire; - - if (to_tensor.is_none()) { - DLManagedTensor* dlmanaged_tensor; - - for (auto it : v) { - ORT_ENFORCE(!IsWebGpuBuffer(it), "DLPack export is not supported for WebGPU OrtValues."); - dlmanaged_tensor = dlpack::OrtValueToDlpack(it); - py::capsule capsule(dlmanaged_tensor, "dltensor", DlpackCapsuleDestructor); - list_dlpacks.append(capsule); - } - } else { - DLManagedTensor* dlmanaged_tensor; - PyObject* capsule = NULL; - PyObject* handle = to_tensor.ptr(); - - for (auto it : v) { - // A new instance of dlpack needs to be created. The object which consumes it - // is responsible for its deletion. - ORT_ENFORCE(!IsWebGpuBuffer(it), "DLPack export is not supported for WebGPU OrtValues."); - dlmanaged_tensor = dlpack::OrtValueToDlpack(it); - if (capsule == NULL) { - capsule = PyCapsule_New(dlmanaged_tensor, "dltensor", NULL); - if (capsule == NULL) - throw std::runtime_error("Unexpected error: empty capsule returned."); - } else { - // The same capsule is reused but FromDLPack rename the capsule into used_dltensor. - PyCapsule_SetName(capsule, "dltensor"); - PyCapsule_SetPointer(capsule, dlmanaged_tensor); - } - obj = PyObject_CallFunctionObjArgs(handle, capsule, NULL); - if (obj == NULL) - throw std::runtime_error("to_tensor returned a null pointer. This may be caused by the data conversion."); - list_dlpacks.append(obj); - Py_DECREF(obj); - } - if (capsule != NULL) { - // This test is never wrong because v is not empty if the execution goes through that path. - // If not present, Guardian detects a potential failure. - Py_DECREF(capsule); - } + if (!to_tensor.is_none() && !PyCallable_Check(to_tensor.ptr())) { + throw py::type_error("OrtValueVector.to_dlpacks expects a callable consumer or None."); + } + // Callbacks can mutate the original vector, including reallocating its storage. + auto values = v; + py::list results; + for (auto& value : values) { + ORT_ENFORCE(!IsWebGpuBuffer(value), "DLPack export is not supported for WebGPU OrtValues."); + auto capsule = py::reinterpret_steal(ToDlpack(value)); + results.append(to_tensor.is_none() ? capsule : to_tensor(capsule)); } - return list_dlpacks; }, - R"pbdoc(Converts all OrtValue into tensors through DLPack protocol, the method creates -a DLPack structure for every tensors, then calls python function `to_tensor` to a new object -consuming the DLPack structure or return a list of capsule if this function is None. + return results; }, + R"pbdoc(Exports a snapshot of the vector as independently owned DLPack capsules. +If `to_tensor` is callable, returns the result of calling it with each capsule. +Consumer exceptions propagate unchanged. Unconsumed capsules release their tensor +references when destroyed, including when conversion fails. :param to_tensor: this function takes a capsule holding a pointer onto a DLPack structure and returns a new tensor which becomes the new owner of the data. This function takes one python object and returns a new python object. It fits the same signature as `torch.from_dlpack`, if None, the method returns a capsule for every new DLPack structure. -:return: a list containing the new tensors or a the new capsules if *to_tensor* is None - -This method is used to replace `tuple(torch.from_dlpack(ov) for ov in ort_values)` -by a faster instruction `tuple(ort_values.to_dlpacks(torch.from_dlpack))`. This loop -is difficult to parallelize as it goes through the GIL many times. -It creates many tensors acquiring ownership of existing OrtValue. -This method saves one object creation and an C++ allocation -for every transferred tensor. +:return: a list containing the new tensors or the new capsules if *to_tensor* is None + +Use `ort_values.to_dlpacks(torch.from_dlpack)` to avoid creating a Python OrtValue +wrapper for every tensor. Each capsule may be consumed only once. CUDA exports +synchronize the tensor's device before invoking the consumer. )pbdoc", py::arg("to_tensor")) #endif diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.cc b/onnxruntime/python/onnxruntime_pybind_state_common.cc index 287c1dedb8a3a..038908efeaf2f 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.cc +++ b/onnxruntime/python/onnxruntime_pybind_state_common.cc @@ -45,20 +45,64 @@ void* empty_cache_fn{nullptr}; #if defined(ENABLE_DLPACK) void DlpackCapsuleDestructor(PyObject* data) { - DLManagedTensor* dlmanaged_tensor = reinterpret_cast(PyCapsule_GetPointer(data, "dltensor")); - if (dlmanaged_tensor) { - // The dlmanaged_tensor has not been consumed, call deleter ourselves. - dlmanaged_tensor->deleter(const_cast(dlmanaged_tensor)); - } else { - // The dlmanaged_tensor has been consumed, - // PyCapsule_GetPointer has set an error indicator. - PyErr_Clear(); + // Do not disturb a pending exception while releasing an unconsumed capsule. + if (PyCapsule_IsValid(data, "dltensor")) { + auto* tensor = static_cast(PyCapsule_GetPointer(data, "dltensor")); + tensor->deleter(tensor); } } -PyObject* ToDlpack(OrtValue& ort_value) { +namespace { +void SynchronizeDlpackDevice(const DLDevice& device) { + if (device.device_type != kDLCUDA) { + return; + } +#if defined(USE_CUDA) || defined(USE_CUDA_PROVIDER_INTERFACE) + if (auto* provider = TryGetProviderInfo_CUDA()) { + ORT_THROW_IF_ERROR(provider->SynchronizeDevice(device.device_id)); + return; + } +#endif + throw py::buffer_error( + "CUDA DLPack interchange requires an available CUDA provider to synchronize the tensor's device. " + "Install a CUDA-enabled ONNX Runtime package and load its CUDA provider. For an already synchronized " + "export, call __dlpack__(stream=-1) explicitly."); +} +} // namespace + +PyObject* ToDlpack(OrtValue& ort_value, py::object stream) { + const auto device = dlpack::GetDlpackDevice(ort_value, ort_value.Get().Location().device.Id()); + bool synchronize = true; + if (!stream.is_none()) { + if (device.device_type != kDLCUDA) { + throw py::buffer_error("CPU DLPack exports require stream=None; omit the stream argument."); + } + if (!PyLong_Check(stream.ptr()) || PyBool_Check(stream.ptr())) { + throw py::type_error("A CUDA DLPack stream must be an integer: -1, 1, 2, or a positive stream handle."); + } + const auto stream_id = PyLong_AsLongLong(stream.ptr()); + if (PyErr_Occurred()) { + throw py::error_already_set(); + } + if (stream_id == 0 || stream_id < -1) { + throw py::value_error( + "Invalid CUDA DLPack stream: use 1 for the legacy default, 2 for the per-thread default, " + "a positive stream handle, or -1 to opt out of synchronization."); + } + synchronize = stream_id != -1; + } + // OrtValue does not record its producer stream, so a device-wide fence is + // necessary before handing its storage to any consumer stream. + if (synchronize) { + SynchronizeDlpackDevice(device); + } DLManagedTensor* dlmanaged_tensor = dlpack::OrtValueToDlpack(ort_value); - return PyCapsule_New(dlmanaged_tensor, "dltensor", DlpackCapsuleDestructor); + PyObject* capsule = PyCapsule_New(dlmanaged_tensor, "dltensor", DlpackCapsuleDestructor); + if (!capsule) { + dlmanaged_tensor->deleter(dlmanaged_tensor); + throw py::error_already_set(); + } + return capsule; } OrtValue FromDlpack(py::object data) { @@ -79,6 +123,9 @@ OrtValue FromDlpack(py::object data) { // Extract DLPack tensor pointer from the capsule carrier. DLManagedTensor* dlmanaged_tensor = (DLManagedTensor*)PyCapsule_GetPointer(dlpack_tensor, "dltensor"); + // The producer's default-stream handoff alone is insufficient for ORT's + // session-specific streams. Fence before importing both objects and raw capsules. + SynchronizeDlpackDevice(dlmanaged_tensor->dl_tensor.device); OrtValue ort_value = dlpack::DlpackToOrtValue(dlmanaged_tensor); // Make sure this capsule will never be used again. PyCapsule_SetName(dlpack_tensor, "used_dltensor"); diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.h b/onnxruntime/python/onnxruntime_pybind_state_common.h index 9574ed33d8738..9ea5f4890d0b5 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.h +++ b/onnxruntime/python/onnxruntime_pybind_state_common.h @@ -543,7 +543,7 @@ bool CheckIfTensor(const std::vector& def_list, // Allocate a capsule retaining an OrtValue reference via OrtValueToDlpack. // Caller is responsible for releasing the capsule. -PyObject* ToDlpack(OrtValue& ort_value); +PyObject* ToDlpack(OrtValue& ort_value, pybind11::object stream = pybind11::none()); // Import a DLPack protocol object or legacy capsule, transferring tensor ownership. OrtValue FromDlpack(pybind11::object data); diff --git a/onnxruntime/test/providers/cuda/test_cases/cuda_test_provider.cc b/onnxruntime/test/providers/cuda/test_cases/cuda_test_provider.cc index 12e78834987c5..5019cdc50b793 100644 --- a/onnxruntime/test/providers/cuda/test_cases/cuda_test_provider.cc +++ b/onnxruntime/test/providers/cuda/test_cases/cuda_test_provider.cc @@ -82,6 +82,11 @@ struct ProviderInfo_CUDA_TestImpl : ProviderInfo_CUDA { int cudaGetDeviceCount() override { return 0; } + Status SynchronizeDevice(int) override { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "The CUDA test provider does not support DLPack device synchronization."); + } + void CUDAExecutionProviderInfo__FromProviderOptions(const ProviderOptions&, CUDAExecutionProviderInfo&) override {} #if defined(USE_CUDA) && defined(ORT_USE_NCCL) && defined(USE_NCCL_P2P) && defined(ENABLE_TRAINING) diff --git a/onnxruntime/test/python/onnxruntime_test_python.py b/onnxruntime/test/python/onnxruntime_test_python.py index c4068224efc6d..800f5d397deec 100644 --- a/onnxruntime/test/python/onnxruntime_test_python.py +++ b/onnxruntime/test/python/onnxruntime_test_python.py @@ -2442,7 +2442,7 @@ def run_and_copy_output(current_session, current_io_binding, run_options=None): np.testing.assert_allclose(run_and_copy_output(session, io_binding), expected) raw_outputs = io_binding._iobinding.get_outputs() with self.assertRaisesRegex(RuntimeError, "DLPack export"): - raw_outputs.dlpack_at(0) + raw_outputs[0].__dlpack__() with self.assertRaisesRegex(RuntimeError, "DLPack export"): raw_outputs.to_dlpacks(None) indexed_raw_output = raw_outputs[0] diff --git a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py index 5cd375f3c0740..25e6a92662ad1 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py +++ b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py @@ -12,6 +12,7 @@ import unittest import weakref +import numpy as np import torch from onnx import TensorProto @@ -164,6 +165,196 @@ def test_ortvalue_dlpack_surface(self): self.assertFalse(hasattr(value_type, "to_dlpack")) for method in ("from_dlpack", "__dlpack__", "__dlpack_device__"): self.assertTrue(callable(getattr(value_type, method))) + self.assertFalse(hasattr(C.OrtValueVector, "dlpack_at")) + + def test_numpy_owner_survives_dlpack_export(self): + for native in (False, True): + for reinterpret in (False, True): + for capsule in (False, True): + with self.subTest(native=native, reinterpret=reinterpret, capsule=capsule): + array = np.arange(6, dtype=np.float32) + owner_ref = weakref.ref(array) + pointer = array.ctypes.data + if reinterpret: + value_type = C.OrtValue if native else ort.OrtValue + value = value_type.ortvalue_from_numpy_with_onnx_type(array, TensorProto.FLOAT) + elif native: + value = C.OrtValue.ortvalue_from_numpy(array, C.OrtDevice(C.OrtDevice.cpu(), 0, 0)) + else: + value = ort.OrtValue.ortvalue_from_numpy(array) + exported = value.__dlpack__() if capsule else torch.from_dlpack(value) + del array, value + gc.collect() + self.assertIsNotNone(owner_ref()) + result = torch.from_dlpack(exported) if capsule else exported + self.assertEqual(result.data_ptr(), pointer) + torch.testing.assert_close(result, torch.arange(6, dtype=torch.float32)) + result[0] = 7 + self.assertEqual(owner_ref()[0], 7) + del exported, result + gc.collect() + self.assertIsNone(owner_ref()) + + def test_numpy_owner_survives_native_copy_and_vector_export(self): + array = np.arange(6, dtype=np.float32) + owner_ref = weakref.ref(array) + value = ort.OrtValue.ortvalue_from_numpy(array) + imported = C.OrtValue.from_dlpack(value) + values = C.OrtValueVector() + values.push_back(imported) + del array, value, imported + gc.collect() + self.assertIsNotNone(owner_ref()) + result = values.to_dlpacks(torch.from_dlpack)[0] + del values + gc.collect() + torch.testing.assert_close(result, torch.arange(6, dtype=torch.float32)) + del result + gc.collect() + self.assertIsNone(owner_ref()) + + def test_numpy_noncontiguous_copy_survives_export(self): + array = np.arange(6, dtype=np.float32)[::-1] + owner_ref = weakref.ref(array) + value = ort.OrtValue.ortvalue_from_numpy(array) + result = torch.from_dlpack(value) + self.assertNotEqual(result.data_ptr(), array.ctypes.data) + del array, value + gc.collect() + self.assertIsNone(owner_ref()) + torch.testing.assert_close(result, torch.arange(6, dtype=torch.float32).flip(0)) + + def test_numpy_reinterpret_rejects_noncontiguous_storage(self): + for value_type in (ort.OrtValue, C.OrtValue): + for array in (np.arange(6, dtype=np.float32)[::-1], np.zeros((2, 3), dtype=np.float32).T): + with ( + self.subTest(value_type=value_type, strides=array.strides), + self.assertRaisesRegex(ValueError, "C-contiguous NumPy array.*ascontiguousarray"), + ): + value_type.ortvalue_from_numpy_with_onnx_type(array, TensorProto.FLOAT) + + def test_cpu_stream_argument(self): + for value in ( + ort.OrtValue.ortvalue_from_shape_and_type([1], TensorProto.FLOAT), + C.OrtValue.from_dlpack(torch.ones(1)), + ): + with self.subTest(value_type=type(value)): + torch.from_dlpack(value.__dlpack__(stream=None)) + for stream in (-1, 0, 1, 2, False, "invalid"): + with self.subTest(stream=stream), self.assertRaisesRegex(BufferError, "require stream=None"): + value.__dlpack__(stream=stream) + + def test_vector_capsules_are_independent(self): + values = C.OrtValueVector() + for offset in (0, 10): + values.push_back(torch.arange(6, dtype=torch.float32) + offset) + for consumer in (None, lambda capsule: capsule): + with self.subTest(consumer=consumer): + capsules = values.to_dlpacks(consumer) + self.assertIsNot(capsules[0], capsules[1]) + results = [torch.from_dlpack(capsule) for capsule in capsules] + for offset, result in zip((0, 10), results, strict=True): + torch.testing.assert_close(result, torch.arange(6, dtype=torch.float32) + offset) + for capsule in capsules: + with self.assertRaisesRegex(TypeError, "valid DLPack capsule"): + C.OrtValue.from_dlpack(capsule) + + def test_vector_unconsumed_capsules_release_owners(self): + for consumer in (None, lambda capsule: None): + with self.subTest(consumer=consumer): + array = np.arange(6, dtype=np.float32) + owner_ref = weakref.ref(array) + values = C.OrtValueVector() + values.push_back(ort.OrtValue.ortvalue_from_numpy(array)._get_c_value()) + del array + results = values.to_dlpacks(consumer) + del values + gc.collect() + if consumer is None: + self.assertIsNotNone(owner_ref()) + del results + gc.collect() + self.assertIsNone(owner_ref()) + + def test_vector_consumer_errors_preserve_exception_and_release_owners(self): + for fail_at in (0, 1): + for consume_before_error in (False, True): + with self.subTest(fail_at=fail_at, consume_before_error=consume_before_error): + arrays = [np.arange(6, dtype=np.float32) + offset for offset in (0, 10)] + owner_refs = [weakref.ref(array) for array in arrays] + values = C.OrtValueVector() + for array in arrays: + values.push_back(ort.OrtValue.ortvalue_from_numpy(array)._get_c_value()) + del array, arrays + error = ValueError("consumer failed") + calls = 0 + + def consumer(capsule, fail_at=fail_at, consume_before_error=consume_before_error, error=error): + nonlocal calls + index = calls + calls += 1 + if index == fail_at: + if consume_before_error: + tensor = torch.from_dlpack(capsule) + self.assertEqual(tensor.numel(), 6) + raise error + return torch.from_dlpack(capsule) + + with self.assertRaises(ValueError) as caught: + values.to_dlpacks(consumer) + self.assertIs(caught.exception, error) + self.assertEqual(calls, fail_at + 1) + # A traceback may legitimately retain the callback's capsule or tensor. + error.__traceback__ = None + del values + gc.collect() + self.assertTrue(all(owner_ref() is None for owner_ref in owner_refs)) + + def test_vector_consumer_can_retain_capsule_after_error(self): + array = np.arange(6, dtype=np.float32) + owner_ref = weakref.ref(array) + values = C.OrtValueVector() + values.push_back(ort.OrtValue.ortvalue_from_numpy(array)._get_c_value()) + retained = [] + + def consumer(capsule): + retained.append(capsule) + raise ValueError("deferred consumption") + + with self.assertRaisesRegex(ValueError, "deferred consumption"): + values.to_dlpacks(consumer) + del array, values + gc.collect() + self.assertIsNotNone(owner_ref()) + result = torch.from_dlpack(retained.pop()) + torch.testing.assert_close(result, torch.arange(6, dtype=torch.float32)) + del result + gc.collect() + self.assertIsNone(owner_ref()) + + def test_vector_consumer_can_mutate_original_vector(self): + values = C.OrtValueVector() + values.push_back(torch.ones(1)) + values.push_back(torch.full((1,), 2.0)) + + def consumer(capsule): + for _ in range(20): + values.push_back(torch.zeros(1)) + return torch.from_dlpack(capsule) + + results = values.to_dlpacks(consumer) + self.assertEqual(len(results), 2) + self.assertEqual(len(values), 42) + self.assertEqual([result.item() for result in results], [1, 2]) + + def test_vector_invalid_consumer(self): + for size in (0, 1): + values = C.OrtValueVector() + if size: + values.push_back(torch.ones(1)) + with self.subTest(size=size), self.assertRaisesRegex(TypeError, "callable consumer or None"): + values.to_dlpacks(42) + self.assertEqual(C.OrtValueVector().to_dlpacks(None), []) def test_native_ortvalue_protocol(self): for dtype, onnx_type in INTERCHANGE_DTYPES: @@ -255,5 +446,81 @@ def test_noncontiguous_input_is_rejected(self): ort.OrtValue.from_dlpack(source) +@unittest.skipUnless( + torch.cuda.is_available() and "CUDAExecutionProvider" in ort.get_available_providers(), + "requires PyTorch CUDA and the ORT CUDA execution provider", +) +class TestDlpackCUDA(unittest.TestCase): + def test_import_waits_for_producer_stream(self): + tensor_device = torch.cuda.device_count() - 1 + producer = torch.cuda.Stream(device=tensor_device) + for value_type in (ort.OrtValue, C.OrtValue): + for capsule in (False, True): + with self.subTest(value_type=value_type, capsule=capsule): + with torch.cuda.stream(producer): + source = torch.zeros(6, device=f"cuda:{tensor_device}") + torch.cuda._sleep(50_000_000) + source.fill_(7) + ready = torch.cuda.Event() + ready.record() + exported = source.__dlpack__() if capsule else source + with torch.cuda.device(0): + imported = value_type.from_dlpack(exported) + self.assertEqual(torch.cuda.current_device(), 0) + self.assertTrue(ready.query()) + torch.testing.assert_close( + torch.from_dlpack(imported), torch.full((6,), 7.0, device=f"cuda:{tensor_device}") + ) + + def test_export_waits_for_producer_and_restores_current_device(self): + tensor_device = torch.cuda.device_count() - 1 + source = torch.zeros(6, device=f"cuda:{tensor_device}") + public = ort.OrtValue.from_dlpack(source) + values = C.OrtValueVector() + values.push_back(public._get_c_value()) + producer = torch.cuda.Stream(device=tensor_device) + consumer = torch.cuda.Stream(device=tensor_device) + for export in ( + lambda: public.__dlpack__(stream=1), + lambda: public._get_c_value().__dlpack__(stream=2), + lambda: public.__dlpack__(stream=consumer.cuda_stream), + lambda: public.__dlpack__(), + lambda: values.to_dlpacks(None)[0], + ): + with torch.cuda.stream(producer): + torch.cuda._sleep(50_000_000) + source.add_(1) + ready = torch.cuda.Event() + ready.record() + with torch.cuda.device(0): + capsule = export() + self.assertEqual(torch.cuda.current_device(), 0) + self.assertTrue(ready.query()) + with torch.cuda.stream(consumer): + result = torch.from_dlpack(capsule).clone() + consumer.synchronize() + torch.testing.assert_close(result, source) + + def test_no_sync_export_is_capture_safe(self): + source = torch.ones(6, device="cuda") + value = ort.OrtValue.from_dlpack(source) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + source.add_(1) + capsule = value.__dlpack__(stream=-1) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(torch.from_dlpack(capsule), source) + + def test_invalid_cuda_stream(self): + value = ort.OrtValue.from_dlpack(torch.ones(1, device="cuda")) + for stream in (0, -2): + with self.subTest(stream=stream), self.assertRaisesRegex(ValueError, "Invalid CUDA DLPack stream"): + value.__dlpack__(stream=stream) + for stream in (False, 1.5, "invalid"): + with self.subTest(stream=stream), self.assertRaisesRegex(TypeError, "must be an integer"): + value.__dlpack__(stream=stream) + + if __name__ == "__main__": unittest.main() From b6ea5d39d79fe325d0a9ab7b0af733be22aa9ef8 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 5 Nov 2026 07:22:41 -0800 Subject: [PATCH 04/11] Address DLPack review findings Defer NumPy owner destruction without blocking native workers, preserve GPU vendor metadata and matching synchronization, and batch stream waits at ORTModule execution boundaries. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/dlpack/dlpack_converter.cc | 12 +- .../migraphx/migraphx_provider_factory.cc | 13 ++ .../migraphx/migraphx_provider_factory.h | 2 + .../onnxruntime_inference_collection.py | 17 +- .../python/onnxruntime_pybind_ortvalue.cc | 50 +++-- .../python/onnxruntime_pybind_state_common.cc | 54 ++++-- .../python/onnxruntime_pybind_state_common.h | 6 +- .../python/onnxruntime_test_python_dlpack.py | 174 +++++++++++++++++- .../training/ortmodule/_training_manager.py | 6 +- .../python/training/ortmodule/_utils.py | 17 +- 10 files changed, 303 insertions(+), 48 deletions(-) diff --git a/onnxruntime/core/dlpack/dlpack_converter.cc b/onnxruntime/core/dlpack/dlpack_converter.cc index 8784c0fe567ed..a57823830860d 100644 --- a/onnxruntime/core/dlpack/dlpack_converter.cc +++ b/onnxruntime/core/dlpack/dlpack_converter.cc @@ -90,7 +90,7 @@ const char* GetOrtDeviceName(const OrtDevice& device) { case OrtDevice::CPU: return CPU; case OrtDevice::GPU: - return CUDA; + return device.Vendor() == OrtDevice::VendorIds::AMD ? HIP : CUDA; default: ORT_THROW("Unknown device type: ", device.Type()); } @@ -136,7 +136,15 @@ DLDevice GetDlpackDevice(const OrtValue& ort_value, const int64_t& device_id) { device.device_type = DLDeviceType::kDLCPU; break; case OrtDevice::GPU: - device.device_type = DLDeviceType::kDLCUDA; + if (location.device.Vendor() == OrtDevice::VendorIds::NVIDIA) { + device.device_type = DLDeviceType::kDLCUDA; + } else if (location.device.Vendor() == OrtDevice::VendorIds::AMD) { + device.device_type = DLDeviceType::kDLROCM; + } else { + ORT_THROW( + "DLPack GPU export supports only NVIDIA CUDA and AMD ROCm tensors. " + "Copy this tensor to CPU before exporting it."); + } break; default: ORT_THROW("Cannot pack tensors on this device."); diff --git a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc index 5d9176d17cc64..51faace1d8caa 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc +++ b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc @@ -69,6 +69,19 @@ struct ProviderInfo_MIGraphX_Impl final : ProviderInfo_MIGraphX { HIP_CALL_THROW(hipMemcpy(dst, src, count, hipMemcpyDeviceToHost)); } + Status SynchronizeDevice(int device_id) override { + int original_device; + HIP_RETURN_IF_ERROR(hipGetDevice(&original_device)); + if (original_device != device_id) { + HIP_RETURN_IF_ERROR(hipSetDevice(device_id)); + } + const auto sync_error = hipDeviceSynchronize(); + const auto restore_error = original_device != device_id ? hipSetDevice(original_device) : hipSuccess; + HIP_RETURN_IF_ERROR(sync_error); + HIP_RETURN_IF_ERROR(restore_error); + return Status::OK(); + } + std::shared_ptr CreateMIGraphXAllocator(OrtDevice::DeviceId device_id, size_t mem_limit, onnxruntime::ArenaExtendStrategy arena_extend_strategy, void* alloc_fn, void* free_fn, void* empty_cache_fn, const OrtArenaCfg* default_memory_arena_cfg) override { if (alloc_fn != nullptr && free_fn != nullptr) { diff --git a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h index c23c9947c8d9b..72045d66d6d3c 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h +++ b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h @@ -5,6 +5,7 @@ #include +#include "core/common/status.h" #include "core/framework/arena_extend_strategy.h" #include "core/framework/ortdevice.h" @@ -18,6 +19,7 @@ struct ProviderInfo_MIGraphX { virtual void MIGraphXMemcpy_DeviceToHost(void* dst, const void* src, size_t count) = 0; virtual std::shared_ptr CreateMIGraphXAllocator(OrtDevice::DeviceId device_id, size_t mem_limit, ArenaExtendStrategy arena_extend_strategy, void* alloc_fn, void* free_fn, void* empty_cache_fn, const OrtArenaCfg* default_memory_arena_cfg) = 0; + virtual Status SynchronizeDevice(int device_id) = 0; protected: ~ProviderInfo_MIGraphX() = default; // Can only be destroyed through a subclass instance diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index c2d538e346d8d..a6013937d83a8 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -1538,9 +1538,10 @@ def __dlpack__(self, *, stream=None): WebGPU OrtValues cannot be exported through DLPack because their data is stored in opaque WebGPU buffers rather than CUDA-addressable memory. - :param stream: CUDA consumer stream: ``None`` or ``1`` for the legacy + :param stream: GPU consumer stream: for CUDA, ``None`` or ``1`` for the legacy default, ``2`` for the per-thread default, or a positive stream handle. - Since OrtValues do not track their producer stream, CUDA exports + ROCm also accepts ``0`` for its default stream. + Since OrtValues do not track their producer stream, GPU exports synchronize the entire tensor device before returning. ``-1`` skips synchronization; the caller must ensure the data is ready. CPU exports require ``None``. @@ -1562,7 +1563,7 @@ def __dlpack_device__(self) -> tuple[int, int]: return self._ortvalue.__dlpack_device__() @classmethod - def from_dlpack(cls, data, /) -> OrtValue: + def from_dlpack(cls, data, /, *, synchronize: bool = True) -> OrtValue: """ Construct an OrtValue from an object that implements the DLPack protocol. @@ -1585,15 +1586,19 @@ def from_dlpack(cls, data, /) -> OrtValue: PyTorch cannot consume ORT's packed float4, int4/uint4, or int2/uint2. No data is copied; the new OrtValue shares memory with the source. - CUDA imports synchronize the tensor's device before returning, so ORT - can safely use the storage on its own streams. + GPU imports synchronize the tensor's device before returning, so ORT + can safely use the storage on its own streams. CUDA and ROCm require + their corresponding CUDA or MIGraphX provider for synchronization. :param data: A tensor object supporting the DLPack protocol, or a raw DLPack PyCapsule. + :param synchronize: Set False only when the caller arranges producer/consumer + stream ordering before using the imported tensor, for example during + graph capture or when preparing a batch followed by one stream wait. :return: An OrtValue wrapping the tensor data. :raises TypeError: If the object does not provide a valid DLPack capsule. """ - return cls(C.OrtValue.from_dlpack(data)) + return cls(C.OrtValue.from_dlpack(data, synchronize=synchronize)) def update_inplace(self, data) -> None: """ diff --git a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc index 378bff1b1eabc..e021d02167302 100644 --- a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc @@ -13,6 +13,7 @@ #include "core/framework/tensor.h" #include "core/framework/sparse_tensor.h" #include "core/framework/TensorSeq.h" +#include "core/common/inlined_containers.h" namespace onnxruntime { namespace python { @@ -28,15 +29,12 @@ void RetainNumpyOwner(OrtValue& ort_value, const py::object& array) { OrtValue tensor_owner = ort_value; auto* tensor = tensor_owner.GetMutable(); auto type = tensor_owner.Type(); - auto array_owner = tensor->OwnsBuffer() ? py::object{} : array; - // Retain the owner in the native value so copies, vectors and DLPack consumers - // remain safe independently of the Python wrapper, including destruction without the GIL. - ort_value.Init(tensor, type, - [tensor_owner = std::move(tensor_owner), array_owner = std::move(array_owner)](void*) mutable { - py::gil_scoped_acquire acquire; - tensor_owner = OrtValue{}; - array_owner = py::object{}; - }); + // The original tensor may itself own a copied NumPy array. Release both + // owners under the GIL, without blocking an ORT worker that another thread is joining. + auto owners = py::make_tuple(py::cast(tensor_owner, py::return_value_policy::copy), + tensor->OwnsBuffer() ? py::none() : array); + auto owner = std::shared_ptr(owners.release().ptr(), ReleasePythonObject); + ort_value.Init(tensor, type, [owner = std::move(owner)](void*) mutable { owner.reset(); }); } std::unique_ptr OrtValueFromShapeAndType(const std::vector& shape, @@ -526,13 +524,15 @@ void addOrtValueMethods(pybind11::module& m) { #endif }) #if defined(ENABLE_DLPACK) - .def_static("from_dlpack", &FromDlpack, py::arg("data"), "Imports a tensor from a DLPack protocol object or legacy capsule without copying.") + .def_static("from_dlpack", &FromDlpack, py::arg("data"), py::kw_only(), py::arg("synchronize") = true, + "Imports a DLPack object or capsule without copying. Set synchronize=False only when the caller " + "arranges producer/consumer stream ordering.") .def("__dlpack__", [](OrtValue* ort_value, py::object stream) -> py::object { ORT_ENFORCE(!IsWebGpuBuffer(*ort_value), "DLPack export is not supported for WebGPU OrtValues."); return py::reinterpret_steal(ToDlpack(*ort_value, std::move(stream))); }, py::arg("stream") = py::none(), "Returns a DLPack representing the tensor (part of __dlpack__ protocol). " "The capsule retains the tensor and its NumPy owner without copying. " - "CUDA exports synchronize the tensor's device unless stream=-1 requests no synchronization.") + "GPU exports synchronize the tensor's device unless stream=-1 requests no synchronization.") .def("__dlpack_device__", [](const OrtValue* ort_value) -> py::tuple { ORT_ENFORCE(ort_value->IsTensor(), "Only tensor type OrtValues are supported"); ORT_ENFORCE(!IsWebGpuBuffer(*ort_value), "DLPack export is not supported for WebGPU OrtValues."); @@ -558,7 +558,7 @@ void addOrtValueMethods(pybind11::module& m) { v->push_back(ortvalue); }) #if defined(ENABLE_DLPACK) - .def("push_back", [](std::vector* v, py::object data) { v->push_back(FromDlpack(std::move(data))); }, "Add an OrtValue imported from a DLPack protocol object or legacy capsule.", py::arg("dlpack_tensor")) + .def("push_back", [](std::vector* v, py::object data, bool synchronize) { v->push_back(FromDlpack(std::move(data), synchronize)); }, "Add an OrtValue imported from a DLPack object or capsule. Set synchronize=False only when arranging stream ordering.", py::arg("dlpack_tensor"), py::kw_only(), py::arg("synchronize") = true) .def("push_back_batch", [](std::vector* v, std::vector& torch_tensors, std::vector& data_ptrs, std::vector& element_types, const std::vector>& shapes, const std::vector& devices) { for (size_t i = 0; i < torch_tensors.size(); ++i) { py::object& element_type = element_types.at(i); @@ -602,16 +602,28 @@ void addOrtValueMethods(pybind11::module& m) { "Raises an exception in any other case.", py::arg("idx")) #if defined(ENABLE_DLPACK) - .def("to_dlpacks", [](const std::vector& v, py::object to_tensor) -> py::list { + .def("to_dlpacks", [](const std::vector& v, py::object to_tensor, bool synchronize) -> py::list { if (!to_tensor.is_none() && !PyCallable_Check(to_tensor.ptr())) { throw py::type_error("OrtValueVector.to_dlpacks expects a callable consumer or None."); } // Callbacks can mutate the original vector, including reallocating its storage. auto values = v; - py::list results; + InlinedHashSet> synchronized_devices; + py::list capsules; for (auto& value : values) { + ORT_ENFORCE(value.IsTensor(), "DLPack export requires a dense tensor OrtValue; initialize the value before exporting."); ORT_ENFORCE(!IsWebGpuBuffer(value), "DLPack export is not supported for WebGPU OrtValues."); - auto capsule = py::reinterpret_steal(ToDlpack(value)); + const auto device = dlpack::GetDlpackDevice(value, value.Get().Location().device.Id()); + py::object stream = py::none(); + if (device.device_type == kDLCUDA || device.device_type == kDLROCM) { + if (!synchronize || !synchronized_devices.emplace(device.device_type, device.device_id).second) { + stream = py::int_(-1); + } + } + capsules.append(py::reinterpret_steal(ToDlpack(value, std::move(stream)))); + } + py::list results; + for (auto capsule : capsules) { results.append(to_tensor.is_none() ? capsule : to_tensor(capsule)); } return results; }, @@ -624,13 +636,15 @@ references when destroyed, including when conversion fails. a new tensor which becomes the new owner of the data. This function takes one python object and returns a new python object. It fits the same signature as `torch.from_dlpack`, if None, the method returns a capsule for every new DLPack structure. +:param synchronize: synchronize each distinct GPU device once before preparing its + capsules. Set False only when the caller arranges producer/consumer stream ordering. :return: a list containing the new tensors or the new capsules if *to_tensor* is None Use `ort_values.to_dlpacks(torch.from_dlpack)` to avoid creating a Python OrtValue -wrapper for every tensor. Each capsule may be consumed only once. CUDA exports -synchronize the tensor's device before invoking the consumer. +wrapper for every tensor. Each capsule may be consumed only once. All capsules +are prepared before any consumer callback runs. )pbdoc", - py::arg("to_tensor")) + py::arg("to_tensor"), py::kw_only(), py::arg("synchronize") = true) #endif ; } diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.cc b/onnxruntime/python/onnxruntime_pybind_state_common.cc index 038908efeaf2f..2770ceee73d2c 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.cc +++ b/onnxruntime/python/onnxruntime_pybind_state_common.cc @@ -9,6 +9,18 @@ namespace py = pybind11; const std::string onnxruntime::python::SessionObjectInitializer::default_logger_id = "Default"; +void ReleasePythonObject(PyObject* object) noexcept { + // After interpreter shutdown, retaining the reference is safer than touching Python. + if (object == nullptr || !Py_IsInitialized()) { + return; + } + if (PyGILState_Check()) { + Py_DECREF(object); + } else { + DeferPythonRelease(object); + } +} + #if defined(USE_OPENVINO) || defined(USE_OPENVINO_PROVIDER_INTERFACE) // TODO remove deprecated global config std::string openvino_device_type; @@ -54,40 +66,50 @@ void DlpackCapsuleDestructor(PyObject* data) { namespace { void SynchronizeDlpackDevice(const DLDevice& device) { - if (device.device_type != kDLCUDA) { + if (device.device_type == kDLCPU) { return; } + if (device.device_type == kDLCUDA) { #if defined(USE_CUDA) || defined(USE_CUDA_PROVIDER_INTERFACE) - if (auto* provider = TryGetProviderInfo_CUDA()) { - ORT_THROW_IF_ERROR(provider->SynchronizeDevice(device.device_id)); - return; - } + if (auto* provider = TryGetProviderInfo_CUDA()) { + ORT_THROW_IF_ERROR(provider->SynchronizeDevice(device.device_id)); + return; + } +#endif + } else if (device.device_type == kDLROCM) { +#ifdef USE_MIGRAPHX + if (auto* provider = TryGetProviderInfo_MIGraphX()) { + ORT_THROW_IF_ERROR(provider->SynchronizeDevice(device.device_id)); + return; + } #endif + } throw py::buffer_error( - "CUDA DLPack interchange requires an available CUDA provider to synchronize the tensor's device. " - "Install a CUDA-enabled ONNX Runtime package and load its CUDA provider. For an already synchronized " - "export, call __dlpack__(stream=-1) explicitly."); + "DLPack synchronization requires a provider matching the tensor's device (CUDA for NVIDIA, MIGraphX for AMD). " + "Install and load the matching provider, or explicitly opt out only after arranging stream ordering: " + "__dlpack__(stream=-1), from_dlpack(..., synchronize=False), or to_dlpacks(..., synchronize=False)."); } } // namespace PyObject* ToDlpack(OrtValue& ort_value, py::object stream) { + ORT_ENFORCE(ort_value.IsTensor(), "DLPack export requires a dense tensor OrtValue; initialize the value before exporting."); const auto device = dlpack::GetDlpackDevice(ort_value, ort_value.Get().Location().device.Id()); bool synchronize = true; if (!stream.is_none()) { - if (device.device_type != kDLCUDA) { + if (device.device_type != kDLCUDA && device.device_type != kDLROCM) { throw py::buffer_error("CPU DLPack exports require stream=None; omit the stream argument."); } if (!PyLong_Check(stream.ptr()) || PyBool_Check(stream.ptr())) { - throw py::type_error("A CUDA DLPack stream must be an integer: -1, 1, 2, or a positive stream handle."); + throw py::type_error("A GPU DLPack stream must be an integer: -1, a default-stream identifier, or a stream handle."); } const auto stream_id = PyLong_AsLongLong(stream.ptr()); if (PyErr_Occurred()) { throw py::error_already_set(); } - if (stream_id == 0 || stream_id < -1) { + if ((device.device_type == kDLCUDA && stream_id == 0) || stream_id < -1) { throw py::value_error( - "Invalid CUDA DLPack stream: use 1 for the legacy default, 2 for the per-thread default, " - "a positive stream handle, or -1 to opt out of synchronization."); + "Invalid GPU DLPack stream: CUDA accepts 1 or 2 for default streams; ROCm accepts 0 or 1. " + "Use a positive stream handle, or -1 to opt out of synchronization."); } synchronize = stream_id != -1; } @@ -105,7 +127,7 @@ PyObject* ToDlpack(OrtValue& ort_value, py::object stream) { return capsule; } -OrtValue FromDlpack(py::object data) { +OrtValue FromDlpack(py::object data, bool synchronize) { if (!PyCapsule_CheckExact(data.ptr())) { auto method = py::reinterpret_steal(PyObject_GetAttrString(data.ptr(), "__dlpack__")); if (method) { @@ -125,7 +147,9 @@ OrtValue FromDlpack(py::object data) { DLManagedTensor* dlmanaged_tensor = (DLManagedTensor*)PyCapsule_GetPointer(dlpack_tensor, "dltensor"); // The producer's default-stream handoff alone is insufficient for ORT's // session-specific streams. Fence before importing both objects and raw capsules. - SynchronizeDlpackDevice(dlmanaged_tensor->dl_tensor.device); + if (synchronize) { + SynchronizeDlpackDevice(dlmanaged_tensor->dl_tensor.device); + } OrtValue ort_value = dlpack::DlpackToOrtValue(dlmanaged_tensor); // Make sure this capsule will never be used again. PyCapsule_SetName(dlpack_tensor, "used_dltensor"); diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.h b/onnxruntime/python/onnxruntime_pybind_state_common.h index 9ea5f4890d0b5..7d0d38d3c2961 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.h +++ b/onnxruntime/python/onnxruntime_pybind_state_common.h @@ -539,6 +539,10 @@ bool CheckIfTensor(const std::vector& def_list, const std::string& name, /*out*/ ONNX_NAMESPACE::TypeProto& type_proto); +// Transfer an owned reference to the existing deferred cleanup queue. +void DeferPythonRelease(PyObject* object) noexcept; +void ReleasePythonObject(PyObject* object) noexcept; + #if defined(ENABLE_DLPACK) // Allocate a capsule retaining an OrtValue reference via OrtValueToDlpack. @@ -546,7 +550,7 @@ bool CheckIfTensor(const std::vector& def_list, PyObject* ToDlpack(OrtValue& ort_value, pybind11::object stream = pybind11::none()); // Import a DLPack protocol object or legacy capsule, transferring tensor ownership. -OrtValue FromDlpack(pybind11::object data); +OrtValue FromDlpack(pybind11::object data, bool synchronize = true); // Destructor for Capsule object holding a DLPack structure. void DlpackCapsuleDestructor(PyObject* data); diff --git a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py index 25e6a92662ad1..75e3882bea97c 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py +++ b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py @@ -52,6 +52,37 @@ SHAPES = ((), (2, 3), (2, 1, 3), (0,), (2, 0, 3)) +class DLDevice(ctypes.Structure): + _fields_ = (("device_type", ctypes.c_int), ("device_id", ctypes.c_int)) + + +class DLDataType(ctypes.Structure): + _fields_ = (("code", ctypes.c_uint8), ("bits", ctypes.c_uint8), ("lanes", ctypes.c_uint16)) + + +class DLTensor(ctypes.Structure): + _fields_ = ( + ("data", ctypes.c_void_p), + ("device", DLDevice), + ("ndim", ctypes.c_int), + ("dtype", DLDataType), + ("shape", ctypes.POINTER(ctypes.c_int64)), + ("strides", ctypes.POINTER(ctypes.c_int64)), + ("byte_offset", ctypes.c_uint64), + ) + + +class DLManagedTensor(ctypes.Structure): + _fields_ = (("dl_tensor", DLTensor), ("manager_ctx", ctypes.c_void_p), ("deleter", ctypes.c_void_p)) + + +def managed_tensor_from_capsule(capsule): + get_pointer = ctypes.pythonapi.PyCapsule_GetPointer + get_pointer.argtypes = [ctypes.py_object, ctypes.c_char_p] + get_pointer.restype = ctypes.c_void_p + return ctypes.cast(get_pointer(capsule, b"dltensor"), ctypes.POINTER(DLManagedTensor)) + + def tensor_bytes(tensor): # Reshape first because PyTorch cannot reinterpret a scalar as a smaller dtype. return tensor.reshape(-1).view(torch.uint8) @@ -224,6 +255,81 @@ def test_numpy_noncontiguous_copy_survives_export(self): self.assertIsNone(owner_ref()) torch.testing.assert_close(result, torch.arange(6, dtype=torch.float32).flip(0)) + def test_numpy_owner_release_without_gil(self): + for contiguous in (False, True): + with self.subTest(contiguous=contiguous): + array = np.arange(6, dtype=np.float32) + if not contiguous: + array = array[::-1] + owner_ref = weakref.ref(array) + value = ort.OrtValue.ortvalue_from_numpy(array) + capsule = value.__dlpack__() + managed = managed_tensor_from_capsule(capsule) + deleter = ctypes.CFUNCTYPE(None, ctypes.POINTER(DLManagedTensor))(managed.contents.deleter) + set_name = ctypes.pythonapi.PyCapsule_SetName + set_name.argtypes = [ctypes.py_object, ctypes.c_char_p] + set_name.restype = ctypes.c_int + self.assertEqual(set_name(capsule, b"used_dltensor"), 0) + del value, array + gc.collect() + self.assertEqual(owner_ref() is not None, contiguous) + # CFUNCTYPE releases the GIL while calling the actual C++ deleter. + deleter(managed) + self.assertEqual(ctypes.pythonapi.Py_MakePendingCalls(), 0) + gc.collect() + self.assertIsNone(owner_ref()) + + def test_gpu_vendor_metadata_without_gpu_allocation(self): + new_capsule = ctypes.pythonapi.PyCapsule_New + new_capsule.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] + new_capsule.restype = ctypes.py_object + # Empty tensors exercise device metadata without allocating or reading GPU memory. + shape = (ctypes.c_int64 * 1)(0) + for device_type in (2, 10): # kDLCUDA, kDLROCM + with self.subTest(device_type=device_type): + managed = DLManagedTensor( + DLTensor(None, DLDevice(device_type, 3), 1, DLDataType(2, 32, 1), shape, None, 0), None, None + ) + capsule = new_capsule(ctypes.addressof(managed), b"dltensor", None) + imported = ort.OrtValue.from_dlpack(capsule, synchronize=False) + exported = None + result = None + values = C.OrtValueVector() + try: + self.assertEqual(imported.__dlpack_device__(), (device_type, 3)) + values.push_back(imported._get_c_value()) + exported = values.to_dlpacks(None, synchronize=False) + for result in (imported.__dlpack__(stream=-1), *exported): + device = managed_tensor_from_capsule(result).contents.dl_tensor.device + self.assertEqual((device.device_type, device.device_id), (device_type, 3)) + finally: + # The synthetic producer's descriptor must outlive every native owner. + del result, imported, exported, values + + def test_import_synchronization_opt_out(self): + source = torch.arange(6, dtype=torch.float32) + for value_type in (ort.OrtValue, C.OrtValue): + with self.subTest(value_type=value_type): + imported = value_type.from_dlpack(source, synchronize=False) + self.assert_tensor(torch.from_dlpack(imported), source) + values = C.OrtValueVector() + values.push_back(source, synchronize=False) + self.assert_tensor(values.to_dlpacks(torch.from_dlpack, synchronize=False)[0], source) + + def test_unsupported_gpu_vendor_is_not_exported_as_cuda(self): + array = np.empty(1, dtype=np.float32) + # A zero-element DirectML-labelled tensor requires no actual GPU memory. + device = C.OrtDevice(C.OrtDevice.gpu(), C.OrtDevice.default_memory(), 0x1414, 0) + values = C.OrtValueVector() + values.push_back_batch([array], [array.ctypes.data], [array.dtype], [[0]], [device]) + value = values[0] + for export in (value.__dlpack_device__, value.__dlpack__, lambda: values.to_dlpacks(None)): + with ( + self.subTest(export=export), + self.assertRaisesRegex(RuntimeError, "supports only NVIDIA CUDA and AMD ROCm"), + ): + export() + def test_numpy_reinterpret_rejects_noncontiguous_storage(self): for value_type in (ort.OrtValue, C.OrtValue): for array in (np.arange(6, dtype=np.float32)[::-1], np.zeros((2, 3), dtype=np.float32).T): @@ -347,6 +453,26 @@ def consumer(capsule): self.assertEqual(len(values), 42) self.assertEqual([result.item() for result in results], [1, 2]) + def test_vector_preparation_failure_does_not_invoke_consumer(self): + array = np.arange(6, dtype=np.float32) + owner_ref = weakref.ref(array) + values = C.OrtValueVector() + values.push_back(ort.OrtValue.ortvalue_from_numpy(array)._get_c_value()) + device = C.OrtDevice(C.OrtDevice.gpu(), C.OrtDevice.default_memory(), 0x1414, 0) + values.push_back_batch([array], [array.ctypes.data], [array.dtype], [[0]], [device]) + calls = [] + + def consumer(capsule): + calls.append(capsule) + return torch.from_dlpack(capsule) + + with self.assertRaisesRegex(RuntimeError, "supports only NVIDIA CUDA and AMD ROCm"): + values.to_dlpacks(consumer) + self.assertEqual(calls, []) + del array, values + gc.collect() + self.assertIsNone(owner_ref()) + def test_vector_invalid_consumer(self): for size in (0, 1): values = C.OrtValueVector() @@ -512,10 +638,56 @@ def test_no_sync_export_is_capture_safe(self): torch.cuda.synchronize() torch.testing.assert_close(torch.from_dlpack(capsule), source) + def test_no_sync_import_and_batch_export_are_capture_safe(self): + source = torch.ones(6, device="cuda") + for value_type in (ort.OrtValue, C.OrtValue): + with self.subTest(value_type=value_type): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + source.add_(1) + imported = value_type.from_dlpack(source, synchronize=False) + values = C.OrtValueVector() + values.push_back(source, synchronize=False) + capsules = values.to_dlpacks(None, synchronize=False) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(torch.from_dlpack(imported), source) + torch.testing.assert_close(torch.from_dlpack(capsules[0]), source) + + def test_batch_prepares_all_capsules_before_callbacks(self): + source = torch.ones(6, device="cuda") + values = C.OrtValueVector() + values.push_back(source, synchronize=False) + values.push_back(source, synchronize=False) + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + calls = 0 + + def consumer(capsule): + nonlocal calls + if calls == 0: + # Any later device-wide fence would invalidate this capture. + graph.capture_begin() + source.add_(1) + calls += 1 + return torch.from_dlpack(capsule) + + with torch.cuda.stream(stream): + try: + results = values.to_dlpacks(consumer) + finally: + if torch.cuda.is_current_stream_capturing(): + graph.capture_end() + graph.replay() + stream.synchronize() + self.assertEqual(calls, 2) + for result in results: + torch.testing.assert_close(result, source) + def test_invalid_cuda_stream(self): value = ort.OrtValue.from_dlpack(torch.ones(1, device="cuda")) for stream in (0, -2): - with self.subTest(stream=stream), self.assertRaisesRegex(ValueError, "Invalid CUDA DLPack stream"): + with self.subTest(stream=stream), self.assertRaisesRegex(ValueError, "Invalid GPU DLPack stream"): value.__dlpack__(stream=stream) for stream in (False, 1.5, "invalid"): with self.subTest(stream=stream), self.assertRaisesRegex(TypeError, "must be an integer"): diff --git a/orttraining/orttraining/python/training/ortmodule/_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py index 5845eac10ec64..2e9711950e54d 100644 --- a/orttraining/orttraining/python/training/ortmodule/_training_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py @@ -78,8 +78,9 @@ def execution_session_run_forward( forward_inputs.push_back(C.aten_ort_tensor_to_ort_value(input)) else: valid_ort_tensor = _utils._torch_tensor_to_dlpack(input) - forward_inputs.push_back(valid_ort_tensor) + forward_inputs.push_back(valid_ort_tensor, synchronize=False) + _utils._synchronize_torch_stream(device) forward_outputs = C.OrtValueVector() # Run and return module outputs. execution_session.run_forward(forward_inputs, forward_outputs, state, gradient_accumulation_manager.cache) @@ -186,9 +187,10 @@ def backward(ctx, *grad_outputs): if grad_output.device.type == "ort": backward_inputs.push_back(C.aten_ort_tensor_to_ort_value(grad_output)) else: - backward_inputs.push_back(_utils._torch_tensor_to_dlpack(grad_output)) + backward_inputs.push_back(_utils._torch_tensor_to_dlpack(grad_output), synchronize=False) backward_inputs.shrink_to_fit() + _utils._synchronize_torch_stream(self._device) # Run and get results backward_outputs = C.OrtValueVector() try: diff --git a/orttraining/orttraining/python/training/ortmodule/_utils.py b/orttraining/orttraining/python/training/ortmodule/_utils.py index 3aeada12de7aa..84a2b635642bc 100644 --- a/orttraining/orttraining/python/training/ortmodule/_utils.py +++ b/orttraining/orttraining/python/training/ortmodule/_utils.py @@ -46,12 +46,20 @@ def set_random_states(states): torch.cuda.set_rng_state(torch_cuda_state) -def _ortvalue_from_torch_tensor(torch_tensor: torch.Tensor) -> C.OrtValue: +def _ortvalue_from_torch_tensor(torch_tensor: torch.Tensor, *, synchronize: bool = True) -> C.OrtValue: """Converts a torch tensor to an OrtValue.""" if torch_tensor.device.type == "ort": return C.aten_ort_tensor_to_ort_value(torch_tensor) # The protocol rejects requires_grad tensors; ORTModule manages autograd separately. - return C.OrtValue.from_dlpack(torch_tensor.detach() if torch_tensor.requires_grad else torch_tensor) + return C.OrtValue.from_dlpack( + torch_tensor.detach() if torch_tensor.requires_grad else torch_tensor, synchronize=synchronize + ) + + +def _synchronize_torch_stream(device: torch.device): + # ORT may use a different, nonblocking stream. Wait once after preparing all inputs. + if device.type == "cuda": + torch.cuda.current_stream(device).synchronize() def _ortvalues_to_torch_tensor( @@ -181,10 +189,13 @@ def _create_iobinding(io_binding, inputs, model, device: torch.device): io_binding.bind_ortvalue_input( value_info.name, OrtValue( - _ortvalue_from_torch_tensor(inputs[idx] if inputs[idx].is_contiguous() else inputs[idx].contiguous()) + _ortvalue_from_torch_tensor( + inputs[idx] if inputs[idx].is_contiguous() else inputs[idx].contiguous(), synchronize=False + ) ), ) + _synchronize_torch_stream(device) device_id = get_device_index(device) for value_info in model.graph.output: io_binding.bind_output(value_info.name, device.type, device_id=device_id) From 3466ddd54e4accaafa55204857b4e368688e6876 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 5 Nov 2026 07:54:31 -0800 Subject: [PATCH 05/11] Fix CUDA synchronization error handling compilation Use ORT_RETURN_IF_ERROR with CUDA_CALL because the file-local CUDA_RETURN_IF_ERROR macro is declared below the provider implementation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/providers/cuda/cuda_provider_factory.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index 2be9b3d2920d5..a9dcce6f9de9b 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -230,15 +230,15 @@ struct ProviderInfo_CUDA_Impl final : ProviderInfo_CUDA { Status SynchronizeDevice(int device_id) override { int original_device; - CUDA_RETURN_IF_ERROR(cudaGetDevice(&original_device)); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetDevice(&original_device))); if (original_device != device_id) { - CUDA_RETURN_IF_ERROR(cudaSetDevice(device_id)); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaSetDevice(device_id))); } const auto sync_error = cudaDeviceSynchronize(); // Restore the caller's device even when synchronization reports an error. const auto restore_error = original_device != device_id ? cudaSetDevice(original_device) : cudaSuccess; - CUDA_RETURN_IF_ERROR(sync_error); - CUDA_RETURN_IF_ERROR(restore_error); + ORT_RETURN_IF_ERROR(CUDA_CALL(sync_error)); + ORT_RETURN_IF_ERROR(CUDA_CALL(restore_error)); return Status::OK(); } From 5daf905bd60effef0d0dd97cac0a8cc6035cae53 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 5 Nov 2026 08:09:19 -0800 Subject: [PATCH 06/11] Address DLPack protocol review feedback Request no-sync GPU exports when importing with synchronization disabled, preserve descriptor AttributeError, and include MIGraphX interface-only provider lookup. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../onnxruntime_inference_collection.py | 1 + .../python/onnxruntime_pybind_state_common.cc | 38 +++++++++++++++---- .../python/onnxruntime_pybind_state_common.h | 4 +- .../python/onnxruntime_test_python_dlpack.py | 36 ++++++++++++++++++ 4 files changed, 70 insertions(+), 9 deletions(-) diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index a6013937d83a8..dc2f5dd21c7a2 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -1595,6 +1595,7 @@ def from_dlpack(cls, data, /, *, synchronize: bool = True) -> OrtValue: :param synchronize: Set False only when the caller arranges producer/consumer stream ordering before using the imported tensor, for example during graph capture or when preparing a batch followed by one stream wait. + GPU protocol producers are asked to export with ``stream=-1`` in this mode. :return: An OrtValue wrapping the tensor data. :raises TypeError: If the object does not provide a valid DLPack capsule. """ diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.cc b/onnxruntime/python/onnxruntime_pybind_state_common.cc index 2770ceee73d2c..32d9df6feb69c 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.cc +++ b/onnxruntime/python/onnxruntime_pybind_state_common.cc @@ -65,6 +65,24 @@ void DlpackCapsuleDestructor(PyObject* data) { } namespace { +py::object GetOptionalDlpackAttribute(const py::object& data, const char* name) { + auto attribute = py::reinterpret_steal(PyObject_GetAttrString(data.ptr(), name)); + if (attribute) { + return attribute; + } + py::error_already_set error; + if (!error.matches(PyExc_AttributeError)) { + throw error; + } + // AttributeError from an existing descriptor is a producer failure, not an + // absent protocol method. Static lookup does not invoke that descriptor again. + auto declared = py::module_::import("inspect").attr("getattr_static")(data, name, py::none()); + if (!declared.is_none()) { + throw error; + } + return py::none(); +} + void SynchronizeDlpackDevice(const DLDevice& device) { if (device.device_type == kDLCPU) { return; @@ -77,7 +95,7 @@ void SynchronizeDlpackDevice(const DLDevice& device) { } #endif } else if (device.device_type == kDLROCM) { -#ifdef USE_MIGRAPHX +#if defined(USE_MIGRAPHX) || defined(USE_MIGRAPHX_PROVIDER_INTERFACE) if (auto* provider = TryGetProviderInfo_MIGraphX()) { ORT_THROW_IF_ERROR(provider->SynchronizeDevice(device.device_id)); return; @@ -129,13 +147,17 @@ PyObject* ToDlpack(OrtValue& ort_value, py::object stream) { OrtValue FromDlpack(py::object data, bool synchronize) { if (!PyCapsule_CheckExact(data.ptr())) { - auto method = py::reinterpret_steal(PyObject_GetAttrString(data.ptr(), "__dlpack__")); - if (method) { - data = method(); - } else if (PyErr_ExceptionMatches(PyExc_AttributeError)) { - PyErr_Clear(); - } else { - throw py::error_already_set(); + auto method = GetOptionalDlpackAttribute(data, "__dlpack__"); + if (!method.is_none()) { + bool request_no_sync = false; + if (!synchronize) { + auto device_method = GetOptionalDlpackAttribute(data, "__dlpack_device__"); + if (!device_method.is_none()) { + const auto device = device_method().cast>(); + request_no_sync = device.first == kDLCUDA || device.first == kDLROCM; + } + } + data = request_no_sync ? method(py::arg("stream") = -1) : method(); } } PyObject* dlpack_tensor = data.ptr(); diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.h b/onnxruntime/python/onnxruntime_pybind_state_common.h index 7d0d38d3c2961..42087ed6d43e8 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.h +++ b/onnxruntime/python/onnxruntime_pybind_state_common.h @@ -182,10 +182,11 @@ extern onnxruntime::ArenaExtendStrategy arena_extend_strategy; } // namespace onnxruntime #endif -#ifdef USE_MIGRAPHX +#if defined(USE_MIGRAPHX) || defined(USE_MIGRAPHX_PROVIDER_INTERFACE) namespace onnxruntime { ProviderInfo_MIGraphX* TryGetProviderInfo_MIGraphX(); ProviderInfo_MIGraphX& GetProviderInfo_MIGraphX(); +#ifdef USE_MIGRAPHX namespace python::migraphx::external { extern void* alloc_fn; extern void* free_fn; @@ -194,6 +195,7 @@ inline bool UseExternalAllocator() { return alloc_fn != nullptr && free_fn != nullptr; } } // namespace python::migraphx::external +#endif } // namespace onnxruntime #endif diff --git a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py index 75e3882bea97c..b5780f71e86fe 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py +++ b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py @@ -297,6 +297,12 @@ def test_gpu_vendor_metadata_without_gpu_allocation(self): values = C.OrtValueVector() try: self.assertEqual(imported.__dlpack_device__(), (device_type, 3)) + for value_type in (ort.OrtValue, C.OrtValue): + for producer in (imported, imported._get_c_value()): + copied = value_type.from_dlpack(producer, synchronize=False) + self.assertEqual(copied.__dlpack_device__(), (device_type, 3)) + del copied + del producer values.push_back(imported._get_c_value()) exported = values.to_dlpacks(None, synchronize=False) for result in (imported.__dlpack__(stream=-1), *exported): @@ -316,6 +322,23 @@ def test_import_synchronization_opt_out(self): values.push_back(source, synchronize=False) self.assert_tensor(values.to_dlpacks(torch.from_dlpack, synchronize=False)[0], source) + def test_import_opt_out_without_gpu_metadata_omits_stream(self): + source = torch.arange(6, dtype=torch.float32) + + class Producer: + def __dlpack__(self): + return source.__dlpack__() + + class CpuProducer(Producer): + def __dlpack_device__(self): + return (1, 0) + + for value_type in (ort.OrtValue, C.OrtValue): + for producer in (Producer(), CpuProducer()): + with self.subTest(value_type=value_type, producer=producer): + imported = value_type.from_dlpack(producer, synchronize=False) + self.assert_tensor(torch.from_dlpack(imported), source) + def test_unsupported_gpu_vendor_is_not_exported_as_cuda(self): array = np.empty(1, dtype=np.float32) # A zero-element DirectML-labelled tensor requires no actual GPU memory. @@ -516,6 +539,11 @@ class InvalidProtocolAttribute: def __dlpack__(self): raise ValueError("DLPack lookup failed") + class InvalidAttributeDescriptor: + @property + def __dlpack__(self): + raise AttributeError("DLPack descriptor failed") + for value_type in (ort.OrtValue, C.OrtValue): for producer, error in ( (InvalidProducer(), "DLPack producer failed"), @@ -523,6 +551,12 @@ def __dlpack__(self): ): with self.subTest(value_type=value_type, producer=producer), self.assertRaisesRegex(ValueError, error): value_type.from_dlpack(producer) + for synchronize in (False, True): + with ( + self.subTest(value_type=value_type, synchronize=synchronize), + self.assertRaisesRegex(AttributeError, "DLPack descriptor failed"), + ): + value_type.from_dlpack(InvalidAttributeDescriptor(), synchronize=synchronize) def test_packed_dtype_boundaries(self): for dtype, onnx_type, storage_type, bits, error in PACKED_DTYPES: @@ -646,12 +680,14 @@ def test_no_sync_import_and_batch_export_are_capture_safe(self): with torch.cuda.graph(graph): source.add_(1) imported = value_type.from_dlpack(source, synchronize=False) + copied = value_type.from_dlpack(imported, synchronize=False) values = C.OrtValueVector() values.push_back(source, synchronize=False) capsules = values.to_dlpacks(None, synchronize=False) graph.replay() torch.cuda.synchronize() torch.testing.assert_close(torch.from_dlpack(imported), source) + torch.testing.assert_close(torch.from_dlpack(copied), source) torch.testing.assert_close(torch.from_dlpack(capsules[0]), source) def test_batch_prepares_all_capsules_before_callbacks(self): From 2cdd2976de057f7eb5eb587bdee5a2edb713d8d0 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 5 Nov 2026 11:37:47 -0800 Subject: [PATCH 07/11] Release the GIL safely during DLPack GPU waits Claim capsule ownership before synchronization, retain export storage across the unlocked wait, and cover concurrent consumption and failure cleanup. Pin SentencePiece to the version compatible with the protobuf runtime used by transformer CI. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../onnxruntime_inference_collection.py | 2 + .../python/onnxruntime_pybind_state_common.cc | 16 ++-- .../python/onnxruntime_test_python_dlpack.py | 94 +++++++++++++++++++ .../transformers-test/requirements.txt | 3 +- 4 files changed, 108 insertions(+), 7 deletions(-) diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index dc2f5dd21c7a2..9b30cfd5f5644 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -1589,6 +1589,8 @@ def from_dlpack(cls, data, /, *, synchronize: bool = True) -> OrtValue: GPU imports synchronize the tensor's device before returning, so ORT can safely use the storage on its own streams. CUDA and ROCm require their corresponding CUDA or MIGraphX provider for synchronization. + The capsule is consumed before synchronization; if that wait fails, the + imported storage is released and the capsule cannot be retried. :param data: A tensor object supporting the DLPack protocol, or a raw DLPack PyCapsule. diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.cc b/onnxruntime/python/onnxruntime_pybind_state_common.cc index 32d9df6feb69c..99fe130ff19c8 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.cc +++ b/onnxruntime/python/onnxruntime_pybind_state_common.cc @@ -90,6 +90,7 @@ void SynchronizeDlpackDevice(const DLDevice& device) { if (device.device_type == kDLCUDA) { #if defined(USE_CUDA) || defined(USE_CUDA_PROVIDER_INTERFACE) if (auto* provider = TryGetProviderInfo_CUDA()) { + py::gil_scoped_release release; ORT_THROW_IF_ERROR(provider->SynchronizeDevice(device.device_id)); return; } @@ -97,6 +98,7 @@ void SynchronizeDlpackDevice(const DLDevice& device) { } else if (device.device_type == kDLROCM) { #if defined(USE_MIGRAPHX) || defined(USE_MIGRAPHX_PROVIDER_INTERFACE) if (auto* provider = TryGetProviderInfo_MIGraphX()) { + py::gil_scoped_release release; ORT_THROW_IF_ERROR(provider->SynchronizeDevice(device.device_id)); return; } @@ -110,8 +112,9 @@ void SynchronizeDlpackDevice(const DLDevice& device) { } // namespace PyObject* ToDlpack(OrtValue& ort_value, py::object stream) { - ORT_ENFORCE(ort_value.IsTensor(), "DLPack export requires a dense tensor OrtValue; initialize the value before exporting."); - const auto device = dlpack::GetDlpackDevice(ort_value, ort_value.Get().Location().device.Id()); + OrtValue owner = ort_value; + ORT_ENFORCE(owner.IsTensor(), "DLPack export requires a dense tensor OrtValue; initialize the value before exporting."); + const auto device = dlpack::GetDlpackDevice(owner, owner.Get().Location().device.Id()); bool synchronize = true; if (!stream.is_none()) { if (device.device_type != kDLCUDA && device.device_type != kDLROCM) { @@ -136,7 +139,7 @@ PyObject* ToDlpack(OrtValue& ort_value, py::object stream) { if (synchronize) { SynchronizeDlpackDevice(device); } - DLManagedTensor* dlmanaged_tensor = dlpack::OrtValueToDlpack(ort_value); + DLManagedTensor* dlmanaged_tensor = dlpack::OrtValueToDlpack(owner); PyObject* capsule = PyCapsule_New(dlmanaged_tensor, "dltensor", DlpackCapsuleDestructor); if (!capsule) { dlmanaged_tensor->deleter(dlmanaged_tensor); @@ -167,14 +170,15 @@ OrtValue FromDlpack(py::object data, bool synchronize) { // Extract DLPack tensor pointer from the capsule carrier. DLManagedTensor* dlmanaged_tensor = (DLManagedTensor*)PyCapsule_GetPointer(dlpack_tensor, "dltensor"); + OrtValue ort_value = dlpack::DlpackToOrtValue(dlmanaged_tensor); + // Claim exclusive ownership under the GIL before the native wait releases it. + // The OrtValue also releases the producer if synchronization subsequently fails. + PyCapsule_SetName(dlpack_tensor, "used_dltensor"); // The producer's default-stream handoff alone is insufficient for ORT's // session-specific streams. Fence before importing both objects and raw capsules. if (synchronize) { SynchronizeDlpackDevice(dlmanaged_tensor->dl_tensor.device); } - OrtValue ort_value = dlpack::DlpackToOrtValue(dlmanaged_tensor); - // Make sure this capsule will never be used again. - PyCapsule_SetName(dlpack_tensor, "used_dltensor"); return ort_value; } diff --git a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py index b5780f71e86fe..b8e0c47c34c1e 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py +++ b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py @@ -9,6 +9,8 @@ import ctypes import gc import math +import threading +import time import unittest import weakref @@ -322,6 +324,35 @@ def test_import_synchronization_opt_out(self): values.push_back(source, synchronize=False) self.assert_tensor(values.to_dlpacks(torch.from_dlpack, synchronize=False)[0], source) + def test_failed_gpu_synchronization_consumes_capsule_once(self): + devices = ((2, "CUDAExecutionProvider"), (10, "MIGraphXExecutionProvider")) + device_type = next((kind for kind, provider in devices if provider not in ort.get_available_providers()), None) + if device_type is None: + self.skipTest("requires a GPU backend without its matching provider") + new_capsule = ctypes.pythonapi.PyCapsule_New + new_capsule.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] + new_capsule.restype = ctypes.py_object + shape = (ctypes.c_int64 * 1)(0) + for value_type in (ort.OrtValue, C.OrtValue): + with self.subTest(value_type=value_type): + releases = [] + deleter = ctypes.CFUNCTYPE(None, ctypes.c_void_p)( + lambda pointer, releases=releases: releases.append(pointer) + ) + managed = DLManagedTensor( + DLTensor(None, DLDevice(device_type, 0), 1, DLDataType(2, 32, 1), shape, None, 0), + None, + ctypes.cast(deleter, ctypes.c_void_p), + ) + capsule = new_capsule(ctypes.addressof(managed), b"dltensor", None) + with self.assertRaisesRegex(BufferError, "provider matching the tensor's device"): + value_type.from_dlpack(capsule) + self.assertEqual(len(releases), 1) + with self.assertRaisesRegex(TypeError, "valid DLPack capsule"): + value_type.from_dlpack(capsule, synchronize=False) + del capsule + self.assertEqual(len(releases), 1) + def test_import_opt_out_without_gpu_metadata_omits_stream(self): source = torch.arange(6, dtype=torch.float32) @@ -611,6 +642,69 @@ def test_noncontiguous_input_is_rejected(self): "requires PyTorch CUDA and the ORT CUDA execution provider", ) class TestDlpackCUDA(unittest.TestCase): + def test_import_claims_capsule_before_releasing_gil(self): + source = torch.zeros(6, device="cuda") + torch.cuda.synchronize() + is_valid = ctypes.pythonapi.PyCapsule_IsValid + is_valid.argtypes = [ctypes.py_object, ctypes.c_char_p] + is_valid.restype = ctypes.c_int + for value_type in (ort.OrtValue, C.OrtValue): + with self.subTest(value_type=value_type): + capsule = source.__dlpack__() + producer = torch.cuda.Stream() + with torch.cuda.stream(producer): + torch.cuda._sleep(1_000_000_000) + source.fill_(7) + ready = torch.cuda.Event() + ready.record() + observed = [] + stop = threading.Event() + + def observer(stop=stop, capsule=capsule, observed=observed, ready=ready): + while not stop.is_set(): + if is_valid(capsule, b"used_dltensor"): + try: + C.OrtValue.from_dlpack(capsule, synchronize=False) + except TypeError: + observed.append(not ready.query()) + return + time.sleep(0.001) + + thread = threading.Thread(target=observer) + thread.start() + try: + imported = value_type.from_dlpack(capsule) + thread.join(timeout=5) + finally: + stop.set() + thread.join() + self.assertEqual(observed, [True], "Another Python thread must run while the import waits for CUDA.") + torch.testing.assert_close(torch.from_dlpack(imported), source) + + def test_export_releases_gil_during_device_wait(self): + source = torch.ones(6, device="cuda") + value = ort.OrtValue.from_dlpack(source) + producer = torch.cuda.Stream() + with torch.cuda.stream(producer): + torch.cuda._sleep(1_000_000_000) + source.add_(1) + ready = torch.cuda.Event() + ready.record() + observed = [] + + def observer(): + time.sleep(0.05) + observed.append(not ready.query()) + + thread = threading.Thread(target=observer) + thread.start() + try: + capsule = value.__dlpack__() + finally: + thread.join() + self.assertEqual(observed, [True], "Another Python thread must run while the export waits for CUDA.") + torch.testing.assert_close(torch.from_dlpack(capsule), source) + def test_import_waits_for_producer_stream(self): tensor_device = torch.cuda.device_count() - 1 producer = torch.cuda.Stream(device=tensor_device) diff --git a/tools/ci_build/requirements/transformers-test/requirements.txt b/tools/ci_build/requirements/transformers-test/requirements.txt index 0a3468093dde9..f4b6c36ad5deb 100644 --- a/tools/ci_build/requirements/transformers-test/requirements.txt +++ b/tools/ci_build/requirements/transformers-test/requirements.txt @@ -9,7 +9,8 @@ torchvision==0.25.0 pillow>=13.0.0 transformers==5.10.1 parameterized>=0.8.1 -sentencepiece +# Newer generated SentencePiece protos require protobuf 7, above the runtime pinned here. +sentencepiece==0.2.1 psutil einops onnxscript==0.6.2 From 8834a7e794bb3a4566419f97d3dac919068f3118 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 5 Nov 2026 12:26:09 -0800 Subject: [PATCH 08/11] Select producer device during DLPack protocol exports Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../migraphx/migraphx_provider_factory.cc | 10 +++ .../migraphx/migraphx_provider_factory.h | 2 + .../onnxruntime_inference_collection.py | 2 + .../python/onnxruntime_pybind_state_common.cc | 61 ++++++++++++++++--- .../python/onnxruntime_test_python_dlpack.py | 38 ++++++++++-- 5 files changed, 101 insertions(+), 12 deletions(-) diff --git a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc index 51faace1d8caa..a85332c7fb66b 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc +++ b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc @@ -82,6 +82,16 @@ struct ProviderInfo_MIGraphX_Impl final : ProviderInfo_MIGraphX { return Status::OK(); } + int GetCurrentGpuDeviceId() override { + int device_id; + HIP_CALL_THROW(hipGetDevice(&device_id)); + return device_id; + } + + void SetCurrentGpuDeviceId(int device_id) override { + HIP_CALL_THROW(hipSetDevice(device_id)); + } + std::shared_ptr CreateMIGraphXAllocator(OrtDevice::DeviceId device_id, size_t mem_limit, onnxruntime::ArenaExtendStrategy arena_extend_strategy, void* alloc_fn, void* free_fn, void* empty_cache_fn, const OrtArenaCfg* default_memory_arena_cfg) override { if (alloc_fn != nullptr && free_fn != nullptr) { diff --git a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h index 72045d66d6d3c..43635d2e79def 100644 --- a/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h +++ b/onnxruntime/core/providers/migraphx/migraphx_provider_factory.h @@ -20,6 +20,8 @@ struct ProviderInfo_MIGraphX { virtual std::shared_ptr CreateMIGraphXAllocator(OrtDevice::DeviceId device_id, size_t mem_limit, ArenaExtendStrategy arena_extend_strategy, void* alloc_fn, void* free_fn, void* empty_cache_fn, const OrtArenaCfg* default_memory_arena_cfg) = 0; virtual Status SynchronizeDevice(int device_id) = 0; + virtual int GetCurrentGpuDeviceId() = 0; + virtual void SetCurrentGpuDeviceId(int device_id) = 0; protected: ~ProviderInfo_MIGraphX() = default; // Can only be destroyed through a subclass instance diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index 9b30cfd5f5644..b1b7d09e1b3f7 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -1589,6 +1589,8 @@ def from_dlpack(cls, data, /, *, synchronize: bool = True) -> OrtValue: GPU imports synchronize the tensor's device before returning, so ORT can safely use the storage on its own streams. CUDA and ROCm require their corresponding CUDA or MIGraphX provider for synchronization. + When that provider is available, GPU protocol exports run with the producer's + device selected and restore the caller's current device, including on failure. The capsule is consumed before synchronization; if that wait fails, the imported storage is released and the capsule cannot be retried. diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.cc b/onnxruntime/python/onnxruntime_pybind_state_common.cc index 99fe130ff19c8..5726c248d4955 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.cc +++ b/onnxruntime/python/onnxruntime_pybind_state_common.cc @@ -1,3 +1,5 @@ +#include + #include "onnxruntime_pybind_exceptions.h" #include "onnxruntime_pybind_state_common.h" @@ -65,6 +67,50 @@ void DlpackCapsuleDestructor(PyObject* data) { } namespace { +template +py::object ExportDlpackOnDevice(const py::object& method, bool no_sync, int device_id, + int original_device, SetDevice set_device) { + if (original_device == device_id) { + return no_sync ? method(py::arg("stream") = -1) : method(); + } + set_device(device_id); + py::object capsule; + std::exception_ptr error; + ORT_TRY { + capsule = no_sync ? method(py::arg("stream") = -1) : method(); + } + ORT_CATCH(...) { + error = std::current_exception(); + } + // Restore on both success and producer failure, and surface restoration errors. + set_device(original_device); + if (error) { + std::rethrow_exception(error); + } + return capsule; +} + +py::object ExportDlpack(const py::object& method, bool no_sync, const std::pair& device) { + if (device.first == kDLCUDA) { +#if defined(USE_CUDA) || defined(USE_CUDA_PROVIDER_INTERFACE) + if (auto* provider = TryGetProviderInfo_CUDA()) { + int original_device; + Ort::ThrowOnError(provider->GetCurrentGpuDeviceId(&original_device)); + return ExportDlpackOnDevice(method, no_sync, device.second, original_device, + [provider](int id) { Ort::ThrowOnError(provider->SetCurrentGpuDeviceId(id)); }); + } +#endif + } else if (device.first == kDLROCM) { +#if defined(USE_MIGRAPHX) || defined(USE_MIGRAPHX_PROVIDER_INTERFACE) + if (auto* provider = TryGetProviderInfo_MIGraphX()) { + return ExportDlpackOnDevice(method, no_sync, device.second, provider->GetCurrentGpuDeviceId(), + [provider](int id) { provider->SetCurrentGpuDeviceId(id); }); + } +#endif + } + return no_sync ? method(py::arg("stream") = -1) : method(); +} + py::object GetOptionalDlpackAttribute(const py::object& data, const char* name) { auto attribute = py::reinterpret_steal(PyObject_GetAttrString(data.ptr(), name)); if (attribute) { @@ -152,15 +198,14 @@ OrtValue FromDlpack(py::object data, bool synchronize) { if (!PyCapsule_CheckExact(data.ptr())) { auto method = GetOptionalDlpackAttribute(data, "__dlpack__"); if (!method.is_none()) { - bool request_no_sync = false; - if (!synchronize) { - auto device_method = GetOptionalDlpackAttribute(data, "__dlpack_device__"); - if (!device_method.is_none()) { - const auto device = device_method().cast>(); - request_no_sync = device.first == kDLCUDA || device.first == kDLROCM; - } + auto device_method = GetOptionalDlpackAttribute(data, "__dlpack_device__"); + if (!device_method.is_none()) { + const auto device = device_method().cast>(); + const bool no_sync = !synchronize && (device.first == kDLCUDA || device.first == kDLROCM); + data = ExportDlpack(method, no_sync, device); + } else { + data = method(); } - data = request_no_sync ? method(py::arg("stream") = -1) : method(); } } PyObject* dlpack_tensor = data.ptr(); diff --git a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py index b8e0c47c34c1e..bb758cb16b4c1 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py +++ b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py @@ -289,8 +289,12 @@ def test_gpu_vendor_metadata_without_gpu_allocation(self): shape = (ctypes.c_int64 * 1)(0) for device_type in (2, 10): # kDLCUDA, kDLROCM with self.subTest(device_type=device_type): + provider = "CUDAExecutionProvider" if device_type == 2 else "MIGraphXExecutionProvider" + device_id = 0 if provider in ort.get_available_providers() else 3 managed = DLManagedTensor( - DLTensor(None, DLDevice(device_type, 3), 1, DLDataType(2, 32, 1), shape, None, 0), None, None + DLTensor(None, DLDevice(device_type, device_id), 1, DLDataType(2, 32, 1), shape, None, 0), + None, + None, ) capsule = new_capsule(ctypes.addressof(managed), b"dltensor", None) imported = ort.OrtValue.from_dlpack(capsule, synchronize=False) @@ -298,18 +302,18 @@ def test_gpu_vendor_metadata_without_gpu_allocation(self): result = None values = C.OrtValueVector() try: - self.assertEqual(imported.__dlpack_device__(), (device_type, 3)) + self.assertEqual(imported.__dlpack_device__(), (device_type, device_id)) for value_type in (ort.OrtValue, C.OrtValue): for producer in (imported, imported._get_c_value()): copied = value_type.from_dlpack(producer, synchronize=False) - self.assertEqual(copied.__dlpack_device__(), (device_type, 3)) + self.assertEqual(copied.__dlpack_device__(), (device_type, device_id)) del copied del producer values.push_back(imported._get_c_value()) exported = values.to_dlpacks(None, synchronize=False) for result in (imported.__dlpack__(stream=-1), *exported): device = managed_tensor_from_capsule(result).contents.dl_tensor.device - self.assertEqual((device.device_type, device.device_id), (device_type, 3)) + self.assertEqual((device.device_type, device.device_id), (device_type, device_id)) finally: # The synthetic producer's descriptor must outlive every native owner. del result, imported, exported, values @@ -642,6 +646,32 @@ def test_noncontiguous_input_is_rejected(self): "requires PyTorch CUDA and the ORT CUDA execution provider", ) class TestDlpackCUDA(unittest.TestCase): + @unittest.skipUnless(torch.cuda.device_count() >= 2, "requires two CUDA devices") + def test_protocol_import_restores_current_device(self): + tensor_device = torch.cuda.device_count() - 1 + with torch.cuda.device(tensor_device): + source = torch.arange(6, dtype=torch.float32, device=f"cuda:{tensor_device}") + torch.cuda.current_stream().synchronize() + + class InvalidProducer: + def __dlpack_device__(self): + return (2, tensor_device) + + def __dlpack__(self, *, stream=None): + if torch.cuda.current_device() != tensor_device: + raise AssertionError("Producer's device must be current") + raise ValueError("Producer export failed") + + for value_type in (ort.OrtValue, C.OrtValue): + for synchronize in (False, True): + with self.subTest(value_type=value_type, synchronize=synchronize), torch.cuda.device(0): + imported = value_type.from_dlpack(source, synchronize=synchronize) + self.assertEqual(torch.cuda.current_device(), 0) + torch.testing.assert_close(torch.from_dlpack(imported), source) + with self.assertRaisesRegex(ValueError, "Producer export failed"): + value_type.from_dlpack(InvalidProducer(), synchronize=synchronize) + self.assertEqual(torch.cuda.current_device(), 0) + def test_import_claims_capsule_before_releasing_gil(self): source = torch.zeros(6, device="cuda") torch.cuda.synchronize() From 24431ef76daa97869282d8850f3dbec87cf638aa Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 5 Nov 2026 12:41:38 -0800 Subject: [PATCH 09/11] Harden DLPack cleanup and portable GPU coverage Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../python/onnxruntime_pybind_ortvalue.cc | 4 +- .../python/onnxruntime_pybind_state_common.cc | 2 +- .../python/onnxruntime_pybind_state_common.h | 4 +- .../python/onnxruntime_test_python_dlpack.py | 49 ++++++++++++------- 4 files changed, 37 insertions(+), 22 deletions(-) diff --git a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc index e021d02167302..9cefb28d4ef01 100644 --- a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc @@ -608,7 +608,7 @@ void addOrtValueMethods(pybind11::module& m) { } // Callbacks can mutate the original vector, including reallocating its storage. auto values = v; - InlinedHashSet> synchronized_devices; + InlinedHashMap> synchronized_devices; py::list capsules; for (auto& value : values) { ORT_ENFORCE(value.IsTensor(), "DLPack export requires a dense tensor OrtValue; initialize the value before exporting."); @@ -616,7 +616,7 @@ void addOrtValueMethods(pybind11::module& m) { const auto device = dlpack::GetDlpackDevice(value, value.Get().Location().device.Id()); py::object stream = py::none(); if (device.device_type == kDLCUDA || device.device_type == kDLROCM) { - if (!synchronize || !synchronized_devices.emplace(device.device_type, device.device_id).second) { + if (!synchronize || !synchronized_devices[device.device_type].insert(device.device_id).second) { stream = py::int_(-1); } } diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.cc b/onnxruntime/python/onnxruntime_pybind_state_common.cc index 5726c248d4955..51cf35776e336 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.cc +++ b/onnxruntime/python/onnxruntime_pybind_state_common.cc @@ -19,7 +19,7 @@ void ReleasePythonObject(PyObject* object) noexcept { if (PyGILState_Check()) { Py_DECREF(object); } else { - DeferPythonRelease(object); + ReleasePythonObjectOnCleanupThread(object); } } diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.h b/onnxruntime/python/onnxruntime_pybind_state_common.h index 855579050e282..bc8c6e2a9fcfc 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.h +++ b/onnxruntime/python/onnxruntime_pybind_state_common.h @@ -532,8 +532,8 @@ bool CheckIfTensor(const std::vector& def_list, const std::string& name, /*out*/ ONNX_NAMESPACE::TypeProto& type_proto); -// Transfer an owned reference to the existing deferred cleanup queue. -void DeferPythonRelease(PyObject* object) noexcept; +// Transfer an owned reference without waiting for the GIL or the main Python thread. +void ReleasePythonObjectOnCleanupThread(PyObject* object) noexcept; void ReleasePythonObject(PyObject* object) noexcept; #if defined(ENABLE_DLPACK) diff --git a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py index bb758cb16b4c1..53e03dd878c7e 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py +++ b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py @@ -257,13 +257,14 @@ def test_numpy_noncontiguous_copy_survives_export(self): self.assertIsNone(owner_ref()) torch.testing.assert_close(result, torch.arange(6, dtype=torch.float32).flip(0)) - def test_numpy_owner_release_without_gil(self): + def test_numpy_owner_release_without_gil_while_main_thread_is_blocked(self): for contiguous in (False, True): with self.subTest(contiguous=contiguous): array = np.arange(6, dtype=np.float32) if not contiguous: array = array[::-1] - owner_ref = weakref.ref(array) + released = threading.Event() + owner_ref = weakref.ref(array, lambda _, released=released: released.set()) value = ort.OrtValue.ortvalue_from_numpy(array) capsule = value.__dlpack__() managed = managed_tensor_from_capsule(capsule) @@ -275,10 +276,20 @@ def test_numpy_owner_release_without_gil(self): del value, array gc.collect() self.assertEqual(owner_ref() is not None, contiguous) - # CFUNCTYPE releases the GIL while calling the actual C++ deleter. - deleter(managed) - self.assertEqual(ctypes.pythonapi.Py_MakePendingCalls(), 0) - gc.collect() + observed = [] + + def release(deleter=deleter, managed=managed, observed=observed, released=released): + # Let the main thread block in join before releasing native storage. + time.sleep(0.05) + # CFUNCTYPE releases the GIL while calling the actual C++ deleter. + deleter(managed) + observed.append(released.wait(timeout=5)) + + worker = threading.Thread(target=release) + worker.start() + worker.join(timeout=10) + self.assertFalse(worker.is_alive(), "Native destruction must not wait for the main thread.") + self.assertEqual(observed, [True], "NumPy storage must be released while the main thread is blocked.") self.assertIsNone(owner_ref()) def test_gpu_vendor_metadata_without_gpu_allocation(self): @@ -642,20 +653,23 @@ def test_noncontiguous_input_is_rejected(self): @unittest.skipUnless( - torch.cuda.is_available() and "CUDAExecutionProvider" in ort.get_available_providers(), - "requires PyTorch CUDA and the ORT CUDA execution provider", + torch.cuda.is_available() + and ("MIGraphXExecutionProvider" if torch.version.hip else "CUDAExecutionProvider") + in ort.get_available_providers(), + "requires PyTorch CUDA/ROCm and the corresponding ORT CUDA/MIGraphX execution provider", ) -class TestDlpackCUDA(unittest.TestCase): - @unittest.skipUnless(torch.cuda.device_count() >= 2, "requires two CUDA devices") +class TestDlpackGPU(unittest.TestCase): + @unittest.skipUnless(torch.cuda.device_count() >= 2, "requires two GPU devices") def test_protocol_import_restores_current_device(self): tensor_device = torch.cuda.device_count() - 1 + device_type = 10 if torch.version.hip else 2 with torch.cuda.device(tensor_device): source = torch.arange(6, dtype=torch.float32, device=f"cuda:{tensor_device}") torch.cuda.current_stream().synchronize() class InvalidProducer: def __dlpack_device__(self): - return (2, tensor_device) + return (device_type, tensor_device) def __dlpack__(self, *, stream=None): if torch.cuda.current_device() != tensor_device: @@ -667,6 +681,7 @@ def __dlpack__(self, *, stream=None): with self.subTest(value_type=value_type, synchronize=synchronize), torch.cuda.device(0): imported = value_type.from_dlpack(source, synchronize=synchronize) self.assertEqual(torch.cuda.current_device(), 0) + self.assertEqual(imported.__dlpack_device__(), (device_type, tensor_device)) torch.testing.assert_close(torch.from_dlpack(imported), source) with self.assertRaisesRegex(ValueError, "Producer export failed"): value_type.from_dlpack(InvalidProducer(), synchronize=synchronize) @@ -708,7 +723,7 @@ def observer(stop=stop, capsule=capsule, observed=observed, ready=ready): finally: stop.set() thread.join() - self.assertEqual(observed, [True], "Another Python thread must run while the import waits for CUDA.") + self.assertEqual(observed, [True], "Another Python thread must run while the import waits for the GPU.") torch.testing.assert_close(torch.from_dlpack(imported), source) def test_export_releases_gil_during_device_wait(self): @@ -732,7 +747,7 @@ def observer(): capsule = value.__dlpack__() finally: thread.join() - self.assertEqual(observed, [True], "Another Python thread must run while the export waits for CUDA.") + self.assertEqual(observed, [True], "Another Python thread must run while the export waits for the GPU.") torch.testing.assert_close(torch.from_dlpack(capsule), source) def test_import_waits_for_producer_stream(self): @@ -765,8 +780,8 @@ def test_export_waits_for_producer_and_restores_current_device(self): producer = torch.cuda.Stream(device=tensor_device) consumer = torch.cuda.Stream(device=tensor_device) for export in ( - lambda: public.__dlpack__(stream=1), - lambda: public._get_c_value().__dlpack__(stream=2), + lambda: public.__dlpack__(stream=0 if torch.version.hip else 1), + lambda: public._get_c_value().__dlpack__(stream=1 if torch.version.hip else 2), lambda: public.__dlpack__(stream=consumer.cuda_stream), lambda: public.__dlpack__(), lambda: values.to_dlpacks(None)[0], @@ -844,9 +859,9 @@ def consumer(capsule): for result in results: torch.testing.assert_close(result, source) - def test_invalid_cuda_stream(self): + def test_invalid_gpu_stream(self): value = ort.OrtValue.from_dlpack(torch.ones(1, device="cuda")) - for stream in (0, -2): + for stream in (-2,) if torch.version.hip else (0, -2): with self.subTest(stream=stream), self.assertRaisesRegex(ValueError, "Invalid GPU DLPack stream"): value.__dlpack__(stream=stream) for stream in (False, 1.5, "invalid"): From 3af546633554b6da3875acd8acf6b0489be4c4b6 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 6 Nov 2026 14:26:08 -0800 Subject: [PATCH 10/11] Support plugin-factory DLPack device handoffs Add optional device synchronization and context callbacks with capability-aware lookup and retained library ownership. Forward callbacks through provider bridges and cover callback lifetime, version guards, and CUDA plugin handoffs. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/cuda_plugin_ep/cuda_plugin_ep_design.md | 10 ++ .../onnxruntime/core/session/environment.h | 13 +- .../core/session/onnxruntime_ep_c_api.h | 41 ++++++ .../providers/cuda/plugin/cuda_ep_factory.cc | 35 +++++ .../providers/cuda/plugin/cuda_ep_factory.h | 4 + onnxruntime/core/session/environment.cc | 27 +++- .../session/plugin_ep/ep_factory_internal.cc | 18 +++ .../session/plugin_ep/ep_factory_internal.h | 2 + .../plugin_ep/ep_factory_internal_impl.h | 3 + .../plugin_ep/ep_factory_provider_bridge.h | 2 + .../onnxruntime_inference_collection.py | 3 +- .../python/onnxruntime_pybind_state_common.cc | 38 ++++- .../library/example_plugin_ep/ep_factory.cc | 65 ++++++++- .../library/example_plugin_ep/ep_test_hooks.h | 5 + .../example_plugin_ep_library.lds | 5 + .../test/framework/ep_device_interop_test.cc | 91 ++++++++++++ .../python/onnxruntime_test_python_dlpack.py | 133 +++++++++++++++++- .../transformers/test_cuda_plugin_ep.py | 71 ++++++++++ 18 files changed, 555 insertions(+), 11 deletions(-) create mode 100644 onnxruntime/test/framework/ep_device_interop_test.cc diff --git a/docs/cuda_plugin_ep/cuda_plugin_ep_design.md b/docs/cuda_plugin_ep/cuda_plugin_ep_design.md index 13d8b59b79bd0..8f5a2d4645906 100644 --- a/docs/cuda_plugin_ep/cuda_plugin_ep_design.md +++ b/docs/cuda_plugin_ep/cuda_plugin_ep_design.md @@ -346,6 +346,16 @@ Implementation details: This is intentionally conservative and correct for the plugin EP's first sync integration. A narrower stream-scoped synchronization strategy can be considered later if profiling shows a need. +For session-independent DLPack handoffs, `CudaEpFactory` also implements the optional +`SynchronizeDevice`, `GetCurrentDevice`, and `SetCurrentDevice` callbacks. Python imports +select the producer's device around its protocol callback and restore the original device, +including on exceptions. Default imports and exports use a device-wide barrier; explicit +no-sync handoffs remain available when the caller arranges stream ordering. +The environment pins the factory's library throughout each handoff callback or wait, so +unregistering the library concurrently cannot unload executing code. +These are provider implementation callbacks rather than application APIs; no C++ wrapper +is needed. + #### 5.1.2 Python Host/Device Tensor Copies The Python `OrtValue` helpers (`update_inplace()` for host-to-device and `numpy()` for device-to-host) historically reached CUDA copies through the legacy provider bridge (`GetProviderInfo_CUDA()`). That bridge requires the provider shared library to export `GetProvider()`, which the CUDA plugin intentionally does not export. diff --git a/include/onnxruntime/core/session/environment.h b/include/onnxruntime/core/session/environment.h index 503c3d2571b5c..cfe4c28950910 100644 --- a/include/onnxruntime/core/session/environment.h +++ b/include/onnxruntime/core/session/environment.h @@ -173,6 +173,17 @@ class Environment { return execution_devices_; } + struct DeviceInterop { + std::shared_ptr library_owner; + OrtEpFactory* factory = nullptr; + }; + + enum class DeviceInteropCapability { Synchronize, + CurrentDevice }; + + // Pins the library across native waits and arbitrary external producer callbacks. + DeviceInterop GetDeviceInterop(const OrtDevice& device, DeviceInteropCapability capability) const; + /// Get hardware device incompatibility details for a specific EP. /// @param ep_name The name of the execution provider to check. /// @param hw The hardware device to check for incompatibility. @@ -308,7 +319,7 @@ class Environment { }; // registration name to EpInfo for library - std::unordered_map> ep_libraries_; + std::unordered_map> ep_libraries_; // combined set of OrtEpDevices for all registered OrtEpFactory instances // std::vector so we can use directly in GetEpDevices. diff --git a/include/onnxruntime/core/session/onnxruntime_ep_c_api.h b/include/onnxruntime/core/session/onnxruntime_ep_c_api.h index f47d5cd8775b7..bb26df1477f82 100644 --- a/include/onnxruntime/core/session/onnxruntime_ep_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_ep_c_api.h @@ -3358,6 +3358,47 @@ struct OrtEpFactory { _In_ const OrtEpDevice* ep_device, _In_ OrtDeviceMemoryType memory_type, _Out_opt_ uint64_t* total_memory_bytes, _Out_opt_ uint64_t* available_memory_bytes); + + /** \brief Wait for all outstanding work on a device without creating a session. + * + * Optional callback used for external tensor handoffs, including DLPack. This must wait for + * all streams on the specified device, not merely an EP-owned stream. Preserve the calling + * thread's current device on success and failure; report synchronization or restoration errors. + * Calls may occur concurrently on different threads. ORT retains the factory and its library + * for the duration of the call. Leave null if device-wide synchronization is unsupported. + * + * \param[in] this_ptr The factory instance. + * \param[in] device_id The device ID from the EP device's DEFAULT memory information. + * \snippet{doc} snippets.dox OrtStatus Return Value + * \since Version 1.32. + */ + ORT_API2_STATUS(SynchronizeDevice, _In_ OrtEpFactory* this_ptr, _In_ int device_id); + + /** \brief Get the calling thread's current device for this factory's runtime. + * + * Optional callback, paired with SetCurrentDevice. Device IDs use the same namespace as + * DEFAULT memory information. Do not modify the output on failure. ORT retains the factory + * and its library during a handoff, including intervening producer callbacks. + * + * \param[in] this_ptr The factory instance. + * \param[out] device_id The current device ID. + * \snippet{doc} snippets.dox OrtStatus Return Value + * \since Version 1.32. + */ + ORT_API2_STATUS(GetCurrentDevice, _In_ OrtEpFactory* this_ptr, _Out_ int* device_id); + + /** \brief Select the calling thread's current device for this factory's runtime. + * + * Optional callback, paired with GetCurrentDevice. Do not change the current device on + * failure. ORT uses this pair to select a producer's device and restore the original device + * after its callback, including on exception. Calls may occur concurrently on different threads. + * + * \param[in] this_ptr The factory instance. + * \param[in] device_id The device ID to select. + * \snippet{doc} snippets.dox OrtStatus Return Value + * \since Version 1.32. + */ + ORT_API2_STATUS(SetCurrentDevice, _In_ OrtEpFactory* this_ptr, _In_ int device_id); }; #ifdef __cplusplus diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc index d2b2d4e639175..bc9f274eadf47 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc +++ b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.cc @@ -62,6 +62,41 @@ CudaEpFactory::CudaEpFactory(const OrtApi& ort_api, const OrtEpApi& ep_api, IsStreamAware = IsStreamAwareImpl; CreateSyncStreamForDevice = CreateSyncStreamForDeviceImpl; GetDeviceMemoryCapacity = GetDeviceMemoryCapacityImpl; + SynchronizeDevice = SynchronizeDeviceImpl; + GetCurrentDevice = GetCurrentDeviceImpl; + SetCurrentDevice = SetCurrentDeviceImpl; +} + +OrtStatus* ORT_API_CALL CudaEpFactory::SynchronizeDeviceImpl(OrtEpFactory*, int device_id) noexcept { + EXCEPTION_TO_STATUS_BEGIN + int original_device; + PL_CUDA_RETURN_IF_ERROR(cudaGetDevice(&original_device)); + Ort::Status status = StatusFromCudaError(cudaSetDevice(device_id)); + if (status.IsOK()) { + status = StatusFromCudaError(cudaDeviceSynchronize()); + } + Ort::Status restore_status = StatusFromCudaError(cudaSetDevice(original_device)); + if (!restore_status.IsOK()) { + return restore_status.release(); + } + return status.release(); + EXCEPTION_TO_STATUS_END +} + +OrtStatus* ORT_API_CALL CudaEpFactory::GetCurrentDeviceImpl(OrtEpFactory*, int* device_id) noexcept { + EXCEPTION_TO_STATUS_BEGIN + int current_device; + PL_CUDA_RETURN_IF_ERROR(cudaGetDevice(¤t_device)); + *device_id = current_device; + return nullptr; + EXCEPTION_TO_STATUS_END +} + +OrtStatus* ORT_API_CALL CudaEpFactory::SetCurrentDeviceImpl(OrtEpFactory*, int device_id) noexcept { + EXCEPTION_TO_STATUS_BEGIN + PL_CUDA_RETURN_IF_ERROR(cudaSetDevice(device_id)); + return nullptr; + EXCEPTION_TO_STATUS_END } CudaEpFactory::~CudaEpFactory() { diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.h b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.h index 0beb6b9def3ac..672787294710a 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_ep_factory.h @@ -125,6 +125,10 @@ class CudaEpFactory : public OrtEpFactory { uint64_t* total_memory_bytes, uint64_t* available_memory_bytes) noexcept; + static OrtStatus* ORT_API_CALL SynchronizeDeviceImpl(OrtEpFactory* this_ptr, int device_id) noexcept; + static OrtStatus* ORT_API_CALL GetCurrentDeviceImpl(OrtEpFactory* this_ptr, int* device_id) noexcept; + static OrtStatus* ORT_API_CALL SetCurrentDeviceImpl(OrtEpFactory* this_ptr, int device_id) noexcept; + const OrtApi& ort_api_; const OrtEpApi& ep_api_; const OrtLogger& default_logger_; diff --git a/onnxruntime/core/session/environment.cc b/onnxruntime/core/session/environment.cc index bc8009b05c9e7..ab0c9f92c539b 100644 --- a/onnxruntime/core/session/environment.cc +++ b/onnxruntime/core/session/environment.cc @@ -598,8 +598,9 @@ Status Environment::RegisterExecutionProviderLibrary(const std::string& registra return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "EP library is already registered under ", registration_name); } - std::unique_ptr ep_info; - ORT_RETURN_IF_ERROR(EpInfo::Create(std::move(ep_library), ep_info, internal_factories)); + std::unique_ptr created_ep_info; + ORT_RETURN_IF_ERROR(EpInfo::Create(std::move(ep_library), created_ep_info, internal_factories)); + std::shared_ptr ep_info{std::move(created_ep_info)}; // Register shared allocators const size_t max_allocator_count = SafeInt(ep_info->execution_devices.size()) * 2; @@ -703,7 +704,7 @@ Status Environment::RegisterExecutionProviderLibrary(const std::string& registra // Notes: // `ep_info`'s target must stay alive for any rollback logic that may need to run. // Directly moving into `ep_libraries_.emplace()` could cause it to be prematurely destroyed if emplace throws - // after ownership has transferred from `ep_info`. This is why we first emplace an empty unique_ptr and then + // after ownership has transferred from `ep_info`. This is why we first emplace an empty shared_ptr and then // swap into it. // // We ignore `emplace_result.second` since we already checked for an existing entry in `ep_libraries_` above. @@ -797,6 +798,26 @@ Status Environment::RegisterExecutionProviderLibrary(const std::string& registra return RegisterExecutionProviderLibrary(registration_name, std::move(ep_library), internal_factories); } +Environment::DeviceInterop Environment::GetDeviceInterop(const OrtDevice& device, + DeviceInteropCapability capability) const { + std::lock_guard lock{mutex_}; + for (const auto& [name, info] : ep_libraries_) { + ORT_UNUSED_PARAMETER(name); + for (const auto& ep_device : info->execution_devices) { + auto* factory = ep_device->GetMutableFactory(); + if (ep_device->device_memory_info != nullptr && + ep_device->device_memory_info->device == device && + factory->ort_version_supported >= 32 && + (capability == DeviceInteropCapability::Synchronize + ? factory->SynchronizeDevice != nullptr + : factory->GetCurrentDevice != nullptr && factory->SetCurrentDevice != nullptr)) { + return {info, factory}; + } + } + } + return {}; +} + Status Environment::UnregisterExecutionProviderLibrary(const std::string& registration_name) { std::lock_guard lock{mutex_}; diff --git a/onnxruntime/core/session/plugin_ep/ep_factory_internal.cc b/onnxruntime/core/session/plugin_ep/ep_factory_internal.cc index 86a6f6df95f19..d6eb8cafb106c 100644 --- a/onnxruntime/core/session/plugin_ep/ep_factory_internal.cc +++ b/onnxruntime/core/session/plugin_ep/ep_factory_internal.cc @@ -38,6 +38,24 @@ EpFactoryInternal::EpFactoryInternal(std::unique_ptr impl OrtEpFactory::InitGraphicsInterop = Forward::InitGraphicsInterop; OrtEpFactory::DeinitGraphicsInterop = Forward::DeinitGraphicsInterop; OrtEpFactory::SelectBestModelCandidate = Forward::SelectBestModelCandidate; + if (auto* factory = GetDeviceInteropFactory(); factory != nullptr && factory->ort_version_supported >= 32) { + if (factory->SynchronizeDevice != nullptr) { + OrtEpFactory::SynchronizeDevice = [](OrtEpFactory* self, int device_id) noexcept -> OrtStatus* { + auto* inner = static_cast(self)->GetDeviceInteropFactory(); + return inner->SynchronizeDevice(inner, device_id); + }; + } + if (factory->GetCurrentDevice != nullptr && factory->SetCurrentDevice != nullptr) { + OrtEpFactory::GetCurrentDevice = [](OrtEpFactory* self, int* device_id) noexcept -> OrtStatus* { + auto* inner = static_cast(self)->GetDeviceInteropFactory(); + return inner->GetCurrentDevice(inner, device_id); + }; + OrtEpFactory::SetCurrentDevice = [](OrtEpFactory* self, int device_id) noexcept -> OrtStatus* { + auto* inner = static_cast(self)->GetDeviceInteropFactory(); + return inner->SetCurrentDevice(inner, device_id); + }; + } + } } InternalExecutionProviderFactory::InternalExecutionProviderFactory(EpFactoryInternal& ep_factory, diff --git a/onnxruntime/core/session/plugin_ep/ep_factory_internal.h b/onnxruntime/core/session/plugin_ep/ep_factory_internal.h index 78a45d4615062..6b6373f0318e5 100644 --- a/onnxruntime/core/session/plugin_ep/ep_factory_internal.h +++ b/onnxruntime/core/session/plugin_ep/ep_factory_internal.h @@ -23,6 +23,8 @@ class EpFactoryInternal : public OrtEpFactory { public: EpFactoryInternal(std::unique_ptr impl); + OrtEpFactory* GetDeviceInteropFactory() const noexcept { return impl_->GetDeviceInteropFactory(); } + const char* GetName() const noexcept { return impl_->GetName(); } const char* GetVendor() const noexcept { return impl_->GetVendor(); } uint32_t GetVendorId() const noexcept { return impl_->GetVendorId(); } diff --git a/onnxruntime/core/session/plugin_ep/ep_factory_internal_impl.h b/onnxruntime/core/session/plugin_ep/ep_factory_internal_impl.h index fa11ae7b52ae0..56180f019447d 100644 --- a/onnxruntime/core/session/plugin_ep/ep_factory_internal_impl.h +++ b/onnxruntime/core/session/plugin_ep/ep_factory_internal_impl.h @@ -92,6 +92,9 @@ class EpFactoryInternalImpl { "GetDeviceMemoryCapacity is not implemented for this EP factory."); } + // Provider-bridge wrappers expose only the callbacks supported by the underlying factory. + virtual OrtEpFactory* GetDeviceInteropFactory() const noexcept { return nullptr; } + virtual OrtStatus* CreateExternalResourceImporterForDevice( _In_ const OrtEpDevice* /*ep_device*/, _Outptr_result_maybenull_ OrtExternalResourceImporterImpl** importer) noexcept { diff --git a/onnxruntime/core/session/plugin_ep/ep_factory_provider_bridge.h b/onnxruntime/core/session/plugin_ep/ep_factory_provider_bridge.h index fe2208603284c..76455efaa36d2 100644 --- a/onnxruntime/core/session/plugin_ep/ep_factory_provider_bridge.h +++ b/onnxruntime/core/session/plugin_ep/ep_factory_provider_bridge.h @@ -27,6 +27,8 @@ class ProviderBridgeEpFactory : public EpFactoryInternalImpl { } private: + OrtEpFactory* GetDeviceInteropFactory() const noexcept override { return &ep_factory_; } + OrtStatus* GetSupportedDevices(EpFactoryInternal& ep_factory, const OrtHardwareDevice* const* devices, size_t num_devices, diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index b1b7d09e1b3f7..3c233601efb1f 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -1588,7 +1588,8 @@ def from_dlpack(cls, data, /, *, synchronize: bool = True) -> OrtValue: No data is copied; the new OrtValue shares memory with the source. GPU imports synchronize the tensor's device before returning, so ORT can safely use the storage on its own streams. CUDA and ROCm require - their corresponding CUDA or MIGraphX provider for synchronization. + their corresponding CUDA or MIGraphX provider for synchronization, + including registered plugin providers with device-wide synchronization support. When that provider is available, GPU protocol exports run with the producer's device selected and restore the caller's current device, including on failure. The capsule is consumed before synchronization; if that wait fails, the diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.cc b/onnxruntime/python/onnxruntime_pybind_state_common.cc index 51cf35776e336..78fd5c7cb8892 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.cc +++ b/onnxruntime/python/onnxruntime_pybind_state_common.cc @@ -67,6 +67,19 @@ void DlpackCapsuleDestructor(PyObject* data) { } namespace { +#if !defined(ORT_MINIMAL_BUILD) +Environment::DeviceInterop GetPluginDeviceInterop(int device_type, int device_id, + Environment::DeviceInteropCapability capability) { + if (device_type != kDLCUDA && device_type != kDLROCM) { + return {}; + } + const auto vendor = device_type == kDLCUDA ? OrtDevice::VendorIds::NVIDIA : OrtDevice::VendorIds::AMD; + return GetOrtEnv()->GetEnvironment().GetDeviceInterop( + OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, vendor, gsl::narrow(device_id)), + capability); +} +#endif + template py::object ExportDlpackOnDevice(const py::object& method, bool no_sync, int device_id, int original_device, SetDevice set_device) { @@ -91,6 +104,17 @@ py::object ExportDlpackOnDevice(const py::object& method, bool no_sync, int devi } py::object ExportDlpack(const py::object& method, bool no_sync, const std::pair& device) { +#if !defined(ORT_MINIMAL_BUILD) + if (auto interop = GetPluginDeviceInterop(device.first, device.second, + Environment::DeviceInteropCapability::CurrentDevice); + interop.factory != nullptr) { + auto* factory = interop.factory; + int original_device; + Ort::ThrowOnError(factory->GetCurrentDevice(factory, &original_device)); + return ExportDlpackOnDevice(method, no_sync, device.second, original_device, + [factory](int id) { Ort::ThrowOnError(factory->SetCurrentDevice(factory, id)); }); + } +#endif if (device.first == kDLCUDA) { #if defined(USE_CUDA) || defined(USE_CUDA_PROVIDER_INTERFACE) if (auto* provider = TryGetProviderInfo_CUDA()) { @@ -133,6 +157,15 @@ void SynchronizeDlpackDevice(const DLDevice& device) { if (device.device_type == kDLCPU) { return; } +#if !defined(ORT_MINIMAL_BUILD) + if (auto interop = GetPluginDeviceInterop(device.device_type, device.device_id, + Environment::DeviceInteropCapability::Synchronize); + interop.factory != nullptr) { + py::gil_scoped_release release; + Ort::ThrowOnError(interop.factory->SynchronizeDevice(interop.factory, device.device_id)); + return; + } +#endif if (device.device_type == kDLCUDA) { #if defined(USE_CUDA) || defined(USE_CUDA_PROVIDER_INTERFACE) if (auto* provider = TryGetProviderInfo_CUDA()) { @@ -151,8 +184,9 @@ void SynchronizeDlpackDevice(const DLDevice& device) { #endif } throw py::buffer_error( - "DLPack synchronization requires a provider matching the tensor's device (CUDA for NVIDIA, MIGraphX for AMD). " - "Install and load the matching provider, or explicitly opt out only after arranging stream ordering: " + "DLPack synchronization requires device-wide synchronization from a matching provider " + "(CUDA for NVIDIA, MIGraphX for AMD). Install and load a provider implementing this capability, " + "or explicitly opt out only after arranging stream ordering: " "__dlpack__(stream=-1), from_dlpack(..., synchronize=False), or to_dlpacks(..., synchronize=False)."); } } // namespace diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc index c362d00fd1dae..76c419acafa96 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc @@ -4,21 +4,32 @@ #include "ep_factory.h" #include +#include +#include #include #include #include +#include #include "ep.h" #include "ep_allocator.h" #include "ep_arena.h" #include "ep_data_transfer.h" #include "ep_stream_support.h" +#define EXAMPLE_PLUGIN_EP_BUILD #include "ep_test_hooks.h" +#undef EXAMPLE_PLUGIN_EP_BUILD #include "core/session/onnxruntime_ep_device_ep_metadata_keys.h" #include "core/session/onnxruntime_session_options_config_keys.h" namespace { +thread_local int dlpack_current_device = 7; +constexpr int kDlpackTestDeviceId = 32760; +std::atomic dlpack_sync_count{0}; +std::atomic dlpack_release_count{0}; +std::atomic dlpack_sync_blocked{false}; + int CompatibilityRank(OrtCompiledModelCompatibility c) { switch (c) { case OrtCompiledModelCompatibility_EP_SUPPORTED_OPTIMAL: @@ -57,6 +68,19 @@ std::optional GetCompatibilityField(const std::string& info, std::s } } // namespace +extern "C" { +EXPORT_SYMBOL int ExampleEpTestHooks_GetDlpackCurrentDevice() { return dlpack_current_device; } +EXPORT_SYMBOL int ExampleEpTestHooks_GetDlpackSyncCount() { return dlpack_sync_count.load(); } +EXPORT_SYMBOL int ExampleEpTestHooks_GetDlpackReleaseCount() { return dlpack_release_count.load(); } +EXPORT_SYMBOL void ExampleEpTestHooks_ResetDlpack() { + dlpack_current_device = 7; + dlpack_sync_count = 0; + dlpack_release_count = 0; + dlpack_sync_blocked = false; +} +EXPORT_SYMBOL void ExampleEpTestHooks_SetDlpackSyncBlocked(int blocked) { dlpack_sync_blocked = blocked != 0; } +} + ExampleEpFactory::ExampleEpFactory(const char* ep_name, ApiPtrs apis, const OrtLogger& default_logger, bool create_duplicate_virtual_devices) : OrtEpFactory{}, @@ -98,11 +122,43 @@ ExampleEpFactory::ExampleEpFactory(const char* ep_name, ApiPtrs apis, const OrtL SelectBestModelCandidate = SelectBestModelCandidateImpl; GetDeviceMemoryCapacity = GetDeviceMemoryCapacityImpl; + const bool dlpack_test = ep_name_.starts_with("example_ep_dlpack"); + if (dlpack_test) { + if (ep_name_.find("context") == std::string::npos) { + SynchronizeDevice = [](OrtEpFactory* self, int device_id) noexcept -> OrtStatus* { + auto* factory = static_cast(self); + if (device_id != kDlpackTestDeviceId) { + return factory->ort_api.CreateStatus(ORT_INVALID_ARGUMENT, "Expected DLPack test device 32760."); + } + ++dlpack_sync_count; + while (dlpack_sync_blocked.load()) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + return nullptr; + }; + } + if (ep_name_.find("sync") == std::string::npos) { + GetCurrentDevice = [](OrtEpFactory*, int* device_id) noexcept -> OrtStatus* { + *device_id = dlpack_current_device; + return nullptr; + }; + SetCurrentDevice = [](OrtEpFactory*, int device_id) noexcept -> OrtStatus* { + dlpack_current_device = device_id; + return nullptr; + }; + } + if (ep_name_.find("old") != std::string::npos) { + ort_version_supported = 31; + } + } + const uint32_t memory_vendor = dlpack_test ? 0x10DE : 0xBE57; + const uint32_t memory_device_id = dlpack_test ? static_cast(kDlpackTestDeviceId) : 0; + // setup the OrtMemoryInfo instances required by the EP. // We pretend the device the EP is running on is GPU. default_memory_info_ = Ort::MemoryInfo{"ExampleEP GPU", OrtMemoryInfoDeviceType_GPU, - /*vendor*/ 0xBE57, /* device_id */ 0, + /*vendor*/ memory_vendor, /* device_id */ memory_device_id, OrtDeviceMemoryType_DEFAULT, /*alignment*/ 0, // it is invalid to use OrtArenaAllocator as that is reserved for the internal @@ -116,7 +172,7 @@ ExampleEpFactory::ExampleEpFactory(const char* ep_name, ApiPtrs apis, const OrtL // create read-only allocator for use with initializers. same info as DEFAULT memory apart from the allocator type. readonly_memory_info_ = Ort::MemoryInfo{"ExampleEP GPU readonly", OrtMemoryInfoDeviceType_GPU, - /*vendor*/ 0xBE57, /* device_id */ 0, + /*vendor*/ memory_vendor, /* device_id */ memory_device_id, OrtDeviceMemoryType_DEFAULT, /*alignment*/ 0, OrtAllocatorType::OrtReadOnlyAllocator}; @@ -125,7 +181,7 @@ ExampleEpFactory::ExampleEpFactory(const char* ep_name, ApiPtrs apis, const OrtL // accessible from. we infer from the type of HOST_ACCESSIBLE that it's CPU accessible. host_accessible_memory_info_ = Ort::MemoryInfo{"ExampleEP GPU pinned", OrtMemoryInfoDeviceType_GPU, - /*vendor*/ 0xBE57, /* device_id */ 0, + /*vendor*/ memory_vendor, /* device_id */ memory_device_id, OrtDeviceMemoryType_HOST_ACCESSIBLE, /*alignment*/ 0, OrtAllocatorType::OrtDeviceAllocator}; @@ -148,6 +204,9 @@ ExampleEpFactory::ExampleEpFactory(const char* ep_name, ApiPtrs apis, const OrtL } ExampleEpFactory::~ExampleEpFactory() { + if (ep_name_.starts_with("example_ep_dlpack")) { + ++dlpack_release_count; + } for (auto* device : virtual_hardware_devices_) { ep_api.ReleaseHardwareDevice(device); } diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep_test_hooks.h b/onnxruntime/test/autoep/library/example_plugin_ep/ep_test_hooks.h index 5969b315cf884..d3a92211299f3 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep_test_hooks.h +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep_test_hooks.h @@ -28,6 +28,11 @@ EXPORT_SYMBOL void ExampleEpTestHooks_ResetPreallocatedOutputQuery(); EXPORT_SYMBOL int ExampleEpTestHooks_GetPreallocatedOutputQueryResult(); EXPORT_SYMBOL int ExampleEpTestHooks_GetPreallocatedOutputBadIndexRejected(); EXPORT_SYMBOL void ExampleEpTestHooks_SetCreateDataTransferFailure(int enabled); +EXPORT_SYMBOL int ExampleEpTestHooks_GetDlpackCurrentDevice(); +EXPORT_SYMBOL int ExampleEpTestHooks_GetDlpackSyncCount(); +EXPORT_SYMBOL int ExampleEpTestHooks_GetDlpackReleaseCount(); +EXPORT_SYMBOL void ExampleEpTestHooks_ResetDlpack(); +EXPORT_SYMBOL void ExampleEpTestHooks_SetDlpackSyncBlocked(int blocked); } // Internal to the library; not exported. diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/example_plugin_ep_library.lds b/onnxruntime/test/autoep/library/example_plugin_ep/example_plugin_ep_library.lds index bed5d9956be7b..b62d047531a57 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/example_plugin_ep_library.lds +++ b/onnxruntime/test/autoep/library/example_plugin_ep/example_plugin_ep_library.lds @@ -8,6 +8,11 @@ VERS_1.0.0 { ExampleEpTestHooks_GetPreallocatedOutputQueryResult; ExampleEpTestHooks_GetPreallocatedOutputBadIndexRejected; ExampleEpTestHooks_SetCreateDataTransferFailure; + ExampleEpTestHooks_GetDlpackCurrentDevice; + ExampleEpTestHooks_GetDlpackSyncCount; + ExampleEpTestHooks_GetDlpackReleaseCount; + ExampleEpTestHooks_ResetDlpack; + ExampleEpTestHooks_SetDlpackSyncBlocked; local: *; }; diff --git a/onnxruntime/test/framework/ep_device_interop_test.cc b/onnxruntime/test/framework/ep_device_interop_test.cc new file mode 100644 index 0000000000000..7be6fd186193e --- /dev/null +++ b/onnxruntime/test/framework/ep_device_interop_test.cc @@ -0,0 +1,91 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(ORT_MINIMAL_BUILD) + +#include "core/session/plugin_ep/ep_factory_internal.h" +#include "gtest/gtest.h" + +namespace onnxruntime::test { +namespace { + +struct InteropFactory : OrtEpFactory { + InteropFactory() : OrtEpFactory{} { + ort_version_supported = ORT_API_VERSION; + SynchronizeDevice = [](OrtEpFactory* self, int device_id) noexcept -> OrtStatus* { + static_cast(self)->last_sync_device = device_id; + return nullptr; + }; + GetCurrentDevice = [](OrtEpFactory* self, int* device_id) noexcept -> OrtStatus* { + *device_id = static_cast(self)->current_device; + return nullptr; + }; + SetCurrentDevice = [](OrtEpFactory* self, int device_id) noexcept -> OrtStatus* { + static_cast(self)->current_device = device_id; + return nullptr; + }; + } + int current_device = 7; + int last_sync_device = -1; +}; + +class InteropFactoryWrapper final : public EpFactoryInternalImpl { + public: + explicit InteropFactoryWrapper(OrtEpFactory& factory) + : EpFactoryInternalImpl("InteropTest", "Test", 0), factory_(factory) {} + + OrtEpFactory* GetDeviceInteropFactory() const noexcept override { return &factory_; } + + OrtStatus* GetSupportedDevices(EpFactoryInternal&, const OrtHardwareDevice* const*, size_t, + OrtEpDevice**, size_t, size_t* count) noexcept override { + *count = 0; + return nullptr; + } + + OrtStatus* CreateIExecutionProvider(const OrtHardwareDevice* const*, const OrtKeyValuePairs* const*, size_t, + const OrtSessionOptions*, const OrtLogger*, + std::unique_ptr*) override { + return OrtApis::CreateStatus(ORT_NOT_IMPLEMENTED, "Interop test factory does not create execution providers."); + } + + private: + OrtEpFactory& factory_; +}; + +TEST(EpDeviceInteropTest, InternalFactoryForwardsSupportedCallbacks) { + InteropFactory inner; + EpFactoryInternal wrapper(std::make_unique(inner)); + auto& factory = static_cast(wrapper); + ASSERT_NE(factory.SynchronizeDevice, nullptr); + ASSERT_NE(factory.GetCurrentDevice, nullptr); + ASSERT_NE(factory.SetCurrentDevice, nullptr); + int current_device = -1; + EXPECT_EQ(factory.GetCurrentDevice(&factory, ¤t_device), nullptr); + EXPECT_EQ(current_device, 7); + EXPECT_EQ(factory.SetCurrentDevice(&factory, 2), nullptr); + EXPECT_EQ(inner.current_device, 2); + EXPECT_EQ(factory.SynchronizeDevice(&factory, 2), nullptr); + EXPECT_EQ(inner.last_sync_device, 2); +} + +TEST(EpDeviceInteropTest, InternalFactoryPreservesOptionalAndVersionGuards) { + for (const bool old_version : {false, true}) { + InteropFactory inner; + if (old_version) { + inner.ort_version_supported = 31; + } else { + inner.SetCurrentDevice = nullptr; + inner.SynchronizeDevice = nullptr; + } + EpFactoryInternal wrapper(std::make_unique(inner)); + auto& factory = static_cast(wrapper); + EXPECT_EQ(factory.SynchronizeDevice, nullptr); + EXPECT_EQ(factory.GetCurrentDevice, nullptr); + EXPECT_EQ(factory.SetCurrentDevice, nullptr); + } +} + +} // namespace +} // namespace onnxruntime::test + +#endif diff --git a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py index 53e03dd878c7e..b1ec613e7da4d 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_dlpack.py +++ b/onnxruntime/test/python/onnxruntime_test_python_dlpack.py @@ -9,10 +9,12 @@ import ctypes import gc import math +import os import threading import time import unittest import weakref +from pathlib import Path import numpy as np import torch @@ -360,7 +362,7 @@ def test_failed_gpu_synchronization_consumes_capsule_once(self): ctypes.cast(deleter, ctypes.c_void_p), ) capsule = new_capsule(ctypes.addressof(managed), b"dltensor", None) - with self.assertRaisesRegex(BufferError, "provider matching the tensor's device"): + with self.assertRaisesRegex(BufferError, "matching provider"): value_type.from_dlpack(capsule) self.assertEqual(len(releases), 1) with self.assertRaisesRegex(TypeError, "valid DLPack capsule"): @@ -652,6 +654,135 @@ def test_noncontiguous_input_is_rejected(self): ort.OrtValue.from_dlpack(source) +class TestDlpackPluginCallbacks(unittest.TestCase): + device_id = 32760 # Avoid colliding with a real GPU provider registered by the build. + + def setUp(self): + library_name = "example_plugin_ep.dll" if os.name == "nt" else "libexample_plugin_ep.so" + self.library_path = Path.cwd() / library_name + if not self.library_path.is_file(): + self.skipTest("requires the example_plugin_ep test library in the build output directory") + self.library = ctypes.CDLL(str(self.library_path)) + self.current_device = self.library.ExampleEpTestHooks_GetDlpackCurrentDevice + self.sync_count = self.library.ExampleEpTestHooks_GetDlpackSyncCount + self.release_count = self.library.ExampleEpTestHooks_GetDlpackReleaseCount + for method in (self.current_device, self.sync_count, self.release_count): + method.argtypes = [] + method.restype = ctypes.c_int + self.library.ExampleEpTestHooks_ResetDlpack() + self.block_sync = self.library.ExampleEpTestHooks_SetDlpackSyncBlocked + self.block_sync.argtypes = [ctypes.c_int] + self.block_sync.restype = None + self.registrations = [] + self.descriptors = [] + + def tearDown(self): + for registration in reversed(getattr(self, "registrations", [])): + ort.unregister_execution_provider_library(registration) + + def register(self, suffix): + name = f"example_ep_dlpack_{suffix}" + ort.register_execution_provider_library(name, str(self.library_path)) + self.registrations.append(name) + return name + + def unregister(self, name): + ort.unregister_execution_provider_library(name) + self.registrations.remove(name) + + def capsule(self): + shape = (ctypes.c_int64 * 1)(0) + managed = DLManagedTensor( + DLTensor(None, DLDevice(2, self.device_id), 1, DLDataType(2, 32, 1), shape, None, 0), None, None + ) + self.descriptors.append((shape, managed)) + new_capsule = ctypes.pythonapi.PyCapsule_New + new_capsule.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] + new_capsule.restype = ctypes.py_object + return new_capsule(ctypes.addressof(managed), b"dltensor", None) + + def test_complementary_factory_capabilities(self): + self.register("sync") + self.register("context") + test = self + + class Producer: + def __dlpack_device__(self): + return (2, test.device_id) + + def __dlpack__(self, *, stream=None): + test.assertEqual(test.current_device(), test.device_id) + return test.capsule() + + for value_type in (ort.OrtValue, C.OrtValue): + value = value_type.from_dlpack(Producer()) + self.assertEqual(self.current_device(), 7) + before = self.sync_count() + capsule = value.__dlpack__() + self.assertEqual(self.sync_count(), before + 1) + self.assertEqual(value.__dlpack_device__(), (2, self.device_id)) + del capsule, value + self.assertEqual(self.sync_count(), 4) + + def test_unregister_inside_producer_pins_factory_until_restoration(self): + name = self.register("both") + test = self + + class Producer: + def __dlpack_device__(self): + return (2, test.device_id) + + def __dlpack__(self, *, stream=None): + test.assertEqual(test.current_device(), test.device_id) + test.unregister(name) + test.assertEqual(test.release_count(), 0) + raise ValueError("Producer export failed after unregister") + + with self.assertRaisesRegex(ValueError, "Producer export failed after unregister"): + ort.OrtValue.from_dlpack(Producer(), synchronize=False) + self.assertEqual(self.current_device(), 7) + self.assertEqual(self.release_count(), 1) + + def test_unregister_during_native_wait_pins_factory(self): + name = self.register("sync") + capsule = self.capsule() + results = [] + errors = [] + + def import_value(): + try: + results.append(ort.OrtValue.from_dlpack(capsule)) + except Exception as error: + errors.append(error) + + worker = threading.Thread(target=import_value) + self.block_sync(1) + worker.start() + try: + deadline = time.monotonic() + 5 + while self.sync_count() == 0 and time.monotonic() < deadline: + time.sleep(0.001) + self.assertEqual(self.sync_count(), 1) + self.unregister(name) + self.assertEqual(self.release_count(), 0) + finally: + self.block_sync(0) + worker.join(timeout=5) + self.assertFalse(worker.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(len(results), 1) + self.assertEqual(self.release_count(), 1) + + def test_old_factory_callbacks_are_not_read(self): + self.register("old") + value = ort.OrtValue.from_dlpack(self.capsule(), synchronize=False) + with self.assertRaises((BufferError, RuntimeError)): + value.__dlpack__() + self.assertEqual(self.sync_count(), 0) + self.assertEqual(self.current_device(), 7) + del value + + @unittest.skipUnless( torch.cuda.is_available() and ("MIGraphXExecutionProvider" if torch.version.hip else "CUDAExecutionProvider") diff --git a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py index 9bea5d4f95831..7889bc4530398 100644 --- a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py +++ b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py @@ -14,6 +14,7 @@ from onnx import OperatorSetIdProto, TensorProto, helper, save import onnxruntime as onnxrt +from onnxruntime.capi import _pybind_state as C try: import faulthandler @@ -34,6 +35,76 @@ DEFAULT_ONNX_OPSET = 27 +class TestCudaPluginDlpack(unittest.TestCase): + def setUp(self): + require_cuda_plugin_ep() + if not torch.cuda.is_available(): + self.skipTest("requires PyTorch CUDA") + self.device_id = get_cuda_plugin_device_id(get_cuda_plugin_device()) + + def test_default_import_export_waits_for_all_streams(self): + device = f"cuda:{self.device_id}" + producer = torch.cuda.Stream(device=self.device_id) + for value_type in (onnxrt.OrtValue, C.OrtValue): + for raw_capsule in (False, True): + with self.subTest(value_type=value_type, raw_capsule=raw_capsule): + with torch.cuda.stream(producer): + source = torch.zeros(6, device=device) + torch.cuda._sleep(50_000_000) + source.fill_(7) + ready = torch.cuda.Event() + ready.record() + exported = source.__dlpack__(stream=-1) if raw_capsule else source + imported = value_type.from_dlpack(exported) + self.assertTrue(ready.query()) + self.assertEqual(imported.__dlpack_device__(), (2, self.device_id)) + with torch.cuda.stream(producer): + torch.cuda._sleep(50_000_000) + source.add_(1) + ready.record() + result = torch.from_dlpack(imported) + self.assertTrue(ready.query()) + torch.testing.assert_close(result, torch.full((6,), 8.0, device=device)) + + values = C.OrtValueVector() + values.push_back(source, synchronize=False) + values.push_back(source, synchronize=False) + with torch.cuda.stream(producer): + torch.cuda._sleep(50_000_000) + source.add_(1) + ready.record() + results = values.to_dlpacks(torch.from_dlpack) + self.assertTrue(ready.query()) + for result in results: + torch.testing.assert_close(result, source) + + @unittest.skipUnless(torch.cuda.device_count() >= 2, "requires two CUDA devices") + def test_protocol_import_restores_device_on_success_and_exception(self): + device_id = get_cuda_plugin_device_id(get_cuda_plugin_devices()[-1]) + original_device = (device_id + 1) % torch.cuda.device_count() + source = torch.arange(6, dtype=torch.float32, device=f"cuda:{device_id}") + torch.cuda.synchronize(device_id) + + class InvalidProducer: + def __dlpack_device__(self): + return (2, device_id) + + def __dlpack__(self, *, stream=None): + if torch.cuda.current_device() != device_id: + raise AssertionError("Producer device was not selected") + raise ValueError("Producer export failed") + + for value_type in (onnxrt.OrtValue, C.OrtValue): + for synchronize in (False, True): + with self.subTest(value_type=value_type, synchronize=synchronize), torch.cuda.device(original_device): + imported = value_type.from_dlpack(source, synchronize=synchronize) + self.assertEqual(torch.cuda.current_device(), original_device) + torch.testing.assert_close(torch.from_dlpack(imported), source) + with self.assertRaisesRegex(ValueError, "Producer export failed"): + value_type.from_dlpack(InvalidProducer(), synchronize=synchronize) + self.assertEqual(torch.cuda.current_device(), original_device) + + def _make_released_opset_model(graph, producer_name="onnx-example"): opset = OperatorSetIdProto() opset.version = DEFAULT_ONNX_OPSET From 0c1e7941688acac8984d0c47a78e00f1f7869a90 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 6 Nov 2026 14:26:09 -0800 Subject: [PATCH 11/11] Avoid math-library initialization in CUDA copy tests Use a native CUDA stream for data-transfer self-copy tests so cuBLAS and cuDNN handle initialization cannot fail before copy behavior is exercised. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../providers/cuda/test_cases/gpu_data_transfer_test.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/providers/cuda/test_cases/gpu_data_transfer_test.cc b/onnxruntime/test/providers/cuda/test_cases/gpu_data_transfer_test.cc index 261c8fb0a56bc..4de480448e963 100644 --- a/onnxruntime/test/providers/cuda/test_cases/gpu_data_transfer_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/gpu_data_transfer_test.cc @@ -68,7 +68,8 @@ TEST(GPUDataTransferSameBufferTest, CopyTensorAsyncSelfCopyIsNoOp) { cudaStream_t cuda_stream = nullptr; ASSERT_EQ(cudaStreamCreate(&cuda_stream), cudaSuccess); - CudaStream stream(cuda_stream, gpu_allocator->Info().device, pinned_alloc, false, true, nullptr, nullptr, info); + auto release_stream = gsl::finally([&]() { EXPECT_EQ(cudaStreamDestroy(cuda_stream), cudaSuccess); }); + Stream stream(cuda_stream, gpu_allocator->Info().device); auto tensor = MakePinnedTensor(pinned_alloc); @@ -93,7 +94,8 @@ TEST(GPUDataTransferSameBufferTest, CopyTensorAsyncSelfCopyDuringCaptureKeepsCap cudaStream_t cuda_stream = nullptr; ASSERT_EQ(cudaStreamCreate(&cuda_stream), cudaSuccess); - CudaStream stream(cuda_stream, gpu_allocator->Info().device, pinned_alloc, false, true, nullptr, nullptr, info); + auto release_stream = gsl::finally([&]() { EXPECT_EQ(cudaStreamDestroy(cuda_stream), cudaSuccess); }); + Stream stream(cuda_stream, gpu_allocator->Info().device); auto tensor = MakePinnedTensor(pinned_alloc);