Skip to content

[Compiler][gfx120] Add modC/reuseA/reuseB to MmaOpGFX1250_WMMAType - #1001

Open
jli-melchior wants to merge 2 commits into
ROCm:mainfrom
jli-melchior:jli/gfx1250/fix-mma-atom
Open

[Compiler][gfx120] Add modC/reuseA/reuseB to MmaOpGFX1250_WMMAType#1001
jli-melchior wants to merge 2 commits into
ROCm:mainfrom
jli-melchior:jli/gfx1250/fix-mma-atom

Conversation

@jli-melchior

Copy link
Copy Markdown
Collaborator

The non-scale WMMA path was missing modC/reuseA/reuseB parameters that its sibling WMMAScaleType already had. The lowering hardcoded WMMACModifier::none and reuseA/reuseB=false, causing mma_atom_call to emit rocdl.wmma ops without an explicit modC attribute. Since the ROCDL intrinsic distinguishes absent modC from modC=0, this produced incorrect ISA encodings and numerical precision regression for bf16 WMMA kernels.

Add modC (int32_t), reuseA (bool), reuseB (bool) to the type with back-compat builders defaulting to 0/false/false, forward them through the Python binding and wrapper, and pass them to the ROCDL ops at emission time.

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

@jli-melchior
jli-melchior force-pushed the jli/gfx1250/fix-mma-atom branch 7 times, most recently from e6cbc6c to 27056e1 Compare August 13, 2026 02:13
@jli-melchior jli-melchior changed the title [WIP][GFX1250] Add modC/reuseA/reuseB to MmaOpGFX1250_WMMAType [GFX1250] Add modC/reuseA/reuseB to MmaOpGFX1250_WMMAType Aug 13, 2026
@jli-melchior jli-melchior changed the title [GFX1250] Add modC/reuseA/reuseB to MmaOpGFX1250_WMMAType [Compiler][gfx120] Add modC/reuseA/reuseB to MmaOpGFX1250_WMMAType Aug 13, 2026
Comment on lines +88 to +94
"bool":$clamp,
// Intrinsic attributes forwarded to the ROCDL WMMA op: modC (I16
// C-operand modifier) and reuseA/reuseB (I1 operand-reuse scheduler
// hints). Default 0/false.
"int32_t":$modC,
"bool":$reuseA,
"bool":$reuseB

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Update to DefaultValuedParam to simplify assembly format.

The non-scale WMMA path was missing modC/reuseA/reuseB parameters that
its sibling WMMAScaleType already had. The lowering hardcoded
WMMACModifier::none and reuseA/reuseB=false, causing mma_atom_call to
emit rocdl.wmma ops without an explicit modC attribute. Since the ROCDL
intrinsic distinguishes absent modC from modC=0, this produced incorrect
ISA encodings and numerical precision regression for bf16 WMMA kernels.

Add modC (int32_t), reuseA (bool), reuseB (bool) to the type with
back-compat builders defaulting to 0/false/false, forward them through
the Python binding and wrapper, and pass them to the ROCDL ops at
emission time.

Co-Authored-By: Claude <noreply@anthropic.com>
@jli-melchior
jli-melchior force-pushed the jli/gfx1250/fix-mma-atom branch from 44bfe88 to 2b691b8 Compare August 14, 2026 08:04
@sjfeng1999

Copy link
Copy Markdown
Collaborator

reuse flags should not be modeled as static fields of the WMMA type, since programmers cannot conveniently vary them across a sequence of WMMA calls.

@sjfeng1999

Copy link
Copy Markdown
Collaborator

Maybe a better way is to add a separate pass to attach these reuse attrs automatically.

@jli-melchior

Copy link
Copy Markdown
Collaborator Author

reuse flags should not be modeled as static fields of the WMMA type, since programmers cannot conveniently vary them across a sequence of WMMA calls.

reuseA/reuseB are bit fields of the wmma instruction that hint the XDL engine to reuse operand data. Exposing them gives users more flexibility for different scheduling scenarios ?

@jli-melchior

Copy link
Copy Markdown
Collaborator Author

Maybe a better way is to add a separate pass to attach these reuse attrs automatically.

Yes it sounds a good idea for a follow-up optimization which programmers may not set these flags, And it also need these fields on the type to construct the updated atom. This PR provides the foundation whether the flags are set manually or by a future automatic pass.

@sjfeng1999

Copy link
Copy Markdown
Collaborator

I can't see how callers could conveniently control reuse flags within a single fx.gemm invocation, which may expand into multiple atom calls. Could you provide a concrete example?

@jli-melchior

Copy link
Copy Markdown
Collaborator Author

I can't see how callers could conveniently control reuse flags within a single fx.gemm invocation, which may expand into multiple atom calls. Could you provide a concrete example?

In mha kernel i haven‘t see use the fx.gemm to call mma atom, the gemm kernel can use this to get a perf improvement, but this need a special pattern for reuseA/reusesB. This may need different pattern according to different cases.

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.

2 participants