Revert "[AMD] Enable General Swizzling ConvertLayoutOp (#7482)"

This reverts commit 318ff2c.

Author: anmyachev

lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp

@@ -278,8 +278,9 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
     assert(cvtNeedsSharedMemory(op.getSrc().getType(), op.getType()));
 
     // Try to use swizzling to implement the conversion
-    if (succeeded(
-            transferWithinBlockSwizzling(op, adaptor.getSrc(), rewriter))) {
+    // HACK Remove once AMD tests pass for the swizzling path
+    if (targetInfo.isCuda() && succeeded(transferWithinBlockSwizzling(
+                                   op, adaptor.getSrc(), rewriter))) {
       return success();
     }
 

test/Conversion/amd/convert_layout.mlir

@@ -1,14 +1,11 @@
-#include "Analysis/AMDGPUAllocation.h"
 #include "PatternTritonGPUOpToLLVM.h"
 #include "Utility.h"
 #include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h"
 #include "triton/Conversion/TritonGPUToLLVM/Utility.h"
 #include "triton/Dialect/TritonGPU/Transforms/Utility.h"
 
-using ::mlir::transferWithinBlockPadding;
 using ::mlir::triton::gpu::AMDMfmaEncodingAttr;
 using ::mlir::triton::gpu::AMDWmmaEncodingAttr;
-using ::mlir::triton::gpu::ConvertLayoutOp;
 using ::mlir::triton::gpu::DotOperandEncodingAttr;
 using ::mlir::triton::gpu::MemDescType;
 using ::triton::gpu::LinearEncodingAttr;
@@ -290,36 +287,6 @@ struct ConvertLayoutOpMFMAToLinearConversion
 protected:
   const TargetInfoBase &targetInfo;
 };
-
-struct ConvertLayoutForcedPadding
-    : public ConvertOpToLLVMPattern<ConvertLayoutOp> {
-
-  explicit ConvertLayoutForcedPadding(LLVMTypeConverter &typeConverter,
-                                      const TargetInfoBase &targetInfo,
-                                      PatternBenefit benefit)
-      : ConvertOpToLLVMPattern<ConvertLayoutOp>(typeConverter, benefit),
-        targetInfo(targetInfo) {}
-
-  LogicalResult
-  matchAndRewrite(ConvertLayoutOp op, OpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const override {
-    if (!op->hasAttr(mlir::triton::AMD::AttrSharedMemPadded))
-      return failure();
-    auto srcType = op.getSrc().getType();
-    auto dstType = op.getType();
-    if (!cvtNeedsSharedMemory(srcType, dstType))
-      return failure();
-
-    auto result = transferWithinBlockPadding(op, adaptor.getSrc(), targetInfo,
-                                             getTypeConverter(), rewriter);
-    rewriter.replaceOp(op, result);
-    return success();
-  }
-
-protected:
-  const TargetInfoBase &targetInfo;
-};
-
 } // namespace
 
 void mlir::triton::AMD::populateConvertLayoutOpToLLVMPatterns(
@@ -329,5 +296,4 @@ void mlir::triton::AMD::populateConvertLayoutOpToLLVMPatterns(
                                                      benefit);
   patterns.add<ConvertLayoutOpMFMAToLinearConversion>(typeConverter, targetInfo,
                                                       benefit);
-  patterns.add<ConvertLayoutForcedPadding>(typeConverter, targetInfo, benefit);
 }

third_party/amd/lib/TritonAMDGPUToLLVM/ConvertLayoutOpToLLVM.cpp

@@ -94,26 +94,15 @@ class OptimizeAMDLDSUsage
     LDBG("Trying fit " << cvtOp << " into " << targetLDSSize << " bytes");
     OpBuilder builder(cvtOp);
 
-    auto ctx = builder.getContext();
     auto srcType = cvtOp.getSrc().getType();
     auto dstType = cvtOp.getType();
 
-    if (!cvtOp->hasAttr(triton::AMD::AttrSharedMemPadded)) {
-      auto emptyAttribute = UnitAttr::get(ctx);
-      // Padded conversion seems more friendly with this optimization
-      // use it instead of general swizzling.
-      cvtOp->setAttr(triton::AMD::AttrSharedMemPadded, emptyAttribute);
-      // if padded layout drops LDS usage on itself, we are done, return
-      if (triton::AMD::getConvertLayoutScratchInBytes(
-              srcType, dstType, /*usePadding*/ true) <= targetLDSSize)
-        return;
-    }
-
     auto srcEnc =
         cast<triton::gpu::DistributedEncodingTrait>(srcType.getEncoding());
     auto dstEnc =
         cast<triton::gpu::DistributedEncodingTrait>(dstType.getEncoding());
 
+    auto ctx = srcEnc.getContext();
     auto rank = srcType.getRank();
 
     unsigned numWarps = triton::gpu::lookupNumWarps(cvtOp);
@@ -255,6 +244,13 @@ class OptimizeAMDLDSUsage
       LDSLimit = targetInfo.getSharedMemorySize();
     }
 
+    auto context = mod.getContext();
+    auto emptyAttribute = UnitAttr::get(context);
+    // TODO choose between padded and swizzled memory patterns
+    mod.walk([emptyAttribute](triton::gpu::ConvertLayoutOp op) -> void {
+      op->setAttr(mlir::triton::AMD::AttrSharedMemPadded, emptyAttribute);
+    });
+
     ModuleAllocation allocAnalysis(
         mod, mlir::triton::AMD::AMDAllocationAnalysisScratchSizeFn);
     if (allocAnalysis.getSharedMemorySize() <= LDSLimit)