Allocator and data transfer support for plugin EP API #25070

skottmckay · 2025-06-14T23:06:13Z

Description

Add allocator and data transfer infrastructure for plugin EP API

Allocators are created via the OrtEpFactory using OrtMemoryInfo that as added to the OrtEpDevice instances the factory returns. This allows allocators to be created outside of an inference session and shared.

When a library is loaded a default instance of each allocator is added to the shared allocators if there is no existing allocator (e.g. user provided custom allocator).
CreateSharedAllocator can be used to replace this default instance with a user configured one. e.g. add an arena or provide other configuration options that are passed through to the OrtEpFactory's CreateAllocator function.

Similarly IDataTransfer is supported by the factory implementing OrtDataTransferImpl, which will also enable data transfer outside of a session. That will be added in a future PR as the synchronization requirements need to be figured out and will affect the public API.

Motivation and Context

github-actions

You can commit the suggested changes from lintrunner.

onnxruntime/core/framework/plugin_data_transfer.cc

skottmckay · 2025-06-14T23:06:41Z

include/onnxruntime/core/framework/allocator.h

@@ -37,6 +44,26 @@ struct OrtArenaCfg {
  int max_dead_bytes_per_chunk;           // use -1 to allow ORT to choose the default
  int initial_growth_chunk_size_bytes;    // use -1 to allow ORT to choose the default
  int64_t max_power_of_two_extend_bytes;  // use -1 to allow ORT to choose the default
+
+  bool IsValid() {


Moving validity check that was scattered around to be here.

skottmckay · 2025-06-14T23:07:19Z

include/onnxruntime/core/session/environment.h

@@ -75,7 +78,7 @@ class Environment {
   * Registers an allocator for sharing between multiple sessions.
   * Return an error if an allocator with the same OrtMemoryInfo is already registered.
   */
-  Status RegisterAllocator(AllocatorPtr allocator);
+  Status RegisterAllocator(OrtAllocator* allocator);


Separating external from internal usage by making the public method take OrtAllocator*. This enables GetSharedAllocator to return any of the OrtAllocator based instances.

skottmckay · 2025-06-14T23:17:03Z

include/onnxruntime/core/session/onnxruntime_c_api.h

+  ORT_API2_STATUS(CopyTensors, _In_ void* this_ptr,
+                  _In_reads_(num_tensors) const OrtValue** src_tensors,
+                  _In_reads_(num_tensors) OrtValue** dst_tensors,
+                  _In_reads_(num_tensors) OrtSyncStream** streams,
+                  _In_ size_t num_tensors);


Design choice: Minimize the EP API by having a single function that can do asynchronous batched copies. The library is free to implement that as synchronous copies. This avoids having additional functions for copying a single tensor.

skottmckay · 2025-06-14T23:19:27Z

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -6096,7 +6185,109 @@ struct OrtCompileApi {

 ORT_RUNTIME_CLASS(Ep);
 ORT_RUNTIME_CLASS(EpFactory);
+ORT_RUNTIME_CLASS(MemoryDevice);  // opaque class to wrap onnxruntime::OrtDevice


I couldn't avoid adding OrtDevice to the API as IDataTransfer uses it. To try and disambiguate/clarify I called it OrtMemoryDevice as we have OrtHardwareDevice, and calling it OrtDevice in the API felt too vague.

skottmckay · 2025-06-14T23:21:49Z

include/onnxruntime/core/session/onnxruntime_c_api.h

+#define ORT_API_T(RETURN_TYPE, NAME, ...) \
+  RETURN_TYPE(ORT_API_CALL* NAME)(__VA_ARGS__) NO_EXCEPTION
+


We were missing the NO_EXCEPTION on a number of functions in the initial plugin EP setup. This macro should help keep things consistent by adding NO_EXCEPTION like ORT_API2_STATUS does.

skottmckay · 2025-06-14T23:25:37Z

onnxruntime/core/framework/allocator.cc

@@ -15,6 +17,61 @@

 #include "core/framework/bfc_arena.h"

+using Status = onnxruntime::common::Status;
+
+Status OrtArenaCfg::FromKeyValuePairs(const OrtKeyValuePairs& kvps, OrtArenaCfg& cfg) {


Support creation from OrtKeyValuePairs to keep the public API minimal. We use OrtKeyValuePairs in general for configuration/options/metadata and this avoids having to add OrtArenaCfg as an additional parameter in some functions like CreateSharedAllocator.

skottmckay · 2025-06-14T23:26:21Z

onnxruntime/core/framework/data_transfer_manager.cc

@@ -16,6 +16,20 @@ Status DataTransferManager::RegisterDataTransfer(std::unique_ptr<IDataTransfer>
  return Status::OK();
 }

+Status DataTransferManager::UnregisterDataTransfer(IDataTransfer* data_transfer) {


Need to unregister if the plugin EP library is unloaded.

skottmckay · 2025-06-14T23:30:16Z

onnxruntime/core/session/allocator_adapters.cc

-  auto st = env->CreateAndRegisterAllocator(*mem_info, arena_cfg);
+  auto& env = ort_env->GetEnvironment();


Removed unnecessary call from here -> OrtEnv.CreateAndRegisterAllocator -> Environment.CreateAndRegisterAllocator.

As the Environment instance is available from OrtEnv there's no need to it to have a CreateAndRegisterAllocator method that forwards.

…locator

Fix MIGraphX build.

edgchen1 · 2025-06-17T00:38:02Z

include/onnxruntime/core/framework/allocator.h

 // some enums are defined in session/onnxruntime_c_api.h but used in ortdevice.h/ortmemory.h
 #include "core/session/onnxruntime_c_api.h"
 #include "core/framework/ortdevice.h"
 #include "core/framework/ortmemoryinfo.h"

+namespace onnxruntime {
+namespace common {
+class Status;


nit: common.h already includes status.h, so we don't need a forward declaration. we could explicitly include status.h too for clarity.

a user of this file would need a complete Status definition anyway to call FromKeyValuePairs().

edgchen1 · 2025-06-17T00:43:25Z

include/onnxruntime/core/session/onnxruntime_c_api.h

+   *
+   * OrtEpDevice maps to the EP factory, and the factory provides the allocator implementation.
+   *
+   * OrtDeviceMemoryType_DEFAULT is always supported for non-CPU based devices.


is the device memory allocator always required? does this also apply to the QNN EP, which only has the HtpSharedMemoryAllocator?

edgchen1 · 2025-06-17T01:04:25Z

onnxruntime/core/framework/allocator.cc

+    ORT_RETURN_IF_ERROR(from_string(it->first, it->second, value));
+    cfg.arena_extend_strategy = onnxruntime::narrow<int32_t>(value);


consider using ParseStringWithClassicLocale() in parse_string.h. it handles conversions to specific int types directly.

Allocator and data transfer support for plugin EP API

57900b5

github-actions bot reviewed Jun 14, 2025

View reviewed changes

onnxruntime/core/framework/plugin_data_transfer.cc Outdated Show resolved Hide resolved

github-advanced-security bot found potential problems Jun 14, 2025

View reviewed changes

onnxruntime/core/framework/plugin_data_transfer.cc Fixed Show fixed Hide fixed

skottmckay commented Jun 14, 2025

View reviewed changes

skottmckay added 12 commits June 15, 2025 09:37

Cleanups

06ca619

Add _Outptr_opt_ SAL define for non-windows platforms

2077b37

*nix build fixes

261f38d

Merge remote-tracking branch 'origin/main' into skottmckay/PluginEpAl…

ed0c3c7

…locator

Fix QNN OrtEpFactory

05aac95

Fix minimal build.

584c2dd

Fix MIGraphX build.

Fix header

2b377b9

Fix unused variable

a8dd28d

Add ORT_API_CALL for x86 build

55d3a8a

Fix cuda and python tests.

5666e6c

lint

cd50bf3

Update python test

dc09763

edgchen1 reviewed Jun 17, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Allocator and data transfer support for plugin EP API #25070

Allocator and data transfer support for plugin EP API #25070

skottmckay commented Jun 14, 2025

Uh oh!

github-actions bot left a comment

Uh oh!

Uh oh!

Uh oh!

skottmckay Jun 14, 2025

Uh oh!

skottmckay Jun 14, 2025

Uh oh!

skottmckay Jun 14, 2025

Uh oh!

skottmckay Jun 14, 2025

Uh oh!

skottmckay Jun 14, 2025

Uh oh!

skottmckay Jun 14, 2025

Uh oh!

skottmckay Jun 14, 2025

Uh oh!

skottmckay Jun 14, 2025

Uh oh!

edgchen1 Jun 17, 2025

Uh oh!

edgchen1 Jun 17, 2025

Uh oh!

edgchen1 Jun 17, 2025

Uh oh!

Uh oh!

		#define ORT_API_T(RETURN_TYPE, NAME, ...) \
		RETURN_TYPE(ORT_API_CALL* NAME)(__VA_ARGS__) NO_EXCEPTION

		auto st = env->CreateAndRegisterAllocator(*mem_info, arena_cfg);
		auto& env = ort_env->GetEnvironment();

		ORT_RETURN_IF_ERROR(from_string(it->first, it->second, value));
		cfg.arena_extend_strategy = onnxruntime::narrow<int32_t>(value);

Allocator and data transfer support for plugin EP API #25070

Are you sure you want to change the base?

Allocator and data transfer support for plugin EP API #25070

Conversation

skottmckay commented Jun 14, 2025

Description

Motivation and Context

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!