Fix NHWC second-pass EP assignment cleanup

[tianleiwu]

Description

This change fixes two related bugs in the NHWC two-pass partitioning flow used by layout-preferring execution providers (QNN, CUDA with prefer_nhwc, etc.):

1. Stale EP ownership. Nodes claimed during the first GetCapability pass could remain assigned to the EP even when the second pass (run after layout transformation) no longer selected them. That stale ownership stranded CPU-kernel nodes on the EP, causing initialization failures such as Kernel not found or invalid NHWC graph state.
2. Phantom resource-accountant budget. When resource-aware partitioning is enabled, the tentative first-pass tags committed memory budget that was never rolled back for nodes dropped in the second pass. That phantom budget leaked into later accounting decisions and could incorrectly starve subsequent partitions.

The fix makes first-pass EP tags tentative: they exist only so the layout transformer and the second GetCapability pass can recognize candidate nodes, and they commit no budget. Budget is committed only for the nodes that survive the second pass.

Summary of Changes

Graph partitioning fix

File	Change
onnxruntime/core/framework/graph_partitioner.cc	Track nodes that were only temporarily assigned during the first NHWC capability pass and clear those assignments for nodes the second pass no longer claims.
onnxruntime/core/framework/graph_partitioner.cc	Clear temporary assignments on early exits from layout transformation failure and load cancellation so the graph is never left half-assigned.
onnxruntime/core/framework/graph_partitioner.cc	Skip resource-accountant accounting during the tentative first pass (TryAssignNodes); capture per-node first-pass costs before capabilities.clear(), then commit budget only for survivors via a deferred-commit step.
include/onnxruntime/core/graph/indexed_sub_graph.h	Add a read-only GetNodeCost(size_t) accessor so the deferred-commit step can reuse the costs computed in the first pass.

Regression coverage

File	Change
onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc	NhwcSecondPassDropFallsBackFromCpuKernelNode: custom NHWC test EP claims Conv and LogSoftmax on pass 1, drops LogSoftmax on pass 2, and verifies the dropped node falls back to a CPU kernel instead of staying owned by the EP.
onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc	NhwcTwoPassAccountingCommitsOnlySurvivors: accounting-aware NHWC test EP (registered as kCudaExecutionProvider) verifies budget is committed once for the surviving Conv and never for the dropped LogSoftmax (no phantom budget, no double-count).
onnxruntime/test/python/transformers/test_cuda_plugin_ep.py	test_nhwc_conv_with_resource_accounting: CUDA plugin-EP smoke test combining NHWC opt-in with session.resource_cuda_partitioning_settings.

Documentation

File	Change
docs/annotated_partitioning/PartitioningWithAnnotationsAndMemoryConstraints.md	Describe how tentative first-pass NHWC tags interact with resource budgeting (budget committed only for second-pass survivors).
docs/cuda_plugin_ep/cuda_plugin_ep_design.md	Note that first-pass tags are tentative and that plugin EPs should attach costs only on first-pass capabilities, mirroring the in-tree CUDA EP.

NHWC opt-in note

NHWC-preferred layout is opt-in:

• In-tree CUDA EP: provider option {"prefer_nhwc": "1"} (alias prefer_nhwc_layout).
• CUDA plugin EP: session option {"ep.cuda.prefer_nhwc_layout": "1"}.

Testing

Build the unit test target, then run:

./onnxruntime_test_all --gtest_filter="InternalTestingEP.Nhwc*:ResourceAccountantTest.*:SessionStateTest.TestResourceAwarePartitioning*"

Before this fix, the second-pass drop scenario failed with:

[ONNXRuntimeError] : 9 : NOT_IMPLEMENTED : Failed to find kernel for LogSoftmax(13) (node:'node_token_0' ep:'TwoPassNhwcTestExecutionProvider'). Kernel not found

Verification performed:

• CPU build: target builds, lint clean, targeted partitioning/accounting tests pass.
• CUDA 13 build (NHWC ops enabled): builds clean; the filtered suite reports 11 tests passed, including the new NhwcTwoPassAccountingCommitsOnlySurvivors and the CUDA-only SessionStateTest.TestResourceAwarePartitioning* cases.

Motivation and Context

This was found during analysis of a "Windows x64 QNN CI Pipeline (static_lib)" test failure.

QNN and other NHWC-preferred EPs use a two-pass partitioning flow around layout transformation. The first pass marks candidate nodes so layout transformation knows what to rewrite, but those assignments were not being rolled back before the second capability query. If the second pass rejected a previously claimed node, the node could remain incorrectly owned by the EP even though it was no longer part of a compiled partition. With resource-aware partitioning enabled, the dropped node's budget was also leaked.

The new regression tests cover both control-flow paths without depending on end-to-end QNN behavior, keeping the failure focused on GraphPartitioner ownership and accounting semantics rather than backend-specific compilation details.

[Copilot]

Pull request overview

This PR addresses a bug in NHWC two-pass EP partitioning where nodes tagged/claimed in the first GetCapability pass could remain assigned to the EP even if they are dropped in the second pass after layout transformation, potentially causing initialization failures (e.g., “Kernel not found”).

Changes:

• Update NHWC partitioning flow to track and clear temporary first-pass EP assignments around layout transformation / second capability query.
• Add a deterministic regression test with a custom two-pass NHWC EP that drops a previously claimed node on the second pass and verifies CPU fallback behavior.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
onnxruntime/core/framework/graph_partitioner.cc	Adds tracking of temporarily assigned nodes and clears EP assignments around NHWC layout transformation + second capability query.
onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc	Adds a new regression test EP + test to ensure nodes dropped on NHWC pass 2 don’t remain incorrectly assigned.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

[yuslepukhin]

The accounting inconsistency is real and the PR does change the situation, even though the direct commit/no-rollback behavior is pre-existing:

Before PR: node committed to accountant + node tagged to EP + EP has no kernel → Initialize fails with "Kernel not found". Phantom accounting never matters because the session never gets further.

After PR: node committed to accountant + node untagged + CPU picks it up → Initialize succeeds. The phantom accounting is now observable on subsequent accounting decisions in the same partitioning operation (next subgraph, control-flow body, or another EP using a shared accountant view).

So the PR doesn't introduce the phantom commit, but it does promote it from "irrelevant because we crash" to "live state that influences subsequent budget decisions". That's a behavioral change in the accounting feature, even if it's a strict improvement on the failure mode.

Currently the only in-tree EP that wires up IResourceAccountant is CUDA, which is not an NHWC-preferred-layout EP — so the buggy intersection (acc_enabled == true and preferred layout NHWC and pass-2 drop) doesn't exist in production today. This is why the bug hasn't surfaced. But:

The fix is built specifically into the NHWC path.
Plugin EPs and future NHWC EPs that opt into accounting (e.g., the CUDA-plugin design referenced in cuda_plugin_ep_design.md) would immediately expose it.
A graph with multiple subgraphs already shares a per-EP accountant across subgraph partitioning calls, so the leak compounds.

I think the right move would be:
Block on accounting: have TryAssignNodes skip AccountForNode when called from the NHWC pass-1 (i.e., when the out-vector is provided), and instead only commit accounting on nodes that survive pass 2. This is cleaner conceptually — the pass-1 tag really is tentative, and tentative tags shouldn't commit budget. It also matches the existing ResetForNewPass philosophy of "don't commit until confirmed". Implementation would move the AccountForNode call out of TryAssignNodes and into a deferred commit after the pass-2 drop loop.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

[tianleiwu]

@yuslepukhin thanks for the detailed analysis — I agree, and I've implemented the suggested approach (commit 6080aeb, with a follow-up doc/comment refinement in f17b183):

• TryAssignNodes now skips AccountForNode during the NHWC pass-1 tentative tagging (i.e. when the out-vector is provided), so tentative tags no longer commit any budget. The pass-1 tag really is tentative, matching the existing ResetForNewPass "don't commit until confirmed" philosophy.
• The per-node pass-1 costs are captured before capabilities.clear(), and budget is committed (AddConsumedAmount) only for nodes that survive pass 2, in a deferred-commit loop after the pass-2 drop loop. Nodes dropped in pass 2 therefore never leak phantom budget.
• New deterministic regression test NhwcTwoPassAccountingCommitsOnlySurvivors wires a SizeBasedStatsAccountant under kCudaExecutionProvider (mirroring the in-tree CUDA EP), drives GraphPartitioner::Partition directly, and asserts the surviving Conv is committed exactly once while the dropped LogSoftmax contributes no budget.

Note on CommitWeightsForNode: the deferred commit only adjusts the consumed total and intentionally does not replay weight commits — the pass-1 pending-weight state is discarded by ResetForNewPass before pass 2, and leaving weights uncommitted is the conservative/safe (over-count, never under-count) direction. This is now documented in a code comment.

NHWC opt-in is prefer_nhwc=1 for the in-tree CUDA EP and ep.cuda.prefer_nhwc_layout=1 for the CUDA plugin EP; the partitioning and plugin-EP design docs were updated accordingly.

Verified on a CUDA build with NHWC ops enabled: the partitioning + resource-accountant suite (including the new test) passes. Could you take another look when you get a chance?