Skip to content

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

Open
YWHyuk wants to merge 1 commit into
refactor/tile-axisfrom
refactor/fusion-frame
Open

[Frontend] Derive the fusion aliasing from the tile axes#299
YWHyuk wants to merge 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.

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 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
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