WA for incorrect strides in subview

Author: AlexAUT

include/triton/Conversion/TritonGPUToLLVM/Utility.h

@@ -376,9 +376,13 @@ class SharedMemoryObject {
     return types;
   }
 
-  SmallVector<Value> getStrides(triton::gpu::MemDescType memDesc, Location loc,
-                                RewriterBase &rewriter) const {
+  SmallVector<Value>
+  getStrides(triton::gpu::MemDescType memDesc, Location loc,
+             RewriterBase &rewriter,
+             ArrayRef<int64_t> overwriteAllocSize = {}) const {
     auto allocShape = memDesc.getAllocShape();
+    if (!overwriteAllocSize.empty())
+      allocShape = overwriteAllocSize;
     auto allocShapePerCTA = triton::gpu::getAllocationShapePerCTA(
         memDesc.getEncoding(), allocShape);
     auto layoutOrder = triton::gpu::getOrder(memDesc);
@@ -699,14 +703,14 @@ emitIndices(Location loc, RewriterBase &rewriter, const TargetInfoBase &target,
     const SharedMemoryObject &smemObj, Location loc, RewriterBase &rewriter,
     const TargetInfoBase &target,
     std::function<void(VectorType, Value /*shmemAddr*/)> perVectorCallback,
-    bool forceLane0 = false);
+    bool forceLane0 = false, ArrayRef<int64_t> overwriteAllocSize = {});
 
 [[nodiscard]] bool emitTransferBetweenRegistersAndShared(
     LinearLayout &regLayout, triton::gpu::MemDescType sharedTy, Type elemLlvmTy,
     std::optional<int32_t> maxVecElems, const SharedMemoryObject &smemObj,
     Location loc, RewriterBase &rewriter, const TargetInfoBase &target,
     std::function<void(VectorType, Value /*shmemAddr*/)> perVectorCallback,
-    bool forceLane0 = false);
+    bool forceLane0 = false, ArrayRef<int64_t> overwriteAllocSize = {});
 
 SmallVector<Value> loadSharedToDistributed(triton::gpu::LocalLoadOp localLoadOp,
                                            Type elemLlvmTy,

lib/Conversion/TritonGPUToLLVM/Utility.cpp

@@ -277,7 +277,8 @@ Value getSmemVecAddr(const LinearLayout &regLayout,
                      const SharedMemoryObject &smemObj,
                      triton::gpu::MemDescType sharedTy, Type elemLlvmTy,
                      Value regId, Value laneId, Value warpId, Value blockId,
-                     Location loc, RewriterBase &rewriter) {
+                     Location loc, RewriterBase &rewriter,
+                     ArrayRef<int64_t> overwriteAllocSize) {
   auto b = TritonLLVMOpBuilder(loc, rewriter);
   MLIRContext *ctx = rewriter.getContext();
   StringAttr kBlock = str_attr("block");
@@ -292,7 +293,8 @@ Value getSmemVecAddr(const LinearLayout &regLayout,
 
   auto smemBase = smemObj.getBase();
   auto smemOffsets = smemObj.getOffsets();
-  auto smemStrides = smemObj.getStrides(sharedTy, loc, rewriter);
+  auto smemStrides =
+      smemObj.getStrides(sharedTy, loc, rewriter, overwriteAllocSize);
   Value smemOffset;
   // When loading or storing to shared memory, we consider two cases for
   // performance reasons:
@@ -410,7 +412,7 @@ bool emitTransferBetweenRegistersAndShared(
     std::optional<int32_t> maxVecElems, const SharedMemoryObject &smemObj,
     Location loc, RewriterBase &rewriter, const TargetInfoBase &target,
     std::function<void(VectorType, Value /*shmemAddr*/)> perVectorCallback,
-    bool forceLane0) {
+    bool forceLane0, ArrayRef<int64_t> overwriteAllocSize) {
   MLIRContext *ctx = rewriter.getContext();
   auto b = TritonLLVMOpBuilder(loc, rewriter);
 
@@ -485,9 +487,10 @@ bool emitTransferBetweenRegistersAndShared(
   SmallVector<Value> ret;
   for (int i = 0; i < numElems / vecElems; i++) {
     auto regId = b.i32_val(i * vecElems);
-    auto vecAddr = getSmemVecAddr(
-        regLayout, regToSharedLayout, invertAllocSharedLayout, smemObj,
-        sharedTy, elemLlvmTy, regId, laneId, warpId, blockId, loc, rewriter);
+    auto vecAddr =
+        getSmemVecAddr(regLayout, regToSharedLayout, invertAllocSharedLayout,
+                       smemObj, sharedTy, elemLlvmTy, regId, laneId, warpId,
+                       blockId, loc, rewriter, overwriteAllocSize);
     perVectorCallback(vecTy, vecAddr);
   }
   return true;
@@ -499,12 +502,12 @@ bool emitTransferBetweenRegistersAndShared(
     const SharedMemoryObject &smemObj, Location loc, RewriterBase &rewriter,
     const TargetInfoBase &target,
     std::function<void(VectorType, Value /*shmemAddr*/)> perVectorCallback,
-    bool forceLane0) {
+    bool forceLane0, ArrayRef<int64_t> overwriteAllocSize) {
   auto regLayout = triton::gpu::toLinearLayout(registerTy.getShape(),
                                                registerTy.getEncoding());
   return emitTransferBetweenRegistersAndShared(
       regLayout, sharedTy, elemLlvmTy, maxVecElems, smemObj, loc, rewriter,
-      target, perVectorCallback, forceLane0);
+      target, perVectorCallback, forceLane0, overwriteAllocSize);
 }
 
 SmallVector<Value> loadSharedToDistributed(triton::gpu::LocalLoadOp localLoadOp,
@@ -515,11 +518,28 @@ SmallVector<Value> loadSharedToDistributed(triton::gpu::LocalLoadOp localLoadOp,
   auto srcTy = localLoadOp.getSrc().getType();
   auto dstTy = localLoadOp.getResult().getType();
 
+  // We overwrite the alloc size if we are a subview to fix subviews in the
+  // fastest dim
+  SmallVector<int64_t> overwriteSmemAllocSize;
+  auto src = localLoadOp.getSrc();
+  if (auto subView = src.getDefiningOp<triton::gpu::MemDescSubviewOp>()) {
+    auto subViewSrcTy =
+        dyn_cast<triton::gpu::MemDescType>(subView.getSrc().getType());
+    if (subViewSrcTy) {
+      auto origAllocSize = subViewSrcTy.getAllocShape();
+      auto srcAllocSize = srcTy.getAllocShape();
+      if (origAllocSize.size() == 3 && srcAllocSize.size() == 2) {
+        overwriteSmemAllocSize = to_vector(origAllocSize.drop_front());
+      }
+    }
+  }
+
   auto b = TritonLLVMOpBuilder(loc, rewriter);
   SmallVector<Value> ret;
   bool success = emitTransferBetweenRegistersAndShared(
       dstTy, srcTy, elemLlvmTy, /*maxVecElems=*/std::nullopt, smemObj, loc,
-      rewriter, target, [&](VectorType vecTy, Value vecAddr) {
+      rewriter, target,
+      [&](VectorType vecTy, Value vecAddr) {
         auto vecVal = b.load(vecTy, vecAddr);
         target.localLoadOpAnnotation(localLoadOp, vecVal);
         vecVal.setAlignment(vecTy.getNumElements() *
@@ -528,7 +548,8 @@ SmallVector<Value> loadSharedToDistributed(triton::gpu::LocalLoadOp localLoadOp,
         for (int v = 0; v < vecTy.getNumElements(); v++) {
           ret.push_back(b.extract_element(elemLlvmTy, vecVal, b.i32_val(v)));
         }
-      });
+      },
+      false, overwriteSmemAllocSize);
   if (!success)
     llvm::report_fatal_error("Failed to emit transfer from shared to register");