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

[AMD] Use Linear Layout convertions for AMDWmma #5255

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

Conversation

joviliast
Copy link
Contributor

Enable LL conwertions for WMMA as well as for MFMA layouts.

See also: #5210

@@ -389,7 +389,7 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
return true;
}
}
if (isa<AMDMfmaEncodingAttr>(parent)) {
if (isa<AMDMfmaEncodingAttr, AMDWmmaEncodingAttr>(parent)) {
Copy link
Contributor

@binarman binarman Nov 25, 2024

Choose a reason for hiding this comment

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

Are you sure this will work?
This is a dot operand branch.

Copy link
Contributor

@binarman binarman left a comment

Choose a reason for hiding this comment

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

LGTM!

@@ -376,7 +376,7 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
// completed before we can remove the layoutIsOK check:
// 1. Support for AMD's WMMA
std::function<bool(Attribute)> layoutIsOK = [&](Attribute layout) {
if (isa<NvidiaMmaEncodingAttr, AMDMfmaEncodingAttr>(layout)) {
if (isa<MmaEncodingTrait>(layout)) {
Copy link
Contributor

Choose a reason for hiding this comment

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

There is another Mfma case below (parent of DotOperand). Does it also need considered?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

It have not been implemented yet. I guess @binarman is assigned.

@sjw36
Copy link
Contributor

sjw36 commented Nov 26, 2024

Also please add a testcase.

Enable LL conwertions for WMMA as well as for MFMA layouts.

See also: triton-lang#5210

Signed-off-by: Ilya Veselov <[email protected]>
std::function<bool(Attribute)> layoutIsOK = [&](Attribute layout) {
if (auto dotOperand = dyn_cast<DotOperandEncodingAttr>(layout)) {
layout = dotOperand.getParent();
if (isa<AMDWmmaEncodingAttr>(layout)) {
Copy link
Contributor

Choose a reason for hiding this comment

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

Please add a comment why this is special cased.

Copy link
Contributor

Choose a reason for hiding this comment

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

Never mind, I see the adjusted comment above.

Copy link
Contributor

@sjw36 sjw36 left a comment

Choose a reason for hiding this comment

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

LGTM, thanks. Please just add the comment for special case.

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.

3 participants