[BACKEND] Don't accumulate the offsets from memdesc_subview on the base (#7515)

Layouts coming from `memdesc_subview` are at heart affine layouts.
This is because partial evaluation of a linear map on some variables
gives you an affine map. More concretely, if `c` is a constant
and `x` is a variable, we have that `A(x ^ c) = A(x) ^ A(c)` where
`A(c)` is a constant. In `memdesc_subview`, `A` is the map from the
matrix into shared memory (the inverse of the shared memory linear
layout)
and `c` are the offsets given in the IR of `memdesc_subview`.

Previously `memdesc_subview` would advance the pointer as `ptr += A(c)`.
This is incorrect, as the actual formula for the address is given by
`ptr + (A(c) ^ A(x))`, so we compensated for this when computing the
address by substracting the offsets and adding them as an xor

https://github.com/triton-lang/triton/blob/8e52b2e483eb072149801443e2e33b0b72d32bc5/lib/Conversion/TritonGPUToLLVM/Utility.cpp#L604-L605

In this PR we untangle these issues by:

1. Just advancing the base ptr when we index over the pipelining
   dimensions (this will be split out into its own op in a future PR)
2. Exposing two methods on SharedMemoryObject to compute the affine part
   of the layout (what we called `A(c)` above) and compute the bits that
   `A(c)` may have on. This second part is useful to perform
   optimisations.
3. Decreeing that shared layout will always return the layout of the
   full allocShape (minus the pipelining dimension). This means that
   `toLinearLayout(MemDescType)` may return a layout of a different
   output shape as the tensor it represents. This should not be a
   problem because of the next poitn
4. To account for the previous point, we generalise `invertAndCompose`
to solve systems `AX = B` where `A` may have an output dimension larger
   than `B`. In this case we still compute the linear map compositon
   `A^{-1}B`, which is well defined as `Im(B) \subset Im(A)`

In the future we could consider always describing shared layouts as maps
from the tensor into the offsets. This would allow us to represent the
linear part of the affine map above via linear layouts, and we could
simply return a layout of the correct shape from
`toLinearLayout(MemDescType)`, rather than an oversized one.
This makes also intuitive sense, as we will never want to store the same
element in two different parts of the shared memory, so this map will
always be well-defined, while its inverse may not be, as it is the case
here.

We also take the chance to give a good clean-up to
`emitTransferBetweenRegistersAndShared` now that the logic is better
defined. The generated code should be comparable or better than the
previous one.

---------

Co-authored-by: apgoucher <[email protected]>

Author: lezcano

include/triton/Conversion/TritonGPUToLLVM/Utility.h

