You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
feat(rocmlir): fuse pointwise chains into rocMLIR elementwise kernels - #1294
FuseROCMlir gains FusePointwiseChain, which outlines a chain of pointwise ops (stepping through transpose, reshape, expand_shape and collapse_shape the way MIGraphX skips conv_dot_reshaper_names) into one rocMLIR elementwise kernel. A chain that reads a single-use conv/gemm/matmul is left to the anchor pattern (chainConsumesSingleUseAnchor), so conv epilogues stay fused into the conv kernel.
Quick autotune has no tuning space for elementwise kernels and aborts on them, so RocMlirAutotune compiles them with the default elem:v1:256,1,4,1,0 perf config instead of searching.
cmake/rocmlirTriton-elementwise-tuning-interface.patch makes ElementwiseParamsAttr implement the RockTuningParamAttrInterface accessors added after the elementwise-fusions series.
Requires a rocmlirTriton with elementwise kernel support; cmake/deps.txt is not bumped here.
dnt-v5-fp16-512x512 (gfx1151, quick autotune, profiling off): 8.154 -> 7.554 ms, 292 -> 245 launches.
This project follows LLVM's incremental-development and AI-tool-use
guidance. See CONTRIBUTING.md
for the project workflow.
Before requesting review, please check that:
The change is focused. Substantial work links the relevant issue
or design discussion.
The PR documents relevant test results and updates affected
documentation.
If AI tools provided substantial assistance, the description
explains what was assisted and how it was validated, and commit
trailers identify the tool. The contributor has reviewed and
understands the result.
Reviewers are assigned through CODEOWNERS where ownership
is configured.
This PR extends hip-ep’s rocMLIR fusion pass to compile standalone pointwise chains as GPU kernels while preserving convolution and matrix-operation epilogues.
Changes:
Adds pointwise-chain outlining and traversal through layout operations.
Selects a default elementwise configuration during autotuning and adds a dependency compatibility patch.
Adds fusion tests for standalone chains, transposes, and anchor ownership.
File
Description
test/lit/Dialect/hip-fuse-rocmlir.mlir
Adds pointwise and mixed-chain fusion tests.
lib/Dialect/Transforms/FuseROCMlir.cpp
Shares outlining logic and adds standalone pointwise matching.
lib/Compiler/RocMlirAutotune.cpp
Bypasses tuning-space enumeration for elementwise kernels.
FuseROCMlir only outlines ops hip-to-tosa can convert. A chain step must
have static shapes (the layout ops it steps through too), no rank-0
result, and pass isTosaExpressiblePointwiseOp: element types rocMLIR
lowers (f16/bf16/f32, signless i1..i64) and the i1 rules hip-to-tosa
claims compare, logical and sign ops by. The chain stops before any other
op, so the prefix still compiles, and that op no longer keeps its
consumer from heading a chain. tensor.reshape is no longer stepped
through: hip-to-tosa and rocMLIR have no lowering for it.
compileRocMlirDefaultPerfConfig (the path with no compileOne, used by the
EP without autotune) now compiles elementwise kernels with the same
elem:v1:256,1,4,1,0 perf config the autotuner uses, shared as
kRocMlirElementwisePerfConfig.
Also applies clang-format to FuseROCMlir.cpp and RocMlirAutotune.cpp.
isTosaExpressiblePointwiseOp returned true for every op past the element
type allow-list, so an i8/i16 hip.div, a floating-point hip.mod, or a
hip.mod with fmod = 1 was outlined into a rock kernel whose hip-to-tosa
conversion then failed. The predicate now mirrors each fusible op's
converter: Div and Mod reuse isTosaExpressibleDivType and
isTosaExpressibleModType plus Mod's fmod = 0 rule, the float-only unary
ops require a float that matches the input type, the binary ops and Where
require broadcastable operands (with an i1 condition for Where), Gelu its
approximate values, and Cast and Transpose a static input. Any op it does
not know returns false.
FuseROCMlir's head check and anchor check now step with the same helper
as the chain itself, so all three agree on where a chain ends.
The fusion lit test gets i16 div and fmod = 1 mod cases, and a second RUN
line that pipes every outlined kernel through --convert-hip-to-tosa.
Co-Authored-By: Cursor <cursoragent@cursor.com>
Made-with: Cursor
Trace indirect predecessors to preserve convolution epilogue anchors
lib/Dialect/Transforms/FuseROCMlir.cpp:411
[P2] For conv -> exp -> add or conv -> transpose -> add with an independent cast -> add input, this guard sees only the exp/transpose predecessor and misses the anchor. If the cast is rewritten first, it claims the add and splits the convolution epilogue, potentially passing a full activation between kernels instead of a small cast result. Trace backward through eligible predecessors to the anchor, and add regressions for both indirect paths.
A chain made only of identity hip.casts and hip.transposes lowers to no
arith or math op, so rock::isElementwiseKernel does not recognise it and
compilation falls through to the tuning space it does not have.
FusePointwiseChain now leaves such a chain on the HIP path; an identity
cast in front of an op that computes still fuses.
chainJoinsAnchorChain (was chainConsumesSingleUseAnchor) traces back
through the chain ops between an anchor and each op of the head's chain,
so conv -> exp -> add and conv -> transpose -> add stay the conv's
epilogue whichever op the driver visits first.
Co-Authored-By: Cursor <cursoragent@cursor.com>
Made-with: Cursor
Re the indirect conv epilogue (conv -> exp -> add and conv -> transpose -> add with a cast -> add input), fixed in 5846ea6. chainJoinsAnchorChain now asks, for each op of the head's chain, whether an anchor's chain reaches it, tracing back through the chain ops in between (isReachedByAnchorChain) instead of looking only for an anchor directly behind it. Both paths have lit cases: the cast is its own kernel, and the conv kernel takes the exp or transpose and the add.
With the current greedy order the conv is rewritten first, so the previous commit gave the same result on these cases; the change removes the dependence on that order.
Lowering can fold a kernel's whole computation away -- hip.max(%x, %x)
becomes its input -- and createTunableParamSpace then aborts with
"Expected to find GEMM, convolution, ... or elementwise op". After the
high-level pipeline, compileAndEmbedRocMlirKernels now checks
hasRocMlirComputation (a gemm-like op or an elementwise kernel), and for a
kernel with neither puts the hip ops it was outlined from back in place of
its dispatches, from a copy taken before the hip->tosa conversion.
compileRocMlirDefaultPerfConfig returns false on such a kernel instead of
aborting.
chainJoinsAnchorChain walks back once from the last op of the chain, with
a worklist, instead of from every op, which was quadratic in chain length.
Co-Authored-By: Cursor <cursoragent@cursor.com>
Made-with: Cursor
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
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.
No description provided.