CreateAndRegisterAllocatorV2 fails to create CudaPinned allocator when shared env. allocators are activated

[AndreyOrb]

Describe the issue

CreateAndRegisterAllocatorV2 with "CudaPinned" fails if "Cpu" allocator was already registered in the environment.

This happens because "mem_info" is not passed into the AllocatorCreationInfo constructor:

onnxruntime/onnxruntime/core/session/environment.cc

Line 229 in 849eee8

AllocatorCreationInfo alloc_creation_info{[](int) { return std::make_unique<CPUAllocator>(); },

The AllocatorCreationInfo constructor then falls back to "Cpu" memory info.

To reproduce

Run the two functions one after another:

void RegisterCpuEnvAllocator(OrtApi& api, OrtEnv* env)
{
	nvtxRangePush("RegisterCpuEnvAllocator");

	OrtMemoryInfo* cpuMemoryInfo;
	ASSERT_ORT_STATUS(api.CreateMemoryInfo("Cpu", OrtDeviceAllocator, 0, OrtMemTypeDefault, &cpuMemoryInfo));

	OrtArenaCfg* cpuArenaConfig;
	ASSERT_ORT_STATUS(api.CreateArenaCfg(0, ArenaExtendStrategy::kNextPowerOfTwo, -1, -1, &cpuArenaConfig));

	// This creates an ORT-internal allocator instance and registers it in the environment for sharing
	vector<const char*> keys, values;
	ASSERT_ORT_STATUS(api.CreateAndRegisterAllocatorV2(env, "CPUExecutionProvider", cpuMemoryInfo, cpuArenaConfig, keys.data(), values.data(), 0));

	api.ReleaseArenaCfg(cpuArenaConfig);
	api.ReleaseMemoryInfo(cpuMemoryInfo);
	nvtxRangePop();
}

void RegisterCudaPinnedEnvAllocator(OrtApi& api, OrtEnv* env)
{
	nvtxRangePush("RegisterCudaPinnedEnvAllocator");

	OrtMemoryInfo* cudaPinnedMemoryInfo;
	ASSERT_ORT_STATUS(api.CreateMemoryInfo("CudaPinned", OrtDeviceAllocator, 0, OrtMemTypeCPUOutput, &cudaPinnedMemoryInfo));

	OrtArenaCfg* cudaPinnedArenaConfig;
	ASSERT_ORT_STATUS(api.CreateArenaCfg(0, ArenaExtendStrategy::kNextPowerOfTwo, -1, -1, &cudaPinnedArenaConfig));

	// This creates an ORT-internal allocator instance and registers it in the environment for sharing
	vector<const char*> keys, values;
	ASSERT_ORT_STATUS(api.CreateAndRegisterAllocatorV2(env, "CPUExecutionProvider", cudaPinnedMemoryInfo, cudaPinnedArenaConfig, keys.data(), values.data(), 0));

	api.ReleaseArenaCfg(cudaPinnedArenaConfig);
	api.ReleaseMemoryInfo(cudaPinnedMemoryInfo);
	nvtxRangePop();
}

The issue happens only if both Cpu and CudaPinned allocators are registered with OrtDeviceAllocator flag, because the other flag (OrtArenaAllocator) uses the correct constructor:

onnxruntime/onnxruntime/core/session/environment.cc

Line 223 in 849eee8

[mem_info](int) { return std::make_unique<CPUAllocator>(mem_info); },

Urgency

Not urgent

Platform

Windows

OS Version

11

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

1.22

ONNX Runtime API

C++

Architecture

X64

Execution Provider

CUDA

Execution Provider Library Version

CUDA 11.8

[AndreyOrb]

Please review the PR: #25218

[skottmckay]

Possibly the issue is this code is using the CPUExecutionProvider name but should be using CUDAExecutionProvider.

	// This creates an ORT-internal allocator instance and registers it in the environment for sharing
	vector<const char*> keys, values;
	ASSERT_ORT_STATUS(api.CreateAndRegisterAllocatorV2(env, "CPUExecutionProvider", cudaPinnedMemoryInfo, cudaPinnedArenaConfig, keys.data(), values.data(), 0));

[tianleiwu]

I agree, the test code shall use "CUDAExecutionProvider" for cuda pinned memory. However, cuda allocator currently only support gpu memory allocator:

onnxruntime/onnxruntime/core/session/environment.cc

Lines 403 to 407 in 849eee8

	AllocatorPtr allocator_ptr = GetProviderInfo_CUDA().CreateCudaAllocator(
	static_cast<int16_t>(mem_info.device.Id()),
	arena_cfg->max_mem,
	static_cast<ArenaExtendStrategy>(arena_cfg->arena_extend_strategy),
	external_info, arena_cfg);

To support cuda pinned memory, the code need change to call CreateCUDAPinnedAllocator when pinned memory is requested. For example, the logic like

if memory type is OrtMemTypeCPUOutput:
   allocator_ptr = GetProviderInfo_CUDA().CreateCUDAPinnedAllocator(CUDA_PINNED);
else
   allocator_ptr = GetProviderInfo_CUDA().CreateCUDAAllocator(...);

BTW, the PR looks good to me since the cpu allocator accepts memory info parameter, so it is reasonable to pass it to allocator constructor.

[AndreyOrb]

@skottmckay You are right, "CUDAExecutionProvider" should be used.
@tianleiwu Just checked, you are also right. Some adjustments to be made.

There's an issue I've created earlier that seems to be more relevant for the conversation: #25211
I will add a PR for #25211 that will distinguish between the memory types and will call CreateCUDAPinnedAllocator when needed.
Stay tuned.

As for this issue, it can be closed if the tests pass.

[AndreyOrb]

@tianleiwu This issue can also be closed now. Thanks a lot for your help!