Fix stride verification to respect divisibility constraints

_verify_stride_on_original was rejecting valid stride decompositions
for B-data and B-scale reads because it probed with arbitrary values
that violated divisibility assumptions (e.g. K not a multiple of 256).
This made floor/Mod terms produce different results per IV step,
causing the verifier to think the stride was non-constant.

Pass div_fwd substitutions into the verifier so probed values always
satisfy the constraints.  This eliminates all residual affine.apply
ops for B-data and B-scale index offsets in the loop body (12 per
test for 1x4, 28 for 2x2).

The only remaining affine.apply in the loop body is a single
bounds-check computation unrelated to data/scale index offsets.

Made-with: Cursor

Author: willghatch

lit_tests/kernel/wave/linearize_loop_affine_maps.py

@@ -8,12 +8,11 @@
 index computation rather than monolithic ``affine.apply`` ops with many
 operands.
 
-Reads whose flat index is truly affine in the IV (A-data GatherToLDS)
-get their stride annotated and codegen emits ``base + IV * stride``.
-Reads with complex preshuffle mappings that contain floor/Mod terms
-interacting with the IV (B-data, B-scale reads) may still use
-``affine.apply`` because the expression is not decomposable as
-``base + IV * stride`` in the symbolic domain.
+All reads whose flat index involves the loop IV -- A-data, B-data,
+A-scale, and B-scale -- get their stride annotated and codegen emits
+``base + IV * stride``.  A single residual ``affine.apply`` may remain
+for the bounds-check computation, but all data/scale index offsets are
+hoisted.
 
 Two wave shapes are tested to cover different workgroup tiling patterns:
   (1, 4):  64 threads, 1 wave in M, 4 in N -- block (256, 192, 256)
@@ -62,17 +61,13 @@ def test_hoisted_indices_in_loop_1x4():
     #              Inside the pipelined scf.for, GatherToLDS reads should use
     #              hoisted base + IV * stride (arith.muli + arith.addi), not
     #              monolithic affine.apply with 6+ operands.
-    #              B-data reads (complex preshuffle) may still use affine.apply
-    #              because the expression is not decomposable as base + IV * stride.
     # CHECK:       scf.for
     # CHECK:       arith.muli
     # CHECK:       arith.addi
     # CHECK:       gather_to_lds
-    #              KEY CHECK -- do not remove.  The affine.apply count and
-    #              CHECK-NOT verify that hoisting is working: only the
-    #              residual B-data/B-scale preshuffle reads (with floor/Mod
-    #              IV interaction) remain as affine.apply in the loop body.
-    # CHECK-COUNT-12: affine.apply
+    #              KEY CHECK -- do not remove.  After gather_to_lds there
+    #              must be NO affine.apply ops; all data/scale index offsets
+    #              (A-data, B-data, A-scale, B-scale) are fully hoisted.
     # CHECK-NOT:   affine.apply
     # CHECK:       scf.yield
 
@@ -88,10 +83,8 @@ def test_hoisted_indices_in_loop_2x2():
     # CHECK:       arith.muli
     # CHECK:       arith.addi
     # CHECK:       gather_to_lds
-    #              KEY CHECK -- do not remove.  The affine.apply count and
-    #              CHECK-NOT verify that hoisting is working: only the
-    #              residual B-data/B-scale preshuffle reads (with floor/Mod
-    #              IV interaction) remain as affine.apply in the loop body.
-    # CHECK-COUNT-28: affine.apply
+    #              KEY CHECK -- do not remove.  After gather_to_lds there
+    #              must be NO affine.apply ops; all data/scale index offsets
+    #              are fully hoisted.
     # CHECK-NOT:   affine.apply
     # CHECK:       scf.yield

wave_lang/kernel/wave/analysis/annotate_iv_strides.py

@@ -168,23 +168,35 @@ def _eval2(iv_val):
     return base, sympy.Integer(stride_val)
 
 
-def _verify_stride_on_original(flat, iv_sym, step, base, stride):
+def _verify_stride_on_original(flat, iv_sym, step, base, stride, div_fwd=None):
     """Numerically verify ``flat(iv=step*i) == base + i*stride``.
 
     The divisibility-substitution path extracts the stride from a
     simplified expression but the base comes from the original.
     Floor/Mod interactions with the IV can invalidate this
     decomposition, so we probe several concrete IV values.
+
+    When *div_fwd* is given, the divisibility substitutions are applied
+    before probing so that the generated probe values satisfy the
+    constraints (e.g. K is always a multiple of 256).  Without this,
+    expressions with ``floor(IV * C / K)`` would get probed with
+    arbitrary K, giving wrong floor results and rejecting valid strides.
     """
-    free_syms = sorted(flat.free_symbols - {iv_sym}, key=str)
+    fwd_dict = dict(div_fwd) if div_fwd else {}
+
+    flat_p = sympy.sympify(flat).subs(fwd_dict) if fwd_dict else flat
+    base_p = sympy.sympify(base).subs(fwd_dict) if fwd_dict else base
+    stride_p = sympy.sympify(stride).subs(fwd_dict) if fwd_dict else stride
+
+    free_syms = sorted(flat_p.free_symbols - {iv_sym}, key=str)
     probe_map = {s: 137 + i * 31 for i, s in enumerate(free_syms)}
     probe_map2 = {s: 251 + i * 47 for i, s in enumerate(free_syms)}
 
     for pmap in (probe_map, probe_map2):
         for i in range(8):
             iv_val = step * i
-            orig = subs_idxc(flat.subs({iv_sym: iv_val})).subs(pmap)
-            expected = subs_idxc(base + i * stride).subs(pmap)
+            orig = subs_idxc(flat_p.subs({iv_sym: iv_val})).subs(pmap)
+            expected = subs_idxc(base_p + i * stride_p).subs(pmap)
             try:
                 if int(orig) != int(expected):
                     return False
@@ -234,7 +246,9 @@ def _try_with_div_subs(flat, iv_sym, step, div_fwd, div_bwd):
 
     base = safe_subs(flat, {iv_sym: sympy.Integer(0)})
 
-    if not _verify_stride_on_original(flat, iv_sym, step, base, stride):
+    if not _verify_stride_on_original(
+        flat, iv_sym, step, base, stride, div_fwd=div_fwd
+    ):
         return None
 
     return base, stride, method