PyTorch fixes AOTInductor crash on incompatible tile fusion
A missing exception type in `_split_iteration_ranges` caused hard compile failures—now fixed.
PyTorch merged a fix for a bug in `_split_iteration_ranges` within `torch/_inductor/codegen/simd.py` that caused AOTInductor to crash during compilation when attempting to fuse a pointwise epilogue with incompatible tiling groups. The function is responsible for determining whether a pointwise kernel can be fused into an existing template (e.g., matmul) tile by splitting iteration ranges. Under certain conditions—specifically when the epilogue's iteration domain is a strict sub-multiple of the template's tiling—the function would hit a bare `assert all(... == 1 for s in remaining)` at the end. This assertion failed with an `AssertionError` instead of raising `CantSplit`, the exception used elsewhere in the same function to signal incompatible tiling. Crucially, both callers driving epilogue fusion (`SIMDKernel.is_compatible` and `Scheduler.speedup_by_fusion`) wrap the `_split_iteration_ranges` call in `try/except CantSplit` blocks to gracefully fall back to unfused codegen. Because the final assert threw `AssertionError`, it escaped those handlers and caused AOTInductor to hard-fail with `InductorError: AssertionError: failed to set ranges ...`.
The fix converts that final assert to `raise CantSplit(remaining, lengths)`, matching the exception contract of the rest of the function. The success path is byte-for-byte unchanged—models that compiled correctly today remain unaffected. Models that previously crashed now compile successfully, with the incompatible epilogue left unfused as a separate kernel instead of aborting the compile entirely. The fix includes a new regression test, `test_leftover_extent_raises_cant_split`, which constructs groups `[2, 2]` and lengths `[[2], []]`—a scenario where all sizes divide cleanly but group 1 is left with extent 2—and asserts that `CantSplit` is raised rather than `AssertionError`. This patch improves AOTInductor reliability by ensuring that tiling-incompatible fusions are skipped gracefully rather than crashing the compiler.
- Bug: `_split_iteration_ranges` ended with `assert all(... == 1 for s in remaining)` that fired when a pointwise epilogue had iteration domain sub-multiple of template tiling.
- Impact: Caused `InductorError: AssertionError: failed to set ranges` during AOTInductor compile, hard-failing compilation instead of gracefully falling back to unfused codegen.
- Fix: Changed assert to `raise CantSplit(remaining, lengths)`, aligning with exception contract of the function so existing safety nets catch it.
Why It Matters
PyTorch users can now compile models with incompatible epilogue fusions without crashes, improving AOTInductor reliability.