[Performance] CudaPinned allocator used wrong backend device_allocator if shared environment allocators are activated

[AndreyOrb]

Describe the issue

The correct CudaPinned allocator from CUDA EP's CudaPinned (pointing the correct CUDAPinnedAllocator) IS replaced with "buggy" CudaPinned allocator (pointing CPUAllocator)!

onnxruntime/onnxruntime/core/session/inference_session.cc

Line 2106 in 849eee8

session_state_ = std::make_unique<SessionState>(

SessionState before applying env. shared allocators:

SessionState right after applying env. shared allocators:

This happens because the environment holds CudaPinned allocator pointing CPUAllocator backend device_allocator.

Here's my code used to register CudaPinned allocator:

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

	OrtMemoryInfo* cudaPinnedMemoryInfo;
	ASSERT_ORT_STATUS(api.CreateMemoryInfo("CudaPinned", OrtArenaAllocator, 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();
}

This is how I revealed the issue:

This is without shared env. allocators:

To reproduce

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

OrtMemoryInfo* cudaPinnedMemoryInfo;
ASSERT_ORT_STATUS(api.CreateMemoryInfo("CudaPinned", OrtArenaAllocator, 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();

}

I tried to search in ORT repo for any clues or similar issues but found nothing.

Urgency

Urgent! Directly impacts performance when shared environment allocators are activated.

Platform

Windows

OS Version

11

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

849eee8

ONNX Runtime API

C++

Architecture

X64

Execution Provider

CUDA

Execution Provider Library Version

CUDA 11.8

Model File

No response

Is this a quantized model?

Unknown

[AndreyOrb]

In both cases, when CudaPinned allocator is registered with or without arena, it uses CPUAllocator backend device_allocator:

With Arena:

onnxruntime/onnxruntime/core/session/environment.cc

Line 223 in 849eee8

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

Without Arena:

onnxruntime/onnxruntime/core/session/environment.cc

Line 229 in 849eee8

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

When CudaPinned memory type is passed, the environment should use

onnxruntime/onnxruntime/core/providers/cuda/cuda_provider_factory.cc

Line 88 in 849eee8

return std::make_unique<CUDAPinnedAllocator>(name);

directly or by querying the CUDA EP (or whatever other EP is relevant) to do it.

[AndreyOrb]

1. CudaPinned memory must be used with "CUDAExecutionProvider"
2. These is an issue preventing the CUDAPinnedAllocator from being created:
   

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

   
   This code does not distinguish between the desired memory types.