@@ -352,6 +352,21 @@ class SharedMemoryObject {
   SmallVector<Value> getStrides(triton::gpu::MemDescType memDesc, Location loc,
                                 RewriterBase &rewriter) const;
 
+  // Returns a mask representing all the bits of the memdesc offsets that
+  // may be modified by an affine offset coming from a memdesc_subview.
+  // The offsets are considered to be in the type of the memdesc.
+  // For padded layouts, we return the offsets without padding.
+  static uint64_t getMaskSpanOffsets(triton::gpu::MemDescType srcTy);
+
+  // Returns whether the shared memory access had a memdesc_subview
+  // that is rank-preserving (soon to be called memdesc_slice)
+  static bool isAffineSharedMemoryAccess(triton::gpu::MemDescType srcTy) {
+    return getMaskSpanOffsets(srcTy) != 0;
+  }
+
+  Value getShmemOffset(Location loc, RewriterBase &rewriter,
+                       triton::gpu::MemDescType srcTy) const;
+
   // TODO(Keren): deprecate the method once AMD backend has cleaned up
   Value getCSwizzleOffset(int dim) const {
     assert(dim >= 0 && dim < offsets.size());
@@ -462,7 +477,6 @@ std::pair<Value, Value> getLaneAndWarpId(OpBuilder &rewriter, Location loc);
 // -----------------------------------------------------------------------
 using LLVM::SharedMemoryObject;
 using ::mlir::LLVM::delinearize;
-using ::mlir::LLVM::SharedMemoryObject;
 using ::mlir::triton::gpu::AMDMfmaEncodingAttr;
 using ::mlir::triton::gpu::AMDWmmaEncodingAttr;
 using ::mlir::triton::gpu::BlockedEncodingAttr;
@@ -474,24 +488,6 @@ using ::mlir::triton::gpu::SliceEncodingAttr;
 Value dot(RewriterBase &rewriter, Location loc, ArrayRef<Value> offsets,
           ArrayRef<Value> strides);
 
-/// Extend 2d shared object to 3d.
-///
-/// If tensor has 3 dimensions, returns original shared object.
-/// If tensor shape is [M, N], return shared object describing shape [1, M, N]
-///
-/// This Function is used to simplify processing of 2d and 3d dot operands,
-/// particularly in the conversion of local_load operation.
-///
-/// \param rewriter
-/// \param loc
-/// \param smemObj
-/// \param shape shape of a tensor represented by smemObj
-/// \returns shared object describing 3d tensor
-SharedMemoryObject
-getExpandedSharedMemoryObject(ConversionPatternRewriter &rewriter, Location loc,
-                              SharedMemoryObject smemObj,
-                              ArrayRef<int64_t> shape);
-
 // "Applies" the given layout by computing layout(indices) and returning the
 // resulting Values.
 //
@@ -568,7 +564,8 @@ void storeDistributedToShared(triton::gpu::MemDescType dstTy,
 SmallVector<Value>
 lowerLdStShared(Location loc, MLIRContext *ctx, LinearLayout cvt,
                 ArrayRef<Value> valsArray, // Input for store, output for load
-                Type llvmElemTy, Value smemBase,
+                Type llvmElemTy, Value smemBase, Value affineOffset,
+                uint64_t maskSpanAffineOffset,
                 ConversionPatternRewriter &rewriter,
                 const TargetInfoBase &targetInfo);
 
@@ -578,20 +575,21 @@ lowerLdStShared(Location loc, MLIRContext *ctx, LinearLayout cvt,
 SmallVector<Value> lowerLdSt(
     Location loc, MLIRContext *ctx, LinearLayout cvt,
     ArrayRef<Value> valsArray, // Input for store, output for load
-    Type llvmElemTy, Value smemBase, ConversionPatternRewriter &rewriter,
+    Type llvmElemTy, Value smemBase, Value affineOffset,
+    uint64_t maskSpanAffineOffset, ConversionPatternRewriter &rewriter,
     const TargetInfoBase &targetInfo, std::optional<int> maybeMaxVecElems,
     std::function<SmallVector<Value>(ConversionPatternRewriter &, Location,
                                      ArrayRef<Value>, Value, int, VectorType)>
         lowerInst);
 
 // 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>
+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,
+               const TargetInfoBase &targetInfo);
 
 SmallVector<Value> unpackLLElements(Location loc, Value llvmStruct,
                                     RewriterBase &rewriter);

include/triton/Tools/LayoutUtils.h

@@ -126,7 +126,7 @@ std::optional<ColumnAction> regPermForDivide(const LinearLayout &A,
 ColumnAction actionRemoveBroadcastedRegs(const LinearLayout &layout);
 
 std::pair<int64_t, ColumnAction>
-actionAdditiveStrides(const LinearLayout &layout);
+actionAdditiveStrides(const LinearLayout &layout, uint64_t maskSpanOffsets);
 
 // For a layout A with A.hasInDim(kReg), repeat the values so that they have
 // the same broadcasting as layout

lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp

@@ -202,6 +202,8 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
 
     assert(permutedInVals.size() == tileSize * nReps);
     SmallVector<Value> outVals;
+    auto affineOffset = b.i32_val(0);
+    auto maskSpanAffineOffset = 0;
     for (int i = 0; i < nReps; ++i) {
       if (i > 0)
         b.barrier();
@@ -210,11 +212,12 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
           ArrayRef<Value>(permutedInVals).slice(i * tileSize, tileSize);
       // Store
       lowerLdStShared(loc, ctx, storeCvt, tileInVals, llvmElemTy, smemBase,
-                      rewriter, targetInfo);
+                      affineOffset, maskSpanAffineOffset, rewriter, targetInfo);
       b.barrier();
       // Load
       SmallVector<Value> tileOutVals = lowerLdStShared(
-          loc, ctx, loadCvt, {}, llvmElemTy, smemBase, rewriter, targetInfo);
+          loc, ctx, loadCvt, {}, llvmElemTy, smemBase, affineOffset,
+          maskSpanAffineOffset, rewriter, targetInfo);
       llvm::append_range(outVals, tileOutVals);
     }
 

lib/Conversion/TritonGPUToLLVM/MemoryOpToLLVM.cpp

@@ -53,8 +53,8 @@ LogicalResult lowerLocalStore(Location loc, MLIRContext *ctx, Value regVal,
   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);
+  lowerLocalLdSt(loc, ctx, cvt, inVals, llvmElemTy, memDescTy, smemObj,
+                 rewriter, targetInfo);
 
   return success();
 }
@@ -177,10 +177,9 @@ struct LocalLoadOpConversion : public ConvertOpToLLVMPattern<LocalLoadOp> {
     auto regTy = cast<RankedTensorType>(regVal.getType());
     auto typeConverter = getTypeConverter();
 
-    auto smemObj = LLVM::getSharedMemoryObjectFromStruct(
-        loc, adaptor.getSrc(),
-        typeConverter->convertType(memDescTy.getElementType()), rewriter);
-    auto llvmElemTy = typeConverter->convertType(regTy.getElementType());
+    auto llvmElemTy = typeConverter->convertType(memDescTy.getElementType());
+    auto smemObj = LLVM::getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(),
+                                                         llvmElemTy, rewriter);
 
     // See [Legacy local_load/local_store]
     if (!targetInfo.isCuda()) {
@@ -206,8 +205,8 @@ struct LocalLoadOpConversion : public ConvertOpToLLVMPattern<LocalLoadOp> {
     auto kOffset = str_attr("offset");
     cvt = cvt.sublayout({kReg, kLane, kWarp}, {kOffset});
 
-    auto outVals = lowerLocalLdSt(op.getLoc(), ctx, cvt, {}, llvmElemTy,
-                                  smemObj.getBase(), rewriter, targetInfo);
+    auto outVals = lowerLocalLdSt(loc, ctx, cvt, {}, llvmElemTy, memDescTy,
+                                  smemObj, rewriter, targetInfo);
 
     Value result = packLLElements(loc, typeConverter, outVals, rewriter, regTy);
     rewriter.replaceOp(op, result);