Revert "[AMD] Enable lowerLocalLdSt for AMD path (#7355)"

This reverts commit 8a5862d.

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h

@@ -25,8 +25,8 @@ class TargetInfoBase {
                             std::optional<Value> ctaId, Value val,
                             Value pred) const = 0;
   virtual Value loadDShared(RewriterBase &rewriter, Location loc, Value ptr,
-                            std::optional<Value> ctaId, Type elemTy, Value pred,
-                            Operation *localLoadOp = nullptr) const = 0;
+                            std::optional<Value> ctaId, Type elemTy,
+                            Value pred) const = 0;
 
   void storeShared(RewriterBase &rewriter, Location loc, Value ptr, Value val,
                    Value pred) const {

include/triton/Conversion/TritonGPUToLLVM/Utility.h

@@ -515,13 +515,10 @@ SmallVector<SmallVector<Value>>
 emitIndices(Location loc, RewriterBase &rewriter, const TargetInfoBase &target,
             Attribute layout, RankedTensorType type, bool withCTAOffset);
 
-// Emits the required padding given shared memory offset
-// - If `offsetInBytes` is true, smemOffset and padding is assumed in bytes.
-// - If false, smemOffset and padding are assumed to be scaled by element
-// bitwidth, in which case, `bitwidth` is not used.
+// Emits the required padding in elements for the given shared memory offset
 Value emitPadding(Location loc, RewriterBase &rewriter,
                   triton::gpu::PaddedSharedEncodingAttr layout,
-                  unsigned bitwidth, Value smemOffset, bool offsetInBytes);
+                  Value smemOffset);
 
 // Emits IR to load data from shared memory into registers, or to store data
 // from registers into shared memory.
@@ -549,33 +546,39 @@ Value emitPadding(Location loc, RewriterBase &rewriter,
     Value laneId, Value warpId,
     std::function<void(VectorType, Value /*shmemAddr*/)> perVectorCallback);
 
+SmallVector<Value> loadSharedToDistributed(triton::gpu::LocalLoadOp localLoadOp,
+                                           Type elemLlvmTy,
+                                           const SharedMemoryObject &smemObj,
+                                           Location loc, RewriterBase &rewriter,
+                                           const TargetInfoBase &target);
+
+void storeDistributedToShared(triton::gpu::MemDescType dstTy,
+                              RankedTensorType srcTy, Type elemLlvmTy,
+                              ArrayRef<Value> srcVals,
+                              const SharedMemoryObject &smemObj, Location loc,
+                              RewriterBase &rewriter,
+                              const TargetInfoBase &target);
+
 // Close cousin of lowerLdStMatrix in MemoryOpToLLVM.cpp
 // We might want to merge them at some point, but having to support
 // ldmatrix.trans makes the code in lowerLdStMatrix a bit specific
 // Lowers to st when valArrays is empty, and to ld when it is not,
 // and returns the output values.
-// 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, Value affineOffset,
+                uint64_t maskSpanAffineOffset,
+                ConversionPatternRewriter &rewriter,
+                const TargetInfoBase &targetInfo);
 
 // 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
 // not, and returns the output values.
-// 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> lowerLdSt(
     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,
+    Type llvmElemTy, Value smemBase, Value affineOffset,
     uint64_t maskSpanAffineOffset, ConversionPatternRewriter &rewriter,
     const TargetInfoBase &targetInfo, std::optional<int> maybeMaxVecElems,
     std::function<SmallVector<Value>(ConversionPatternRewriter &, Location,
@@ -589,8 +592,7 @@ lowerLocalLdSt(Location loc, MLIRContext *ctx,
                ArrayRef<Value> valsArray, // Input for store, empty for load
                Type llvmElemTy, triton::gpu::MemDescType srcTy,
                SharedMemoryObject smemObj, ConversionPatternRewriter &rewriter,
-               const TargetInfoBase &targetInfo,
-               Operation *localLoadOp = nullptr);
+               const TargetInfoBase &targetInfo);
 
 SmallVector<Value> unpackLLElements(Location loc, Value llvmStruct,
                                     RewriterBase &rewriter);

include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td

@@ -370,7 +370,7 @@ When vec=2, elements are swizzled in pairs of 2.  In other words, the element at
   let genVerifyDecl = 1;
 }
 
