From 2c4b2e34802930e3b8cf544e9345aa811434a449 Mon Sep 17 00:00:00 2001
From: lezcano <lezcano-93@hotmail.com>
Date: Tue, 26 Nov 2024 17:44:29 +0000
Subject: [PATCH] Fix Blocked FMA path in isLayoutOK

Fixes https://github.com/triton-lang/triton/pull/5235/files/de18e21ddf5bf03f17f779fef032d53ea87a53a0#r1858955613
---
 .../TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp         | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp
index aedb18a245e6..02a62d1668e6 100644
--- 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;
       }