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

Fixes #4795

Author: whitneywhtsang

include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h

@@ -97,6 +97,7 @@ class TargetInfoBase {
   virtual bool supportLdMatrix() const { return false; }
   virtual bool supportStMatrix() const { return false; }
   virtual bool isCuda() const { return false; }
+  virtual bool isXpu() const { return false; }
 
   // Annotate target specific information to local load operations during
   // lowering to LLVM. `llLoadOp` is the generated LLVM load op.

lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp

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

test/Conversion/amd/convert_layout.mlir

@@ -0,0 +1,52 @@
+// RUN: triton-opt %s -split-input-file --allocate-shared-memory --convert-triton-amdgpu-to-llvm=arch=gfx942 --cse| FileCheck %s
+
+#blocked0 = #ttg.blocked<{sizePerThread = [1, 4], threadsPerWarp = [16, 4], warpsPerCTA = [2, 2], order = [1, 0], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [1, 0]}>
+#blocked1 = #ttg.blocked<{sizePerThread = [4, 1], threadsPerWarp = [4, 16], warpsPerCTA = [2, 2], order = [0, 1], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [1, 0]}>
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "hip:gfx942", "ttg.threads-per-warp" = 64 : i32} {
+  // CHECK: llvm.mlir.global external @global_smem
+  tt.func @convert_layout_general_swizzling(%arg0: tensor<64x64xf32, #blocked0>, %arg1: tensor<64x64x!tt.ptr<f32>, #blocked1>) {
+
+    // verify that following convert layout uses general swizzling path
+
+    // CHECK: [[CST_128:%.*]] = llvm.mlir.constant(128 : i32) : i32
+
+    // Part of offset computation generated by applyLinearLayout function
+    // CHECK: [[SEL:%.*]]= llvm.select {{.*}}, {{.*}}, [[CST_128]]
+    // CHECK: [[OFFSET_0:%.*]] = llvm.xor {{.*}}, [[SEL]]
+    // CHECK: [[OFFSET_1:%.*]] = llvm.xor {{.*}}, [[OFFSET_0]] : i32
+
+    // Part of offset computation generated by lowerLdSt function after applyLinearLayout
+    // CHECK: [[OFFSET_2:%.*]] = llvm.xor [[OFFSET_1]], {{.*}} : i32
+    // CHECK: [[OFFSET_3:%.*]] = llvm.xor [[OFFSET_2]], {{.*}} : i32
+    // CHECK: [[OFFSET_4:%.*]] = llvm.add [[OFFSET_3]], {{.*}} : i32
+    // CHECK: llvm.getelementptr inbounds {{.*}}{{\[}}[[OFFSET_4]]{{\]}}
+
+    %0 = ttg.convert_layout %arg0 : tensor<64x64xf32, #blocked0> -> tensor<64x64xf32, #blocked1>
+    tt.store %arg1, %0 : tensor<64x64x!tt.ptr<f32>, #blocked1>
+    tt.return
+  }
+}
+
+// -----
+
+#blocked0 = #ttg.blocked<{sizePerThread = [1, 4], threadsPerWarp = [16, 4], warpsPerCTA = [2, 2], order = [1, 0], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [1, 0]}>
+#blocked1 = #ttg.blocked<{sizePerThread = [4, 1], threadsPerWarp = [4, 16], warpsPerCTA = [2, 2], order = [0, 1], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [1, 0]}>
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "hip:gfx942", "ttg.threads-per-warp" = 64 : i32} {
+  // CHECK-LABEL: convert_layout_padding_swizzling
+  tt.func @convert_layout_padding_swizzling(%arg0: tensor<64x64xf32, #blocked0>, %arg1: tensor<64x64x!tt.ptr<f32>, #blocked1>) {
+
+    // verify that following convert layout uses padded path
+    // see getVecAddr lambda in transferWithinBlockImpl function
+
+    // CHECK-DAG: [[CST_0:%.*]] = llvm.mlir.constant(0 : i32) : i32
+    // CHECK-DAG: [[CST_5:%.*]] = llvm.mlir.constant(5 : i32) : i32
+    // CHECK-DAG: [[OFFSET_0:%.*]] = llvm.lshr {{.*}}, [[CST_5]] : i32
+    // CHECK: [[OFFSET_1:%.*]] = llvm.shl [[OFFSET_0]], [[CST_0]] : i32
+    // CHECK: [[OFFSET_2:%.*]] = llvm.add [[OFFSET_1]], {{.*}} : i32
+    // CHECK: llvm.getelementptr inbounds {{.*}}{{\[}}[[OFFSET_2]]{{\]}}
+
+    %0 = ttg.convert_layout %arg0 {amdgpu.use_padded_scratch_shmem} : tensor<64x64xf32, #blocked0> -> tensor<64x64xf32, #blocked1>
+    tt.store %arg1, %0 : tensor<64x64x!tt.ptr<f32>, #blocked1>
+    tt.return
+  }
+}

third_party/amd/lib/TritonAMDGPUToLLVM/ConvertLayoutOpToLLVM.cpp

@@ -1,3 +1,4 @@
+#include "Analysis/AMDGPUAllocation.h"
 #include "PatternTritonGPUOpToLLVM.h"
 #include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h"
 #include "triton/Conversion/TritonGPUToLLVM/Utility.h"
@@ -110,6 +111,35 @@ class ConvertLayoutOpMFMAToLinearConversion
     return success();
   }
 
+private:
+  const TargetInfoBase &targetInfo;
+};
+
+class ConvertLayoutForcedPadding
+    : public ConvertOpToLLVMPattern<ConvertLayoutOp> {
+public:
+  ConvertLayoutForcedPadding(LLVMTypeConverter &typeConverter,
+                             const TargetInfoBase &targetInfo,
+                             PatternBenefit benefit)
+      : ConvertOpToLLVMPattern(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();
+  }
+
 private:
   const TargetInfoBase &targetInfo;
 };
@@ -120,4 +150,7 @@ void mlir::triton::AMD::populateConvertLayoutOpToLLVMPatterns(
     RewritePatternSet &patterns, PatternBenefit benefit) {
   patterns.add<ConvertLayoutOpMFMAToLinearConversion>(typeConverter, targetInfo,
                                                       benefit);
+  patterns.add<ConvertLayoutForcedPadding>(typeConverter, targetInfo, benefit);
+  // No need to convert when ForcedSwizzling as it's already the default
+  // lowering
 }

third_party/amd/lib/TritonAMDGPUToLLVM/OptimizeLDSUsage.cpp

@@ -94,15 +94,26 @@ 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);
@@ -244,13 +255,6 @@ 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)

third_party/intel/lib/TritonIntelGPUToLLVM/TargetInfo.h

@@ -76,6 +76,8 @@ class TargetInfo : public mlir::triton::TargetInfoBase {
                              StringRef name, StringRef value,
                              unsigned addressSpace) const;
 
+  bool isXpu() const override { return true; }
+
 protected:
   virtual bool isSupportedWarpReduceOp(Operation *op, unsigned numLanesToReduce,
                                        unsigned warpSize) const = 0;