[CUDA EP Plugin] ResourceAcountant integration (#28028)

This pull request introduces several enhancements and refactorings to
the resource accounting and execution provider (EP) infrastructure, with
a focus on better support for plugin-based CUDA execution providers. The
most significant changes include the addition of type-erased arithmetic
for resource accounting, improved handling of resource budgets for
plugin EPs, and more robust device matching logic. These updates
increase maintainability, enforce stricter type safety, and ensure
correct resource tracking across both in-tree and plugin-based EPs.

**Resource accounting improvements:**

* Added type-erased arithmetic functions (`AddResourceCounts`,
`ResourceCountExceeds`, `FormatResourceCount`) for `ResourceCount` to
enforce exhaustive handling of variant types and improve type safety.
[[1]](diffhunk://#diff-7b1c9ef14536f9a66ed370cb729b6609d12c5907b460d8f145a7ad5a401e0fb6R29-R40)
[[2]](diffhunk://#diff-03c846683a6d76ded189d6ef24dc545da89ca418d0bce5cf1243d33cf1e2ac06R320-R351)
* Refactored the `IResourceAccountant` interface: replaced
`ResetPendingWeights` with `ResetForNewPass`, which resets both the stop
flag and pending weights, and introduced a protected
`ResetPendingWeightsImpl` for subclass-specific cleanup.
[[1]](diffhunk://#diff-7b1c9ef14536f9a66ed370cb729b6609d12c5907b460d8f145a7ad5a401e0fb6L64-R83)
[[2]](diffhunk://#diff-7b1c9ef14536f9a66ed370cb729b6609d12c5907b460d8f145a7ad5a401e0fb6R92-R96)
[[3]](diffhunk://#diff-03c846683a6d76ded189d6ef24dc545da89ca418d0bce5cf1243d33cf1e2ac06L123-R123)
[[4]](diffhunk://#diff-e2d3910ae7593ee7ba4fd74e53f738fa973ae2fc32c069f1088ba458b91f8d4bL280-R280)
[[5]](diffhunk://#diff-e2d3910ae7593ee7ba4fd74e53f738fa973ae2fc32c069f1088ba458b91f8d4bL351-R351)

**Plugin CUDA EP and resource budget enforcement:**

* Added `kCudaPluginExecutionProvider` constant and updated logic to
ensure plugin EPs correctly map to their in-tree accountant counterparts
and are included in device matching and partitioning.
[[1]](diffhunk://#diff-442c270eea3703252c48e97a7573960e14bf27a45a4443348840ed565330bf70R34)
[[2]](diffhunk://#diff-b20f416b9fe3b85423eea6707c38753351a3f1b8ef7a319858b27794507e0686L102)
[[3]](diffhunk://#diff-a8f614056d63b5b3325eea1d855afc96550c977c16d8fdba641012a79194b7b5L186-R187)
[[4]](diffhunk://#diff-a8f614056d63b5b3325eea1d855afc96550c977c16d8fdba641012a79194b7b5L206-R207)
[[5]](diffhunk://#diff-a8f614056d63b5b3325eea1d855afc96550c977c16d8fdba641012a79194b7b5L228-R229)
[[6]](diffhunk://#diff-e2d3910ae7593ee7ba4fd74e53f738fa973ae2fc32c069f1088ba458b91f8d4bL1192-R1200)
* Updated plugin EP infrastructure to pass and utilize resource
accountant pointers, enabling host-side resource budget enforcement for
plugin EPs and ensuring correct node assignment.
[[1]](diffhunk://#diff-fb00c9a234d8cc889927a22de94acfcfd893b56505e8ed613961b1bf13c0e435R19)
[[2]](diffhunk://#diff-fb00c9a234d8cc889927a22de94acfcfd893b56505e8ed613961b1bf13c0e435R54-R57)
[[3]](diffhunk://#diff-6dac10650c4e1c5a55b95378173b33e95b300bf7c2350d8476088693b98652a5R16-R17)
[[4]](diffhunk://#diff-6dac10650c4e1c5a55b95378173b33e95b300bf7c2350d8476088693b98652a5L239-R259)
[[5]](diffhunk://#diff-6dac10650c4e1c5a55b95378173b33e95b300bf7c2350d8476088693b98652a5R273-R281)
[[6]](diffhunk://#diff-0890d267a71ca02f4173c2ab226e6c5707fcbbf6bbb5f602fa5d92aa82f42a80R14-R22)
[[7]](diffhunk://#diff-0890d267a71ca02f4173c2ab226e6c5707fcbbf6bbb5f602fa5d92aa82f42a80R233-R241)

**Device matching and partitioning:**

* Improved device matching heuristics to consider both in-tree and
plugin CUDA EPs, and updated logic to prefer runtime device ordinals for
more reliable device selection.

Other minor changes include code style cleanups and additional includes
for completeness.

Author: yuslepukhin

include/onnxruntime/core/framework/resource_accountant.h

@@ -26,6 +26,18 @@ struct Node;
 // for different EPs
 using ResourceCount = std::variant<size_t>;
 
+// Type-erased arithmetic for ResourceCount values.
+// Implementations use std::visit so the compiler enforces exhaustive handling
+// of all variant members — adding a new type to ResourceCount will produce
+// build errors at each call site that must be addressed.
+//
+// NOTE: These functions are NOT available through the provider bridge (shared library EPs).
+// Budget enforcement for bridge-based EPs (e.g., in-tree CUDA EP) will be moved to the
+// graph partitioner in a follow-up PR.
+ResourceCount AddResourceCounts(const ResourceCount& a, const ResourceCount& b);
+bool ResourceCountExceeds(const ResourceCount& a, const ResourceCount& b);
+std::string FormatResourceCount(const ResourceCount& rc);
+
 /// <summary>
 /// This class is used for graph partitioning by EPs
 /// It stores the cumulative amount of the resource such as
@@ -61,9 +73,14 @@ class IResourceAccountant {
 
   bool IsStopIssued() const noexcept { return stop_assignment_; }
 
-  // Called before each GetCapability pass to discard pending weight tracking
-  // from a previous (discarded) pass. Default no-op for stats-based accountants.
-  virtual void ResetPendingWeights() {}
+  // Called before each GetCapability pass to reset per-pass state:
+  // clears the stop flag (which only applies to the pass that set it)
+  // and discards pending weight tracking from a previous (discarded) pass.
+  // Subclasses override ResetPendingWeightsImpl for EP-specific cleanup.
+  void ResetForNewPass() {
+    stop_assignment_ = false;
+    ResetPendingWeightsImpl();
+  }
 
   // Called when a node's cost is committed (AccountForNode/AccountForAllNodes).
   // Moves the node's pending weights into the committed set so they persist
@@ -72,6 +89,11 @@ class IResourceAccountant {
 
   static std::string MakeUniqueNodeName(const Node& node);
 
+ protected:
+  // Override to discard EP-specific pending weight tracking.
+  // Default no-op for stats-based accountants.
+  virtual void ResetPendingWeightsImpl() {}
+
  private:
   bool stop_assignment_ = false;
   std::optional<ResourceCount> threshold_;

include/onnxruntime/core/graph/constants.h

@@ -31,6 +31,7 @@ constexpr size_t kMaxExecutionProviderNameLen = 30;
 
 constexpr const char* kCpuExecutionProvider = "CPUExecutionProvider";
 constexpr const char* kCudaExecutionProvider = "CUDAExecutionProvider";
+constexpr const char* kCudaPluginExecutionProvider = "CudaPluginExecutionProvider";
 constexpr const char* kCudaNHWCExecutionProvider = "CUDANHWCExecutionProvider";
 constexpr const char* kDnnlExecutionProvider = "DnnlExecutionProvider";
 constexpr const char* kOpenVINOExecutionProvider = "OpenVINOExecutionProvider";

include/onnxruntime/core/session/onnxruntime_cxx_api.h

@@ -3729,5 +3729,6 @@ using UnownedSharedPrePackedWeightCache =
 
 ///< Wraps OrtEpApi::GetEnvConfigEntries()
 Ort::KeyValuePairs GetEnvConfigEntries();
+
 }  // namespace Ort
 #include "onnxruntime_cxx_inline.h"

include/onnxruntime/core/session/onnxruntime_cxx_inline.h

@@ -4174,4 +4174,5 @@ inline OpSchema GetOpSchema(const char* name, int max_inclusive_version, const c
   ThrowOnError(GetEpApi().GetOpSchema(name, max_inclusive_version, domain, &schema));
   return OpSchema{schema};
 }
+
 }  // namespace Ort

onnxruntime/core/framework/graph_partitioner.cc

@@ -277,7 +277,7 @@ static Status GetCapabilityForEP(const GetCapabilityForEPParams& params, const l
     ORT_RETURN_IF_ERROR(create_graph_viewer(sub_graph_holder, graph_viewer));
 
     if (params.resource_accountant) {
-      params.resource_accountant->ResetPendingWeights();
+      params.resource_accountant->ResetForNewPass();
     }
     capabilities = get_capabilities(current_ep, *graph_viewer, kernel_lookup, params.resource_accountant,
                                     graph_optimizer_registry);
@@ -348,7 +348,7 @@ static Status GetCapabilityForEP(const GetCapabilityForEPParams& params, const l
     ORT_RETURN_IF_ERROR(create_graph_viewer(sub_graph_holder, graph_viewer));
 
     if (params.resource_accountant) {
-      params.resource_accountant->ResetPendingWeights();
+      params.resource_accountant->ResetForNewPass();
     }
     capabilities = get_capabilities(current_ep, *graph_viewer, kernel_lookup, params.resource_accountant,
                                     graph_optimizer_registry);
@@ -1189,7 +1189,15 @@ static Status PartitionOnnxFormatModel(const PartitionParams& partition_params,
     for (const auto& ep : execution_providers) {
       IResourceAccountant* resource_accountant = nullptr;
       if (acc_map.has_value()) {
-        auto hit = acc_map->find(ep->Type());
+        // Plugin EPs have a different Type() than the in-tree EP they replace
+        // (e.g., kCudaPluginExecutionProvider vs kCudaExecutionProvider), but the
+        // accountant is registered under the in-tree EP name. Translate the key
+        // so plugin EPs find the correct accountant.
+        const auto& ep_type = ep->Type();
+        const auto accountant_key = ep_type == kCudaPluginExecutionProvider
+                                        ? std::string{kCudaExecutionProvider}
+                                        : ep_type;
+        auto hit = acc_map->find(accountant_key);
         if (hit != acc_map->end()) {
           resource_accountant = hit->second.get();
         }

onnxruntime/core/framework/layering_annotations.cc

@@ -183,7 +183,8 @@ bool MatchEpDevice(const EpDeviceView& ep,
     if (target_specifier.empty()) {
       if (ep.device_type == OrtDevice::GPU) return true;
       // Heuristic fallback for common GPU EPs if hardware info is missing
-      return ep.ep_name == kCudaExecutionProvider || ep.ep_name == kDmlExecutionProvider;
+      return ep.ep_name == kCudaExecutionProvider || ep.ep_name == kCudaPluginExecutionProvider ||
+             ep.ep_name == kDmlExecutionProvider;
     }
     // "gpu:<vendor>" or "gpu:<index>"
     if (ep.device_type == OrtDevice::GPU) {
@@ -203,7 +204,7 @@ bool MatchEpDevice(const EpDeviceView& ep,
           ep.vendor_id == OrtDevice::VendorIds::INTEL) return true;
       // Heuristic: gpu:nvidia -> CUDA
       if (CaseInsensitiveCompare(target_specifier, "nvidia") &&
-          ep.ep_name == kCudaExecutionProvider) return true;
+          (ep.ep_name == kCudaExecutionProvider || ep.ep_name == kCudaPluginExecutionProvider)) return true;
     }
     return false;
   }
@@ -225,7 +226,7 @@ bool MatchEpDevice(const EpDeviceView& ep,
   }
   // "cuda"
   if (CaseInsensitiveCompare(target_type_str, "cuda")) {
-    return ep.ep_name == kCudaExecutionProvider;
+    return ep.ep_name == kCudaExecutionProvider || ep.ep_name == kCudaPluginExecutionProvider;
   }
   // "dml"
   if (CaseInsensitiveCompare(target_type_str, "dml")) {
@@ -284,7 +285,13 @@ std::optional<std::string> EpLayeringMatcher::Match(gsl::span<const OrtEpDevice*
         ep_device.ep_name,
         device_type,
         has_hw ? ep_device.device->vendor_id : 0u,
-        has_hw ? static_cast<OrtDevice::DeviceId>(ep_device.device->device_id) : OrtDevice::DeviceId{},
+        // Prefer the device ordinal from device_memory_info (set by the EP factory to
+        // a runtime device ordinal such as a CUDA ordinal) over the OrtHardwareDevice::device_id
+        // which is a hardware-type identifier and not guaranteed to be a stable runtime ordinal.
+        ep_device.device_memory_info
+            ? ep_device.device_memory_info->device.Id()
+            : (has_hw ? static_cast<OrtDevice::DeviceId>(ep_device.device->device_id)
+                      : OrtDevice::DeviceId{}),
         has_hw ? std::string_view(ep_device.device->vendor) : std::string_view{}};
 
     if (MatchEpDevice(view, target_type_str, target_specifier, rule.device)) {

onnxruntime/core/framework/resource_accountant.cc

@@ -120,7 +120,7 @@ class SizeBasedStatsAccountant : public IResourceAccountant {
     }
   }
 
-  void ResetPendingWeights() override {
+  void ResetPendingWeightsImpl() override {
     pending_weights_.clear();
     pending_weights_by_node_.clear();
   }
@@ -317,4 +317,36 @@ std::string IResourceAccountant::MakeUniqueNodeName(const Node& node) {
   return result;
 }
 
+ResourceCount AddResourceCounts(const ResourceCount& a, const ResourceCount& b) {
+  return std::visit(
+      [](auto lhs, auto rhs) -> ResourceCount {
+        static_assert(std::is_same_v<decltype(lhs), decltype(rhs)>,
+                      "AddResourceCounts requires both operands to hold the same type. "
+                      "Handle the new ResourceCount variant member.");
+        if constexpr (std::is_integral_v<decltype(lhs)>) {
+          return static_cast<decltype(lhs)>(SafeInt<decltype(lhs)>(lhs) + rhs);
+        } else {
+          return lhs + rhs;
+        }
+      },
+      a, b);
+}
+
+bool ResourceCountExceeds(const ResourceCount& a, const ResourceCount& b) {
+  return std::visit(
+      [](auto lhs, auto rhs) -> bool {
+        static_assert(std::is_same_v<decltype(lhs), decltype(rhs)>,
+                      "ResourceCountExceeds requires both operands to hold the same type. "
+                      "Handle the new ResourceCount variant member.");
+        return lhs > rhs;
+      },
+      a, b);
+}
+
+std::string FormatResourceCount(const ResourceCount& rc) {
+  return std::visit(
+      [](auto val) -> std::string { return std::to_string(val); },
+      rc);
+}
+
 }  // namespace onnxruntime

onnxruntime/core/providers/cuda/plugin/cuda_ep.cc

@@ -11,12 +11,15 @@
 #include "ep/get_capability_utils.h"
 
 #include <cstring>
+#include <limits>
 #include <stdexcept>
 #include <string>
 #include <string_view>
 #include <unordered_map>
 #include <unordered_set>
 
+#include "core/graph/constants.h"
+
 namespace onnxruntime {
 namespace cuda_plugin {
 
@@ -227,12 +230,15 @@ OrtStatus* ORT_API_CALL CudaEp::GetCapabilityImpl(
       cpu_preferred_nodes));
 
   // Phase 3: Add final supported nodes (tentative minus CPU-preferred).
+  // Resource budget enforcement is handled by the host after GetCapability returns.
+
   for (const OrtNode* ort_node : candidate_nodes) {
-    if (cpu_preferred_nodes.count(ort_node) == 0) {
-      Ort::ConstNode node{ort_node};
-      RETURN_IF_ERROR(ep_api.EpGraphSupportInfo_AddSingleNode(
-          graph_support_info, node));
+    if (cpu_preferred_nodes.count(ort_node) != 0) {
+      continue;
     }
+
+    RETURN_IF_ERROR(ep_api.EpGraphSupportInfo_AddSingleNode(
+        graph_support_info, ort_node));
   }
 
   return nullptr;

onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h

@@ -99,7 +99,6 @@ class OrtStreamAdapter {
 #include "core/providers/common.h"
 
 namespace onnxruntime {
-inline constexpr const char* kCudaPluginExecutionProvider = "CudaPluginExecutionProvider";
 
 // Forward declaration of GetEnvironmentVar for plugin builds on Windows.
 // Defined in provider_api_shims.cc; mirrors the provider_api.h declaration

onnxruntime/core/session/abi_ep_types.h

@@ -16,6 +16,7 @@
 namespace onnxruntime {
 struct EpGraph;
 struct EpNode;
+class IResourceAccountant;
 }  // namespace onnxruntime
 
 /// <summary>
@@ -50,4 +51,8 @@ struct OrtEpGraphSupportInfo {
   const onnxruntime::EpGraph& ort_graph;
   std::vector<NodeGrouping> node_groupings;
   const onnxruntime::IExecutionProvider::IKernelLookup& kernel_lookup;
+
+  // Optional resource accountant for capacity-aware partitioning.
+  // Owned by the graph partitioner; lifetime exceeds this struct.
+  onnxruntime::IResourceAccountant* resource_accountant = nullptr;
 };