From d185bf444de1dd5b573da8d089d4a97417cf68d4 Mon Sep 17 00:00:00 2001 From: FFFrog <85333600+FFFrog@users.noreply.github.com> Date: Thu, 1 Jun 2023 18:13:38 +0800 Subject: [PATCH] [CANN] Add IOBinding Support For CANN EP (#15802) ### Description Add IOBinding Support For CANN EP ### Motivation and Context Now, Users can use IOBinding feature to speed up the inference on CANN. --- .../providers/cann/cann_execution_provider.cc | 41 ++++++++------- .../providers/cann/cann_execution_provider.h | 3 ++ .../providers/cann/cann_provider_factory.cc | 22 ++++++++ .../providers/cann/cann_provider_factory.h | 10 ++++ .../onnxruntime_inference_collection.py | 12 +++-- .../python/onnxruntime_pybind_mlvalue.cc | 50 +++++++++++++++++++ .../python/onnxruntime_pybind_mlvalue.h | 14 ++++++ .../python/onnxruntime_pybind_ortvalue.cc | 15 ++++++ .../python/onnxruntime_pybind_state.cc | 5 ++ 9 files changed, 150 insertions(+), 22 deletions(-) diff --git a/onnxruntime/core/providers/cann/cann_execution_provider.cc b/onnxruntime/core/providers/cann/cann_execution_provider.cc index e0a2db0626..75996e809c 100644 --- a/onnxruntime/core/providers/cann/cann_execution_provider.cc +++ b/onnxruntime/core/providers/cann/cann_execution_provider.cc @@ -1434,6 +1434,28 @@ Status CANNExecutionProvider::Compile(const std::vector& fuse return Status::OK(); } +AllocatorPtr CANNExecutionProvider::CreateCannAllocator(OrtDevice::DeviceId device_id, size_t npu_mem_limit, + ArenaExtendStrategy arena_extend_strategy, + OrtArenaCfg* default_memory_arena_cfg) { + AllocatorCreationInfo default_memory_info( + [](OrtDevice::DeviceId id) { + return std::make_unique(id, CANN); + }, + device_id, + true, + {default_memory_arena_cfg ? *default_memory_arena_cfg + : OrtArenaCfg(npu_mem_limit, + static_cast(arena_extend_strategy), + -1, + -1, + -1, + -1L)}, + true, + false); + + return CreateAllocator(default_memory_info); +} + void CANNExecutionProvider::RegisterAllocator(AllocatorManager& allocator_manager) { OrtDevice cann_device{OrtDevice::NPU, OrtDevice::MemType::DEFAULT, info_.device_id}; OrtDevice pinned_device{OrtDevice::CPU, OrtDevice::MemType::CANN_PINNED, DEFAULT_CPU_ALLOCATOR_DEVICE_ID}; @@ -1444,23 +1466,8 @@ void CANNExecutionProvider::RegisterAllocator(AllocatorManager& allocator_manage cann_alloc = allocator_manager.GetAllocator(OrtMemTypeDefault, cann_device); if (!cann_alloc) { - AllocatorCreationInfo default_memory_info( - [](OrtDevice::DeviceId id) { - return std::make_unique(id, CANN); - }, - cann_device.Id(), - true, - {info_.default_memory_arena_cfg ? *info_.default_memory_arena_cfg - : OrtArenaCfg(info_.npu_mem_limit, - static_cast(info_.arena_extend_strategy), - -1, - -1, - -1, - -1)}, - true, - false); - - cann_alloc = CreateAllocator(default_memory_info); + cann_alloc = CreateCannAllocator(info_.device_id, info_.npu_mem_limit, info_.arena_extend_strategy, + info_.default_memory_arena_cfg); allocator_manager.InsertAllocator(cann_alloc); } diff --git a/onnxruntime/core/providers/cann/cann_execution_provider.h b/onnxruntime/core/providers/cann/cann_execution_provider.h index f10ad7af9d..26f863db8c 100644 --- a/onnxruntime/core/providers/cann/cann_execution_provider.h +++ b/onnxruntime/core/providers/cann/cann_execution_provider.h @@ -81,6 +81,9 @@ class CANNExecutionProvider : public IExecutionProvider { return CANNExecutionProviderInfo::ToProviderOptions(info_); } + static AllocatorPtr CreateCannAllocator(OrtDevice::DeviceId device_id, size_t npu_mem_limit, + ArenaExtendStrategy arena_extend_strategy, + OrtArenaCfg* arena_cfg); void RegisterAllocator(AllocatorManager& allocator_manager) override; void RegisterStreamHandlers(IStreamCommandHandleRegistry& stream_handle_registry) const override; diff --git a/onnxruntime/core/providers/cann/cann_provider_factory.cc b/onnxruntime/core/providers/cann/cann_provider_factory.cc index 636ca22591..679a42be86 100644 --- a/onnxruntime/core/providers/cann/cann_provider_factory.cc +++ b/onnxruntime/core/providers/cann/cann_provider_factory.cc @@ -32,6 +32,21 @@ std::unique_ptr CANNProviderFactory::CreateProvider() { } struct ProviderInfo_CANN_Impl : ProviderInfo_CANN { + int cannGetDeviceCount() override { + uint32_t num_devices = 0; + CANN_CALL_THROW(aclrtGetDeviceCount(&num_devices)); + return num_devices; + } + + void cannMemcpy_HostToDevice(void* dst, const void* src, size_t count) override { + CANN_CALL_THROW(aclrtMemcpy(dst, count, src, count, ACL_MEMCPY_HOST_TO_DEVICE)); + CANN_CALL_THROW(aclrtSynchronizeStream(0)); + } + + void cannMemcpy_DeviceToHost(void* dst, const void* src, size_t count) override { + CANN_CALL_THROW(aclrtMemcpy(dst, count, src, count, ACL_MEMCPY_DEVICE_TO_HOST)); + } + void CANNExecutionProviderInfo__FromProviderOptions(const ProviderOptions& options, CANNExecutionProviderInfo& info) override { info = CANNExecutionProviderInfo::FromProviderOptions(options); @@ -41,6 +56,13 @@ struct ProviderInfo_CANN_Impl : ProviderInfo_CANN { CreateExecutionProviderFactory(const CANNExecutionProviderInfo& info) override { return std::make_shared(info); } + + std::shared_ptr CreateCannAllocator(int16_t device_id, size_t npu_mem_limit, + onnxruntime::ArenaExtendStrategy arena_extend_strategy, + OrtArenaCfg* default_memory_arena_cfg) override { + return CANNExecutionProvider::CreateCannAllocator(device_id, npu_mem_limit, arena_extend_strategy, + default_memory_arena_cfg); + } } g_info; struct CANN_Provider : Provider { diff --git a/onnxruntime/core/providers/cann/cann_provider_factory.h b/onnxruntime/core/providers/cann/cann_provider_factory.h index ec40812023..44068760df 100644 --- a/onnxruntime/core/providers/cann/cann_provider_factory.h +++ b/onnxruntime/core/providers/cann/cann_provider_factory.h @@ -10,14 +10,24 @@ #include "core/providers/cann/cann_provider_options.h" namespace onnxruntime { +class IAllocator; +class IDataTransfer; struct IExecutionProviderFactory; struct CANNExecutionProviderInfo; +enum class ArenaExtendStrategy : int32_t; struct ProviderInfo_CANN { + virtual int cannGetDeviceCount() = 0; + + virtual void cannMemcpy_HostToDevice(void* dst, const void* src, size_t count) = 0; + virtual void cannMemcpy_DeviceToHost(void* dst, const void* src, size_t count) = 0; virtual void CANNExecutionProviderInfo__FromProviderOptions(const onnxruntime::ProviderOptions& options, onnxruntime::CANNExecutionProviderInfo& info) = 0; virtual std::shared_ptr CreateExecutionProviderFactory(const onnxruntime::CANNExecutionProviderInfo& info) = 0; + virtual std::shared_ptr + CreateCannAllocator(int16_t device_id, size_t npu_mem_limit, onnxruntime::ArenaExtendStrategy arena_extend_strategy, + OrtArenaCfg* default_memory_arena_cfg) = 0; }; } // namespace onnxruntime diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index d269d7d6f9..b6672958c7 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -20,6 +20,8 @@ if typing.TYPE_CHECKING: def get_ort_device_type(device_type: str, device_index) -> C.OrtDevice: if device_type == "cuda": return C.OrtDevice.cuda() + elif device_type == "cann": + return C.OrtDevice.cann() elif device_type == "cpu": return C.OrtDevice.cpu() elif device_type == "ort": @@ -490,7 +492,7 @@ class IOBinding: def bind_input(self, name, device_type, device_id, element_type, shape, buffer_ptr): """ :param name: input name - :param device_type: e.g. cpu, cuda + :param device_type: e.g. cpu, cuda, cann :param device_id: device id, e.g. 0 :param element_type: input element type :param shape: input shape @@ -529,7 +531,7 @@ class IOBinding: ): """ :param name: output name - :param device_type: e.g. cpu, cuda, cpu by default + :param device_type: e.g. cpu, cuda, cann, cpu by default :param device_id: device id, e.g. 0 :param element_type: output element type :param shape: output shape @@ -629,7 +631,7 @@ class OrtValue: A copy of the data in the Numpy object is held by the OrtValue only if the device is NOT cpu :param numpy_obj: The Numpy object to construct the OrtValue from - :param device_type: e.g. cpu, cuda, cpu by default + :param device_type: e.g. cpu, cuda, cann, cpu by default :param device_id: device id, e.g. 0 """ # Hold a reference to the numpy object (if device_type is 'cpu') as the OrtValue @@ -654,7 +656,7 @@ class OrtValue: :param shape: List of integers indicating the shape of the OrtValue :param element_type: The data type of the elements in the OrtValue (numpy type) - :param device_type: e.g. cpu, cuda, cpu by default + :param device_type: e.g. cpu, cuda, cann, cpu by default :param device_id: device id, e.g. 0 """ if shape is None or element_type is None: @@ -694,7 +696,7 @@ class OrtValue: def device_name(self): """ - Returns the name of the device where the OrtValue's data buffer resides e.g. cpu, cuda + Returns the name of the device where the OrtValue's data buffer resides e.g. cpu, cuda, cann """ return self._ortvalue.device_name().lower() diff --git a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc index ff4c2837a6..46934ef678 100644 --- a/onnxruntime/python/onnxruntime_pybind_mlvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_mlvalue.cc @@ -181,6 +181,56 @@ std::unique_ptr GetGPUDataTransfer() { #endif +#ifdef USE_CANN +void CpuToCannMemCpy(void* dst, const void* src, size_t num_bytes) { + GetProviderInfo_CANN().cannMemcpy_HostToDevice(dst, src, num_bytes); +} + +void CannToCpuMemCpy(void* dst, const void* src, size_t num_bytes) { + GetProviderInfo_CANN().cannMemcpy_DeviceToHost(dst, src, num_bytes); +} + +const std::unordered_map* GetCannToHostMemCpyFunction() { + static std::unordered_map map{ + {OrtDevice::NPU, CannToCpuMemCpy}}; + + return ↦ +} + +bool IsCannDeviceIdValid(const onnxruntime::logging::Logger& logger, int id) { + int num_devices = GetProviderInfo_CANN().cannGetDeviceCount(); + + if (0 == num_devices) { + LOGS(logger, WARNING) << "your system does not have a CANN capable device."; + return false; + } + + if (id < 0 || id >= num_devices) { + LOGS(logger, WARNING) << "cann_device=" << id << " is invalid, must choose device ID between 0 and " + << num_devices - 1; + return false; + } + + return true; +} + +AllocatorPtr GetCannAllocator(OrtDevice::DeviceId id) { + size_t npu_mem_limit = std::numeric_limits::max(); + onnxruntime::ArenaExtendStrategy arena_extend_strategy = onnxruntime::ArenaExtendStrategy::kNextPowerOfTwo; + + static auto* id_to_allocator_map = + std::make_unique>().release(); + auto hit = id_to_allocator_map->find(id); + if (hit == id_to_allocator_map->end()) { + auto cann_allocator = GetProviderInfo_CANN().CreateCannAllocator(id, npu_mem_limit, arena_extend_strategy, nullptr); + hit = id_to_allocator_map->emplace(id, std::move(cann_allocator)).first; + } + + return hit->second; +} + +#endif + #ifdef USE_ROCM void CpuToRocmMemCpy(void* dst, const void* src, size_t num_bytes) { GetProviderInfo_ROCM().rocmMemcpy_HostToDevice(dst, src, num_bytes); diff --git a/onnxruntime/python/onnxruntime_pybind_mlvalue.h b/onnxruntime/python/onnxruntime_pybind_mlvalue.h index 76b8819319..e1eb9aff86 100644 --- a/onnxruntime/python/onnxruntime_pybind_mlvalue.h +++ b/onnxruntime/python/onnxruntime_pybind_mlvalue.h @@ -74,6 +74,20 @@ std::unique_ptr GetGPUDataTransfer(); #endif +#ifdef USE_CANN + +void CpuToCannMemCpy(void* dst, const void* src, size_t num_bytes); + +void CannToCpuMemCpy(void* dst, const void* src, size_t num_bytes); + +const std::unordered_map* GetCannToHostMemCpyFunction(); + +bool IsCannDeviceIdValid(const onnxruntime::logging::Logger& logger, int id); + +AllocatorPtr GetCannAllocator(OrtDevice::DeviceId id); + +#endif + #ifdef USE_ROCM bool IsRocmDeviceIdValid(const onnxruntime::logging::Logger& logger, int id); diff --git a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc index 49bc314de3..367c06ec8b 100644 --- a/onnxruntime/python/onnxruntime_pybind_ortvalue.cc +++ b/onnxruntime/python/onnxruntime_pybind_ortvalue.cc @@ -68,6 +68,19 @@ void addOrtValueMethods(pybind11::module& m) { throw std::runtime_error( "Can't allocate memory on the CUDA device using this package of OnnxRuntime. " "Please use the CUDA package of OnnxRuntime to use this feature."); +#endif + } else if (device.Type() == OrtDevice::NPU) { +#ifdef USE_CANN + if (!IsCannDeviceIdValid(logging::LoggingManager::DefaultLogger(), device.Id())) { + throw std::runtime_error("The provided device id doesn't match any available NPUs on the machine."); + } + + CreateGenericMLValue(nullptr, GetCannAllocator(device.Id()), "", array_on_cpu, ml_value.get(), + true, false, CpuToCannMemCpy); +#else + throw std::runtime_error( + "Can't allocate memory on the CANN device using this package of OnnxRuntime. " + "Please use the CANN package of OnnxRuntime to use this feature."); #endif } else { throw std::runtime_error("Unsupported device: Cannot place the OrtValue on this device"); @@ -276,6 +289,8 @@ void addOrtValueMethods(pybind11::module& m) { GetPyObjFromTensor(ml_value->Get(), obj, nullptr, GetCudaToHostMemCpyFunction()); #elif USE_ROCM GetPyObjFromTensor(ml_value->Get(), obj, nullptr, GetRocmToHostMemCpyFunction()); +#elif USE_CANN + GetPyObjFromTensor(ml_value->Get(), obj, nullptr, GetCannToHostMemCpyFunction()); #else GetPyObjFromTensor(ml_value->Get(), obj, nullptr, nullptr); #endif diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index 0a7f6fb73d..56a1c4e52e 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -151,7 +151,11 @@ const char* GetDeviceName(const OrtDevice& device) { case OrtDevice::FPGA: return "FPGA"; case OrtDevice::NPU: +#ifdef USE_CANN + return CANN; +#else return "NPU"; +#endif default: ORT_THROW("Unknown device type: ", device.Type()); } @@ -1143,6 +1147,7 @@ void addObjectMethods(py::module& m, ExecutionProviderRegistrationFn ep_registra .def("device_type", &OrtDevice::Type, R"pbdoc(Device Type.)pbdoc") .def_static("cpu", []() { return OrtDevice::CPU; }) .def_static("cuda", []() { return OrtDevice::GPU; }) + .def_static("cann", []() { return OrtDevice::NPU; }) .def_static("fpga", []() { return OrtDevice::FPGA; }) .def_static("npu", []() { return OrtDevice::NPU; }) .def_static("default_memory", []() { return OrtDevice::MemType::DEFAULT; });