[AMD] Use Linear Layout convertions for AMDWmma

Enable LL conwertions for WMMA as well as for MFMA layouts. See also: #5210 Signed-off-by: Ilya Veselov <[email protected]>
triton-lang · Nov 25, 2024 · 2a3a84d · 2a3a84d
1 parent 09fcc52
commit 2a3a84d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp
@@ -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)) {
         return !useLegacyMMAConversion;
       }
       if (auto dotOperand = dyn_cast<DotOperandEncodingAttr>(layout)) {
@@ -389,7 +389,7 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
             return true;
           }
         }
-        if (isa<AMDMfmaEncodingAttr>(parent)) {
+        if (isa<AMDMfmaEncodingAttr, AMDWmmaEncodingAttr>(parent)) {
           return true;
         }
         return false;