saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-17 18:40:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[CANN] Add IOBinding Support For CANN EP (#15802)

Browse source 
### Description
Add IOBinding Support For CANN EP

### Motivation and Context
Now, Users can use IOBinding feature to speed up the inference on CANN.
This commit is contained in:
FFFrog 2023-06-01 18:13:38 +08:00 committed by GitHub
parent 8c85d990c2
commit d185bf444d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 150 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
onnxruntime/core/providers/cann/cann_execution_provider.cc
View file

@ -1434,6 +1434,28 @@ Status CANNExecutionProvider::Compile(const std::vector<FusedNodeAndGraph>& 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<CANNAllocator>(id, CANN);
},
device_id,
true,
{default_memory_arena_cfg ? *default_memory_arena_cfg
: OrtArenaCfg(npu_mem_limit,
static_cast<int>(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<CANNAllocator>(id, CANN);
},
cann_device.Id(),
true,
{info_.default_memory_arena_cfg ? *info_.default_memory_arena_cfg
: OrtArenaCfg(info_.npu_mem_limit,
static_cast<int>(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);
}

3
onnxruntime/core/providers/cann/cann_execution_provider.h
View file

@ -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;

22
onnxruntime/core/providers/cann/cann_provider_factory.cc
View file

@ -32,6 +32,21 @@ std::unique_ptr<IExecutionProvider> 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<CANNProviderFactory>(info);
}
std::shared_ptr<IAllocator> 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 {

10
onnxruntime/core/providers/cann/cann_provider_factory.h
View file

@ -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<onnxruntime::IExecutionProviderFactory>
CreateExecutionProviderFactory(const onnxruntime::CANNExecutionProviderInfo& info) = 0;
virtual std::shared_ptr<onnxruntime::IAllocator>
CreateCannAllocator(int16_t device_id, size_t npu_mem_limit, onnxruntime::ArenaExtendStrategy arena_extend_strategy,
OrtArenaCfg* default_memory_arena_cfg) = 0;
};
} // namespace onnxruntime

12
onnxruntime/python/onnxruntime_inference_collection.py
View file

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

50
onnxruntime/python/onnxruntime_pybind_mlvalue.cc
View file

@ -181,6 +181,56 @@ std::unique_ptr<IDataTransfer> 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<OrtDevice::DeviceType, MemCpyFunc>* GetCannToHostMemCpyFunction() {
static std::unordered_map<OrtDevice::DeviceType, MemCpyFunc> map{
{OrtDevice::NPU, CannToCpuMemCpy}};
return &map;
}
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<size_t>::max();
onnxruntime::ArenaExtendStrategy arena_extend_strategy = onnxruntime::ArenaExtendStrategy::kNextPowerOfTwo;
static auto* id_to_allocator_map =
std::make_unique<std::unordered_map<OrtDevice::DeviceId, AllocatorPtr>>().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);

14
onnxruntime/python/onnxruntime_pybind_mlvalue.h
View file

@ -74,6 +74,20 @@ std::unique_ptr<IDataTransfer> 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<OrtDevice::DeviceType, MemCpyFunc>* 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);

15
onnxruntime/python/onnxruntime_pybind_ortvalue.cc
View file

@ -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<Tensor>(), obj, nullptr, GetCudaToHostMemCpyFunction());
#elif USE_ROCM
GetPyObjFromTensor(ml_value->Get<Tensor>(), obj, nullptr, GetRocmToHostMemCpyFunction());
#elif USE_CANN
GetPyObjFromTensor(ml_value->Get<Tensor>(), obj, nullptr, GetCannToHostMemCpyFunction());
#else
GetPyObjFromTensor(ml_value->Get<Tensor>(), obj, nullptr, nullptr);
#endif

5
onnxruntime/python/onnxruntime_pybind_state.cc
View file

@ -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; });