[RFC][MLIR] Enable Dynamic and Tighter Affine Unrolling via ValueBoundsConstraintSet

MLIR
2

Background:

  • When the lower and upper bounds of a loop are symbolic expressions, existing affine infrastructure struggles to support them. Even though we can infer their exact ranges using the Presburger library, these bounds are ultimately not compile-time constants, which prevents current affine transformation layers from processing them effectively.
  • The above scenario serves as an excellent example

Proposed Changes / Core Idea

The core idea of this proposal is to introduce a transformation that rewrites the loop structure to ensure that either the lower bound or the upper bound of the affine loop becomes a strict compile-time constant.

This transformation effectively bridges the gap between constrained symbolic bounds and the rigid constant requirements of the existing affine infrastructure.

Detailed Design: Loop Peeling and Normalization via Range Analysis

The proposed transformation leverages the value ranges inferred by the Presburger library to split a loop with symbolic bounds into two parts: a highly optimized fully-constant main loop and a residual tail loop for edge cases.

Mathematical Formulation

Given an affine loop with symbolic lower bound LB and upper bound UB, and a constant stride S:

  • The Presburger solver provides local bounds such that LB∈[LBmin​,LBmax​] and UB∈[UBmin​,UBmax​].

  • We compute a conservative minimal trip count TCsafe​ that is guaranteed to execute safely without exceeding the actual runtime upper bound.

  • The iteration space covered by the main loop is then defined by TCsafe​. The remaining iterations are handed over to the tail loop starting at LBnew​=LB+TCsafe​×S.

Concrete Example

Consider the following pseudo-IR scenario where the bounds are symbols but have known Presburger ranges:

MLIR

// Original Loop
// Presburger Analysis: %bound  ∈ [1, 3]
//                     %bound1 ∈ [8, 10]
%bound  = ...
%bound1 = ...
for %iv = %bound to %bound1 step 2 {
  use(%iv)
}

The transformation replaces the single dynamic loop with a static-trip-count core loop followed by a tail loop:

MLIR

// Transformed IR
%bound  = ... // Range [1, 3]
%bound1 = ... // Range [8, 10]

// 1. Normalized Main Loop (Exposes constant bounds to Affine infrastructure)
for %iv_norm = 0 to 3 step 1 {
  // Induction Variable Rematerialization
  %iv_reconstructed = %bound + %iv_norm * 2
  use(%iv_reconstructed)
}

// 2. Tail Handover
// Calculates the residual lower bound based on the constant steps already taken
%bound_tail = %bound + 6 // 6 = 3 steps * step_size 2 Range [7, 9]
for %iv_tail = %bound_tail to %bound1 step 2 {
  use(%iv_tail)
}

Benefits to Affine Infrastructure

By presenting a loop bounded strictly from 0 to 3 step 1, the existing Affine optimization pipeline can immediately apply loop unrolling, polyhedral fusion, or SIMD vectorization to the main loop body without any modifications to its core logic.

Design Discussion: Pass vs. Pattern

I am not entirely sure whether this feature should be implemented as a dedicated pass or a rewrite pattern. Currently, I lean towards implementing it as a pass, though I have not yet settled on an ideal name for it. I would like to finalize this architectural choice before diving into the implementation.

End

Thank you to everyone in the community for reviewing this proposal and providing feedback. I would also like to extend my gratitude to Gemini for assisting with the refinement, terminology polishing, and formatting of this RFC.:wink: cc: @krzysz00 @ftynse @bondhugula