[EP ABI] Infer OrtDevice for plugin EP from registered OrtMemoryInfo (microsoft#25308)

adrianlizarraga · web-flow · commit 0ccecf71f434 · 2025-07-07T23:02:15.000-07:00
### Description
- Infer `OrtDevice` for a plugin EP from the registered `OrtMemoryInfo`
for device memory.
- Fix potential `nullptr` dereference when a `PluginExecutionProvider`
tries to log a message without a valid logger. Now, constructing a
`PluginExecutionProvider` requires passing a valid logger.



### Motivation and Context
Address a `TODO` to properly set the `OrtDevice` for a
`PluginExecutionProvider` instance.
diff --git a/include/onnxruntime/core/framework/execution_provider.h b/include/onnxruntime/core/framework/execution_provider.h
@@ -79,6 +79,10 @@ class IExecutionProvider {
       : default_device_(device), type_{type} {
   }
 
+  IExecutionProvider(const std::string& type, OrtDevice device, const logging::Logger& logger)
+      : default_device_(device), type_{type}, logger_{&logger} {
+  }
+
   /*
      default device for this ExecutionProvider
   */
diff --git a/onnxruntime/core/session/ep_plugin_provider_interfaces.cc b/onnxruntime/core/session/ep_plugin_provider_interfaces.cc
@@ -53,11 +53,9 @@ PluginExecutionProviderFactory::CreateProvider(const OrtSessionOptions& session_
     ORT_THROW("Error creating execution provider: ", status.ToString());
   }
 
-  auto ep_wrapper = std::make_unique<PluginExecutionProvider>(UniqueOrtEp(ort_ep, OrtEpDeleter(ep_factory_)),
-                                                              session_options, ep_factory_, devices_);
-  ep_wrapper->SetLogger(session_logger.ToInternal());
-
-  return ep_wrapper;
+  return std::make_unique<PluginExecutionProvider>(UniqueOrtEp(ort_ep, OrtEpDeleter(ep_factory_)),
+                                                   session_options, ep_factory_, devices_,
+                                                   *session_logger.ToInternal());
 }
 
 /// <summary>
@@ -86,10 +84,43 @@ struct PluginEpMetaDefNameFunctor {
 // PluginExecutionProvider
 //
 
+static OrtDevice GetOrtDeviceForPluginEp(gsl::span<const OrtEpDevice* const> ep_devices) {
+  // Get the OrtDevice from OrtEpDevice.device_memory_info if it is set. Otherwise, we set it to CPU.
+  // If there are multiple OrtEpDevice instances, the device_memory_info must be consistent for all.
+
+  ORT_ENFORCE(!ep_devices.empty());  // Should not be possible to create an EP without OrtEpDevices.
+
+  const OrtMemoryInfo* device_memory_info = ep_devices[0]->device_memory_info;
+
+  // Check assertion that all OrtEpDevice instances must have equivalent device_memory_infos
+  bool all_match = std::all_of(ep_devices.begin() + 1, ep_devices.end(),
+                               [mem_a = device_memory_info](const OrtEpDevice* ep_device) {
+                                 const OrtMemoryInfo* mem_b = ep_device->device_memory_info;
+
+                                 if (mem_a == mem_b) {
+                                   return true;  // Point to the same OrtMemoryInfo instance.
+                                 }
+
+                                 if (mem_a == nullptr || mem_b == nullptr) {
+                                   return false;  // One is nullptr and the other is not.
+                                 }
+
+                                 // Both non-null but point to different instances. Use operator==.
+                                 return *mem_a == *mem_b;
+                               });
+  if (!all_match) {
+    ORT_THROW("Error creating execution provider '", ep_devices[0]->ep_name,
+              "': expected all OrtEpDevice instances to use the same device_memory_info.");
+  }
+
+  return device_memory_info != nullptr ? device_memory_info->device : OrtDevice();
+}
+
 PluginExecutionProvider::PluginExecutionProvider(UniqueOrtEp ep, const OrtSessionOptions& session_options,
                                                  OrtEpFactory& ep_factory,
-                                                 gsl::span<const OrtEpDevice* const> ep_devices)
-    : IExecutionProvider(ep->GetName(ep.get()), OrtDevice()),  // TODO: What to do about OrtDevice for plugins?
+                                                 gsl::span<const OrtEpDevice* const> ep_devices,
+                                                 const logging::Logger& logger)
+    : IExecutionProvider(ep->GetName(ep.get()), GetOrtDeviceForPluginEp(ep_devices), logger),
       ort_ep_(std::move(ep)),
       ep_factory_(ep_factory),
       ep_devices_(ep_devices.begin(), ep_devices.end()) {
diff --git a/onnxruntime/core/session/ep_plugin_provider_interfaces.h b/onnxruntime/core/session/ep_plugin_provider_interfaces.h
@@ -65,7 +65,7 @@ class PluginExecutionProvider : public IExecutionProvider {
 
  public:
   explicit PluginExecutionProvider(UniqueOrtEp ep, const OrtSessionOptions& session_options, OrtEpFactory& ep_factory,
-                                   gsl::span<const OrtEpDevice* const> ep_devices);
+                                   gsl::span<const OrtEpDevice* const> ep_devices, const logging::Logger& logger);
   ~PluginExecutionProvider();
 
   std::vector<std::unique_ptr<ComputeCapability>>
diff --git a/onnxruntime/core/session/provider_policy_context.cc b/onnxruntime/core/session/provider_policy_context.cc
@@ -10,6 +10,7 @@
 
 #include "core/framework/error_code_helper.h"
 #include "core/session/abi_devices.h"
+#include "core/session/abi_logger.h"
 #include "core/session/ep_factory_internal.h"
 #include "core/session/ep_plugin_provider_interfaces.h"
 #include "core/session/inference_session.h"
@@ -355,7 +356,7 @@ Status ProviderPolicyContext::CreateExecutionProvider(const Environment& env, Or
         info.ep_factory->CreateEp(info.ep_factory, info.hardware_devices.data(), info.ep_metadata.data(),
                                   info.hardware_devices.size(), &options, &logger, &api_ep)));
     ep = std::make_unique<PluginExecutionProvider>(UniqueOrtEp(api_ep, OrtEpDeleter(*info.ep_factory)), options,
-                                                   *info.ep_factory, info.devices);
+                                                   *info.ep_factory, info.devices, *logger.ToInternal());
   }
 
   return Status::OK();
