From 7f9571fb6536cc7270798c757fea52451cefbccb Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 25 Mar 2026 23:18:38 -0700 Subject: [PATCH 1/6] Fix NHWC second-pass EP assignment cleanup --- .../core/framework/graph_partitioner.cc | 30 +++- .../internal_testing_partitioning_tests.cc | 139 ++++++++++++++++++ 2 files changed, 166 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 0d3a84f30e1fb..029dc520c29be 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -6,6 +6,7 @@ #include #include #include +#include #include "core/common/inlined_containers.h" #include "core/common/string_utils.h" @@ -93,7 +94,8 @@ static void BuildFusedKernelDef(KernelDefBuilder& builder, const IndexedSubGraph /// Indexed subgraph which needs to be assigned /// The EP to assign the Indexed subgraph to static bool TryAssignNodes(Graph& graph, const IndexedSubGraph& capability, - const std::string& provider_type) { + const std::string& provider_type, + std::vector* newly_assigned_nodes = nullptr) { // Before assigning the ep to any node, first walk through all the nodes and ensure // none of the nodes have already been assigned. If a node is assigned, simply return. for (auto node_index : capability.nodes) { @@ -107,7 +109,11 @@ static bool TryAssignNodes(Graph& graph, const IndexedSubGraph& capability, const bool acc_enabled = capability.IsAccountingEnabled(); for (size_t i = 0, limit = capability.nodes.size(); i < limit; ++i) { auto* node = graph.GetNode(capability.nodes[i]); + const bool was_unassigned = node->GetExecutionProviderType().empty(); node->SetExecutionProviderType(provider_type); + if (newly_assigned_nodes != nullptr && was_unassigned) { + newly_assigned_nodes->push_back(node->Index()); + } if (acc_enabled) { capability.AccountForNode(i); } @@ -115,6 +121,17 @@ static bool TryAssignNodes(Graph& graph, const IndexedSubGraph& capability, return true; } +static void ClearExecutionProviderAssignments(Graph& graph, + const std::vector& node_indices, + const std::string& provider_type) { + for (NodeIndex node_index : node_indices) { + auto* node = graph.GetNode(node_index); + if (node != nullptr && node->GetExecutionProviderType() == provider_type) { + node->SetExecutionProviderType(""); + } + } +} + #endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) static bool TryAssignSingleNode(Graph& graph, @@ -299,16 +316,22 @@ static Status GetCapabilityForEP(const GetCapabilityForEPParams& params, const l // CPU EP layout transformation happens later when level 3 transformers are run. if (params.mode != GraphPartitioner::Mode::kAssignOnly && params.transform_layout.get() && current_ep.GetPreferredLayout() == DataLayout::NHWC) { + std::vector nodes_temporarily_assigned_to_ep; for (auto& capability : capabilities) { - TryAssignNodes(graph, *capability->sub_graph, ep_type); + TryAssignNodes(graph, *capability->sub_graph, ep_type, &nodes_temporarily_assigned_to_ep); } const NodeIndex first_new_node = graph.MaxNodeIndex(); // Perform layout transformation on the specific EP assigned graph bool modified = false; - ORT_RETURN_IF_ERROR(params.transform_layout(graph, modified, current_ep, params.debug_graph_fn)); + auto transform_status = params.transform_layout(graph, modified, current_ep, params.debug_graph_fn); + if (!transform_status.IsOK()) { + ClearExecutionProviderAssignments(graph, nodes_temporarily_assigned_to_ep, ep_type); + return transform_status; + } if (params.check_load_cancellation_fn()) { + ClearExecutionProviderAssignments(graph, nodes_temporarily_assigned_to_ep, ep_type); return ORT_MAKE_STATUS(ONNXRUNTIME, MODEL_LOAD_CANCELED, "GetCapabilities was canceled by user request"); } @@ -326,6 +349,7 @@ static Status GetCapabilityForEP(const GetCapabilityForEPParams& params, const l const NodeIndex end_node = graph.MaxNodeIndex(); + ClearExecutionProviderAssignments(graph, nodes_temporarily_assigned_to_ep, ep_type); capabilities.clear(); #if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) diff --git a/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc b/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc index f023cc668a995..93290f43a2c1b 100644 --- a/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc +++ b/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc @@ -5,10 +5,13 @@ #include "core/common/logging/logging.h" #include "core/framework/compute_capability.h" +#include "core/framework/model_metadef_id_generator.h" #include "core/framework/utils.h" +#include "core/providers/partitioning_utils.h" #include "core/session/inference_session.h" #include "test/unittest_util/framework_test_utils.h" +#include "test/unittest_util/graph_transform_test_builder.h" #include "test/internal_testing_ep/internal_testing_execution_provider.h" #include "test/test_environment.h" #include "test/util/include/asserts.h" @@ -32,6 +35,87 @@ using namespace onnxruntime::internal_testing_ep; // it would be possible to use ORT format models but the same partitioning code would run either way #if !defined(ORT_MINIMAL_BUILD) +#if !defined(DISABLE_CONTRIB_OPS) +namespace { + +class TwoPassNhwcTestExecutionProvider : public IExecutionProvider { + public: + TwoPassNhwcTestExecutionProvider() : IExecutionProvider{"TwoPassNhwcTestExecutionProvider"} { + } + + DataLayout GetPreferredLayout() const override { + return DataLayout::NHWC; + } + + std::vector> + GetCapability(const GraphViewer& graph_viewer, + const IKernelLookup&, + const GraphOptimizerRegistry&, + IResourceAccountant*) const override { + ++get_capability_calls_; + const bool second_pass = get_capability_calls_ > 1; + + auto generate_metadef_name = [this, &graph_viewer]() { + HashValue model_hash; + const int metadef_id = metadef_id_generator_.GenerateId(graph_viewer, model_hash); + return std::string(Type()) + "_" + std::to_string(model_hash) + "_" + std::to_string(metadef_id); + }; + + std::vector> capabilities; + for (const auto node_index : graph_viewer.GetNodesInTopologicalOrder()) { + const Node* node = graph_viewer.GetNode(node_index); + if (node == nullptr) { + continue; + } + + const bool is_conv = node->OpType() == "Conv"; + const bool is_log_softmax = node->OpType() == "LogSoftmax"; + if (!is_conv && !is_log_softmax) { + continue; + } + + if (second_pass && is_log_softmax) { + continue; + } + + const auto& assigned_ep = node->GetExecutionProviderType(); + if (!assigned_ep.empty() && assigned_ep != Type()) { + continue; + } + + capabilities.push_back(utils::MakeComputeCapability(graph_viewer, + std::vector{node}, + generate_metadef_name, + Type(), + false)); + } + + return capabilities; + } + + Status Compile(const std::vector& fused_nodes, + std::vector& node_compute_funcs) override { + for (size_t i = 0; i < fused_nodes.size(); ++i) { + NodeComputeInfo compute_info; + compute_info.create_state_func = [](ComputeContext*, FunctionState*) { return 0; }; + compute_info.release_state_func = [](FunctionState) {}; + compute_info.compute_func = [](FunctionState, const OrtApi*, OrtKernelContext*) { + return Status::OK(); + }; + node_compute_funcs.push_back(std::move(compute_info)); + } + + return Status::OK(); + } + + private: + mutable size_t get_capability_calls_{0}; + mutable ModelMetadefIdGenerator metadef_id_generator_; +}; + +} // namespace +#endif // !defined(DISABLE_CONTRIB_OPS) + #define ORT_MODEL_FOLDER ORT_TSTR("testdata/") auto RunTest(const std::string& op, const ORTCHAR_T* model_path) { @@ -129,6 +213,61 @@ TEST(InternalTestingEP, TestDependenciesCorrectlyHandled) { ASSERT_EQ(num_other_nodes, 2); } +#if !defined(DISABLE_CONTRIB_OPS) +TEST(InternalTestingEP, NhwcSecondPassDropFallsBackFromCpuKernelNode) { + std::unordered_map domain_to_version{{kOnnxDomain, 13}, {kMSDomain, 1}}; + Model model("NhwcSecondPassDropFallsBackFromCpuKernelNode", + false, + ModelMetaData(), + PathString(), + IOnnxRuntimeOpSchemaRegistryList(), + domain_to_version, + {}, + DefaultLoggingManager().DefaultLogger()); + + Graph& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + + const std::vector tensor_shape{1, 1, 3, 3}; + auto* input = builder.MakeInput(std::optional>{tensor_shape}); + auto* weights = builder.MakeInitializer(std::vector{1, 1, 1, 1}, std::vector{1.0f}); + auto* conv_output = builder.MakeIntermediate(std::optional>{tensor_shape}); + auto* output = builder.MakeOutput(std::optional>{tensor_shape}); + + builder.AddConvNode(input, weights, conv_output); + builder.AddNode("LogSoftmax", std::vector{conv_output}, std::vector{output}); + builder.SetGraphOutputs(); + + ASSERT_STATUS_OK(graph.Resolve()); + + std::string model_data; + ASSERT_TRUE(model.ToProto().SerializeToString(&model_data)); + + SessionOptions so; + auto session = std::make_unique(so, GetEnvironment()); + ASSERT_STATUS_OK(session->RegisterExecutionProvider(std::make_unique())); + + ASSERT_STATUS_OK(session->Load(model_data.data(), static_cast(model_data.size()))); + ASSERT_STATUS_OK(session->Initialize()); + + bool saw_log_softmax = false; + int num_ep_nodes = 0; + for (const auto& node : session->GetGraph().Nodes()) { + if (node.GetExecutionProviderType() == "TwoPassNhwcTestExecutionProvider") { + ++num_ep_nodes; + } + + if (node.OpType() == "LogSoftmax") { + saw_log_softmax = true; + EXPECT_NE(node.GetExecutionProviderType(), "TwoPassNhwcTestExecutionProvider"); + } + } + + EXPECT_GT(num_ep_nodes, 0); + EXPECT_TRUE(saw_log_softmax); +} +#endif // !defined(DISABLE_CONTRIB_OPS) + // Infrastructure that was used to check NNAPI coverage. // Ideally this could be updated to read the model paths, supported ops and stop ops from input files // and provide info on the partitions so no code changes are required to investigate different scenarios. From 0c6c3752514ace66921531fa4668e66aba7932f6 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 1 Apr 2026 14:45:40 -0700 Subject: [PATCH 2/6] address feedback --- .../core/framework/graph_partitioner.cc | 19 +++++++++++++++++-- .../internal_testing_partitioning_tests.cc | 14 +++++++++++--- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 029dc520c29be..877b4b4670562 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -349,7 +349,8 @@ static Status GetCapabilityForEP(const GetCapabilityForEPParams& params, const l const NodeIndex end_node = graph.MaxNodeIndex(); - ClearExecutionProviderAssignments(graph, nodes_temporarily_assigned_to_ep, ep_type); + // Keep pass-1 EP assignments through the second GetCapability call so that EPs can + // recognize already-tagged nodes (e.g. nodes transformed into kMSInternalNHWCDomain). capabilities.clear(); #if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) @@ -380,20 +381,34 @@ static Status GetCapabilityForEP(const GetCapabilityForEPParams& params, const l reset_assignment_unclaimed_nodes(); if (params.check_load_cancellation_fn()) { + ClearExecutionProviderAssignments(graph, nodes_temporarily_assigned_to_ep, ep_type); return ORT_MAKE_STATUS(ONNXRUNTIME, MODEL_LOAD_CANCELED, "GetCapabilities was canceled by user request"); } - // all nodes with an index >= first_new_node with domain of kMSInternalNHWCDomain should be in the capabilities + // Collect pass-2 node indices and track new nodes for NHWC domain validation. + InlinedHashSet pass2_node_indices; InlinedHashSet new_nodes_in_capabilities; for (const auto& capability : capabilities) { for (auto node_index : capability->sub_graph->nodes) { + pass2_node_indices.insert(node_index); if (node_index >= first_new_node) { new_nodes_in_capabilities.insert(node_index); } } } + // Clear pass-1 temporary assignments for nodes NOT re-claimed in pass 2. + // Nodes present in both passes keep their EP tag for correct downstream assignment. + for (NodeIndex node_index : nodes_temporarily_assigned_to_ep) { + if (pass2_node_indices.count(node_index) == 0) { + auto* node = graph.GetNode(node_index); + if (node != nullptr && node->GetExecutionProviderType() == ep_type) { + node->SetExecutionProviderType(""); + } + } + } + for (NodeIndex idx = first_new_node; idx < end_node; ++idx) { const Node* node = graph.GetNode(idx); if (node != nullptr && node->Domain() == kMSInternalNHWCDomain) { diff --git a/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc b/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc index 93290f43a2c1b..89f73a8d20de0 100644 --- a/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc +++ b/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc @@ -52,8 +52,17 @@ class TwoPassNhwcTestExecutionProvider : public IExecutionProvider { const IKernelLookup&, const GraphOptimizerRegistry&, IResourceAccountant*) const override { - ++get_capability_calls_; - const bool second_pass = get_capability_calls_ > 1; + // Detect second pass by checking if any node already has our EP type assigned + // (set during the first-pass assignment). Real NHWC EPs use this pattern to + // recognize nodes that were transformed into kMSInternalNHWCDomain. + bool second_pass = false; + for (const auto node_index : graph_viewer.GetNodesInTopologicalOrder()) { + const Node* node = graph_viewer.GetNode(node_index); + if (node != nullptr && node->GetExecutionProviderType() == Type()) { + second_pass = true; + break; + } + } auto generate_metadef_name = [this, &graph_viewer]() { HashValue model_hash; @@ -109,7 +118,6 @@ class TwoPassNhwcTestExecutionProvider : public IExecutionProvider { } private: - mutable size_t get_capability_calls_{0}; mutable ModelMetadefIdGenerator metadef_id_generator_; }; From e8a645cb8dab37ac02759ac5b5533540f7a63695 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 2 Jun 2026 23:49:37 +0000 Subject: [PATCH 3/6] use InlinedVector for consistent --- onnxruntime/core/framework/graph_partitioner.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 877b4b4670562..01b0d566c8b7e 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -95,7 +95,7 @@ static void BuildFusedKernelDef(KernelDefBuilder& builder, const IndexedSubGraph /// The EP to assign the Indexed subgraph to static bool TryAssignNodes(Graph& graph, const IndexedSubGraph& capability, const std::string& provider_type, - std::vector* newly_assigned_nodes = nullptr) { + InlinedVector* newly_assigned_nodes = nullptr) { // Before assigning the ep to any node, first walk through all the nodes and ensure // none of the nodes have already been assigned. If a node is assigned, simply return. for (auto node_index : capability.nodes) { @@ -122,7 +122,7 @@ static bool TryAssignNodes(Graph& graph, const IndexedSubGraph& capability, } static void ClearExecutionProviderAssignments(Graph& graph, - const std::vector& node_indices, + const InlinedVector& node_indices, const std::string& provider_type) { for (NodeIndex node_index : node_indices) { auto* node = graph.GetNode(node_index); @@ -316,7 +316,7 @@ static Status GetCapabilityForEP(const GetCapabilityForEPParams& params, const l // CPU EP layout transformation happens later when level 3 transformers are run. if (params.mode != GraphPartitioner::Mode::kAssignOnly && params.transform_layout.get() && current_ep.GetPreferredLayout() == DataLayout::NHWC) { - std::vector nodes_temporarily_assigned_to_ep; + InlinedVector nodes_temporarily_assigned_to_ep; for (auto& capability : capabilities) { TryAssignNodes(graph, *capability->sub_graph, ep_type, &nodes_temporarily_assigned_to_ep); } From 333620db92b33d46a14e8348d238f1fa599e87b3 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 4 Jun 2026 00:12:34 +0000 Subject: [PATCH 4/6] fix minimal build --- .../internal_testing_partitioning_tests.cc | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc b/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc index 89f73a8d20de0..e68212c9a6421 100644 --- a/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc +++ b/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc @@ -5,19 +5,22 @@ #include "core/common/logging/logging.h" #include "core/framework/compute_capability.h" -#include "core/framework/model_metadef_id_generator.h" #include "core/framework/utils.h" -#include "core/providers/partitioning_utils.h" #include "core/session/inference_session.h" #include "test/unittest_util/framework_test_utils.h" -#include "test/unittest_util/graph_transform_test_builder.h" #include "test/internal_testing_ep/internal_testing_execution_provider.h" #include "test/test_environment.h" #include "test/util/include/asserts.h" #include "test/util/include/inference_session_wrapper.h" #include "test/util/include/test_utils.h" +#if !defined(ORT_MINIMAL_BUILD) +#include "core/framework/model_metadef_id_generator.h" +#include "core/providers/partitioning_utils.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#endif // !defined(ORT_MINIMAL_BUILD) + #include "gtest/gtest.h" #include "gmock/gmock.h" From 6080aeb12428e594832b99c912994232ba24ca4f Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 4 Jun 2026 13:07:21 -0700 Subject: [PATCH 5/6] update resource accountant --- ...ningWithAnnotationsAndMemoryConstraints.md | 9 + docs/cuda_plugin_ep/cuda_plugin_ep_design.md | 2 + .../core/graph/indexed_sub_graph.h | 7 + .../core/framework/graph_partitioner.cc | 48 +++- .../internal_testing_partitioning_tests.cc | 258 ++++++++++++++++++ .../transformers/test_cuda_plugin_ep.py | 28 ++ 6 files changed, 351 insertions(+), 1 deletion(-) diff --git a/docs/annotated_partitioning/PartitioningWithAnnotationsAndMemoryConstraints.md b/docs/annotated_partitioning/PartitioningWithAnnotationsAndMemoryConstraints.md index 1131c5cf7dacd..34092fe9a0307 100644 --- a/docs/annotated_partitioning/PartitioningWithAnnotationsAndMemoryConstraints.md +++ b/docs/annotated_partitioning/PartitioningWithAnnotationsAndMemoryConstraints.md @@ -283,6 +283,15 @@ The value of `session.resource_cuda_partitioning_settings` is a comma-separated The stats file path follows the same resolution rules described above: relative paths are resolved against the model's directory, absolute paths are used as-is. +### Interaction with NHWC Layout (`prefer_nhwc`) + +EPs that prefer the NHWC data layout — for example, the CUDA EP when it is created with the provider option `prefer_nhwc=1` (CUDA plugin EP: session option `ep.cuda.prefer_nhwc_layout=1`) — partition layout-sensitive subgraphs in two passes: + +1. **First pass (tentative):** The EP tags the nodes it could claim so the layout transformer can rewrite them into the NHWC (`com.microsoft.nhwc`) form. +2. **Second pass (final):** After the layout transform, the EP runs capability detection again and fuses/optimizes the rewritten nodes. Some first-pass nodes may be dropped here (for example, a node whose NHWC form is not actually supported), in which case they fall back to a later EP. + +Because the first-pass tags are tentative, ONNX Runtime does **not** commit any memory budget for them. The budget is committed only for the nodes that survive the second pass; the cost of a node that is dropped is never counted against the memory limit. This keeps the accumulated memory estimate accurate when `prefer_nhwc` is combined with `session.resource_cuda_partitioning_settings`, so a dropped node does not consume phantom budget that could prematurely halt assignment of later nodes. + ## Combining Both Features Layer annotations and capacity-aware partitioning can be used together. When both are configured: - Layer annotations provide the initial node-to-device mapping. diff --git a/docs/cuda_plugin_ep/cuda_plugin_ep_design.md b/docs/cuda_plugin_ep/cuda_plugin_ep_design.md index 15f8188505b37..af422f7b365ae 100644 --- a/docs/cuda_plugin_ep/cuda_plugin_ep_design.md +++ b/docs/cuda_plugin_ep/cuda_plugin_ep_design.md @@ -390,6 +390,8 @@ That behavior is now implemented by tracking: The final support set is chosen from `candidate_nodes`, with the existing CPU-preferred-node filtering applied only where appropriate. +When resource accounting is also enabled (`session.resource_cuda_partitioning_settings`), this two-pass flow interacts with the partitioner's budget commit in an important way. The first-pass tags are tentative — ORT applies them only so the layout transformer can rewrite the nodes — so the partitioner does **not** commit any accountant budget for them. After the second pass, the partitioner commits budget only for the first-pass nodes that survived (still claimed by the plugin), using the per-node costs captured during the first pass. Nodes dropped on the second pass therefore never consume budget, and surviving nodes are counted exactly once. Plugin EPs that attach accounting costs should do so only on first-pass (newly claimed) capabilities, mirroring the in-tree CUDA EP, which leaves already-assigned second-pass nodes cost-free and relies on the partitioner's deferred commit. + **B. Cache the shim provider pointer at kernel creation** Migrated CUDA kernels expect `info.GetExecutionProvider()` to return the shim `CUDAExecutionProvider`, not the outer `PluginExecutionProvider`. The adapter now resolves that relationship once during kernel creation, captures the shim provider's runtime-config object, and uses `CudaKernel` accessors for later provider-setting reads. diff --git a/include/onnxruntime/core/graph/indexed_sub_graph.h b/include/onnxruntime/core/graph/indexed_sub_graph.h index 54e878761ba87..a00a92ec8eb9b 100644 --- a/include/onnxruntime/core/graph/indexed_sub_graph.h +++ b/include/onnxruntime/core/graph/indexed_sub_graph.h @@ -124,6 +124,13 @@ struct IndexedSubGraph { nodes_costs.emplace_back(cost); } + // Read-only access to the pre-computed resource cost for the node at cost_index. + // Should call IsAccountingEnabled() first. + const ResourceCount& GetNodeCost(size_t cost_index) const { + assert(cost_index < nodes_costs.size()); + return nodes_costs[cost_index]; + } + private: // subgraph meta definition. std::unique_ptr meta_def_; diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 01b0d566c8b7e..8d5f6192b2a5c 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -106,6 +106,13 @@ static bool TryAssignNodes(Graph& graph, const IndexedSubGraph& capability, } } + // When newly_assigned_nodes is provided, this call performs the tentative first-pass + // tagging of the NHWC two-pass partitioning flow. Those tags only exist so the layout + // transformer and the second GetCapability pass can recognize the candidate nodes; they + // may be dropped again before any partition is committed. Tentative tags must therefore + // not commit resource-accountant budget here. The budget for nodes that survive the + // second pass is committed later (see the deferred-commit step in GetCapabilityForEP). + const bool is_tentative_pass = (newly_assigned_nodes != nullptr); const bool acc_enabled = capability.IsAccountingEnabled(); for (size_t i = 0, limit = capability.nodes.size(); i < limit; ++i) { auto* node = graph.GetNode(capability.nodes[i]); @@ -114,7 +121,7 @@ static bool TryAssignNodes(Graph& graph, const IndexedSubGraph& capability, if (newly_assigned_nodes != nullptr && was_unassigned) { newly_assigned_nodes->push_back(node->Index()); } - if (acc_enabled) { + if (acc_enabled && !is_tentative_pass) { capability.AccountForNode(i); } } @@ -349,6 +356,23 @@ static Status GetCapabilityForEP(const GetCapabilityForEPParams& params, const l const NodeIndex end_node = graph.MaxNodeIndex(); + // Pass-1 tags were applied tentatively without committing any resource-accountant + // budget (see TryAssignNodes). Capture the per-node costs computed during pass-1, + // keyed by node index, so the budget for the nodes that survive the second pass can + // be committed after the drop step below. The costs must be captured here because + // capabilities.clear() destroys the pass-1 capabilities (and their costs) next. + InlinedHashMap pass1_node_costs; + if (params.resource_accountant != nullptr) { + for (const auto& capability : capabilities) { + const auto& sub_graph = *capability->sub_graph; + if (sub_graph.IsAccountingEnabled()) { + for (size_t i = 0, limit = sub_graph.nodes.size(); i < limit; ++i) { + pass1_node_costs.insert_or_assign(sub_graph.nodes[i], sub_graph.GetNodeCost(i)); + } + } + } + } + // Keep pass-1 EP assignments through the second GetCapability call so that EPs can // recognize already-tagged nodes (e.g. nodes transformed into kMSInternalNHWCDomain). capabilities.clear(); @@ -409,6 +433,28 @@ static Status GetCapabilityForEP(const GetCapabilityForEPParams& params, const l } } + // Commit resource-accountant budget for pass-1 tentatively-tagged nodes that survived + // the second pass (still claimed by this EP). Pass-1 deliberately deferred this commit + // (TryAssignNodes skipped accounting) so that nodes dropped in the loop above never + // leak phantom budget into later accounting decisions. New nodes introduced for the + // second pass (e.g. NHWC ops) carry their own costs and are accounted normally when + // their partitions are placed, so they are intentionally excluded here. + if (params.resource_accountant != nullptr) { + for (NodeIndex node_index : nodes_temporarily_assigned_to_ep) { + if (pass2_node_indices.count(node_index) == 0) { + continue; + } + const auto* node = graph.GetNode(node_index); + if (node == nullptr || node->GetExecutionProviderType() != ep_type) { + continue; + } + auto cost_it = pass1_node_costs.find(node_index); + if (cost_it != pass1_node_costs.end()) { + params.resource_accountant->AddConsumedAmount(cost_it->second); + } + } + } + for (NodeIndex idx = first_new_node; idx < end_node; ++idx) { const Node* node = graph.GetNode(idx); if (node != nullptr && node->Domain() == kMSInternalNHWCDomain) { diff --git a/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc b/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc index e68212c9a6421..76f2e7fac8630 100644 --- a/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc +++ b/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc @@ -16,8 +16,18 @@ #include "test/util/include/test_utils.h" #if !defined(ORT_MINIMAL_BUILD) +#include "core/framework/config_options.h" +#include "core/framework/execution_providers.h" +#include "core/framework/ep_context_options.h" +#include "core/framework/fuse_nodes_funcs.h" +#include "core/framework/graph_partitioner.h" +#include "core/framework/kernel_registry_manager.h" #include "core/framework/model_metadef_id_generator.h" +#include "core/framework/resource_accountant.h" +#include "core/graph/constants.h" +#include "core/optimizer/graph_optimizer_registry.h" #include "core/providers/partitioning_utils.h" +#include "core/session/onnxruntime_session_options_config_keys.h" #include "test/unittest_util/graph_transform_test_builder.h" #endif // !defined(ORT_MINIMAL_BUILD) @@ -124,6 +134,122 @@ class TwoPassNhwcTestExecutionProvider : public IExecutionProvider { mutable ModelMetadefIdGenerator metadef_id_generator_; }; +// Variant of the two-pass NHWC EP used to validate that the resource accountant +// is updated correctly across the NHWC two-pass partitioning flow. +// +// It reports kCudaExecutionProvider as its type so that the SizeBasedStatsAccountant +// (which CreateAccountants registers under kCudaExecutionProvider) is wired to it, +// mirroring the real in-tree CUDA EP. Like the CUDA EP, it attaches accounting costs +// only to first-pass (newly claimed) capabilities. Second-pass survivors are already +// tagged with this EP and take the "previously assigned" branch, so they carry no +// cost and rely on the partitioner's deferred commit (using the captured first-pass +// costs) for their budget. Dropped nodes therefore never leak budget. +class AccountingNhwcTestExecutionProvider : public IExecutionProvider { + public: + AccountingNhwcTestExecutionProvider() : IExecutionProvider{kCudaExecutionProvider} { + } + + DataLayout GetPreferredLayout() const override { + return DataLayout::NHWC; + } + + std::vector> + GetCapability(const GraphViewer& graph_viewer, + const IKernelLookup&, + const GraphOptimizerRegistry&, + IResourceAccountant* resource_accountant) const override { + if (resource_accountant != nullptr) { + observed_accountant_ = resource_accountant; + } + + bool second_pass = false; + for (const auto node_index : graph_viewer.GetNodesInTopologicalOrder()) { + const Node* node = graph_viewer.GetNode(node_index); + if (node != nullptr && node->GetExecutionProviderType() == Type()) { + second_pass = true; + break; + } + } + + auto generate_metadef_name = [this, &graph_viewer]() { + HashValue model_hash; + const int metadef_id = metadef_id_generator_.GenerateId(graph_viewer, model_hash); + return std::string(Type()) + "_" + std::to_string(model_hash) + "_" + std::to_string(metadef_id); + }; + + std::vector> capabilities; + for (const auto node_index : graph_viewer.GetNodesInTopologicalOrder()) { + const Node* node = graph_viewer.GetNode(node_index); + if (node == nullptr) { + continue; + } + + const bool is_conv = node->OpType() == "Conv"; + const bool is_log_softmax = node->OpType() == "LogSoftmax"; + if (!is_conv && !is_log_softmax) { + continue; + } + + // Drop LogSoftmax on the second pass to model the EP releasing a node that + // it tentatively claimed on the first pass. + if (second_pass && is_log_softmax) { + continue; + } + + const auto& assigned_ep = node->GetExecutionProviderType(); + if (!assigned_ep.empty() && assigned_ep != Type()) { + continue; + } + + const bool already_claimed = (assigned_ep == Type()); + + auto capability = utils::MakeComputeCapability(graph_viewer, + std::vector{node}, + generate_metadef_name, + Type(), + false); + + // Mirror the in-tree CUDA EP: only newly-claimed (first-pass) capabilities carry + // accounting costs. Already-tagged second-pass survivors take the "previously + // assigned" branch and carry no cost. + if (resource_accountant != nullptr && !already_claimed) { + capability->sub_graph->SetAccountant(resource_accountant); + for (auto cost_node_index : capability->sub_graph->nodes) { + const Node* cost_node = graph_viewer.GetNode(cost_node_index); + capability->sub_graph->AppendNodeCost(resource_accountant->ComputeResourceCount(*cost_node)); + } + } + + capabilities.push_back(std::move(capability)); + } + + return capabilities; + } + + Status Compile(const std::vector& fused_nodes, + std::vector& node_compute_funcs) override { + for (size_t i = 0; i < fused_nodes.size(); ++i) { + NodeComputeInfo compute_info; + compute_info.create_state_func = [](ComputeContext*, FunctionState*) { return 0; }; + compute_info.release_state_func = [](FunctionState) {}; + compute_info.compute_func = [](FunctionState, const OrtApi*, OrtKernelContext*) { + return Status::OK(); + }; + node_compute_funcs.push_back(std::move(compute_info)); + } + + return Status::OK(); + } + + IResourceAccountant* observed_accountant() const { + return observed_accountant_; + } + + private: + mutable ModelMetadefIdGenerator metadef_id_generator_; + mutable IResourceAccountant* observed_accountant_ = nullptr; +}; + } // namespace #endif // !defined(DISABLE_CONTRIB_OPS) @@ -277,6 +403,138 @@ TEST(InternalTestingEP, NhwcSecondPassDropFallsBackFromCpuKernelNode) { EXPECT_GT(num_ep_nodes, 0); EXPECT_TRUE(saw_log_softmax); } + +// Validates that the resource accountant is updated correctly across the NHWC two-pass +// partitioning flow: a node tentatively claimed on the first pass but dropped on the +// second pass must NOT consume budget (no phantom), while a node that survives must be +// committed exactly once (no double-count). This guards the fix where first-pass NHWC +// tags are tentative and budget is committed only for second-pass survivors. +TEST(InternalTestingEP, NhwcTwoPassAccountingCommitsOnlySurvivors) { + std::unordered_map domain_to_version{{kOnnxDomain, 13}, {kMSDomain, 1}}; + Model model("NhwcTwoPassAccountingCommitsOnlySurvivors", + false, + ModelMetaData(), + PathString(), + IOnnxRuntimeOpSchemaRegistryList(), + domain_to_version, + {}, + DefaultLoggingManager().DefaultLogger()); + + Graph& graph = model.MainGraph(); + ModelTestBuilder builder(graph); + + const std::vector tensor_shape{1, 1, 3, 3}; + auto* input = builder.MakeInput(std::optional>{tensor_shape}); + auto* weights = builder.MakeInitializer(std::vector{1, 1, 1, 1}, std::vector{1.0f}); + auto* conv_output = builder.MakeIntermediate(std::optional>{tensor_shape}); + auto* output = builder.MakeOutput(std::optional>{tensor_shape}); + + builder.AddConvNode(input, weights, conv_output); + builder.AddNode("LogSoftmax", std::vector{conv_output}, std::vector{output}); + builder.SetGraphOutputs(); + + ASSERT_STATUS_OK(graph.Resolve()); + + // Helper to read the size_t held by a ResourceCount. + auto get_size = [](const ResourceCount& rc) -> size_t { + const auto* value = std::get_if(&rc); + EXPECT_NE(value, nullptr) << "ResourceCount does not hold size_t"; + return value != nullptr ? *value : 0; + }; + + // Build a fresh ad-hoc accountant (no stats file) via the real factory. + auto make_accountant = [](std::optional& acc_map) -> IResourceAccountant* { + ConfigOptions config; + // Large memory limit so nothing is offloaded; empty stats file => ad-hoc cost mode. + EXPECT_STATUS_OK(config.AddConfigEntry(kOrtSessionOptionsResourceCudaPartitioningSettings, "1048576,")); + EXPECT_STATUS_OK(CreateAccountants(config, PathString(), acc_map)); + EXPECT_TRUE(acc_map.has_value()); + auto it = acc_map->find(kCudaExecutionProvider); + return it != acc_map->end() ? it->second.get() : nullptr; + }; + + // Reference per-node costs computed independently on fresh accountants. + const Node* conv_node = nullptr; + const Node* log_softmax_node = nullptr; + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "Conv") { + conv_node = &node; + } else if (node.OpType() == "LogSoftmax") { + log_softmax_node = &node; + } + } + ASSERT_NE(conv_node, nullptr); + ASSERT_NE(log_softmax_node, nullptr); + + std::optional ref_conv_map; + std::optional ref_ls_map; + IResourceAccountant* ref_conv_acc = make_accountant(ref_conv_map); + IResourceAccountant* ref_ls_acc = make_accountant(ref_ls_map); + ASSERT_NE(ref_conv_acc, nullptr); + ASSERT_NE(ref_ls_acc, nullptr); + const size_t expected_conv_cost = get_size(ref_conv_acc->ComputeResourceCount(*conv_node)); + const size_t expected_log_softmax_cost = get_size(ref_ls_acc->ComputeResourceCount(*log_softmax_node)); + ASSERT_GT(expected_conv_cost, 0u); + ASSERT_GT(expected_log_softmax_cost, 0u); + + // Drive partitioning directly with the accounting-aware NHWC EP. The accountant is + // created internally by GraphPartitioner from the config option below (keyed to + // kCudaExecutionProvider, which matches the EP's type). + ExecutionProviders execution_providers; + auto& default_logger = DefaultLoggingManager().DefaultLogger(); + auto ep = std::make_unique(); + auto* ep_raw = ep.get(); + ep->SetLogger(&default_logger); + ASSERT_STATUS_OK(execution_providers.Add(kCudaExecutionProvider, std::move(ep))); + + KernelRegistryManager krm; + ASSERT_STATUS_OK(krm.RegisterKernels(execution_providers)); + + SessionOptions sess_options; + ASSERT_STATUS_OK(sess_options.config_options.AddConfigEntry( + kOrtSessionOptionsResourceCudaPartitioningSettings, "1048576,")); + + // Capture the accountant's consumed amount when the survivor partition is assigned. + // on_partition_assignment_fn runs after GetCapabilityForEP completes (and therefore + // after the deferred commit), but before PlaceNode adds any further cost. + std::optional observed_consumed; + OnPartitionAssignmentFunction on_assignment = + [&](const Graph&, const ComputeCapability&, const std::string& assigned_ep_type) { + if (assigned_ep_type == kCudaExecutionProvider && ep_raw->observed_accountant() != nullptr) { + observed_consumed = get_size(ep_raw->observed_accountant()->GetConsumedAmount()); + } + }; + + auto graph_optimizer_registry = std::make_unique( + &sess_options, nullptr /*cpu_ep*/, &default_logger); + + GraphPartitioner partitioner(krm, execution_providers, std::move(graph_optimizer_registry), []() -> bool { return false; }, on_assignment); + + layout_transformation::TransformLayoutFunction transform_layout_fn = + [](Graph&, bool& modified, const IExecutionProvider&, + const layout_transformation::DebugGraphFn&) -> Status { + modified = false; + return Status::OK(); + }; + layout_transformation::DebugGraphFn debug_graph_fn; + + FuncManager func_mgr; + ASSERT_STATUS_OK( + partitioner.Partition(graph, func_mgr, transform_layout_fn, + sess_options.config_options, default_logger, nullptr /*layering_index*/, + GraphPartitioner::Mode::kNormal, + epctx::ModelGenOptions{}, + debug_graph_fn)); + + ASSERT_TRUE(observed_consumed.has_value()) + << "Expected the surviving Conv partition to be assigned to the EP."; + // Conv survived: committed exactly once. + EXPECT_EQ(*observed_consumed, expected_conv_cost) + << "Survivor Conv should be committed exactly once."; + // LogSoftmax was dropped on the second pass: its cost must not leak (no phantom budget). + EXPECT_NE(*observed_consumed, expected_conv_cost + expected_log_softmax_cost) + << "Dropped LogSoftmax must not consume budget."; +} #endif // !defined(DISABLE_CONTRIB_OPS) // Infrastructure that was used to check NNAPI coverage. diff --git a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py index 99e669f73eb72..3cb1392919398 100644 --- a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py +++ b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py @@ -910,6 +910,34 @@ def expected_fn(feed): result = _run_nhwc_model_test(target_device, "GridSample", model, {"X": x, "grid": grid}, expected_fn) self.assertEqual(result, TEST_PASS, "GridSample (NHWC) plugin test failed") + def test_nhwc_conv_with_resource_accounting(self): + # Smoke test for the NHWC two-pass partitioning flow combined with the resource + # accountant (session.resource_cuda_partitioning_settings). The NHWC layout + # transform makes the CUDA EP claim Conv nodes tentatively on the first pass; the + # budget for surviving nodes is committed only after the second pass. This guards + # against regressions where dropped first-pass tags would leak phantom budget. With + # a large limit, the Conv must still be claimed by the plugin and run correctly. + target_device = get_cuda_plugin_device() + inputs = { + "X": np.random.rand(1, 2, 4, 4).astype(np.float32), + "W": np.random.rand(3, 2, 3, 3).astype(np.float32), + } + # Large ad-hoc memory limit (1 GB in KB) with no stats file (trailing comma) so the + # Conv comfortably fits and remains assigned to the plugin EP. + session_config = { + **_NHWC_CONFIG, + "session.resource_cuda_partitioning_settings": "1048576,", + } + result = run_operator_test( + target_device, + create_conv_model, + inputs, + _expected_conv, + session_config=session_config, + nhwc_ops={"Conv"}, + ) + self.assertTrue(result, "Conv (NHWC + resource accounting) plugin test failed") + # ---- Standard op tests ---- def test_op_reshape(self): From f17b18319925ab31356755693f11d64f864c4246 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 4 Jun 2026 13:59:46 -0700 Subject: [PATCH 6/6] address review feedback: enable NHWC tests in contrib-disabled builds; document weight-commit trade-off - Remove DISABLE_CONTRIB_OPS guards around the NHWC two-pass partitioning test EPs and tests; they only use ONNX-domain ops (Conv, LogSoftmax) and core framework utilities, so they now provide regression coverage in contrib-disabled builds. - Document why the deferred survivor commit only adjusts AddConsumedAmount and intentionally does not replay CommitWeightsForNode (pending weight state is discarded by ResetForNewPass before pass 2; leaving weights uncommitted is the conservative/safe over-count direction). --- onnxruntime/core/framework/graph_partitioner.cc | 8 ++++++++ .../internal_testing_partitioning_tests.cc | 7 +++---- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 8d5f6192b2a5c..e66c16b5b195d 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -439,6 +439,14 @@ static Status GetCapabilityForEP(const GetCapabilityForEPParams& params, const l // leak phantom budget into later accounting decisions. New nodes introduced for the // second pass (e.g. NHWC ops) carry their own costs and are accounted normally when // their partitions are placed, so they are intentionally excluded here. + // + // Only the consumed total is adjusted here (AddConsumedAmount); the per-node initializer + // weight tracking (CommitWeightsForNode) is intentionally not replayed. The pending weight + // state computed in pass 1 is discarded by ResetForNewPass before pass 2 and cannot be + // committed for survivors without re-probing, which pass 2 does not do for already-tagged + // nodes. Leaving those weights uncommitted is the safe direction: in ad-hoc accounting mode + // a shared initializer may be re-counted in a later partitioning iteration (a conservative + // over-estimate) but is never under-counted, so the configured budget can never be exceeded. if (params.resource_accountant != nullptr) { for (NodeIndex node_index : nodes_temporarily_assigned_to_ep) { if (pass2_node_indices.count(node_index) == 0) { diff --git a/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc b/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc index 76f2e7fac8630..9ee38987b2cd2 100644 --- a/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc +++ b/onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc @@ -48,7 +48,9 @@ using namespace onnxruntime::internal_testing_ep; // it would be possible to use ORT format models but the same partitioning code would run either way #if !defined(ORT_MINIMAL_BUILD) -#if !defined(DISABLE_CONTRIB_OPS) +// These NHWC two-pass partitioning helpers and tests only use ONNX-domain ops +// (Conv, LogSoftmax) and core framework utilities, so they are intentionally not +// guarded by DISABLE_CONTRIB_OPS and provide regression coverage in contrib-disabled builds. namespace { class TwoPassNhwcTestExecutionProvider : public IExecutionProvider { @@ -251,7 +253,6 @@ class AccountingNhwcTestExecutionProvider : public IExecutionProvider { }; } // namespace -#endif // !defined(DISABLE_CONTRIB_OPS) #define ORT_MODEL_FOLDER ORT_TSTR("testdata/") @@ -350,7 +351,6 @@ TEST(InternalTestingEP, TestDependenciesCorrectlyHandled) { ASSERT_EQ(num_other_nodes, 2); } -#if !defined(DISABLE_CONTRIB_OPS) TEST(InternalTestingEP, NhwcSecondPassDropFallsBackFromCpuKernelNode) { std::unordered_map domain_to_version{{kOnnxDomain, 13}, {kMSDomain, 1}}; Model model("NhwcSecondPassDropFallsBackFromCpuKernelNode", @@ -535,7 +535,6 @@ TEST(InternalTestingEP, NhwcTwoPassAccountingCommitsOnlySurvivors) { EXPECT_NE(*observed_consumed, expected_conv_cost + expected_log_softmax_cost) << "Dropped LogSoftmax must not consume budget."; } -#endif // !defined(DISABLE_CONTRIB_OPS) // Infrastructure that was used to check NNAPI coverage. // Ideally this could be updated to read the model paths, supported ops and stop ops from input files