-
Notifications
You must be signed in to change notification settings - Fork 1.7k
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
base: main
Are you sure you want to change the base?
Conversation
@@ -389,7 +389,7 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion | |||
return true; | |||
} | |||
} | |||
if (isa<AMDMfmaEncodingAttr>(parent)) { | |||
if (isa<AMDMfmaEncodingAttr, AMDWmmaEncodingAttr>(parent)) { |
There was a problem hiding this comment.
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.
There was a problem hiding this 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)) { |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
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)) { |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
There was a problem hiding this 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.
Enable LL conwertions for WMMA as well as for MFMA layouts.
See also: #5210