Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[GPU] Add col_major optional attribute to MMAAttr #19860

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

qedawkins
Copy link
Contributor

@qedawkins qedawkins commented Jan 30, 2025

For MMA intrinsics with symmetric register layouts we have the freedom
to choose between row major and column major output layouts by simply
swapping the input operands. This adds a flag to MMAAttr that allows
specifying whether the result should be column major.

A follow up PR will turn the column major variant on by default for
mfma as this enables vectorization of writes in non-transpose cases.

@Groverkss
Copy link
Contributor

Nice! You can use the documentation from here if you like: https://github.com/iree-org/iree/blob/main/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureTensorLayouts.cpp#L582

We already do this for attention.

Also, this should really always be the canonical form of MFMA instructions. It has A, B, C all in row major format. The only thing this can messup is the broadcast semantics of MFMA instructions because they are not symmetric for A and B, however we should never really expose the broadcast semantics by themselves anyway, so this is okay, but an additional comment for it would be good

@qedawkins qedawkins force-pushed the swap_mfma_operands branch 2 times, most recently from c68cec0 to 3799a2e Compare February 5, 2025 17:57
For MMA intrinsics with symmetric register layouts we have the freedom
to choose between row major and column major output layouts by simply
swapping the input operands. This adds a flag to `MMAAttr` that allows
specifying whether the result should be column major.

A follow up PR will turn the column major variant on by default for
mfma as this enables vectorization of writes in non-transpose cases.
@qedawkins qedawkins changed the title [Codegen][LLVMGPU] Swap MFMA operands for better store performance [GPU] Add col_major optional attribute to MMAAttr Feb 5, 2025
@qedawkins qedawkins requested review from bjacob and Groverkss February 5, 2025 17:59
@qedawkins qedawkins marked this pull request as ready for review February 5, 2025 17:59
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