Zero-copy I/O for plugin EPs with HOST_ACCESSIBLE memory

include/onnxruntime/core/session/onnxruntime_ep_c_api.h

@@ -2551,6 +2551,32 @@ struct OrtEp {
    * \since Version 1.26.
    */
   ORT_API2_STATUS(GetAvailableResource, _In_ const OrtEp* this_ptr, _Out_ OrtResourceCount* available);
+
+  /** \brief Returns the OrtMemoryInfo the EP wants used for the given OrtMemType.
+   *
+   * Lets an EP declare, per OrtMemType, the memory the runtime should associate with that
+   * role (default device memory, CPU-side inputs, CPU-side outputs). ORT may consult this
+   * any time it needs to resolve placement for the EP.
+   *
+   * Implementations should be deterministic: a given OrtMemType should always map to the
+   * same OrtMemoryInfo for the lifetime of the OrtEp. Caching the answer up front is the
+   * recommended pattern; returned pointers must remain valid while the OrtEp is alive.
+   *
+   * Return nullptr for any OrtMemType to defer to ORT's built-in resolution for that type.
+   * Plugins may opt in selectively.
+   *
+   * \param[in] this_ptr The OrtEp instance.
+   * \param[in] mem_type The memory type to query.
+   * \return The OrtMemoryInfo the EP wants associated with the given mem_type, or nullptr
+   *         to defer to ORT.
+   *
+   * \note Implementation of this function is optional. If set to NULL, ORT applies its
+   *       built-in resolution for every OrtMemType.
+   *
+   * \since Version 1.26.
+   */
+  ORT_API_T(const OrtMemoryInfo*, GetMemoryInfoByMemType, _In_ const OrtEp* this_ptr,
+            _In_ OrtMemType mem_type);
 };
 
 /** \brief The function signature that ORT will call to create OrtEpFactory instances.

onnxruntime/core/framework/allocation_planner.cc

@@ -913,12 +913,7 @@ class PlannerImpl {
           ProcessDef(index, node_output);
           OrtDevice output_device = exec_provider->GetOrtDeviceByMemType(p_kernel_def->OutputMemoryType(i));
 #if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
-          // Downstream nodes of certain providers may require a CPU accessible location override
-          // to make sure the EP does not incur an unnecessary copy.
-          // We only do it for CPU based EPs. We are not likely to encounter
-          // non CPU devices here since they are already taken care of by using MemCpy nodes earlier.
-          // However, we still ignore them.
-          if (output_device.Type() == OrtDevice::CPU) {
+          if (output_device.UsesCpuMemory()) {
             const auto& output_name = node_output->Name();
             const auto consumers = graph_viewer_.GetConsumerNodes(output_name);
             for (const auto* consumer : consumers) {

onnxruntime/core/framework/utils.cc

@@ -50,6 +50,52 @@ bool ProviderIsCpuBased(const IExecutionProvider& provider) {
   return provider.GetDevice().Type() == OrtDevice::CPU;
 }
 
+// Returns true if src memory can satisfy tgt's requirements without a data copy.
+//
+// HOST_ACCESSIBLE → DEFAULT is valid: the device can access HOST_ACCESSIBLE memory directly.
+// DEFAULT → HOST_ACCESSIBLE is NOT valid: HOST_ACCESSIBLE implies CPU consumers, and DEFAULT
+// memory is device-only — the CPU cannot read it.
+//
+// For the mixed case, src alignment must meet tgt's minimum requirement.
+// Alignment 0 means "unspecified" and is treated as compatible with any requirement.
+bool CanSourceSatisfyTarget(const OrtDevice& src, const OrtDevice& tgt) {
+  const bool src_is_cpu_mem = src.UsesCpuMemory();
+  const bool tgt_is_cpu_mem = tgt.UsesCpuMemory();
+
+  // Identical devices are always compatible.
+  if (src == tgt) {
+    return true;
+  }
+
+  // Alignment 0 means "unspecified" — treat tgt as compatible with any alignment requirement.
+  const bool is_alignment_satisfied = tgt.GetAlignment() == 0 ||
+                                      src.GetAlignment() >= tgt.GetAlignment();
+
+  const bool is_same_physical_device = src.Type() == tgt.Type() &&
+                                       src.Vendor() == tgt.Vendor() &&
+                                       src.Id() == tgt.Id();
+
+  // Both are CPU-accessible (CPU type or HOST_ACCESSIBLE memory).
+  if (src_is_cpu_mem && tgt_is_cpu_mem) {
+    // CPU target can read from any CPU or HOST_ACCESSIBLE source, regardless of the source device
+    if (tgt.Type() == OrtDevice::CPU) {
+      return is_alignment_satisfied;
+    }
+    // Both are HOST_ACCESSIBLE on some device: require the same physical device.
+    return is_same_physical_device && is_alignment_satisfied;
+  }
+
+  // HOST_ACCESSIBLE source can serve a DEFAULT target on the same physical device —
+  // the device can DMA from HOST_ACCESSIBLE memory directly.
+  // The reverse (DEFAULT → HOST_ACCESSIBLE) is unsafe: HOST_ACCESSIBLE implies CPU consumers,
+  // and DEFAULT memory is device-only so the CPU cannot read it.
+  if (src_is_cpu_mem && !tgt_is_cpu_mem) {
+    return is_same_physical_device && is_alignment_satisfied;
+  }
+
+  return false;
+}
+
 bool IsMemcpyNode(const Node& node) {
   return node.Domain() == kOnnxDomain &&
          (node.OpType() == "MemcpyFromHost" || node.OpType() == "MemcpyToHost");
@@ -117,6 +163,27 @@ const std::string& GetNodeInputProviderType(const SessionState::NodeInfo& info)
   return required_provider_type;
 }
 
+// Populate device_fetches for the output-copy path.
+// When the user pre-allocates a fetch buffer, reuse it directly as the EP's output buffer if
+// the user's buffer (tgt) can satisfy the EP's output device (src) requirements — i.e.,
+// CanSourceSatisfyTarget(tgt, src). This avoids a post-execution copy.
+// Otherwise inserts an empty placeholder for the EP to allocate into.
+static void PopulateDeviceFetches(gsl::span<const MLValueCopyInfo> fetch_copy_info,
+                                  const std::vector<OrtValue>& fetches,
+                                  std::vector<OrtValue>& device_fetches) {
+  ORT_ENFORCE(fetch_copy_info.size() >= fetches.size());
+  device_fetches.reserve(fetches.size());
+  for (size_t i = 0; i < fetches.size(); ++i) {
+    const auto& src = fetch_copy_info[i].source_device;
+    const auto& tgt = fetch_copy_info[i].target_device;
+    if (CanSourceSatisfyTarget(tgt, src) && fetches[i].IsAllocated()) {
+      device_fetches.push_back(fetches[i]);
+    } else {
+      device_fetches.push_back({});
+    }
+  }
+}
+
 // Copy MLValue. Uses DataTransferManager for device copy if necessary. If copy_tensor_pairs/copy_sparse_pairs is provided,
 // src/dst pairs that need a device copy are added to copy_pairs so copying can be batches by the DataTransferManager
 // implementation for performance reasons.
@@ -132,8 +199,9 @@ static Status BatchOrCopyMLValue(const SessionState& session_state,
                                  std::vector<IDataTransfer::SrcDstPair>* copy_tensor_pairs = nullptr)
 #endif
 {
-  // same device so direct copy
-  if (copy_info.source_device == copy_info.target_device) {
+  // No data transfer needed if devices are identical, or the source can satisfy the target
+  // (HOST_ACCESSIBLE source serving a DEFAULT target on the same physical device).
+  if (CanSourceSatisfyTarget(copy_info.source_device, copy_info.target_device)) {
     target_mlvalue = source_mlvalue;
     return Status::OK();
   }
@@ -324,7 +392,7 @@ static bool FinalizeCopyInfoForFeeds(gsl::span<const OrtDevice> feed_locations,
   for (size_t i = 0, end = feed_locations.size(); i < end; ++i) {
     copy_info[i].source_device = feed_locations[i];
 
-    if (copy_info[i].source_device != copy_info[i].target_device) {
+    if (!CanSourceSatisfyTarget(copy_info[i].source_device, copy_info[i].target_device)) {
       copy_needed = true;
     }
   }
@@ -345,7 +413,7 @@ static bool FinalizeCopyInfoForFetches(gsl::span<const OrtDevice* const>& fetch_
       copy_info[i].target_device = *alloc_info;
     }
 
-    if (copy_info[i].source_device != copy_info[i].target_device) {
+    if (!CanSourceSatisfyTarget(copy_info[i].source_device, copy_info[i].target_device)) {
       copy_needed = true;
     }
   }
@@ -652,22 +720,9 @@ ExecuteGraphImpl(const SessionState& session_state,
       feeds_to_use = device_feeds;
     }
 
-    auto num_outputs = fetches.size();
     const auto& fetch_copy_info = feeds_fetches_manager.GetFetchesDeviceCopyInfo();
-
     if (device_copy_checks.output_copy_needed == DeviceCopyCheck::Copy) {
-      // need intermediate fetches. use pre-allocated fetches where possible.
-      device_fetches.reserve(num_outputs);
-
-      for (size_t i = 0; i < num_outputs; ++i) {
-        if (fetch_copy_info[i].source_device == fetch_copy_info[i].target_device && fetches[i].IsAllocated()) {
-          device_fetches.push_back(fetches[i]);
-        } else {
-          // use temporary value
-          device_fetches.push_back({});
-        }
-      }
-
+      PopulateDeviceFetches(fetch_copy_info, fetches, device_fetches);
       p_fetches = &device_fetches;
     }
 
@@ -808,22 +863,10 @@ common::Status ExecutePartialGraphImpl(const SessionState& session_state, FeedsF
       p_feeds = device_feeds;
     }
 
-    auto num_outputs = fetches.size();
     const auto& fetch_copy_info = feeds_fetches_manager.GetFetchesDeviceCopyInfo();
 
     if (device_copy_checks.output_copy_needed == DeviceCopyCheck::Copy) {
-      // need intermediate fetches. use pre-allocated fetches where possible.
-      device_fetches.reserve(num_outputs);
-
-      for (size_t i = 0; i < num_outputs; ++i) {
-        if (fetch_copy_info[i].source_device == fetch_copy_info[i].target_device && fetches[i].IsAllocated()) {
-          device_fetches.push_back(fetches[i]);
-        } else {
-          // use temporary value
-          device_fetches.push_back({});
-        }
-      }
-
+      PopulateDeviceFetches(fetch_copy_info, fetches, device_fetches);
       p_fetches = &device_fetches;
     }
 

onnxruntime/core/framework/utils.h

@@ -57,6 +57,11 @@ bool ProviderIsCpuBased(const IExecutionProvider& provider);
 
 bool IsMemcpyNode(const Node& node);
 
+// Returns true if src memory can satisfy tgt's requirements without a data copy.
+// HOST_ACCESSIBLE -> DEFAULT is valid (device can access HOST_ACCESSIBLE memory directly).
+// DEFAULT -> HOST_ACCESSIBLE is NOT valid (CPU cannot read device-only memory).
+bool CanSourceSatisfyTarget(const OrtDevice& src, const OrtDevice& tgt);
+
 common::Status CopyOneInputAcrossDevices(const SessionState& session_state, const std::string& input_name,
                                          const OrtValue& orig_mlvalue, OrtValue& new_mlvalue);
 

onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc

@@ -121,12 +121,26 @@ 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.
+// Single source of truth for the OrtEp::GetMemoryInfoByMemType callback (added in EP API
+// version 26): version-gated, null-checked. Returns nullptr if the EP did not opt in or if
+// the EP returned nullptr to defer to ORT's built-in fallback.
+static const OrtMemoryInfo* TryGetEpMemoryInfo(const OrtEp& ep, OrtMemType mem_type) {
+  if (ep.ort_version_supported >= 26 && ep.GetMemoryInfoByMemType != nullptr) {
+    return ep.GetMemoryInfoByMemType(&ep, mem_type);
+  }
+  return nullptr;
+}
+
+static OrtDevice GetOrtDeviceForPluginEp(const OrtEp& ep, gsl::span<const OrtEpDevice* const> ep_devices) {
+  // Get the OrtDevice from the Ep's default memory info. 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.
 
+  if (const OrtMemoryInfo* info = TryGetEpMemoryInfo(ep, OrtMemTypeDefault)) {
+    return info->device;
+  }
+
   const OrtMemoryInfo* device_memory_info = ep_devices[0]->device_memory_info;
 
   // Check assertion that all OrtEpDevice instances must have equivalent device_memory_infos
@@ -169,7 +183,7 @@ PluginExecutionProvider::PluginExecutionProvider(UniqueOrtEp ep, const OrtSessio
                                                  gsl::span<const OrtEpDevice* const> ep_devices,
                                                  std::shared_ptr<KernelRegistry> kernel_registry,
                                                  const logging::Logger& logger)
-    : IExecutionProvider(ep->GetName(ep.get()), GetOrtDeviceForPluginEp(ep_devices),
+    : IExecutionProvider(ep->GetName(ep.get()), GetOrtDeviceForPluginEp(*ep, ep_devices),
                          std::vector<const OrtEpDevice*>(ep_devices.begin(), ep_devices.end()), logger),
       ort_ep_(std::move(ep)),
       ep_factory_(ep_factory),
@@ -219,6 +233,13 @@ PluginExecutionProvider::PluginExecutionProvider(UniqueOrtEp ep, const OrtSessio
   }
 }
 
+OrtDevice PluginExecutionProvider::GetOrtDeviceByMemType(OrtMemType mem_type) const {
+  if (const OrtMemoryInfo* info = TryGetEpMemoryInfo(*ort_ep_, mem_type)) {
+    return info->device;
+  }
+  return IExecutionProvider::GetOrtDeviceByMemType(mem_type);
+}
+
 PluginExecutionProvider::~PluginExecutionProvider() {
   if (ort_ep_ && !api_node_compute_infos_.empty() && ort_ep_->ReleaseNodeComputeInfos != nullptr) {
     ort_ep_->ReleaseNodeComputeInfos(ort_ep_.get(), api_node_compute_infos_.data(),

onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.h

@@ -149,6 +149,8 @@ class PluginExecutionProvider : public IExecutionProvider {
   common::Status ReplayGraph(int graph_annotation_id) override;
   OrtGraphCaptureNodeAssignmentPolicy GetGraphCaptureNodeAssignmentPolicy() const override;
 
+  OrtDevice GetOrtDeviceByMemType(OrtMemType mem_type) const override;
+
  private:
   const logging::Logger& GetEpLoggerOrDefault() const;