[BACKEND] Share ld/st.shared lowering between convert_layout and local_load/store (#7248)

lezcano · web-flow · commit dcf41f6af8b4 · 2025-06-20T10:48:05.000Z
We recently heavily optimised the ld/st.shared lowering for convert_layout. Let's reuse it in other ops for better perf and to have just one source of truth. For now we just activate it for NVIDIA's path. Fixing it for AMD should be easy (and beneficial) but I don't have access to an AMD GPU so... cc @antiagainst
diff --git a/include/triton/Conversion/TritonGPUToLLVM/Utility.h b/include/triton/Conversion/TritonGPUToLLVM/Utility.h
@@ -574,6 +574,15 @@ lowerLdStShared(Location loc, MLIRContext *ctx, LinearLayout cvt,
                 ConversionPatternRewriter &rewriter,
                 const TargetInfoBase &targetInfo);
 
+// Lower local_load/local_store via ld.shared/st.shared
+SmallVector<Value> lowerLocalLdSt(Location loc, MLIRContext *ctx,
+                                  // Map from registers to offset
+                                  LinearLayout cvt, ArrayRef<Value> valsArray,
+                                  // Input for store, output for load
+                                  Type llvmElemTy, Value smemBase,
+                                  ConversionPatternRewriter &rewriter,
+                                  const TargetInfoBase &targetInfo);
+
 SmallVector<Value> unpackLLElements(Location loc, Value llvmStruct,
                                     RewriterBase &rewriter);
 
diff --git a/lib/Conversion/TritonGPUToLLVM/MemoryOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/MemoryOpToLLVM.cpp
@@ -29,6 +29,36 @@ void lowerDistributedToShared(
                            targetInfo, llvmOpCount);
 }
 
