[AMD] Fix kwidth printing and parsing by binarman · Pull Request #3676 · triton-lang/triton

binarman · 2024-04-16T18:14:23Z

This PR fixes:

parsing and printing of kWidth attribute for MFMA and WMMA layouts

binarman · 2024-04-16T18:14:38Z

binarman · 2024-04-16T18:16:17Z

+  auto mmaParent = parent.dyn_cast<NvidiaMmaEncodingAttr>();
  printer << "<{"
          << "opIdx = " << getOpIdx() << ", parent = " << getParent();
-  if (mmaParent && mmaParent.isAmpere())


I have a question related to Nvidia.

This part mismatches parsing:

Parse for Ampere and Hopper

Print only for Ampere

Can anyone clarify, which behavior is correct?

I'm 60% sure that printing for Hopper is correct, but also in the 40% case that it's incorrect then printing it will make it obvious to someone else that it shouldn't be there. :) Thanks for noticing this.

Also this whole thing seems unnecessarily confusing. It seems to me that we shouldn't need any of this custom parsing logic. Instead:

Make kWidth an optional attribute with a default of 0 or -1.

Use a regular MLIR parser and printer, no need for the custom C++ at all. It will not print kWidth when it's 0, because that's the default.

Add a verifier that checks that kWidth is 0 for those encodings that don't support kWidth. Also check that the kWidth is explicitly equal to AMDWmmaEncodingAttr::getMNKDimPerWMMAInstr()[2] for AMDWmma.

Not your sin, I guess you don't need to change it.

I can try to rework this with default value approach without custom parsing.
I suspect this will change formatting, but we can experiment and see if this make sense

It is OK to change formatting. You'll have to fix the tests, but I've done it, it's not too bad.

Turned out declarative format can express same formatting 🙂

Yup. Great to see this getting cleaned up. Thanks @binarman!

jlebar

Requesting changes because I think this needs a unit test.

jlebar · 2024-04-16T19:39:50Z

          << "opIdx = " << getOpIdx() << ", parent = " << getParent();
-  if (mmaParent && mmaParent.isAmpere())
+  bool isAmdWithKWidth =
+      parent.isa<AMDMfmaEncodingAttr>() || parent.isa<AMDWmmaEncodingAttr>();


This seems like a similar inconsistency to Ampere/Hopper. Here we're printing kWidth if AMDWmma, but if AMDWmma we ignore any parsed value and unconditionally do kWidth = AMDWmmaEncodingAttr::getMNKDimPerWMMAInstr()[2];

jlebar · 2024-04-16T19:41:27Z

+  auto mmaParent = parent.dyn_cast<NvidiaMmaEncodingAttr>();
  printer << "<{"
          << "opIdx = " << getOpIdx() << ", parent = " << getParent();
-  if (mmaParent && mmaParent.isAmpere())


Also this whole thing seems unnecessarily confusing. It seems to me that we shouldn't need any of this custom parsing logic. Instead:

Make kWidth an optional attribute with a default of 0 or -1.

Use a regular MLIR parser and printer, no need for the custom C++ at all. It will not print kWidth when it's 0, because that's the default.

Add a verifier that checks that kWidth is 0 for those encodings that don't support kWidth. Also check that the kWidth is explicitly equal to AMDWmmaEncodingAttr::getMNKDimPerWMMAInstr()[2] for AMDWmma.

Not your sin, I guess you don't need to change it.

