[AIRADSW-735] ONNX control flow on CPU with MIGraphX branch execution on GPU

[urpetkov-amd]

Run ONNX control flow on CPU with MIGraphX branch execution on GPU

Summary

This PR adds an opt-in cpu_control_flow mode to the AMDGPU/MIGraphX execution provider.

When enabled, ONNX Runtime executes control-flow operators such as If on the CPU while MIGraphX compiles and executes the branch bodies on the GPU. This follows the same high-level execution model used by DirectML: host-side control-flow orchestration with GPU-accelerated subgraphs.

The original MIGraphX behavior remains the default. With cpu_control_flow=0, MIGraphX continues to compile the control-flow operator and its nested subgraphs as one fused program.

The change also adds the graph serialization and input-binding support required to compile control-flow branches as standalone MIGraphX programs without breaking the legacy fused path.

Motivation

The lmx-v1-fp32-512x512 model contains an If whose condition depends on a runtime scalar value (alpha). Under the original MIGraphX path, the If is lowered into a MIGraphX control-flow operation.

Although the branch computation runs on the GPU, selecting the branch requires MIGraphX to:

1. Copy the condition from GPU memory to host memory.
2. Synchronize the HIP stream.
3. Read the condition on the host.
4. Select and execute the appropriate GPU branch.

The generated program contains operations equivalent to hip::copy_from_gpu and hip::sync_stream around the condition. This introduces a blocking GPU-to-CPU synchronization in the middle of inference. Investigation of the lmx model showed this synchronization dominating execution time, with roughly 6 ms attributable to the condition transfer and stream wait.

Execution model

Default mode: cpu_control_flow=0

• MIGraphX claims the parent graph containing the If.
• Nested branch graphs are preserved as subgraphs of the fused MIGraphX program.
• MIGraphX lowers and executes the If.
• The runtime condition is copied to the host internally by MIGraphX.
• This preserves the existing behavior and does not require CPU fallback when MIGraphX supports the complete graph.

New mode: cpu_control_flow=1

• ONNX Runtime leaves the control-flow boundary on the CPU.
• If, Loop, and Scan are excluded from the top-level MIGraphX partition.
• Greater and Cast are also left on the CPU by the current partitioning policy so that a generated boolean condition remains CPU-resident.
• ONNX Runtime evaluates the condition and invokes the selected branch.
• MIGraphX independently compiles and executes the branch subgraphs on the GPU.
• CPU-resident values consumed by a GPU branch are uploaded to GPU staging buffers before run_async.

For the tested If model, this changes the execution flow from:

MIGraphX fused graph
  -> GPU condition
  -> GPU-to-CPU copy
  -> stream synchronization
  -> host branch selection
  -> GPU branch

to:

ONNX Runtime CPU control flow
  -> CPU condition and branch selection
  -> selected MIGraphX GPU branch

Implementation

Control-flow-aware partitioning

MIGraphXExecutionProvider::GetCapability() now distinguishes between the top-level graph and a graph owned by a control-flow node.

When CPU control flow is disabled:

• MIGraphX does not independently claim nested control-flow branch graphs.
• The parent control-flow operator is serialized and compiled with its original nested branches.

When CPU control flow is enabled:

• Control-flow boundary nodes are excluded from the top-level MIGraphX partition.
• MIGraphX is allowed to claim and compile nodes inside branch subgraphs.
• Supported branch nodes are fused because MIGraphX is a compile-only EP and requires fused nodes to flow through Compile().

This produces separate MIGraphX programs for the selected branch bodies while ONNX Runtime retains ownership of control-flow execution.

Standalone branch input serialization

ONNX control-flow branch graphs can consume values from their parent graph without declaring those values as normal branch inputs. ONNX Runtime exposes these values as implicit or outer-scope inputs.

That representation works when the branch remains nested under an If, but it is insufficient when the branch is extracted and serialized as a standalone model for MIGraphX compilation. A standalone model must declare every non-initializer external value as a graph input.

GraphToProtoImpl() now promotes outer-scope values to graph inputs only for the root graph being serialized. Recursive serialization of nested control-flow graphs disables promotion.

This distinction is required for both modes:

• Standalone branch compilation needs promoted outer-scope inputs.
• Legacy fused If serialization must preserve outer-scope captures.

Complete input-name binding

The previous input mapping only covered the most direct fused-node inputs. Standalone branch programs can also depend on graph inputs and implicit inputs captured from the parent graph.

BuildInputNameIndices() constructs the kernel-context input mapping in this order:

1. Fused-node boundary inputs.
2. Graph inputs, including promoted branch feeds.
3. Fused-node implicit inputs.

GetGpuInputData() now:

• Uses the original pointer when an input is already in GPU memory.
• Allocates or reuses a HIP device buffer for a CPU input.
• Copies the CPU tensor to the device on the active HIP stream.
• Returns the device pointer used to construct the MIGraphX argument.

Stream and HIP graph behavior

CPU control flow creates synchronization and ownership boundaries between ONNX Runtime's CPU execution and independently compiled MIGraphX branches.

To preserve correctness:

• MIGraphX reports itself as not stream-aware when CPU control flow is enabled.
• HIP graph execution is disabled for sessions using CPU control flow.
• CPU input uploads use the branch's active HIP stream.
• Branch execution is synchronized before returning control to ONNX Runtime.

Result

This PR provides an opt-in execution strategy for ONNX control flow that avoids MIGraphX's expensive mid-inference GPU-to-host condition synchronization.

The final behavior is:

• Control-flow decisions execute on the CPU.
• Selected branch bodies continue to execute on the GPU through MIGraphX.
• Runtime parent-scope values are correctly exposed and bound to standalone branch programs.
• CPU branch inputs are safely staged to GPU memory.
• The original fused MIGraphX behavior remains available and functional.