Revert "Revert "[AMD] Reuse explicit swizzling pattern for ConvertLayoutOp (#7636)"" (#4888)

This reverts commit 77c9289.

Fixes #4885

Author: whitneywhtsang

include/triton/Conversion/TritonGPUToLLVM/Utility.h

@@ -557,13 +557,14 @@ Value emitPadding(Location loc, RewriterBase &rewriter,
 // calcPaddedOffset is a lambda that takes a base offset (mlir::Value)
 // and computes a new offset (mlir::Value) by applying padding based on
 // shared memory layout.
-SmallVector<Value> lowerLdStShared(
-    Location loc, MLIRContext *ctx, LinearLayout cvt,
-    ArrayRef<Value> valsArray, // Input for store, output for load
-    Type llvmElemTy, Value smemBase,
-    std::function<Value(Value)> calcPaddedOffset, Value affineOffset,
-    uint64_t maskSpanAffineOffset, ConversionPatternRewriter &rewriter,
-    const TargetInfoBase &targetInfo, Operation *localLoadOp = nullptr);
+SmallVector<Value>
+lowerLdStShared(Location loc, MLIRContext *ctx, LinearLayout cvt,
+                ArrayRef<Value> valsArray, // Input for store, output for load
+                Type llvmElemTy, Value smemBase,
+                std::function<Value(Value)> calcPaddedOffset,
+                Value affineOffset, uint64_t maskSpanAffineOffset,
+                RewriterBase &rewriter, const TargetInfoBase &targetInfo,
+                Operation *localLoadOp = nullptr);
 
 // Lower an ld/st-like operation given a layout and a callback that creates the
 // PTX instruction Lowers to st when valArrays is empty, and to ld when it is
@@ -576,10 +577,10 @@ SmallVector<Value> lowerLdSt(
     ArrayRef<Value> valsArray, // Input for store, output for load
     Type llvmElemTy, Value smemBase,
     std::function<Value(Value)> calcPaddedOffset, Value affineOffset,
-    uint64_t maskSpanAffineOffset, ConversionPatternRewriter &rewriter,
+    uint64_t maskSpanAffineOffset, RewriterBase &rewriter,
     const TargetInfoBase &targetInfo, std::optional<int> maybeMaxVecElems,
-    std::function<SmallVector<Value>(ConversionPatternRewriter &, Location,
-                                     ArrayRef<Value>, Value, int, VectorType)>
+    std::function<SmallVector<Value>(RewriterBase &, Location, ArrayRef<Value>,
+                                     Value, int, VectorType)>
         lowerInst);
 
 // Lower local_load/local_store via ld.shared/st.shared
@@ -588,7 +589,7 @@ lowerLocalLdSt(Location loc, MLIRContext *ctx,
                LinearLayout cvt,          // Map from registers to offset
                ArrayRef<Value> valsArray, // Input for store, empty for load
                Type llvmElemTy, triton::gpu::MemDescType srcTy,
-               SharedMemoryObject smemObj, ConversionPatternRewriter &rewriter,
+               SharedMemoryObject smemObj, RewriterBase &rewriter,
                const TargetInfoBase &targetInfo,
                Operation *localLoadOp = nullptr);
 
@@ -643,6 +644,12 @@ Value transferWithinBlockPadding(triton::gpu::ConvertLayoutOp op, Value src,
                                  const LLVMTypeConverter *typeConverter,
                                  RewriterBase &rewriter);
 
+LogicalResult
+transferWithinBlockSwizzling(triton::gpu::ConvertLayoutOp op, Value src,
+                             const TargetInfoBase &targetInfo,
+                             const LLVMTypeConverter *typeConverter,
+                             RewriterBase &rewriter);
+
 SmallVector<Value> inlineRegionImpl(RewriterBase &rewriter, Region &region,
                                     ArrayRef<Value> args,
                                     mlir::TypeID terminatorTypeId,

lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp

@@ -110,167 +110,6 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
     return success();
   }
 
-  SmallVector<Value> transferWithinBlockSwizzlingImpl(
-      Location loc, ConversionPatternRewriter &rewriter,
-      const LinearLayout &srcLayout, const LinearLayout &dstLayout,
-      ArrayRef<Value> inVals, Type llvmElemTy, Value smemBase) const {
-    auto *ctx = rewriter.getContext();
-    auto b = TritonLLVMOpBuilder(loc, rewriter);
-    // We handle transformations recursively as they all need a preprocessing
-    // and a postprocessing step.
-
-    // Handle pointer types as 64-bit integers
-    if (isa<LLVM::LLVMPointerType>(llvmElemTy)) {
-      auto llvmElemTyPtr = i64_ty;
-      auto newInVals = llvm::to_vector(llvm::map_range(inVals, [&](Value v) {
-        return b.ptrtoint(llvmElemTyPtr, v).getResult();
-      }));
-      auto outVals =
-          transferWithinBlockSwizzlingImpl(loc, rewriter, srcLayout, dstLayout,
-                                           newInVals, llvmElemTyPtr, smemBase);
-      for (auto &v : outVals) {
-        v = b.inttoptr(llvmElemTy, v);
-      }
-      return outVals;
-    }
-
-    // Handle sub-byte elements like i1
-    if (llvmElemTy.getIntOrFloatBitWidth() < 8) {
-      // Upcast to i8
-      auto i8ElemTy = i8_ty;
-      auto newInVals = llvm::to_vector(llvm::map_range(
-          inVals, [&](Value v) { return b.zext(i8ElemTy, v).getResult(); }));
-      auto outVals = transferWithinBlockSwizzlingImpl(
-          loc, rewriter, srcLayout, dstLayout, newInVals, i8ElemTy, smemBase);
-      for (auto &v : outVals) {
-        v = b.trunc(llvmElemTy, v);
-      }
-      return outVals;
-    }
-
-    // Remove broadcasting in src
-    auto removeBroadcastSrc = actionRemoveBroadcastedRegs(srcLayout);
-    if (!removeBroadcastSrc.isIdentity()) {
-      auto prmtSrc = removeBroadcastSrc.apply(srcLayout);
-      auto newInVals = removeBroadcastSrc.apply(inVals);
-      return transferWithinBlockSwizzlingImpl(loc, rewriter, prmtSrc, dstLayout,
-                                              newInVals, llvmElemTy, smemBase);
-    }
-
-    // Remove broadcasting in dst
-    auto removeBroadcastDst = actionRemoveBroadcastedRegs(dstLayout);
-    if (!removeBroadcastDst.isIdentity()) {
-      auto prmtDst = removeBroadcastDst.apply(dstLayout);
-      auto outVals = transferWithinBlockSwizzlingImpl(
-          loc, rewriter, srcLayout, prmtDst, inVals, llvmElemTy, smemBase);
-      return broadcastAs(outVals, dstLayout);
-    }
-
-    // At this point we have a type that's at least 8-bit
-    // and we don't have broadcasting in the registers
-    auto bitwidth = llvmElemTy.getIntOrFloatBitWidth();
-    auto smem = optimalSwizzling(srcLayout, dstLayout, bitwidth);
-
-    // Extract reps from smem
-    auto kReg = str_attr("register");
-    auto kReps = str_attr("reps");
-    auto nReps = smem.getInDimSize(kReps);
-    auto reps = LinearLayout::identity1D(nReps, kReg, kReps);
-
-    auto totalStoreCvt = srcLayout.invertAndCompose(smem);
-    auto totalLoadCvt = dstLayout.invertAndCompose(smem);
-
-    // The permutation exists by construction of the reps dimension in
-    // optimalSwizzling
-    auto permStore =
-        regPermForDivide(totalStoreCvt, reps, /*left=*/false).value();
-    totalStoreCvt = permStore.apply(totalStoreCvt);
-    auto permutedInVals = permStore.apply(inVals);
-    auto permLoad =
-        regPermForDivide(totalLoadCvt, reps, /*left=*/false).value();
-    totalLoadCvt = permLoad.apply(totalLoadCvt);
-
-    // Remove the reps and flatten into offset
-    auto storeCvt = *divideRight(totalStoreCvt, reps);
-    auto loadCvt = *divideRight(totalLoadCvt, reps);
-    auto kOffset = str_attr("offset");
-    storeCvt = storeCvt.reshapeOuts({{kOffset, storeCvt.getTotalOutDimSize()}});
-    loadCvt = loadCvt.reshapeOuts({{kOffset, loadCvt.getTotalOutDimSize()}});
-
-    auto tileSize = storeCvt.getInDimSize(kReg);
-
-    assert(permutedInVals.size() == tileSize * nReps);
-    SmallVector<Value> outVals;
-    auto noPaddingOffset = [](Value v) { return v; };
-    auto affineOffset = b.i32_val(0);
-    auto maskSpanAffineOffset = 0;
-    for (int i = 0; i < nReps; ++i) {
-      if (i > 0)
-        b.barrier();
-
-      auto tileInVals =
-          ArrayRef<Value>(permutedInVals).slice(i * tileSize, tileSize);
-      // Store
-      lowerLdStShared(loc, ctx, storeCvt, tileInVals, llvmElemTy, smemBase,
-                      noPaddingOffset, affineOffset, maskSpanAffineOffset,
-                      rewriter, targetInfo);
-      b.barrier();
-      // Load
-      SmallVector<Value> tileOutVals = lowerLdStShared(
-          loc, ctx, loadCvt, {}, llvmElemTy, smemBase, noPaddingOffset,
-          affineOffset, maskSpanAffineOffset, rewriter, targetInfo);
-      llvm::append_range(outVals, tileOutVals);
-    }
-
-    // Undo the permLoad used to divideRight
-    outVals = permLoad.inverse().apply(outVals);
-    return outVals;
-  }
-
-  LogicalResult
-  transferWithinBlockSwizzling(ConvertLayoutOp op, Value src,
-                               ConversionPatternRewriter &rewriter) const {
-    // Fallback for now to standard lowering if it can use stmatrix
-    auto scratchConfig =
-        getScratchConfigForCvt(op.getSrc().getType(), op.getType());
-    bool isStMatrix = targetInfo.canUseStMatrix(
-        op.getSrc().getType(), scratchConfig.repShape,
-        scratchConfig.paddedRepShape, scratchConfig.order,
-        /*swizzleByteSize=*/0);
-    if (isStMatrix) {
-      return failure();
-    }
-
-    auto loc = op.getLoc();
-    auto *ctx = op.getContext();
-    auto srcTy = op.getSrc().getType();
-    auto dstTy = op.getType();
-
-    // Remove the kBlock dimension from the layout as it's the identity in the
-    // cvt
-    auto srcLayout = toLinearLayout(srcTy);
-    auto dstLayout = toLinearLayout(dstTy);
-    auto kReg = str_attr("register");
-    auto kLane = str_attr("lane");
-    auto kWarp = str_attr("warp");
-    srcLayout = srcLayout.sublayout({kReg, kLane, kWarp},
-                                    to_vector(srcLayout.getOutDimNames()));
-    dstLayout = dstLayout.sublayout({kReg, kLane, kWarp},
-                                    to_vector(dstLayout.getOutDimNames()));
-
-    auto llvmElemTy = getTypeConverter()->convertType(srcTy.getElementType());
-    auto smemBase =
-        LLVM::getSharedMemoryBase(loc, rewriter, targetInfo, op.getOperation());
-    auto inVals = unpackLLElements(loc, src, rewriter);
-    auto outVals = transferWithinBlockSwizzlingImpl(
-        loc, rewriter, srcLayout, dstLayout, inVals, llvmElemTy, smemBase);
-
-    Value result =
-        packLLElements(loc, getTypeConverter(), outVals, rewriter, dstTy);
-    rewriter.replaceOp(op, result);
-    return success();
-  }
-
   LogicalResult transferWithinBlock(ConvertLayoutOp op,
                                     const LinearLayout &srcLayout,
                                     const LinearLayout &dstLayout,
@@ -280,8 +119,9 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
 
     // Try to use swizzling to implement the conversion
     // HACK Remove once XPU tests pass for the swizzling path
-    if (!targetInfo.isXpu() && succeeded(transferWithinBlockSwizzling(
-                                   op, adaptor.getSrc(), rewriter))) {
+    if (!targetInfo.isXpu() &&
+        succeeded(transferWithinBlockSwizzling(op, adaptor.getSrc(), targetInfo,
+                                               getTypeConverter(), rewriter))) {
       return success();
     }