Skip to content

Commit cef03e3

Browse files
willghatchwillghatch
committed
follow-ups
1 parent ab288b3 commit cef03e3

6 files changed

Lines changed: 110 additions & 91 deletions

File tree

‎lit_tests/kernel/wave/bounds_check.py‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -79,17 +79,16 @@ def read_write(
7979
# CHECK-LABEL: test_bounds_check
8080
# CHECK: func.func @read_write
8181
# CHECK: %[[C227:.*]] = arith.constant 227 : index
82-
# M row index (linearized), then N / vector mask setup
83-
# CHECK: %[[D2:.*]] = affine.apply #map{{.*}}()[%thread_id_x, %block_id_x]
84-
# CHECK: %[[D3:.*]] = arith.cmpi slt, %[[D2]], %[[C227]] : index
85-
# CHECK: vector.broadcast %[[D3]] : i1 to vector<16xi1>
86-
# CHECK: %[[D0:.*]] = affine.apply #map1{{.*}}()[%block_id_y]
87-
# CHECK: vector.broadcast %[[D0]] : index to vector<16xindex>
88-
# CHECK: arith.addi {{.*}}, {{.*}} overflow<nsw, nuw> : vector<16xindex>
82+
# CHECK: %[[D0:.*]] = affine.apply #map()[%block_id_y]
8983

9084
# N per-element mask (fastest dim gets iota)
9185
# CHECK: arith.cmpi slt, {{.*}}, {{.*}} : vector<16xindex>
9286

87+
# CHECK: %[[D2:.*]] = affine.apply #map{{.*}}()[%thread_id_x, %block_id_x]
88+
89+
# M scalar mask
90+
# CHECK: %[[D3:.*]] = arith.cmpi slt, %[[D2]], %[[C227]] : index
91+
9392
# Bounds check
9493
# CHECK: %[[D1:.*]] = arith.cmpi slt, %[[D0]], %[[C227]] : index
9594
# CHECK: %[[D4:.*]] = arith.cmpi sge, %[[D0]], %[[C227]] : index

‎lit_tests/kernel/wave/codegen.py‎

Lines changed: 19 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -808,30 +808,27 @@ def dynamic_copy(a: tkl.Memory[M, N, ADDRESS_SPACE, tkl.f16]):
808808
print(dynamic_copy.asm)
809809

810810
# CHECK-LABEL: test_dynamic_copy
811-
# CHECK-DAG: #[[dc_map_ceildiv:.*]] = affine_map<()[s0] -> (s0 ceildiv 16)>
812-
# CHECK-DAG: #[[dc_map_row:.*]] = affine_map<()[s0, s1] -> (s0 + s1 * 16 - (s0 floordiv 64) * 48)>
813-
# CHECK-DAG: #[[dc_map_colbase:.*]] = affine_map<()[s0] -> (s0 * 16)>
814-
# CHECK-DAG: #[[dc_map_linear:.*]] = affine_map<()[s0, s1, s2, s3] -> ((s0 * s1) * 16 + s0 * s2 + s3 * 16 - ((s2 floordiv 64) * s0) * 48)>
811+
# CHECK-DAG: #[[map1:.*]] = affine_map<()[s0] -> (s0 * 16)>
812+
# CHECK-DAG: #[[map2:.*]] = affine_map<()[s0, s1] -> (s0 + s1 * 16 - (s0 floordiv 64) * 48)>
815813
# CHECK: func.func @dynamic_copy
816-
# CHECK-SAME: (%[[ARG0:.*]]: !stream.binding, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
814+
# CHECH-SAME: (%[[ARG0:.*]]: !stream.binding, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
817815
# CHECK-DAG: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<16xf16>
818816
# CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : vector<16xindex>
819817
# CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
820818
# CHECK-DAG: %[[WORKGROUP_ID_0:.*]] = gpu.block_id x
821819
# CHECK-DAG: %[[WORKGROUP_ID_1:.*]] = gpu.block_id y
822820
# CHECK-DAG: %[[THREAD_ID_X:.*]] = gpu.thread_id x
823-
# CHECK: %[[D1:.*]] = affine.apply #[[dc_map_row]]()[%[[THREAD_ID_X]], %[[WORKGROUP_ID_0]]]
824-
# CHECK: %[[D2:.*]] = arith.cmpi slt, %[[D1]], %arg1 : index
825-
# CHECK: %[[D3:.*]] = vector.broadcast %[[D2]] : i1 to vector<16xi1>
826-
# CHECK: %[[D4:.*]] = affine.apply #[[dc_map_colbase]]()[%[[WORKGROUP_ID_1]]]
827-
# CHECK: %[[D5:.*]] = vector.broadcast %[[D4]] : index to vector<16xindex>
828-
# CHECK: %[[D6:.*]] = arith.addi %[[D5]], %[[CST_0]] overflow<nsw, nuw> : vector<16xindex>
829-
# CHECK: %[[D7:.*]] = vector.broadcast %arg2 : index to vector<16xindex>
830-
# CHECK: %[[D8:.*]] = arith.cmpi slt, %[[D6]], %[[D7]] : vector<16xindex>
831-
# CHECK: %[[D9:.*]] = arith.andi %[[D3]], %[[D8]] : vector<16xi1>
832-
# CHECK: %[[D10:.*]] = affine.apply #[[dc_map_linear]]()[%arg2, %[[WORKGROUP_ID_0]], %[[THREAD_ID_X]], %[[WORKGROUP_ID_1]]]
833-
# CHECK: %[[D11:.*]] = vector.maskedload %{{.*}}[%[[D10]]], %[[D9]], %[[CST]] : memref<1073741822xf16, strided<[1]>>, vector<16xi1>, vector<16xf16> into vector<16xf16>
834-
# CHECK: vector.maskedstore {{.*}}[{{.*}}], %[[D9]], %[[D11]] : memref<{{.*}}xf16{{.*}}>, vector<16xi1>, vector<16xf16>
821+
# CHECK: %[[D1:.*]] = affine.apply #[[map1]]()[%[[WORKGROUP_ID_1]]]
822+
# CHECK: %[[D2:.*]] = vector.broadcast %[[D1]] : index to vector<16xindex>
823+
# CHECK: %[[D3:.*]] = arith.addi %[[D2]], %[[CST_0]] overflow<nsw, nuw> : vector<16xindex>
824+
# CHECK: %[[D4:.*]] = vector.broadcast %{{.*}} : index to vector<16xindex>
825+
# CHECK: %[[D5:.*]] = arith.cmpi slt, %[[D3]], %[[D4]] : vector<16xindex>
826+
# CHECK: %[[D6:.*]] = affine.apply #[[map2]]()[%[[THREAD_ID_X]], %[[WORKGROUP_ID_0]]]
827+
# CHECK: %[[D7:.*]] = arith.cmpi slt, %[[D6]], %{{.*}} : index
828+
# CHECK: %[[D8:.*]] = vector.broadcast %[[D7]] : i1 to vector<16xi1>
829+
# CHECK: %[[D9:.*]] = arith.andi %[[D5]], %[[D8]] : vector<16xi1>
830+
# CHECK: %[[D10:.*]] = vector.maskedload %{{.*}}[%{{.*}}], %[[D9]], %[[CST]] : memref<1073741822xf16, strided<[1]>>, vector<16xi1>, vector<16xf16> into vector<16xf16>
831+
# CHECK: vector.maskedstore {{.*}}[{{.*}}], %[[D9]], %[[D10]] : memref<{{.*}}xf16{{.*}}>, vector<16xi1>, vector<16xf16>
835832
# CHECK: func.func @isolated_benchmark$async(%[[ARG:.*]]: !hal.buffer_view, %[[FENCE:.*]]: !hal.fence, %{{.*}}: !hal.fence)
836833
# CHECK: %[[D0:.*]] = hal.buffer_view.dim<%[[ARG]] : !hal.buffer_view>[0] : index
837834
# CHECK: %[[D1:.*]] = hal.buffer_view.dim<%[[ARG]] : !hal.buffer_view>[1] : index
@@ -1635,26 +1632,20 @@ def repeat(
16351632
print(test.asm)
16361633

16371634
# CHECK-LABEL: test_tiled_reduce_min_unaligned
1638-
# CHECK-DAG: #[[tr_map_lane:.*]] = affine_map<()[s0] -> (s0 * 2 - (s0 floordiv 64) * 128)>
1639-
# CHECK-DAG: #[[tr_map_blk:.*]] = affine_map<()[s0, s1] -> (s0 * 527 + s1 * 2 - (s1 floordiv 64) * 128)>
1635+
# CHECK-DAG: #[[map1:.*]] = affine_map<()[s0, s1] -> (s0 * 128 + s1 * 2 - (s1 floordiv 64) * 128)>
16401636
# CHECK: func @tiled_reduce_min_unaligned
16411637
# CHECK-DAG: %[[cst_0:.*]] = arith.constant dense<527> : vector<2xindex>
16421638
# CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[0, 1]> : vector<2xindex>
16431639
# CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
16441640
# CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
16451641
# CHECK-DAG: %[[C5:.*]] = arith.constant 5 : index
1646-
# CHECK-DAG: %[[C128:.*]] = arith.constant 128 : index
1647-
# CHECK-DAG: %[[BLK_Y:.+]] = gpu.block_id y
16481642
# CHECK: %[[THREAD_ID_X:.+]] = gpu.thread_id x
1649-
# CHECK: %{{.*}} = affine.apply #[[tr_map_lane]]()[%[[THREAD_ID_X]]]
1650-
# CHECK: %{{.*}} = affine.apply #[[tr_map_blk]]()[%[[BLK_Y]], %[[THREAD_ID_X]]]
16511643
# Tiled Reduction Loop
16521644
# CHECK: scf.for %[[ITER:.*]] = %[[C0]] to %[[C5]] step %[[C1]]
1653-
# CHECK: %[[MUL_K:.*]] = arith.muli %[[ITER]], %[[C128]] overflow<nsw, nuw> : index
1654-
# CHECK: vector.broadcast %[[MUL_K]] : index to vector<2xindex>
1655-
# CHECK: arith.addi {{.*}}, {{.*}} overflow<nsw, nuw> : vector<2xindex>
1656-
# CHECK: %[[D9:.*]] = arith.cmpi slt, {{.*}}, %[[cst_0]] : vector<2xindex>
1657-
# CHECK: arith.addi {{.*}}, %[[MUL_K]] overflow<nsw, nuw> : index
1645+
# CHECK: %[[D6:.*]] = affine.apply #[[map1]]()[%[[ITER]], %[[THREAD_ID_X]]]
1646+
# CHECK: %[[D7:.*]] = vector.broadcast %[[D6]] : index to vector<2xindex>
1647+
# CHECK: %[[D8:.*]] = arith.addi %[[D7]], %[[cst_1]] overflow<nsw, nuw> : vector<2xindex>
1648+
# CHECK: %[[D9:.*]] = arith.cmpi slt, %[[D8]], %[[cst_0]] : vector<2xindex>
16581649
# CHECK-COUNT-2: vector.maskedload %{{.*}}[%{{.*}}], %[[D9]]
16591650

16601651

‎lit_tests/kernel/wave/linearize_loop_affine_maps.py‎

Lines changed: 38 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,19 @@
11
# RUN: python %s | FileCheck %s
22

33
"""
4-
Test that linearization eliminates affine maps from the scf.for loop body.
4+
Test that linearization hoists loop-invariant index bases out of scf.for.
55
66
When M, N, K are dynamic (not substituted at compile time), the pipelined
7-
loop body should contain no affine.apply operations -- all index
8-
computation should be lowered to plain arith ops by the linearization and
9-
stride-annotation passes.
7+
loop body should use hoisted ``arith.muli``/``arith.addi`` for GatherToLDS
8+
index computation rather than monolithic ``affine.apply`` ops with many
9+
operands.
10+
11+
Reads whose flat index is truly affine in the IV (A-data GatherToLDS)
12+
get their stride annotated and codegen emits ``base + IV * stride``.
13+
Reads with complex preshuffle mappings that contain floor/Mod terms
14+
interacting with the IV (B-data, B-scale reads) may still use
15+
``affine.apply`` because the expression is not decomposable as
16+
``base + IV * stride`` in the symbolic domain.
1017
1118
Two wave shapes are tested to cover different workgroup tiling patterns:
1219
(1, 4): 64 threads, 1 wave in M, 4 in N -- block (256, 192, 256)
@@ -46,26 +53,45 @@ def _compile_dynamic_preshuffle_b(shape, block, wave_shape):
4653

4754

4855
@run_test
49-
def test_no_affine_in_loop_1x4():
50-
"""Dynamic 256x192x256 block, (1,4) wave: no affine.apply inside scf.for."""
56+
def test_hoisted_indices_in_loop_1x4():
57+
"""Dynamic 256x192x256 block, (1,4) wave: GatherToLDS uses hoisted arith."""
5158
_compile_dynamic_preshuffle_b((8192, 3072, 8192), (256, 192, 256), (1, 4))
5259

53-
# CHECK-LABEL: test_no_affine_in_loop_1x4
60+
# CHECK-LABEL: test_hoisted_indices_in_loop_1x4
5461
# CHECK: func.func @
55-
# No affine.apply between the pipelined scf.for and its yield.
56-
# affine.apply may still appear in the epilogue after the loop.
62+
# Inside the pipelined scf.for, GatherToLDS reads should use
63+
# hoisted base + IV * stride (arith.muli + arith.addi), not
64+
# monolithic affine.apply with 6+ operands.
65+
# B-data reads (complex preshuffle) may still use affine.apply
66+
# because the expression is not decomposable as base + IV * stride.
5767
# CHECK: scf.for
68+
# CHECK: arith.muli
69+
# CHECK: arith.addi
70+
# CHECK: gather_to_lds
71+
# KEY CHECK -- do not remove. The affine.apply count and
72+
# CHECK-NOT verify that hoisting is working: only the
73+
# residual B-data/B-scale preshuffle reads (with floor/Mod
74+
# IV interaction) remain as affine.apply in the loop body.
75+
# CHECK-COUNT-12: affine.apply
5876
# CHECK-NOT: affine.apply
5977
# CHECK: scf.yield
6078

6179

6280
@run_test
63-
def test_no_affine_in_loop_2x2():
64-
"""Dynamic 256x224x256 block, (2,2) wave: no affine.apply inside scf.for."""
81+
def test_hoisted_indices_in_loop_2x2():
82+
"""Dynamic 256x224x256 block, (2,2) wave: GatherToLDS uses hoisted arith."""
6583
_compile_dynamic_preshuffle_b((8192, 3584, 8192), (256, 224, 256), (2, 2))
6684

67-
# CHECK-LABEL: test_no_affine_in_loop_2x2
85+
# CHECK-LABEL: test_hoisted_indices_in_loop_2x2
6886
# CHECK: func.func @
6987
# CHECK: scf.for
88+
# CHECK: arith.muli
89+
# CHECK: arith.addi
90+
# CHECK: gather_to_lds
91+
# KEY CHECK -- do not remove. The affine.apply count and
92+
# CHECK-NOT verify that hoisting is working: only the
93+
# residual B-data/B-scale preshuffle reads (with floor/Mod
94+
# IV interaction) remain as affine.apply in the loop body.
95+
# CHECK-COUNT-28: affine.apply
7096
# CHECK-NOT: affine.apply
7197
# CHECK: scf.yield

‎tests/kernel/wave_gemm_mxfp_test.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -360,9 +360,9 @@ def testScaledBatchedGemmMXFP4Codegen(use_water_backend: bool, tmp_path: Path):
360360
# We encode the exact registers and wait counts as we want to know if
361361
# they suddenly change due to backend or upstream MLIR changes.
362362
if use_water_backend:
363-
vgpr_count = 172
363+
vgpr_count = 164
364364
vgpr_spill_count = 0
365-
sgpr_count = 61
365+
sgpr_count = 57
366366
sgpr_spill_count = 0
367367
waitcounts = [
368368
"s_waitcnt lgkmcnt(0)",

‎wave_lang/kernel/compiler/wave_codegen/read_write.py‎

Lines changed: 11 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
# See https://llvm.org/LICENSE.txt for license information.
55
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
66

7+
import functools
78
from typing import Any, Optional
89

910
import math
@@ -74,7 +75,6 @@
7475
cast_vector,
7576
gen_sympy_index,
7677
gen_sympy_index_hoisted,
77-
gen_sympy_index_no_affine,
7878
get_constant_attr,
7979
get_type_or_element_type,
8080
handle_op,
@@ -311,51 +311,24 @@ def _build_mask(
311311
return None
312312

313313
idxc = IndexingContext.current()
314-
subs_map = add_emitter_subs(emitter, dynamic_values)
315314

316-
mask_parts: list[Value] = []
315+
conditions = []
317316
for key, bound in bounds.items():
318317
if isinstance(key, sympy.Symbol) and key in index:
319318
start = _get_start_index(index[key])
320319
fastest_dim = get_fastest_index(index)
321320
last_dim = list(index)[fastest_dim]
322321
if key == last_dim:
323322
start = start + idxc.iota(elements_per_thread)
324-
cond = gen_sympy_index(subs_map, start < bound)
323+
conditions.append(start < bound)
325324
else:
326-
iv_sym, iv_val, all_iv_syms = _iv_context(emitter, key)
327-
in_loop = _get_enclosing_scf_for() is not None
328-
hoist_ip = _hoist_before_loop(emitter) if in_loop else None
329-
key_val = gen_sympy_index_hoisted(
330-
subs_map,
331-
key,
332-
iv_sym=iv_sym,
333-
iv_val=iv_val,
334-
hoist_ip=hoist_ip,
335-
all_iv_syms=all_iv_syms,
336-
iv_offset_cache=emitter.iv_offset_cache,
337-
)
338-
bound_val = gen_sympy_index(subs_map, bound)
339-
if isinstance(key_val.type, VectorType) and not isinstance(
340-
bound_val.type, VectorType
341-
):
342-
bound_val = vector_d.broadcast(key_val.type, bound_val)
343-
elif isinstance(bound_val.type, VectorType) and not isinstance(
344-
key_val.type, VectorType
345-
):
346-
key_val = vector_d.broadcast(bound_val.type, key_val)
347-
cond = arith_d.cmpi(arith_d.CmpIPredicate.slt, key_val, bound_val)
348-
349-
mask_vec_type = VectorType.get(
350-
[elements_per_thread], IntegerType.get_signless(1)
351-
)
352-
if cond.type != mask_vec_type:
353-
cond = vector_d.broadcast(mask_vec_type, cond)
354-
mask_parts.append(cond)
325+
conditions.append(key < bound)
355326

356-
mask = mask_parts[0]
357-
for part in mask_parts[1:]:
358-
mask = arith_d.andi(mask, part)
327+
mask_expr = functools.reduce(
328+
lambda a, b: sympy.And(a, b),
329+
conditions,
330+
)
331+
mask = gen_sympy_index(add_emitter_subs(emitter, dynamic_values), mask_expr)
359332

360333
mask_vec_type = VectorType.get([elements_per_thread], IntegerType.get_signless(1))
361334
if mask.type != mask_vec_type:
@@ -636,11 +609,7 @@ def _compute_branchless_valid_bytes(
636609
)
637610
zero_valid = arith_d.constant(uint64, get_constant_attr(0, uint64))
638611

639-
subs_map = add_emitter_subs(emitter)
640-
if _get_enclosing_scf_for() is not None:
641-
cond_val = gen_sympy_index_no_affine(subs_map, guard_condition)
642-
else:
643-
cond_val = gen_sympy_index(subs_map, guard_condition)
612+
cond_val = gen_sympy_index(add_emitter_subs(emitter), guard_condition)
644613
i1 = IntegerType.get_signless(1)
645614
if cond_val.type != i1:
646615
zero_idx = arith_d.constant(cond_val.type, 0)
@@ -1209,7 +1178,7 @@ def _handle_read_linear_index(
12091178
subs_map = add_emitter_subs(emitter, dynamic_vals_map_start)
12101179
sym_strides = _sym_strides_for_flat_memref(kb_src, input_shape)
12111180
linear_buffer_ops = emitter.options.eliminate_epilogue
1212-
lin_src = _linear_read_linearize_memref_maybe_hoisted(
1181+
lin_src, _ = _linearize_memref_hoisted(
12131182
emitter,
12141183
kb_src,
12151184
sym_strides,

‎wave_lang/kernel/wave/analysis/annotate_iv_strides.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,31 @@ def _eval2(iv_val):
168168
return base, sympy.Integer(stride_val)
169169

170170

171+
def _verify_stride_on_original(flat, iv_sym, step, base, stride):
172+
"""Numerically verify ``flat(iv=step*i) == base + i*stride``.
173+
174+
The divisibility-substitution path extracts the stride from a
175+
simplified expression but the base comes from the original.
176+
Floor/Mod interactions with the IV can invalidate this
177+
decomposition, so we probe several concrete IV values.
178+
"""
179+
free_syms = sorted(flat.free_symbols - {iv_sym}, key=str)
180+
probe_map = {s: 137 + i * 31 for i, s in enumerate(free_syms)}
181+
probe_map2 = {s: 251 + i * 47 for i, s in enumerate(free_syms)}
182+
183+
for pmap in (probe_map, probe_map2):
184+
for i in range(8):
185+
iv_val = step * i
186+
orig = subs_idxc(flat.subs({iv_sym: iv_val})).subs(pmap)
187+
expected = subs_idxc(base + i * stride).subs(pmap)
188+
try:
189+
if int(orig) != int(expected):
190+
return False
191+
except (TypeError, ValueError):
192+
return False
193+
return True
194+
195+
171196
def _try_with_div_subs(flat, iv_sym, step, div_fwd, div_bwd):
172197
"""Try stride extraction after applying divisibility substitutions.
173198
@@ -179,6 +204,11 @@ def _try_with_div_subs(flat, iv_sym, step, div_fwd, div_bwd):
179204
The stride is extracted from the simplified expression but the base
180205
is computed from the ORIGINAL flat to preserve floor/Mod structure
181206
needed for correct MLIR integer lowering.
207+
208+
A numerical verification confirms the decomposition is valid in
209+
the original domain; when floor/Mod terms interact with the IV
210+
(e.g. in B-scale preshuffle mappings), the decomposition may be
211+
invalid and the extraction is rejected.
182212
"""
183213
if not div_fwd:
184214
return None
@@ -203,6 +233,10 @@ def _try_with_div_subs(flat, iv_sym, step, div_fwd, div_bwd):
203233
stride = mem_simplify(sympy.sympify(stride).subs(bwd_dict))
204234

205235
base = safe_subs(flat, {iv_sym: sympy.Integer(0)})
236+
237+
if not _verify_stride_on_original(flat, iv_sym, step, base, stride):
238+
return None
239+
206240
return base, stride, method
207241

208242

0 commit comments

Comments
 (0)