-def PaddedSharedEncodingAttr
+def PaddeddSharedEncodingAttr
     : TritonGPU_Attr<"PaddedSharedEncoding", "padded_shared_encoding",
                      [SharedEncodingTrait, LayoutEncodingTrait]> {
   let mnemonic = "padded_shared";

lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp

@@ -201,7 +201,6 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
 
     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) {
@@ -212,13 +211,12 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
           ArrayRef<Value>(permutedInVals).slice(i * tileSize, tileSize);
       // Store
       lowerLdStShared(loc, ctx, storeCvt, tileInVals, llvmElemTy, smemBase,
-                      noPaddingOffset, affineOffset, maskSpanAffineOffset,
-                      rewriter, targetInfo);
+                      affineOffset, maskSpanAffineOffset, rewriter, targetInfo);
       b.barrier();
       // Load
       SmallVector<Value> tileOutVals = lowerLdStShared(
-          loc, ctx, loadCvt, {}, llvmElemTy, smemBase, noPaddingOffset,
-          affineOffset, maskSpanAffineOffset, rewriter, targetInfo);
+          loc, ctx, loadCvt, {}, llvmElemTy, smemBase, affineOffset,
+          maskSpanAffineOffset, rewriter, targetInfo);
       llvm::append_range(outVals, tileOutVals);
     }
 

lib/Conversion/TritonGPUToLLVM/MemoryOpToLLVM.cpp

@@ -12,6 +12,24 @@ using namespace mlir;
 using namespace mlir::triton;
 using namespace mlir::triton::gpu;
 
