Skip to content

[Frontend] Derive the fusion aliasing from the tile axes#299

Merged
YWHyuk merged 1 commit into
refactor/tile-axisfrom
refactor/fusion-frame
Jul 14, 2026
Merged

[Frontend] Derive the fusion aliasing from the tile axes#299
YWHyuk merged 1 commit into
refactor/tile-axisfrom
refactor/fusion-frame

Conversation

@YWHyuk

@YWHyuk YWHyuk commented Jul 11, 2026

Copy link
Copy Markdown
Collaborator

Stacked on #298. A fused epilogue or prologue renames its loop variables to the template's, in the tile's declared order -- that is what set_ranges consumes, and the order the reduction MVOUT reads the DRAM strides in. Every template spelled that list out by hand as a dim_aliasing dict:

epilogue_dim_aliasing = {"index0": "index1", "index1": "index0"}   # GEMM reduction
prologue_dim_aliasing = {"index0": "index2", "index1": "index1"}   # GEMM weight
dim_aliasing          = {"index0": "c0", "index1": "tile_n", ...}  # conv single-batch

Sixteen of these across gemm, bmm and the four conv templates. Each is exactly the loop name of every axis in the corresponding tile, in declared order -- the same tile the template already builds (#298). So they are all aliasing(tile.axes), a list.

The change

  • Replace the sixteen dicts with aliasing(tile.axes).
  • The keys of these dicts were never read; every consumer took .values(). So dim_aliasing becomes a plain list and the six consumer sites drop .values().
  • input_dim_aliasing and weight_dim_aliasing, carried in prologue_info, were never read at all -- deleted.

Two small facts it surfaced

  • conv single-batch declared its degenerate batch axis with loop=None while its aliasing named it "c0". Since the axis has stride 0, Symbol("c0")*0 and Integer(0) are the same DRAM index term, so naming it "c0" is byte-identical and lets the aliasing derive.
  • sdpa keeps its hand-written aliasing. Its epilogue frame is the kernel output rank (4), not its 3-axis output tile, so it is not the tile's loop order. Left as is and noted.

Verification

Regenerated every kernel from scratch (directories deleted first).

result
gemm (mm, +relu, addmm, reduction, prologue, N==1, var_mean) byte-identical MLIR
bmm (bmm, +relu, reduction, prologue) byte-identical MLIR
conv (single-batch, SB-strided, multi-tile, batched) byte-identical MLIR
functional vs CPU: gemm 6, bmm 4, conv 4 pass, max abs diff 4.2e-05
tests/ops/fusion/ (7 tests) fuse the same kernels

The equivalence of all sixteen dicts to the derivation was first confirmed by asserting aliasing(tile.axes) == list(dict.values()) across the whole suite, then the dicts were removed.

Not in this PR

This makes the aliasing derive from the tile; it does not yet make the fusion gate use the tile. The gate still decides which epilogues are legal by parsing the reduction stride out of str(node) and by coeff == 1 -- the path that produced #255. Replacing that with stride-matching against the tile's axes (which closes the silent M == N transpose hole and the view-epilogue misreject) touches the gate that #295 is already rewriting, so it stacks on #295 separately, not here.

@YWHyuk
YWHyuk force-pushed the refactor/tile-axis branch from a3bf0ed to 90d8b22 Compare July 11, 2026 09:09
@YWHyuk
YWHyuk force-pushed the refactor/fusion-frame branch from 0c47614 to 4fa02b1 Compare July 11, 2026 09:09
A fused epilogue or prologue renames its loop variables to the template's, in the
tile's declared order -- that is what set_ranges consumes, and the order the
reduction MVOUT reads the DRAM strides in. Every template spelled that list out by
hand as a `dim_aliasing` dict:

    epilogue_dim_aliasing = {"index0": "index1", "index1": "index0"}   # GEMM reduction
    prologue_dim_aliasing = {"index0": "index2", "index1": "index1"}   # GEMM weight
    dim_aliasing          = {"index0": "c0", "index1": "tile_n", ...}  # conv single-batch

Sixteen of these across gemm, bmm and the four conv templates. Each is exactly the
loop name of every axis in the corresponding tile, in declared order -- the same
tile the template already builds. So they are all `aliasing(tile.axes)`, a list.

The keys of these dicts were never read; every consumer took `.values()`. So
`dim_aliasing` becomes a plain list, and the six consumer sites drop the
`.values()`. `input_dim_aliasing` and `weight_dim_aliasing`, carried in
prologue_info, were never read at all and are deleted.

conv single-batch declared its degenerate batch axis with loop=None while its
aliasing named it "c0"; since the axis has stride 0, `Symbol("c0")*0` and
`Integer(0)` are the same DRAM index term, so naming it "c0" is byte-identical and
lets the aliasing derive.

sdpa keeps its hand-written aliasing: its epilogue frame is the kernel output rank
(4), not its 3-axis output tile, so it is not the tile's loop order. Left as is.

Verified by regenerating every kernel from scratch: byte-identical MLIR for the
gemm, bmm and conv cases, functional mode matches CPU (max abs diff 4.2e-05), and
tests/ops/fusion fuses the same kernels. The equivalence of all sixteen dicts to
the derivation was first confirmed by asserting it across the suite, then the
dicts were removed.
@YWHyuk
YWHyuk force-pushed the refactor/tile-axis branch from 90d8b22 to 56d26dd Compare July 13, 2026 15:07
@YWHyuk
YWHyuk force-pushed the refactor/fusion-frame branch from 4fa02b1 to 4e0fdfe Compare July 13, 2026 15:07
@YWHyuk
YWHyuk merged commit aaae111 into refactor/tile-axis Jul 14, 2026
@YWHyuk
YWHyuk deleted the refactor/fusion-frame branch July 14, 2026 05:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant