Skip to content

trunk/271250efec976ae0138286284100b157763b2dc5

@eellison eellison tagged this 27 Jul 14:47
Index inversion only analyzes a single logical loop variable. Equal-numel layout consumers can retain a higher-rank iteration domain, causing compatible producer-to-clone paths such as attention transposes and reduced scale swizzles to miss fusion before inversion is attempted.

Speculatively flatten the pointwise consumer to the producer write domain and retry the existing inversion analysis. Require the flattened consumer write to be contiguous, and restore the original loop state when inversion fails. This keeps the transformation generic instead of adding nested-reduction-specific fusion behavior or relying on fusion ordering.

The regression coverage includes the attention layout clone from #188635, failed speculative reindex rollback, and the nested RMSNorm block-scale swizzle integration.

Test Plan:

```
python test/inductor/test_loop_ordering.py -k test_reindex_attention_layout_clone_for_index_inversion
python test/inductor/test_loop_ordering.py -k test_failed_reindex_for_index_inversion_rollback
python test/inductor/test_loop_ordering.py -k test_rejected_index_inversion_rollback
python test/inductor/test_nested_reduction.py -k rmsnorm_block_scale_swizzle
env -u HTTPS_PROXY -u https_proxy -u HTTP_PROXY -u http_proxy lintrunner -a
```

Authored with an AI assistant.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/190402
Approved by: https://github.com/drisspg
Assets 2
Loading