Skip to content

Commit 5b88e4e

Browse files
AndreyOrbCopilot
andauthored
Allow choosing memory type for CudaPinned allocator (microsoft#26892)
### Description CudaPinned allocator uses hardcoded to DeviceAllocator memory type. The PR comes to allow choosing memory type for CudaPinned allocator between DeviceAllocator and ArenaAllocator. ### Motivation and Context Fixed issue microsoft#26887 --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent c1da27c commit 5b88e4e

6 files changed

Lines changed: 38 additions & 2 deletions

File tree

onnxruntime/core/providers/cuda/cuda_execution_provider.cc

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,23 @@ AllocatorPtr CUDAExecutionProvider::CreateCudaAllocator(const CUDAAllocatorParam
203203
}
204204
}
205205

206+
AllocatorPtr CUDAExecutionProvider::CreateCudaPinnedAllocator(const CUDAAllocatorParams& cuda_allocator_params) {
207+
const auto* arena_cfg = cuda_allocator_params.arena_cfg;
208+
AllocatorCreationInfo pinned_memory_info(
209+
[](OrtDevice::DeviceId id) {
210+
return std::make_unique<CUDAPinnedAllocator>(id, CUDA_PINNED);
211+
},
212+
cuda_allocator_params.device_id,
213+
true,
214+
{arena_cfg ? *arena_cfg
215+
: OrtArenaCfg(cuda_allocator_params.cuda_mem_threshold,
216+
static_cast<int>(cuda_allocator_params.arena_extend_strategy), -1, -1, -1, -1L)},
217+
// stream-aware flag (intentionally set to false for this allocator)
218+
false);
219+
220+
return CreateAllocator(pinned_memory_info);
221+
}
222+
206223
CUDAExecutionProvider::PerThreadContext::PerThreadContext(OrtDevice::DeviceId device_id, cudaStream_t stream, size_t /*gpu_mem_limit*/,
207224
ArenaExtendStrategy /*arena_extend_strategy*/, CUDAExecutionProviderExternalAllocatorInfo /*external_allocator_info*/,
208225
OrtArenaCfg* /*default_memory_arena_cfg*/) {

onnxruntime/core/providers/cuda/cuda_execution_provider.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,8 @@ class CUDAExecutionProvider : public IExecutionProvider {
115115

116116
static AllocatorPtr CreateCudaAllocator(const CUDAAllocatorParams& cuda_allocator_params);
117117

118+
static AllocatorPtr CreateCudaPinnedAllocator(const CUDAAllocatorParams& cuda_allocator_params);
119+
118120
ITuningContext* GetTuningContext() const override;
119121

120122
std::unique_ptr<profiling::EpProfiler> GetProfiler() override;

onnxruntime/core/providers/cuda/cuda_provider_factory.cc

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,15 @@ struct ProviderInfo_CUDA_Impl final : ProviderInfo_CUDA {
189189
params.arena_cfg = default_memory_arena_cfg;
190190
return CUDAExecutionProvider::CreateCudaAllocator(params);
191191
}
192+
193+
std::shared_ptr<IAllocator> CreateCudaPinnedAllocator(int16_t device_id, size_t gpu_mem_limit, onnxruntime::ArenaExtendStrategy arena_extend_strategy, const OrtArenaCfg* default_memory_arena_cfg) override {
194+
CUDAExecutionProvider::CUDAAllocatorParams params{};
195+
params.device_id = device_id;
196+
params.cuda_mem_threshold = gpu_mem_limit;
197+
params.arena_extend_strategy = arena_extend_strategy;
198+
params.arena_cfg = default_memory_arena_cfg;
199+
return CUDAExecutionProvider::CreateCudaPinnedAllocator(params);
200+
}
192201
} g_info;
193202

194203
struct CUDA_Provider : Provider {

onnxruntime/core/providers/cuda/cuda_provider_factory.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ struct ProviderInfo_CUDA {
5353

5454
virtual std::shared_ptr<onnxruntime::IExecutionProviderFactory> CreateExecutionProviderFactory(const onnxruntime::CUDAExecutionProviderInfo& info) = 0;
5555
virtual std::shared_ptr<onnxruntime::IAllocator> CreateCudaAllocator(int16_t device_id, size_t gpu_mem_limit, onnxruntime::ArenaExtendStrategy arena_extend_strategy, onnxruntime::CUDAExecutionProviderExternalAllocatorInfo& external_allocator_info, const OrtArenaCfg* default_memory_arena_cfg) = 0;
56+
virtual std::shared_ptr<onnxruntime::IAllocator> CreateCudaPinnedAllocator(int16_t device_id, size_t gpu_mem_limit, onnxruntime::ArenaExtendStrategy arena_extend_strategy, const OrtArenaCfg* default_memory_arena_cfg) = 0;
5657

5758
// This function is the entry point to CUDA EP's UT cases.
5859
// All tests are only called from onnxruntime_provider_test.

onnxruntime/core/session/environment.cc

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -403,9 +403,11 @@ Status Environment::CreateAndRegisterAllocatorV2(const std::string& provider_typ
403403
#if defined(USE_CUDA) || defined(USE_CUDA_PROVIDER_INTERFACE)
404404
if (provider_type == onnxruntime::kCudaExecutionProvider) {
405405
if (mem_info.device.MemType() == OrtDevice::MemType::HOST_ACCESSIBLE) {
406-
AllocatorPtr allocator_ptr = GetProviderInfo_CUDA().CreateCUDAPinnedAllocator(
406+
AllocatorPtr allocator_ptr = GetProviderInfo_CUDA().CreateCudaPinnedAllocator(
407407
static_cast<int16_t>(mem_info.device.Id()),
408-
onnxruntime::CUDA_PINNED);
408+
arena_cfg->max_mem,
409+
static_cast<ArenaExtendStrategy>(arena_cfg->arena_extend_strategy),
410+
arena_cfg);
409411
return RegisterAllocatorImpl(allocator_ptr);
410412
} else {
411413
CUDAExecutionProviderInfo cuda_ep_info;

onnxruntime/test/providers/cuda/test_cases/cuda_test_provider.cc

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,11 @@ struct ProviderInfo_CUDA_TestImpl : ProviderInfo_CUDA {
105105
return nullptr;
106106
}
107107

108+
std::shared_ptr<onnxruntime::IAllocator> CreateCudaPinnedAllocator(int16_t, size_t, onnxruntime::ArenaExtendStrategy,
109+
const OrtArenaCfg*) override {
110+
return nullptr;
111+
}
112+
108113
void TestAll() override {
109114
// TestAll is the entry point of CUDA EP's internal tests.
110115
// Those internal tests are not directly callable from onnxruntime_provider_test

0 commit comments

Comments
 (0)