[Frontend] Derive the fusion aliasing from the tile axes#299
Open
YWHyuk wants to merge 1 commit into
Open
Conversation
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.
a3bf0ed to
90d8b22
Compare
0c47614 to
4fa02b1
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_rangesconsumes, and the order the reduction MVOUT reads the DRAM strides in. Every template spelled that list out by hand as adim_aliasingdict: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
aliasing(tile.axes)..values(). Sodim_aliasingbecomes a plain list and the six consumer sites drop.values().input_dim_aliasingandweight_dim_aliasing, carried inprologue_info, were never read at all -- deleted.Two small facts it surfaced
loop=Nonewhile its aliasing named it"c0". Since the axis has stride 0,Symbol("c0")*0andInteger(0)are the same DRAM index term, so naming it"c0"is byte-identical and lets the aliasing derive.Verification
Regenerated every kernel from scratch (directories deleted first).
tests/ops/fusion/(7 tests)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 bycoeff == 1-- the path that produced #255. Replacing that with stride-matching against the tile's axes (which closes the silentM == Ntranspose hole and the view-epilogue misreject) touches the gate that #295 is already rewriting, so it stacks on #295 separately, not here.