From 2a3a84da91ed6761b10e39ae2af7559f285f0a09 Mon Sep 17 00:00:00 2001
From: Ilya Veselov <iveselov.nn@gmail.com>
Date: Mon, 25 Nov 2024 15:32:29 +0100
Subject: [PATCH] [AMD] Use Linear Layout convertions for AMDWmma

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

See also: https://github.com/triton-lang/triton/pull/5210

Signed-off-by: Ilya Veselov <iveselov.nn@gmail.com>
---
 lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp
index aab97c7dd2b3..f15021e3c3fe 100644
--- 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;