From 417358fdc0a116e6c7c8a5d38fe305bd40c70cc4 Mon Sep 17 00:00:00 2001
From: Hongtao Yu <hoy@fb.com>
Date: Thu, 18 Jan 2024 10:42:56 -0800
Subject: [PATCH] Enable swizzling SMEM for transposed operand

---
 include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td b/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
index 68797e11eb23..5a008aae0f20 100644
--- a/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
+++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td
@@ -131,6 +131,8 @@ compared to 1*64 when the hasLeadingOffset is false.
 
         if (mfmaEnc) {
           int kDimNum = dotOpEnc.getOpIdx() == 0 ? 1 : 0;
+          if (needTrans)
+            kDimNum = 1 - kDimNum;
           bool isKDimInner = (order[0] == kDimNum);
           if (isKDimInner) {
             const int numBanks = 32;