Fix Blocked FMA path in isLayoutOK (#5260)

Fixes https://github.com/triton-lang/triton/pull/5235/files/de18e21ddf5bf03f17f779fef032d53ea87a53a0#r1858955613
triton-lang · Nov 26, 2024 · 3f1d70f · 3f1d70f
1 parent 68a08dd
commit 3f1d70f
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp
@@ -376,13 +376,16 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
     // completed before we can remove the layoutIsOK check:
     // 1. Support for AMD's WMMA
     std::function<bool(Attribute)> layoutIsOK = [&](Attribute layout) {
-      if (auto dotOperand = dyn_cast<DotOperandEncodingAttr>(layout)) {
-        layout = dotOperand.getParent();
-      }
-
       if (isa<NvidiaMmaEncodingAttr, AMDMfmaEncodingAttr>(layout)) {
         return !useLegacyMMAConversion;
       }
+      if (auto dotOperand = dyn_cast<DotOperandEncodingAttr>(layout)) {
+        if (isa<NvidiaMmaEncodingAttr, AMDMfmaEncodingAttr>(
+                dotOperand.getParent())) {
+          return !useLegacyMMAConversion;
+        }
+        return false;
+      }
       if (isa<BlockedEncodingAttr, LinearEncodingAttr>(layout)) {
         return true;
       }