From 7e2408880e963bcfdd2b898c7b6464506545cec2 Mon Sep 17 00:00:00 2001 From: Takeshi Watanabe Date: Fri, 31 Jan 2025 07:23:56 +0900 Subject: [PATCH] Enable dlpack by default (#23110) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description This PR will enable python dlpack interface by default. ### Motivation and Context dlpack python interface is useful in inference mode not only training mode. Since some inference result preprocess may be written in torch and making unnecessary device transfer should be reduced in those cases. closes https://github.com/microsoft/onnxruntime/issues/15963 closes https://github.com/microsoft/onnxruntime/issues/22061 TODOs: - [x] Add tests like https://github.com/microsoft/onnxruntime/blob/5407c69028ae6dd4e87521aea147c22153d8e6c7/orttraining/orttraining/test/python/orttraining_test_ortvalue.py that's unrelated to training feature --------- Co-authored-by: Xavier Dupré Co-authored-by: Justin Chu --- cmake/CMakeLists.txt | 7 +++++ .../external/onnxruntime_external_deps.cmake | 4 +-- cmake/onnxruntime_providers_cpu.cmake | 6 +++- cmake/onnxruntime_python.cmake | 3 ++ .../python/onnxruntime_pybind_ortvalue.cc | 13 ++++---- .../python/onnxruntime_pybind_state_common.cc | 2 +- .../python/onnxruntime_pybind_state_common.h | 4 +-- .../test/python/onnxruntime_test_python.py | 30 ++++++++++++++++--- 8 files changed, 51 insertions(+), 18 deletions(-) diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 8650cc53d9..962b42c190 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -241,6 +241,9 @@ option(onnxruntime_ENABLE_CPUINFO "Enable cpuinfo" ON) # ATen fallback support option(onnxruntime_ENABLE_ATEN "Enable ATen fallback" OFF) +# dlpack support +cmake_dependent_option(onnxruntime_ENABLE_DLPACK "Enable dlpack" ON "onnxruntime_ENABLE_TRAINING OR onnxruntime_ENABLE_ATEN OR onnxruntime_ENABLE_PYTHON" OFF) + # Triton support option(onnxruntime_ENABLE_TRITON "Enable Triton" OFF) @@ -1603,6 +1606,10 @@ if (onnxruntime_ENABLE_TRAINING) list(APPEND onnxruntime_EXTERNAL_LIBRARIES tensorboard) endif() +if (onnxruntime_ENABLE_DLPACK) + add_compile_definitions(ENABLE_DLPACK) +endif() + if (UNIX OR onnxruntime_USE_NCCL) # MPI is INDEPENDENT of NCCL for now. You can build NCLL without MPI and launch multi-GPU with your own launcher. if (onnxruntime_USE_MPI) diff --git a/cmake/external/onnxruntime_external_deps.cmake b/cmake/external/onnxruntime_external_deps.cmake index e956569698..4e5875f969 100644 --- a/cmake/external/onnxruntime_external_deps.cmake +++ b/cmake/external/onnxruntime_external_deps.cmake @@ -576,8 +576,8 @@ if (onnxruntime_RUN_ONNX_TESTS) endif() -if(onnxruntime_ENABLE_ATEN) - message(STATUS "Aten fallback is enabled.") +if(onnxruntime_ENABLE_DLPACK) + message(STATUS "dlpack is enabled.") FetchContent_Declare( dlpack URL ${DEP_URL_dlpack} diff --git a/cmake/onnxruntime_providers_cpu.cmake b/cmake/onnxruntime_providers_cpu.cmake index 4ae89a3922..3e1a9edbd1 100644 --- a/cmake/onnxruntime_providers_cpu.cmake +++ b/cmake/onnxruntime_providers_cpu.cmake @@ -103,7 +103,7 @@ if (onnxruntime_ENABLE_TRAINING_OPS AND NOT onnxruntime_ENABLE_TRAINING) list(REMOVE_ITEM onnxruntime_providers_src ${onnxruntime_cpu_full_training_only_srcs}) endif() -if (onnxruntime_ENABLE_ATEN) +if (onnxruntime_ENABLE_DLPACK) file(GLOB_RECURSE onnxruntime_providers_dlpack_srcs CONFIGURE_DEPENDS "${ONNXRUNTIME_ROOT}/core/dlpack/dlpack_converter.cc" "${ONNXRUNTIME_ROOT}/core/dlpack/dlpack_converter.h" @@ -191,6 +191,10 @@ endif() if (onnxruntime_ENABLE_ATEN) target_compile_definitions(onnxruntime_providers PRIVATE ENABLE_ATEN) +endif() + +if (onnxruntime_ENABLE_DLPACK) + target_compile_definitions(onnxruntime_providers PRIVATE ENABLE_DLPACK) # DLPack is a header-only dependency set(DLPACK_INCLUDE_DIR ${dlpack_SOURCE_DIR}/include) target_include_directories(onnxruntime_providers PRIVATE ${DLPACK_INCLUDE_DIR}) diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 15a2862ced..b4a251dc6e 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -128,6 +128,9 @@ endif() if (onnxruntime_ENABLE_ATEN) target_compile_definitions(onnxruntime_pybind11_state PRIVATE ENABLE_ATEN) +endif() + +if (onnxruntime_ENABLE_DLPACK) target_include_directories(onnxruntime_pybind11_state PRIVATE ${dlpack_SOURCE_DIR}/include) endif() diff --git a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc index 6a57fc5f90..66ceacda75 100644 --- a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc @@ -13,9 +13,6 @@ #include "core/framework/tensor.h" #include "core/framework/sparse_tensor.h" #include "core/framework/TensorSeq.h" -#ifdef ENABLE_TRAINING -#include "core/dlpack/dlpack_converter.h" -#endif namespace onnxruntime { namespace python { @@ -350,7 +347,7 @@ void addOrtValueMethods(pybind11::module& m) { py::object obj = GetPyObjFromTensor(*ml_value, nullptr, nullptr); #endif return obj; }) -#ifdef ENABLE_TRAINING +#if defined(ENABLE_DLPACK) .def("to_dlpack", [](OrtValue* ort_value) -> py::object { 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 " @@ -373,7 +370,7 @@ void addOrtValueMethods(pybind11::module& m) { .def("push_back", [](std::vector* v, const OrtValue& ortvalue) { v->push_back(ortvalue); }) -#ifdef ENABLE_TRAINING +#if defined(ENABLE_DLPACK) .def("push_back", [](std::vector* v, py::object dlpack_tensor, const bool is_bool_tensor) { v->push_back(FromDlpack(dlpack_tensor.ptr(), is_bool_tensor)); }, "Add a new OrtValue after being ownership was transferred from the DLPack structure.", py::arg("dlpack_tensor"), py::arg("is_bool_tensor") = false) .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) { @@ -415,7 +412,7 @@ void addOrtValueMethods(pybind11::module& m) { "In case of a boolean tensor, method to_dlpacks returns a uint8 tensor instead of a boolean tensor. " "If torch consumes the dlpack structure, `.to(torch.bool)` must be applied to the torch tensor " "to get a boolean tensor.") -#ifdef ENABLE_TRAINING +#if defined(ENABLE_DLPACK) .def("dlpack_at", [](std::vector* v, const size_t idx) { 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)); }, @@ -424,7 +421,7 @@ void addOrtValueMethods(pybind11::module& m) { "(such as onnx.TensorProto.FLOAT)." "Raises an exception in any other case.", py::arg("idx")) -#ifdef ENABLE_TRAINING +#if defined(ENABLE_DLPACK) .def("to_dlpacks", [](const std::vector& v, py::object to_tensor) -> py::list { if (v.size() == 0) return py::list(); @@ -494,7 +491,7 @@ for every transferred tensor. #endif ; -#ifdef ENABLE_TRAINING +#if defined(ENABLE_DLPACK) m.def( "is_dlpack_uint8_tensor", [](py::capsule cap) -> bool { // case ONNX_NAMESPACE::TensorProto_DataType_BOOL: diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.cc b/onnxruntime/python/onnxruntime_pybind_state_common.cc index cec4dfc141..55ea264571 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.cc +++ b/onnxruntime/python/onnxruntime_pybind_state_common.cc @@ -43,7 +43,7 @@ onnxruntime::ROCMExecutionProviderExternalAllocatorInfo external_allocator_info{ onnxruntime::ArenaExtendStrategy arena_extend_strategy = onnxruntime::ArenaExtendStrategy::kNextPowerOfTwo; #endif -#ifdef ENABLE_TRAINING +#if defined(ENABLE_DLPACK) void DlpackCapsuleDestructor(PyObject* data) { DLManagedTensor* dlmanaged_tensor = reinterpret_cast(PyCapsule_GetPointer(data, "dltensor")); diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.h b/onnxruntime/python/onnxruntime_pybind_state_common.h index 995341b0f8..8d4a882b14 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.h +++ b/onnxruntime/python/onnxruntime_pybind_state_common.h @@ -12,7 +12,7 @@ #include "core/session/environment.h" #include "core/session/abi_session_options_impl.h" #include "core/session/inference_session.h" -#ifdef ENABLE_TRAINING +#if defined(ENABLE_DLPACK) #include "core/dlpack/dlpack_converter.h" #endif @@ -410,7 +410,7 @@ bool CheckIfTensor(const std::vector& def_list, const std::string& name, /*out*/ ONNX_NAMESPACE::TypeProto& type_proto); -#ifdef ENABLE_TRAINING +#if defined(ENABLE_DLPACK) // Allocate a new Capsule object, which takes the ownership of OrtValue. // Caller is responsible for releasing. diff --git a/onnxruntime/test/python/onnxruntime_test_python.py b/onnxruntime/test/python/onnxruntime_test_python.py index 91310cfc2a..3af6e8ccac 100644 --- a/onnxruntime/test/python/onnxruntime_test_python.py +++ b/onnxruntime/test/python/onnxruntime_test_python.py @@ -17,6 +17,7 @@ import numpy as np from helper import get_name import onnxruntime as onnxrt +from onnxruntime.capi import _pybind_state as C from onnxruntime.capi.onnxruntime_pybind11_state import Fail, OrtValueVector, RunOptions # handle change from python 3.8 and on where loading a dll from the current directory needs to be explicitly allowed. @@ -325,8 +326,6 @@ class TestInferenceSession(unittest.TestCase): self.assertEqual(option["user_compute_stream"], "1") self.assertEqual(option["has_user_compute_stream"], "1") - from onnxruntime.capi import _pybind_state as C - session_options = C.get_default_session_options() # TRT plugins registered as custom op domain should only be added once in session option regardless of number of session creation @@ -1421,6 +1420,31 @@ class TestInferenceSession(unittest.TestCase): outs = session.run(output_names=["output"], input_feed=upstreams_onnxrt)[0] self.assertTrue(np.allclose(inps, outs)) + @unittest.skipIf(not hasattr(C.OrtValue, "from_dlpack"), "dlpack not enabled in this build") + def test_ort_value_dlpack(self): + # Tests originally from orttraining/orttraining/test/python/orttraining_test_ortvalue.py testOrtValueDlPack_float32 + numpy_arr_input = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32) + ortvalue = onnxrt.OrtValue.ortvalue_from_numpy(numpy_arr_input) + self.assertEqual(numpy_arr_input.shape, tuple(ortvalue.shape())) + ptr = ortvalue._ortvalue.data_ptr() + + dlp = ortvalue._ortvalue.to_dlpack() + self.assertFalse(C.is_dlpack_uint8_tensor(dlp)) + ortvalue2 = C.OrtValue.from_dlpack(dlp, False) + self.assertEqual(ptr, ortvalue2.data_ptr()) + new_array = ortvalue2.numpy() + np.testing.assert_equal(numpy_arr_input, new_array) + + dlp = ortvalue._ortvalue.__dlpack__() + self.assertFalse(C.is_dlpack_uint8_tensor(dlp)) + ortvalue2 = C.OrtValue.from_dlpack(dlp, False) + self.assertEqual(ptr, ortvalue2.data_ptr()) + new_array = ortvalue2.numpy() + np.testing.assert_equal(numpy_arr_input, new_array) + + device = ortvalue._ortvalue.__dlpack_device__() + self.assertEqual((1, 0), device) + def test_sparse_tensor_coo_format(self): cpu_device = onnxrt.OrtDevice.make("cpu", 0) shape = [9, 9] @@ -1694,8 +1718,6 @@ class TestInferenceSession(unittest.TestCase): check_failure([("a", {1: 2})], [{3: 4}]) def test_register_custom_e_ps_library(self): - from onnxruntime.capi import _pybind_state as C - available_eps = C.get_available_providers() # skip amd gpu build if "ROCMExecutionProvider" in available_eps: