fix(partitioning_utils): include Loop/If/Scan implicit inputs in MetaDef (#28608)
## Summary
`utils::MakeComputeCapability` is the shared helper used by
`utils::CreateSupportedPartitions` to build an
`IndexedSubGraph::MetaDef` from a group of supported nodes. When a
supported group contains a control-flow op (`Loop`, `If`, `Scan`),
`MakeComputeCapability` currently walks only `node->InputDefs()` and
silently drops the outer-scope captures (`node->ImplicitInputDefs()`).
The captures never enter `meta_def->inputs`, so after
`Graph::FinalizeFuseSubGraph` the fused node's `InputDefs()` is missing
them — the EP that owns the fused subgraph has no boundary value-info
for the captured tensors and cannot bind them at Compute time.
This PR adds a second loop in `MakeComputeCapability` that walks
`node->ImplicitInputDefs()` with the same "produced inside the partition
→ skip, otherwise add to subgraph inputs" semantics already applied to
`InputDefs()`.
## Why this is the right fix
`onnxruntime::Node` partitions inputs into two arrays by design:
- `InputDefs()` — formal operand list as declared in the op's ONNX
schema.
- `ImplicitInputDefs()` — outer-scope SSA values referenced from inside
body subgraphs of `Loop` / `If` / `Scan`. These are real boundary inputs
at runtime (the body kernel reads them) but they don't appear in the
op's formal operand list.
`Graph::FinalizeFuseSubGraph` consumes only `meta_def->inputs` to
populate the fused node's `InputDefs()` and rewire outer-scope edges. So
whatever `MakeComputeCapability` puts in `meta_def->inputs` is what ends
up at the fused-node boundary. Omitting `ImplicitInputDefs()` here means
the captures are unreachable downstream — there is no other place that
can patch them back in.
The fix is intentionally a mirror of the existing `InputDefs()` loop
(same `Contains(node_outputs, ...)` produced-inside check, same
`ordered_subgraph_inputs.push_back` ordering). The new loop runs after
the explicit loop so explicit-operand index ordering in
`meta_def->inputs` is preserved (EPs that have implicitly relied on
`meta_def->inputs[i].name == node.InputDefs()[i].name` for
non-control-flow op groups are not perturbed).
## Scope of impact
Only EPs that consume `utils::MakeComputeCapability` /
`utils::CreateSupportedPartitions` are affected. A quick audit:
| EP | Uses `partitioning_utils::MakeComputeCapability`? | Affected by
bug? |
|---|---|---|
| Plugin EPs (`EpGraphSupportInfo_AddNodesToFuse` →
`PluginExecutionProvider::GetCapability`) | yes, in
`onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc` |
**yes** |
| `internal_testing_ep` (used by ORT's own unit tests) | yes, in
`onnxruntime/test/internal_testing_ep/internal_testing_execution_provider.cc`
| **yes** |
| TensorRT, MIGraphX, NV-TRT-RTX, VitisAI | no — they build
`MetaDef::inputs` themselves and already walk `ImplicitInputDefs()`
(e.g. `tensorrt_execution_provider.cc:2084`,
`migraphx_execution_provider.cc:735`) | no |
| DML / CPU / CUDA / ROCm / OpenVINO / QNN / CANN / WebGPU / CoreML |
don't use it for Loop/If/Scan fusion paths | no |
So the impact is bounded to the plugin EP architecture (ORT 1.23+) and
the in-tree testing EP — both of which delegate boundary calculation to
this shared helper.
## Reproduction
The bug is reproducible against this repo's `internal_testing_ep`. No
external code required.
A minimal repro model with a Loop body that captures an outer-scope
tensor `B`:
```python
# build_repro.py — produces a ~1.5 KB onnx
import numpy as np, onnx
from onnx import TensorProto, helper as h, numpy_helper as nph
A = h.make_tensor_value_info("A", TensorProto.FLOAT, ["N", 2, 2])
B = h.make_tensor_value_info("B", TensorProto.FLOAT, [2, 2])
out = h.make_tensor_value_info("v_final", TensorProto.FLOAT, [2, 2])
acc_init = nph.from_array(np.zeros((2, 2), np.float32), name="acc_init")
cond_init = nph.from_array(np.array([1], np.bool_), name="cond_init")
sq_ax = nph.from_array(np.array([0], np.int64), name="sq_ax")
body = h.make_graph(
nodes=[
h.make_node("Gather", ["A", "iter"], ["slice"], axis=0),
h.make_node("Add", ["slice", "B"], ["tmp"]), # captures outer B
h.make_node("Add", ["acc_in", "tmp"], ["acc_out"]),
h.make_node("Identity", ["cond_in"], ["cond_out"]),
],
name="loop_body",
inputs=[h.make_tensor_value_info("iter", TensorProto.INT64, []),
h.make_tensor_value_info("cond_in", TensorProto.BOOL, []),
h.make_tensor_value_info("acc_in", TensorProto.FLOAT, [2, 2])],
outputs=[h.make_tensor_value_info("cond_out", TensorProto.BOOL, []),
h.make_tensor_value_info("acc_out", TensorProto.FLOAT, [2, 2])],
)
g = h.make_graph(
nodes=[
h.make_node("Shape", ["A"], ["M_1d"], start=0, end=1),
h.make_node("Squeeze", ["M_1d", "sq_ax"], ["M"]),
h.make_node("Loop", ["M", "cond_init", "acc_init"], ["v_final"], body=body),
],
name="loop_with_outer_capture",
inputs=[A, B], outputs=[out],
initializer=[acc_init, cond_init, sq_ax],
)
onnx.save(h.make_model(g, opset_imports=[h.make_opsetid("", 16)]),
"loop_with_outer_capture.onnx")
```
Observable bug path (against any EP using `CreateSupportedPartitions`,
e.g. `InternalTestingExecutionProvider`):
```cpp
// Claim every node (Shape/Squeeze/Constant/Loop) as compiled.
SessionOptions so;
InferenceSession session(so, env);
session.RegisterExecutionProvider(
std::make_unique<InternalTestingExecutionProvider>(/*supported=*/{...}));
session.Load("loop_with_outer_capture.onnx");
session.Initialize();
// In EP::Compile, iterate fused_node.InputDefs():
// for (const auto* in : fused_node.InputDefs()) std::cerr << in->Name() << "\n";
// BEFORE this fix: only "A" is printed (Shape(A) makes A explicit;
// B is consumed only via Loop's ImplicitInputDefs and gets dropped).
// AFTER this fix: both "A" and "B" are printed.
```
A small unit-test fixture exercising the same path can be added to
`onnxruntime/test/providers/partitioning_utils_test.cc` following the
existing `CheckAllNodesProcessed` pattern, asserting that
`result[0]->sub_graph->GetMetaDef()->inputs` contains `B` when the
supported group includes the Loop.
## What this PR changes
A single hunk in
`onnxruntime/core/providers/partitioning_utils.cc::MakeComputeCapability`,
immediately after the existing `for (const auto* input :
node->InputDefs()) { ... }`:
```cpp
// Region-bearing ops (Loop/If/Scan) reference outer-scope SSA values via
// ImplicitInputDefs rather than InputDefs. When an EP claims the whole
// control-flow op, those implicit captures must also be in MetaDef::inputs
// so FinalizeFuseSubGraph can rewire the outer-scope edges onto the fused
// node's InputDefs. Without this, plugin EPs that fuse Loop/If/Scan lose
// the captures at the fused-node boundary and cannot resolve them at
// Compute time.
for (const auto* input : node->ImplicitInputDefs()) {
if (!input->Exists()) {
continue;
}
if (!Contains(node_outputs, input)) {
if (!Contains(subgraph_inputs, input)) {
subgraph_inputs.insert(input);
ordered_subgraph_inputs.push_back(input);
}
}
}
```
## Risks / migration
- **No ABI change.** `MakeComputeCapability` signature unchanged.
`IndexedSubGraph::MetaDef` schema unchanged.
- **No semantic regression for op groups without control flow.** The new
loop only adds elements; for partitions that contain no `Loop` / `If` /
`Scan`, `ImplicitInputDefs()` is empty on every node and the new loop is
a no-op.
- **Behavior change for plugin EPs that fuse Loop/If/Scan.** Their fused
node's `InputDefs()` gains the captures. EPs that were silently fishing
out captures via a workaround (e.g. walking the original Loop node's
`ImplicitInputDefs()` themselves at Compile time) would see those names
show up via the standard fused-node `InputDefs()` API. Audit above shows
no in-tree EP that uses `partitioning_utils` had such a workaround — TRT
/ MIGraphX / etc. roll their own MetaDef without calling
`MakeComputeCapability`.
## Validation
- Verified the fix end-to-end against a downstream plugin EP that claims
a `Loop` node as part of a fused partition (Loop body captures an
outer-scope tensor): without this fix, the EP cannot resolve the
captured tensor name at the fused-node boundary; with the fix the
captured tensor appears in `fused_node.InputDefs()` and session
initialization + the EP's Compile both succeed.
- No `partitioning_utils.cc` changes between `origin/main` and the patch
base, so it applies cleanly.
- Existing `onnxruntime_test_all --gtest_filter=PartitioningUtilsTest.*`
cases still pass (the fix only adds behavior for control-flow ops;
non-control-flow partitions are byte-for-byte identical to before).
onnxruntime
eefc22f2 - fix(partitioning_utils): include Loop/If/Scan implicit inputs in MetaDef (#28608)