diff --git a/onnxruntime/test/framework/ep_plugin_provider_test.cc b/onnxruntime/test/framework/ep_plugin_provider_test.cc
@@ -9,6 +9,7 @@
 #include "core/session/abi_devices.h"
 #include "core/session/onnxruntime_cxx_api.h"
 #include "test/util/include/asserts.h"
+#include "test/util/include/test_environment.h"
 
 namespace onnxruntime::test {
 
@@ -56,24 +57,58 @@ struct TestOrtEpFactory : ::OrtEpFactory {
 
 static TestOrtEpFactory g_test_ort_ep_factory{};
 
+std::unique_ptr<OrtHardwareDevice> MakeTestOrtHardwareDevice(OrtHardwareDeviceType type) {
+  auto hw_device = std::make_unique<OrtHardwareDevice>();
+  hw_device->type = type;
+  hw_device->vendor_id = 0xBE57;
+  hw_device->device_id = 0;
+  hw_device->vendor = "Contoso";
+  return hw_device;
+}
+
+std::unique_ptr<OrtEpDevice> MakeTestOrtEpDevice(const OrtHardwareDevice* hardware_device,
+                                                 const OrtMemoryInfo* device_memory_info = nullptr,
+                                                 const OrtMemoryInfo* host_accessible_memory_info = nullptr) {
+  auto ep_device = std::make_unique<OrtEpDevice>();
+  ep_device->ep_name = "TestOrtEp";
+  ep_device->ep_vendor = "Contoso";
+  ep_device->device = hardware_device;
+  ep_device->ep_factory = &g_test_ort_ep_factory;
+  ep_device->device_memory_info = device_memory_info;
+  ep_device->host_accessible_memory_info = host_accessible_memory_info;
+  return ep_device;
+}
+
+OrtDevice MakeTestOrtDevice(OrtDevice::DeviceType device_type, OrtDevice::MemoryType memory_type) {
+  return OrtDevice(device_type, memory_type, /*vendor_id*/ 0xBE57, /*device_id*/ 0, /*alignment*/ 16);
+}
+
 struct MakeTestOrtEpResult {
   std::unique_ptr<IExecutionProvider> ep;  // the IExecutionProvider wrapping the TestOrtEp
   gsl::not_null<TestOrtEp*> ort_ep;        // the wrapped TestOrtEp, owned by `ep`
 };
 
 // Creates an IExecutionProvider that wraps a TestOrtEp.
 // The TestOrtEp is also exposed so that tests can manipulate its function pointers directly.
-MakeTestOrtEpResult MakeTestOrtEp() {
+MakeTestOrtEpResult MakeTestOrtEp(std::vector<const OrtEpDevice*> ep_devices = {}) {
+  // Default OrtHardwareDevice and OrtEpDevice used if the caller does not explicitly provide ep_devices.
+  static std::unique_ptr<OrtHardwareDevice> ort_hw_device = MakeTestOrtHardwareDevice(OrtHardwareDeviceType_CPU);
+  static std::unique_ptr<OrtEpDevice> ort_ep_device = MakeTestOrtEpDevice(ort_hw_device.get());
+
   auto ort_ep_raw = std::make_unique<TestOrtEp>().release();
   auto ort_ep = UniqueOrtEp(ort_ep_raw, OrtEpDeleter{g_test_ort_ep_factory});
   auto ort_session_options = Ort::SessionOptions{};
-  auto ort_ep_device = OrtEpDevice{};
-  std::vector<const OrtEpDevice*> ep_devices{&ort_ep_device};
 
+  if (ep_devices.empty()) {
+    ep_devices.push_back(ort_ep_device.get());
+  }
+
+  auto& logging_manager = DefaultLoggingManager();
   auto ep = std::make_unique<PluginExecutionProvider>(std::move(ort_ep),
                                                       *static_cast<const OrtSessionOptions*>(ort_session_options),
                                                       g_test_ort_ep_factory,
-                                                      ep_devices);
+                                                      ep_devices,
+                                                      logging_manager.DefaultLogger());
 
   auto result = MakeTestOrtEpResult{std::move(ep), ort_ep_raw};
   return result;
@@ -177,4 +212,109 @@ TEST(PluginExecutionProviderTest, ShouldConvertDataLayoutForOp) {
 #endif  // !defined(ORT_NO_EXCEPTIONS)
 }
 
+TEST(PluginExecutionProviderTest, InferOrtDeviceFromDeviceMemoryInfo) {
+  // 1 OrtEpDevice without a device_memory_info.
+  // PluginExecutionProvider should decide to use a default OrtDevice.
+  {
+    auto ort_hw_device = test_plugin_ep::MakeTestOrtHardwareDevice(OrtHardwareDeviceType_CPU);
+    auto ort_ep_device = test_plugin_ep::MakeTestOrtEpDevice(ort_hw_device.get());
+    std::vector<const OrtEpDevice*> ep_devices{ort_ep_device.get()};
+
+    auto [ep, ort_ep] = test_plugin_ep::MakeTestOrtEp(ep_devices);
+    ASSERT_EQ(ep->GetOrtDeviceByMemType(OrtMemTypeDefault), OrtDevice());
+  }
+
+  // 1 OrtEpDevice with a device_memory_info.
+  // PluginExecutionProvider should decide to use the OrtDevice from the device_memory_info.
+  {
+    auto ort_device = test_plugin_ep::MakeTestOrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT);
+    auto ort_memory_info = std::make_unique<OrtMemoryInfo>("TestOrtEp GPU", OrtAllocatorType::OrtDeviceAllocator,
+                                                           ort_device, OrtMemTypeDefault);
+
+    auto ort_hw_device = test_plugin_ep::MakeTestOrtHardwareDevice(OrtHardwareDeviceType_GPU);
+    auto ort_ep_device = test_plugin_ep::MakeTestOrtEpDevice(ort_hw_device.get(),
+                                                             /*device_memory_info*/ ort_memory_info.get());
+    std::vector<const OrtEpDevice*> ep_devices{ort_ep_device.get()};
+
+    auto [ep, ort_ep] = test_plugin_ep::MakeTestOrtEp(ep_devices);
+    ASSERT_EQ(ep->GetOrtDeviceByMemType(OrtMemTypeDefault), ort_device);
+  }
+
+  // 2 OrtEpDevice instances with the same device_memory_info.
+  // PluginExecutionProvider should decide to use the OrtDevice from the device_memory_info.
+  {
+    auto ort_device = test_plugin_ep::MakeTestOrtDevice(OrtDevice::CPU, OrtDevice::MemType::DEFAULT);
+    auto ort_memory_info = std::make_unique<OrtMemoryInfo>("TestOrtEp CPU", OrtAllocatorType::OrtDeviceAllocator,
+                                                           ort_device, OrtMemTypeDefault);
+
+    auto ort_hw_device_gpu = test_plugin_ep::MakeTestOrtHardwareDevice(OrtHardwareDeviceType_GPU);
+    auto ort_hw_device_npu = test_plugin_ep::MakeTestOrtHardwareDevice(OrtHardwareDeviceType_NPU);
+    auto ort_ep_device_gpu = test_plugin_ep::MakeTestOrtEpDevice(ort_hw_device_gpu.get(), ort_memory_info.get());
+    auto ort_ep_device_npu = test_plugin_ep::MakeTestOrtEpDevice(ort_hw_device_npu.get(), ort_memory_info.get());
+    std::vector<const OrtEpDevice*> ep_devices{ort_ep_device_gpu.get(), ort_ep_device_npu.get()};
+
+    auto [ep, ort_ep] = test_plugin_ep::MakeTestOrtEp(ep_devices);
+    ASSERT_EQ(ep->GetOrtDeviceByMemType(OrtMemTypeDefault), ort_device);
+  }
+
+  // 2 OrtEpDevice instances with the different (but equivalent) device_memory_info pointers.
+  // PluginExecutionProvider should decide to use a OrtDevice that is equal to the devices used by both
+  // device_memory_info pointers.
+  {
+    auto ort_device = test_plugin_ep::MakeTestOrtDevice(OrtDevice::CPU, OrtDevice::MemType::DEFAULT);
+    auto ort_memory_info_0 = std::make_unique<OrtMemoryInfo>("TestOrtEp CPU", OrtAllocatorType::OrtDeviceAllocator,
+                                                             ort_device, OrtMemTypeDefault);
+    auto ort_memory_info_1 = std::make_unique<OrtMemoryInfo>("TestOrtEp CPU", OrtAllocatorType::OrtDeviceAllocator,
+                                                             ort_device, OrtMemTypeDefault);
+
+    auto ort_hw_device_gpu = test_plugin_ep::MakeTestOrtHardwareDevice(OrtHardwareDeviceType_GPU);
+    auto ort_hw_device_npu = test_plugin_ep::MakeTestOrtHardwareDevice(OrtHardwareDeviceType_NPU);
+    auto ort_ep_device_gpu = test_plugin_ep::MakeTestOrtEpDevice(ort_hw_device_gpu.get(), ort_memory_info_0.get());
+    auto ort_ep_device_npu = test_plugin_ep::MakeTestOrtEpDevice(ort_hw_device_npu.get(), ort_memory_info_1.get());
+    std::vector<const OrtEpDevice*> ep_devices{ort_ep_device_gpu.get(), ort_ep_device_npu.get()};
+
+    auto [ep, ort_ep] = test_plugin_ep::MakeTestOrtEp(ep_devices);
+    ASSERT_EQ(ep->GetOrtDeviceByMemType(OrtMemTypeDefault), ort_device);
+  }
+
+  // 1 OrtEpDevice with only a host_accessible_memory_info.
+  // PluginExecutionProvider should decide to use a default OrtDevice (cpu).
+  {
+    auto ort_device = test_plugin_ep::MakeTestOrtDevice(OrtDevice::GPU, OrtDevice::MemType::HOST_ACCESSIBLE);
+    auto ort_memory_info = std::make_unique<OrtMemoryInfo>("TestOrtEp GPU", OrtAllocatorType::OrtDeviceAllocator,
+                                                           ort_device, OrtMemTypeDefault);
+
+    auto ort_hw_device = test_plugin_ep::MakeTestOrtHardwareDevice(OrtHardwareDeviceType_GPU);
+    auto ort_ep_device = test_plugin_ep::MakeTestOrtEpDevice(ort_hw_device.get(),
+                                                             /*device_memory_info*/ nullptr,
+                                                             /*host_accessible_memory_info*/ ort_memory_info.get());
+    std::vector<const OrtEpDevice*> ep_devices{ort_ep_device.get()};
+
+    auto [ep, ort_ep] = test_plugin_ep::MakeTestOrtEp(ep_devices);
+    ASSERT_EQ(ep->GetOrtDeviceByMemType(OrtMemTypeDefault), OrtDevice());
+  }
+
+#if !defined(ORT_NO_EXCEPTIONS)
+  // 2 OrtEpDevice instances with DIFFERENT device_memory_info instances.
+  // Should throw an exception on construction of PluginExecutionProvider.
+  {
+    auto ort_device_gpu = test_plugin_ep::MakeTestOrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT);
+    auto ort_memory_info_gpu = std::make_unique<OrtMemoryInfo>("TestOrtEp GPU", OrtAllocatorType::OrtDeviceAllocator,
+                                                               ort_device_gpu, OrtMemTypeDefault);
+
+    auto ort_device_npu = test_plugin_ep::MakeTestOrtDevice(OrtDevice::NPU, OrtDevice::MemType::DEFAULT);
+    auto ort_memory_info_npu = std::make_unique<OrtMemoryInfo>("TestOrtEp NPU", OrtAllocatorType::OrtDeviceAllocator,
+                                                               ort_device_npu, OrtMemTypeDefault);
+
+    auto ort_hw_device_gpu = test_plugin_ep::MakeTestOrtHardwareDevice(OrtHardwareDeviceType_GPU);
+    auto ort_hw_device_npu = test_plugin_ep::MakeTestOrtHardwareDevice(OrtHardwareDeviceType_NPU);
+    auto ort_ep_device_gpu = test_plugin_ep::MakeTestOrtEpDevice(ort_hw_device_gpu.get(), ort_memory_info_gpu.get());
+    auto ort_ep_device_npu = test_plugin_ep::MakeTestOrtEpDevice(ort_hw_device_npu.get(), ort_memory_info_npu.get());
+    std::vector<const OrtEpDevice*> ep_devices{ort_ep_device_gpu.get(), ort_ep_device_npu.get()};
+
+    ASSERT_THROW(test_plugin_ep::MakeTestOrtEp(ep_devices), OnnxRuntimeException);
+  }
+#endif  // !defined(ORT_NO_EXCEPTIONS)
+}
+
 }  // namespace onnxruntime::test
