From 550293d9adcabd0739f862d62c5f6e8d8da00c70 Mon Sep 17 00:00:00 2001 From: Yulong Wang <7679871+fs-eire@users.noreply.github.com> Date: Fri, 8 Sep 2023 16:37:35 -0700 Subject: [PATCH] OrtMemoryInfo: support new name "WebGPU_Buffer" (#17469) ### Description Add new name "WebGPU_Buffer" to OrtMemoryInfo. This is one of the prerequisites for supporting IO binding for WebGPU buffer in onnxruntime-web. list of prerequisites PRs: #17465 #17469 (this one) --- .../onnxruntime/core/framework/allocator.h | 1 + onnxruntime/core/framework/allocator.cc | 20 ++++++------------- 2 files changed, 7 insertions(+), 14 deletions(-) diff --git a/include/onnxruntime/core/framework/allocator.h b/include/onnxruntime/core/framework/allocator.h index e2a90852e5..cbc2208b6b 100644 --- a/include/onnxruntime/core/framework/allocator.h +++ b/include/onnxruntime/core/framework/allocator.h @@ -48,6 +48,7 @@ constexpr const char* HIP = "Hip"; constexpr const char* HIP_PINNED = "HipPinned"; constexpr const char* OpenVINO_CPU = "OpenVINO_CPU"; constexpr const char* OpenVINO_GPU = "OpenVINO_GPU"; +constexpr const char* WEBGPU_BUFFER = "WebGPU_Buffer"; constexpr size_t kAllocAlignment = 256; diff --git a/onnxruntime/core/framework/allocator.cc b/onnxruntime/core/framework/allocator.cc index ceb3b3bbd7..2499ead9ef 100644 --- a/onnxruntime/core/framework/allocator.cc +++ b/onnxruntime/core/framework/allocator.cc @@ -137,26 +137,18 @@ ORT_API_STATUS_IMPL(OrtApis::CreateMemoryInfo, _In_ const char* name1, enum OrtA enum OrtMemType mem_type1, _Outptr_ OrtMemoryInfo** out) { if (strcmp(name1, onnxruntime::CPU) == 0) { *out = new OrtMemoryInfo(onnxruntime::CPU, type, OrtDevice(), id1, mem_type1); - } else if (strcmp(name1, onnxruntime::CUDA) == 0) { + } else if (strcmp(name1, onnxruntime::CUDA) == 0 || + strcmp(name1, onnxruntime::OpenVINO_GPU) == 0 || + strcmp(name1, onnxruntime::DML) == 0 || + strcmp(name1, onnxruntime::HIP) == 0 || + strcmp(name1, onnxruntime::WEBGPU_BUFFER) == 0) { *out = new OrtMemoryInfo( - onnxruntime::CUDA, type, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, static_cast(id1)), id1, + name1, type, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, static_cast(id1)), id1, mem_type1); } else if (strcmp(name1, onnxruntime::CUDA_PINNED) == 0) { *out = new OrtMemoryInfo( onnxruntime::CUDA_PINNED, type, OrtDevice(OrtDevice::CPU, OrtDevice::MemType::CUDA_PINNED, static_cast(id1)), id1, mem_type1); - } else if (strcmp(name1, onnxruntime::OpenVINO_GPU) == 0) { - *out = new OrtMemoryInfo( - onnxruntime::OpenVINO_GPU, type, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, static_cast(id1)), - id1, mem_type1); - } else if (strcmp(name1, onnxruntime::DML) == 0) { - *out = new OrtMemoryInfo( - onnxruntime::DML, type, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, static_cast(id1)), - id1, mem_type1); - } else if (strcmp(name1, onnxruntime::HIP) == 0) { - *out = new OrtMemoryInfo( - onnxruntime::HIP, type, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, static_cast(id1)), id1, - mem_type1); } else if (strcmp(name1, onnxruntime::HIP_PINNED) == 0) { *out = new OrtMemoryInfo( onnxruntime::HIP_PINNED, type, OrtDevice(OrtDevice::CPU, OrtDevice::MemType::HIP_PINNED, static_cast(id1)),