+// blocked -> shared.
+// Swizzling in shared memory to avoid bank conflict. Normally used for
+// A/B operands of dots.
+void lowerDistributedToShared(Location loc, Value src, Value dst,
+                              Value adaptorSrc,
+                              const SharedMemoryObject &smemObj,
+                              const LLVMTypeConverter *typeConverter,
+                              ConversionPatternRewriter &rewriter,
+                              const TargetInfoBase &targetInfo) {
+  auto srcTy = cast<RankedTensorType>(src.getType());
+  auto dstTy = cast<MemDescType>(dst.getType());
+  auto elemTy = typeConverter->convertType(srcTy.getElementType());
+
+  auto inVals = unpackLLElements(loc, adaptorSrc, rewriter);
+  storeDistributedToShared(dstTy, srcTy, elemTy, inVals, smemObj, loc, rewriter,
+                           targetInfo);
+}
+
 LogicalResult lowerLocalStore(Location loc, MLIRContext *ctx, Value regVal,
                               MemDescType memDescTy, SharedMemoryObject smemObj,
                               ArrayRef<Value> inVals,
@@ -21,25 +39,19 @@ LogicalResult lowerLocalStore(Location loc, MLIRContext *ctx, Value regVal,
   auto regTy = cast<RankedTensorType>(regVal.getType());
   auto llvmElemTy = typeConverter->convertType(memDescTy.getElementType());
 
+  auto regLayout = toLinearLayout(regTy);
+  auto sharedLayout = toLinearLayout(memDescTy);
+  auto cvt = regLayout.invertAndCompose(sharedLayout);
+
+  auto kBlock = str_attr("block");
+  // NYI. We would need to emit a map.shared::cluster instruction.
+  if (!cvt.isTrivialOver({kBlock})) {
+    return failure();
+  }
   auto kReg = str_attr("register");
   auto kLane = str_attr("lane");
   auto kWarp = str_attr("warp");
   auto kOffset = str_attr("offset");
-  auto regLayout = toLinearLayout(regTy);
-  auto paddedLayout =
-      dyn_cast<triton::gpu::PaddedSharedEncodingAttr>(memDescTy.getEncoding());
-  LinearLayout cvt = LinearLayout::empty();
-  if (paddedLayout) {
-    cvt = regLayout.reshapeOuts({{kOffset, regLayout.getTotalOutDimSize()}});
-  } else {
-    auto sharedLayout = toLinearLayout(memDescTy);
-    cvt = regLayout.invertAndCompose(sharedLayout);
-    auto kBlock = str_attr("block");
-    // NYI. We would need to emit a map.shared::cluster instruction.
-    if (!cvt.isTrivialOver({kBlock})) {
-      return failure();
-    }
-  }
   cvt = cvt.sublayout({kReg, kLane, kWarp}, {kOffset});
   lowerLocalLdSt(loc, ctx, cvt, inVals, llvmElemTy, memDescTy, smemObj,
                  rewriter, targetInfo);
@@ -103,12 +115,25 @@ struct LocalAllocOpConversion
                                       loc, rewriter);
     // If there is an initial tensor, store it into the shared memory.
     if (op.getSrc()) {
-      auto *ctx = op.getContext();
-      auto inVals = unpackLLElements(loc, adaptor.getSrc(), rewriter);
-      if (failed(lowerLocalStore(loc, ctx, op.getSrc(), memDescTy, smemObj,
-                                 inVals, typeConverter, rewriter,
-                                 targetInfo))) {
-        return failure();
+      // [Legacy local_load/local_store]
+      // TODO(Lezcano) We should activate this path for other targets as it's
+      // more efficient. AFAIK The main blockers are:
+      // - The legacy path calls localLoadOpAnnotation
+      // - The legacy path calls llvm.load/llvm.store unconditionally, while
+      //   the AMD lowering of storeDShared does not, even when the predicate
+      //   is constant true.
+      if (targetInfo.isCuda()) {
+        auto *ctx = op.getContext();
+        auto inVals = unpackLLElements(loc, adaptor.getSrc(), rewriter);
+        if (failed(lowerLocalStore(loc, ctx, op.getSrc(), memDescTy, smemObj,
+                                   inVals, typeConverter, rewriter,
+                                   targetInfo))) {
+          return failure();
+        }
+      } else {
+        lowerDistributedToShared(loc, op.getSrc(), op.getResult(),
+                                 adaptor.getSrc(), smemObj, typeConverter,
+                                 rewriter, targetInfo);
       }
     }
     auto retVal = getStructFromSharedMemoryObject(loc, smemObj, rewriter);
@@ -156,31 +181,32 @@ struct LocalLoadOpConversion : public ConvertOpToLLVMPattern<LocalLoadOp> {
     auto smemObj = LLVM::getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(),
                                                          llvmElemTy, rewriter);
 
-    auto sharedEnc =
-        cast<triton::gpu::SharedEncodingTrait>(memDescTy.getEncoding());
+    // See [Legacy local_load/local_store]
+    if (!targetInfo.isCuda()) {
+      SmallVector<Value> outVals = loadSharedToDistributed(
+          op, llvmElemTy, smemObj, loc, rewriter, targetInfo);
+      Value result =
+          packLLElements(loc, typeConverter, outVals, rewriter, regTy);
+      rewriter.replaceOp(op, result);
+      return success();
+    }
+
+    auto regLayout = toLinearLayout(regTy);
+    auto sharedLayout = toLinearLayout(memDescTy);
+    auto cvt = regLayout.invertAndCompose(sharedLayout);
+    auto kBlock = str_attr("block");
+    // NYI. We would need to emit a map.shared::cluster instruction.
+    if (!cvt.isTrivialOver({kBlock})) {
+      return failure();
+    }
     auto kReg = str_attr("register");
     auto kLane = str_attr("lane");
     auto kWarp = str_attr("warp");
     auto kOffset = str_attr("offset");
-    auto regLayout = toLinearLayout(regTy);
-    auto paddedLayout =
-        dyn_cast<triton::gpu::PaddedSharedEncodingAttr>(sharedEnc);
-    LinearLayout cvt = LinearLayout::empty();
-    if (paddedLayout) {
-      cvt = regLayout.reshapeOuts({{kOffset, regLayout.getTotalOutDimSize()}});
-    } else {
-      auto sharedLayout = toLinearLayout(memDescTy);
-      cvt = regLayout.invertAndCompose(sharedLayout);
-      auto kBlock = str_attr("block");
-      // NYI. We would need to emit a map.shared::cluster instruction.
-      if (!cvt.isTrivialOver({kBlock})) {
-        return failure();
-      }
-    }
     cvt = cvt.sublayout({kReg, kLane, kWarp}, {kOffset});
 
     auto outVals = lowerLocalLdSt(loc, ctx, cvt, {}, llvmElemTy, memDescTy,
-                                  smemObj, rewriter, targetInfo, op);
+                                  smemObj, rewriter, targetInfo);
 
     Value result = packLLElements(loc, typeConverter, outVals, rewriter, regTy);
     rewriter.replaceOp(op, result);
@@ -217,9 +243,14 @@ struct LocalStoreOpConversion
     auto smemObj = LLVM::getSharedMemoryObjectFromStruct(loc, adaptor.getDst(),
                                                          llvmElemTy, rewriter);
     auto inVals = unpackLLElements(loc, adaptor.getSrc(), rewriter);
-    if (failed(lowerLocalStore(loc, ctx, regVal, memDescTy, smemObj, inVals,
-                               typeConverter, rewriter, targetInfo))) {
-      return failure();
+    if (targetInfo.isCuda()) {
+      if (failed(lowerLocalStore(loc, ctx, regVal, memDescTy, smemObj, inVals,
+                                 typeConverter, rewriter, targetInfo))) {
+        return failure();
+      }
+    } else {
+      lowerDistributedToShared(loc, regVal, memDescVal, adaptor.getSrc(),
+                               smemObj, typeConverter, rewriter, targetInfo);
     }
 
     rewriter.eraseOp(op);