ciflow/trunk/186104: [Inductor] Allow output-input aliasing in torch.cond branches (#186104)

Summary:

Allow output-input aliasing in torch.cond branches.

torch.cond branches are mutually exclusive — only one branch executes at runtime. When the true branch returns an operand directly (e.g., identity/skip path), the output aliases the input. This is safe because the non-taken branch never runs, so there is no concurrent access to aliased buffers.

Previously, two checks rejected this pattern:

1. Dynamo (higher_order_ops.py): CondHigherOrderVariable.supports_aliasing = False rejected aliasing during tracing with a graph break.


2. Functionalization (cond.py): _check_alias_and_mutation rejected aliasing during AOT functionalization.

Note: Dynamo already allowed aliasing in inference (supports_aliasing = not torch.is_grad_enabled()), but torch.export traces with grad enabled even for inference models, so the check was ineffective for the AOTI export path.

Fix:

• Set supports_aliasing = True unconditionally on CondHigherOrderVariable. Aliasing is safe for cond regardless of grad context because branches are exclusive — even with autograd, only the taken branch's gradient flows.


• Skip the alias check in the functionalization impl (keep mutation check).

Motivation: PrismNet's _clone_layer_outputs uses lo.clone() to break output-input aliasing required by torch.cond. Removing this constraint saves ~0.046ms per inference (2 DHEN layers at B=2048).

Adds test_cond_output_input_aliasing regression test.

Test Plan:

buck2 test fbcode//mode/opt -c 'fbcode.nvcc_arch=a100,h100' fbcode//caffe2/test/inductor:control_flow -- --exact 'fbcode//caffe2/test/inductor:control_flow - test_cond_output_input_aliasing_device_cpu (caffe2.test.inductor.test_control_flow.CondTests)' 'fbcode//caffe2/test/inductor:control_flow - test_cond_output_input_aliasing_device_cuda (caffe2.test.inductor.test_control_flow.CondTests)' 'fbcode//caffe2/test/inductor:control_flow - test_cond_output_input_aliasing_multi_output_device_cpu (caffe2.test.inductor.test_control_flow.CondTests)' 'fbcode//caffe2/test/inductor:control_flow - test_cond_output_input_aliasing_multi_output_device_cuda (caffe2.test.inductor.test_control_flow.CondTests)' 'fbcode//caffe2/test/inductor:control_flow - test_cond_output_input_aliasing_both_branches_device_cpu (caffe2.test.inductor.test_control_flow.CondTests)' 'fbcode//caffe2/test/inductor:control_flow - test_cond_output_input_aliasing_both_branches_device_cuda (caffe2.test.inductor.test_control_flow.CondTests)' 'fbcode//caffe2/test/inductor:control_flow - test_cond_input_mutation_still_rejected (caffe2.test.inductor.test_control_flow.CondTests)'

Test session: https://www.internalfb.com/intern/testinfra/testrun/10414574313882076

PrismNet repro — all 11 tests pass, including CondOutputInputAlias: without clone which previously failed with "Replace return input with return input.clone()". Confirms lo.clone() workaround in _clone_layer_outputs can be removed.

buck2 run fbcode//mode/opt fbcode//aps_models/ads/gmp/models/ads_mtml_prism_net_dedicated_model/experimental/benchmarks:repro_aoti_inductor_bugs

Reviewed By: kqfu, atalman

Differential Revision: D106106380