jlebar · 2024-04-16T19:41:37Z

  Attribute _kWidth = attrs.get("kWidth");
  if (_kWidth) {
-    if (!mmaParent || mmaParent.isVolta()) {
+    bool isAmdWithKWidth =


Write a unit test?

antiagainst · 2024-04-16T23:16:40Z

  Attribute _kWidth = attrs.get("kWidth");
  if (_kWidth) {
-    if (!mmaParent || mmaParent.isVolta()) {
+    bool isAmdWithKWidth =


antiagainst · 2024-04-16T23:17:56Z

Also, @binarman, please make sure all tests pass before marking pull requests as ready for review. As addressing bots happy needs changing the impl. (You can ping me if you need to trigger CI runs.)

binarman · 2024-05-01T13:30:43Z

@antiagainst @jlebar
I've updated this PR, can you take a look?

antiagainst

Cool, looks better. I still have a few comments. Major one is actuall the tests.

antiagainst · 2024-05-14T06:24:50Z

+  auto mmaParent = parent.dyn_cast<NvidiaMmaEncodingAttr>();
  printer << "<{"
          << "opIdx = " << getOpIdx() << ", parent = " << getParent();
-  if (mmaParent && mmaParent.isAmpere())


Yup. Great to see this getting cleaned up. Thanks @binarman!

antiagainst · 2024-05-14T21:06:52Z

(converting to draft as we said internally. will mark it open after approval from me)

This PR fixes: - parsing and printing of kWidth attribute for MFMA and WMMA layouts - comment in 06-fused-attention.p tutorial

antiagainst

Awesome, thanks for bearing with my nitpicking! :)

jlebar

Thanks!

binarman requested a review from ptillet as a code owner April 16, 2024 18:14

binarman commented Apr 16, 2024

View reviewed changes

jlebar suggested changes Apr 16, 2024

View reviewed changes

antiagainst requested changes Apr 16, 2024

View reviewed changes

binarman mentioned this pull request Apr 29, 2024

[AMD] Move MFMA shortcut check to not compute scratch buffer shape if it is not needed #3790

Merged

binarman force-pushed the fix_kwidth_printing branch from c936414 to 5017d00 Compare May 1, 2024 13:00

jlebar requested a review from antiagainst May 1, 2024 20:50

alefimov-amd force-pushed the fix_kwidth_printing branch 2 times, most recently from 794e451 to dff9246 Compare May 9, 2024 14:50

antiagainst requested changes May 14, 2024

View reviewed changes

antiagainst marked this pull request as draft May 14, 2024 21:05

binarman added 4 commits May 15, 2024 15:17

[AMD] Fix kwidth printing and parsing

f17d5a3

This PR fixes: - parsing and printing of kWidth attribute for MFMA and WMMA layouts - comment in 06-fused-attention.p tutorial

addressing review comments

4277e19

rebase fix

e252415

review fixes: rework verifier function, fix comments

09680cf

binarman force-pushed the fix_kwidth_printing branch from dff9246 to 09680cf Compare May 15, 2024 21:21

binarman and others added 5 commits May 16, 2024 15:19

change error messages, add lit test

f8fbbd2

Make error messages clear

5151536

Merge remote-tracking branch 'origin/main' into pr-3676

e21089b

Fix new failures after merge

b4dc772

Drop unit tests for attribute

1cc5650

antiagainst approved these changes May 17, 2024

View reviewed changes

antiagainst marked this pull request as ready for review May 17, 2024 00:00

jlebar approved these changes May 17, 2024

View reviewed changes

Comment thread include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td Resolved

Comment thread include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td Resolved

Comment thread lib/Dialect/TritonGPU/IR/Dialect.cpp Resolved

Address comments

018e928

antiagainst enabled auto-merge (squash) May 17, 2024 00:47

antiagainst merged commit 7e9a8e1 into triton-lang:main May 17, 2024

jlebar mentioned this pull request May 19, 2024

[CI] Rename check_cuda_only() to check_cuda_or_hip(). #3946

Merged

ethanwee1 mentioned this pull request Sep 22, 2026

[release 3.8] Support Python 3.15 ROCm/triton#971

Merged

4 tasks done

Conversation

binarman commented Apr 16, 2024 • edited by antiagainst Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

binarman commented Apr 16, 2024

Uh oh!

binarman Apr 16, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jlebar Apr 16, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jlebar left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jlebar Apr 16, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

antiagainst commented Apr 16, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

binarman commented May 1, 2024

Uh oh!

antiagainst left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

antiagainst commented May 14, 2024

Uh oh!

antiagainst left a comment

Choose a reason for hiding this comment

Uh oh!

jlebar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

binarman commented Apr 16, 2024 •

edited by antiagainst

Loading

binarman Apr 16, 2024 •

edited

Loading

jlebar Apr 16, 2024 •

edited

Loading

jlebar Apr 16, 2024 •

edited

Loading

antiagainst commented Apr 16, 2024 •

edited

Loading