RFC: Polyhedral Optimization Pass for PyTorch Inductor

compiler
1

RFC: Polyhedral Optimization Pass for PyTorch Inductor

This RFC proposes adding an optional polyhedral optimization pass to PyTorch Inductor to enable fusion of operations that current fusion heuristics cannot handle. The initial test case demonstrates ~1.4x speedup on RMSNorm + chunking + gating patterns (found in SwiGLU) used in architectures like Llama.

MOTIVATION

PyTorch Inductor currently relies on pattern matching and heuristic-based fusion strategies. While effective for many cases, these approaches fail to fuse certain operation sequences where fusion is mathematically valid and beneficial. For a case study demonstrating the speedup for SwiGLU, see [Example] Benefits of Polyhedral Optimization by morrison-turnansky · Pull Request #3 · morrison-turnansky/pytorch · GitHub

High-Level Design

Add a new opt-in compilation pass that users can enable via compilation flags. This optimization pass would target the loop level IR. We would use a subset of polyhedral optimization specifically targeting tensor workflows.

Scope Constraints (Initial Release)

  1. Inference Only: Training support deferred to avoid complexity with backward pass.

    1. Ultimately these techniques can be extended to support training and inference, but we want to defer the potential complexities of training such as synchronization issues.

  2. Subset of Polyhedral: Use a simplified version of polyhedral analysis focusing on identifying fusion opportunities and therefore minimizing overhead during compilation.

    1. We will focus on identifying fusion opportunities via loop level dependency analysis. Initially, we will follow a general heuristic that fusion will be profitable (see testing strategy).

    2. We can later expand this to use optimization functions standard in polyhedral analysis to rigorously determine if fusion will be profitable.

    3. Refer to reference for an example of lightweight polyhedral analysis.

  3. Opt-In: Default disabled to ensure zero impact on existing workflows.

Dynamic vs Static Shapes

A primary benefit of polyhedral analysis is that it handles dynamic shapes by design. Loop level dependency analysis depends on the rank of the tensor, not the specific shape.The end goal would be full dynamic shape support.

Breaking Changes

None. This is an additive feature:

  1. Default behavior unchanged

  2. Requires explicit opt-in via compilation flag

  3. Falls back to standard Inductor fusion when polyhedral analysis is unavailable or unprofitable

Testing Strategy

  1. There are two areas to test at a function level

    1. Does our pass result in the generated kernel being fused?

      1. This can be verified by a test similar to test_fusion_codegen in the example. This follows the pattern of tests in pytorch/test/inductor/test_loop_ordering.py.

    2. Numerical Accuracy

      1. This should be treated similarly to eager/inductor tests. Eager will be the source of truth.

  2. End to End Performance Analysis

    1. Verify a set of desired/common models actually see a performance increase. Monitor performance regressions.

We propose that for a set of graphs where we expect polyhedral to make substantial changes, i.e. Swiglu + others, we add a test verifying that fusion occurred and a numerical accuracy test.

References

  1. See simplified polyhedral form as a precedent for using a subset of polyhedral optimizations on tensor workflows: https://mlir.llvm.org/docs/Rationale/RationaleSimplifiedPolyhedralForm/

  2. For a high level background on polyhedral optimization and a discussion on its wide usage in compilers: http://polyhedral.info/

  3. For a detailed discussion on polyhedral see https://pliss2019.github.io/albert_cohen_slides.pdf