+LogicalResult lowerLocalStore(Location loc, MLIRContext *ctx, Value regVal,
+                              MemDescType memDescTy, SharedMemoryObject smemObj,
+                              ArrayRef<Value> inVals,
+                              const LLVMTypeConverter *typeConverter,
+                              ConversionPatternRewriter &rewriter,
+                              const TargetInfoBase &targetInfo) {
+  auto regTy = cast<RankedTensorType>(regVal.getType());
+  auto llvmElemTy = typeConverter->convertType(memDescTy.getElementType());
+
+  auto regLayout = toLinearLayout(regTy.getShape(), regTy.getEncoding());
+  auto sharedLayout =
+      toLinearLayout(memDescTy.getShape(), memDescTy.getEncoding());
+  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");
+  cvt = cvt.sublayout({kReg, kLane, kWarp}, {kOffset});
+  lowerLocalLdSt(loc, ctx, cvt, inVals, llvmElemTy, smemObj.getBase(), rewriter,
+                 targetInfo);
+
+  return success();
+}
+
 struct GlobalScratchAllocOpConversion
     : public ConvertOpToLLVMPattern<triton::gpu::GlobalScratchAllocOp> {
   const TargetInfoBase *targetInfo;
@@ -77,17 +107,34 @@ struct LocalAllocOpConversion
     Location loc = op->getLoc();
     Value smemBase =
         LLVM::getSharedMemoryBase(loc, rewriter, targetInfo, op.getOperation());
-    auto resultTy = cast<MemDescType>(op.getType());
+    auto memDescTy = cast<MemDescType>(op.getType());
     auto typeConverter = getTypeConverter();
 
-    auto llvmElemTy = typeConverter->convertType(resultTy.getElementType());
-    auto smemObj = SharedMemoryObject(smemBase, llvmElemTy, resultTy.getRank(),
+    auto llvmElemTy = typeConverter->convertType(memDescTy.getElementType());
+    auto smemObj = SharedMemoryObject(smemBase, llvmElemTy, memDescTy.getRank(),
                                       loc, rewriter);
     // If there is an initial tensor, store it into the shared memory.
     if (op.getSrc()) {
-      lowerDistributedToShared(loc, op.getSrc(), op.getResult(),
-                               adaptor.getSrc(), smemObj, typeConverter,
-                               rewriter, targetInfo);
+      // [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);
     rewriter.replaceOp(op, retVal);
@@ -122,27 +169,48 @@ struct LocalLoadOpConversion : public ConvertOpToLLVMPattern<LocalLoadOp> {
   LogicalResult
   matchAndRewrite(LocalLoadOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    return lowerSharedToDistributed(op, adaptor, getTypeConverter(), rewriter);
-  }
-
-private:
-  LogicalResult
-  lowerSharedToDistributed(LocalLoadOp op, LocalLoadOpAdaptor adaptor,
-                           const LLVMTypeConverter *typeConverter,
-                           ConversionPatternRewriter &rewriter) const {
     auto loc = op.getLoc();
-    auto srcTy = op.getSrc().getType();
-    auto dstTy = op.getResult().getType();
+    auto *ctx = op.getContext();
+    auto memDescVal = op.getSrc();
+    auto regVal = op.getResult();
+    auto memDescTy = cast<MemDescType>(memDescVal.getType());
+    auto regTy = cast<RankedTensorType>(regVal.getType());
+    auto typeConverter = getTypeConverter();
 
     auto smemObj = LLVM::getSharedMemoryObjectFromStruct(
         loc, adaptor.getSrc(),
-        typeConverter->convertType(srcTy.getElementType()), rewriter);
-    auto elemLlvmTy = typeConverter->convertType(dstTy.getElementType());
+        typeConverter->convertType(memDescTy.getElementType()), rewriter);
+    auto llvmElemTy = typeConverter->convertType(regTy.getElementType());
+
+    // 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();
+    }
 
-    SmallVector<Value> outVals = loadSharedToDistributed(
-        op, elemLlvmTy, smemObj, loc, rewriter, targetInfo);
+    auto regLayout = toLinearLayout(regTy.getShape(), regTy.getEncoding());
+    auto sharedLayout =
+        toLinearLayout(memDescTy.getShape(), memDescTy.getEncoding());
+    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");
+    cvt = cvt.sublayout({kReg, kLane, kWarp}, {kOffset});
+
+    auto outVals = lowerLocalLdSt(op.getLoc(), ctx, cvt, {}, llvmElemTy,
+                                  smemObj.getBase(), rewriter, targetInfo);
 
-    Value result = packLLElements(loc, typeConverter, outVals, rewriter, dstTy);
+    Value result = packLLElements(loc, typeConverter, outVals, rewriter, regTy);
     rewriter.replaceOp(op, result);
 
     return success();
@@ -167,20 +235,30 @@ struct LocalStoreOpConversion
   LogicalResult
   matchAndRewrite(triton::gpu::LocalStoreOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    auto *ctx = op.getContext();
+    Value regVal = op.getSrc();
     Value memDescVal = op.getDst();
-    auto llvmElemTy =
-        getTypeConverter()->convertType(op.getDst().getType().getElementType());
-    auto smemObj = LLVM::getSharedMemoryObjectFromStruct(
-        op.getLoc(), adaptor.getDst(), llvmElemTy, rewriter);
-
+    auto typeConverter = getTypeConverter();
+    auto memDescTy = cast<MemDescType>(memDescVal.getType());
+    auto llvmElemTy = typeConverter->convertType(memDescTy.getElementType());
+    auto smemObj = LLVM::getSharedMemoryObjectFromStruct(loc, adaptor.getDst(),
+                                                         llvmElemTy, rewriter);
+    auto inVals = unpackLLElements(loc, adaptor.getSrc(), rewriter);
     std::pair<size_t, Type> llvmOpCount;
-    lowerDistributedToShared(op.getLoc(), op.getSrc(), op.getDst(),
-                             adaptor.getSrc(), smemObj, getTypeConverter(),
-                             rewriter, targetInfo, &llvmOpCount);
+    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,
+                               &llvmOpCount);
+    }
 
     targetInfo.localStoreOpAnnotation(op, llvmOpCount.first,
                                       llvmOpCount.second);
-
     rewriter.eraseOp(op);
     return success();
   }
diff --git a/lib/Conversion/TritonGPUToLLVM/Utility.cpp b/lib/Conversion/TritonGPUToLLVM/Utility.cpp
@@ -588,6 +588,35 @@ lowerLdStShared(Location loc, MLIRContext *ctx, LinearLayout cvt,
   return outVals;
 }
 
+SmallVector<Value> lowerLocalLdSt(Location loc, MLIRContext *ctx,
+                                  LinearLayout cvt, ArrayRef<Value> valsArray,
+                                  // Input for store, output for load
+                                  Type llvmElemTy, Value smemBase,
+                                  ConversionPatternRewriter &rewriter,
+                                  const TargetInfoBase &targetInfo) {
+  assert(cvt.getNumOutDims() == 1);
+  assert(*cvt.getOutDimNames().begin() == str_attr("offset"));
+  auto isStore = !valsArray.empty();
+  // Remove broadcasting in the registers
+  auto removeBroadcastSrc = actionRemoveBroadcastedRegs(cvt);
+  if (!removeBroadcastSrc.isIdentity()) {
+    auto prmtCvt = removeBroadcastSrc.apply(cvt);
+    auto inVals = to_vector(valsArray);
+    if (isStore) {
+      inVals = removeBroadcastSrc.apply(inVals);
+    }
+    auto outVals = lowerLdStShared(loc, ctx, prmtCvt, inVals, llvmElemTy,
+                                   smemBase, rewriter, targetInfo);
+    if (!isStore) {
+      outVals = broadcastAs(outVals, cvt);
+    }
+    return outVals;
+  }
+
+  return lowerLdStShared(loc, ctx, cvt, valsArray, llvmElemTy, smemBase,
+                         rewriter, targetInfo);
+}
+
 bool emitTransferBetweenRegistersAndShared(
     LinearLayout &regLayout, triton::gpu::MemDescType sharedTy, Type elemLlvmTy,
     std::optional<int32_t> maxVecElems, const SharedMemoryObject &smemObj,
diff --git a/test/Conversion/tritongpu_to_llvm.mlir b/test/Conversion/tritongpu_to_llvm.mlir
@@ -1918,7 +1918,7 @@ module attributes {"ttg.target" = "cuda:75", "ttg.num-ctas" = 1 : i32, "ttg.num-
 module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
   // CHECK-LABEL: @vectorize_shmem_load
   // CHECK: llvm.load
-  // CHECK-SAME: {alignment = 8 : i64} : !llvm.ptr<3> -> vector<8xi8>
+  // CHECK-SAME: {alignment = 8 : i64} : !llvm.ptr<3> -> vector<2xi32>
   // CHECK-NOT: llvm.load
   tt.func public @vectorize_shmem_load(%shmem : !ttg.memdesc<16x16xi8, #shared, #smem>) {
     %0 = ttg.local_load %shmem : !ttg.memdesc<16x16xi8, #shared, #smem> -> tensor<16x16xi8, #blocked>
@@ -1933,9 +1933,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
 #smem = #ttg.shared_memory
 module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, "ttg.threads-per-warp" = 32 : i32} {
   // CHECK-LABEL: @vectorize_shmem_store
-  // CHECK: llvm.store
-  // CHECK-SAME: {alignment = 64 : i64} : vector<16xi32>, !llvm.ptr<3>
-  // CHECK-NOT: llvm.store
+  // CHECK-COUNT-4:  llvm.store {{.*}} {alignment = 16 : i64} : vector<4xi32>, !llvm.ptr<3>
   tt.func public @vectorize_shmem_store(%block : tensor<64x64xi32, #blocked>) {
     %0 = ttg.local_alloc %block : (tensor<64x64xi32, #blocked>) -> !ttg.memdesc<64x64xi32, #shared, #smem>
     tt.return
