[TTGIR] Disallow memdesc_index 1D -> 1D (#7673)

We make memdesc_index always be rank-reducing.

Author: lezcano

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

@@ -214,8 +214,6 @@ def TTG_MemDescIndexOp : TTG_Op<"memdesc_index", [Pure, MemDescViewTrait]> {
      - the output shape is 4x16xf16, and
      - index = 1.
     Then the output descriptor is equivalent to input[1], where input is the logical tensor.
-
-    When the input is of rank 1 (i.e, shape=[k]), the output will have shape=[1].
   }];
 
   let arguments = (ins TTG_MemDescType:$src, I32:$index);

lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp

@@ -477,17 +477,21 @@ struct MemDescIndexOpConversion
     auto *ctx = op->getContext();
     auto b = TritonLLVMOpBuilder(loc, rewriter);
     auto srcTy = op.getSrc().getType();
-    auto destTy = op.getResult().getType();
+    auto dstTy = op.getResult().getType();
     auto llvmElemTy = getTypeConverter()->convertType(srcTy.getElementType());
 
+    // getAllocationShapePerCTA returns the correct number fp4 elements that we
+    // need to skip when we have fp4Padded=True. getShapePerCTA does not account
+    // for this
+    auto stride = product(
+        getAllocationShapePerCTA(dstTy.getEncoding(), dstTy.getShape()));
+    Value offset = b.mul(op.getIndex(), b.i32_val(stride));
     auto smemObj = getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(),
                                                    llvmElemTy, rewriter);
     auto base = smemObj.getBase();
     auto elemPtrTy = base.getType();
-    Value stride = smemObj.getStrides(srcTy, loc, rewriter).front();
-    Value offset = b.mul(op.getIndex(), stride);
     auto prevOffsets = smemObj.getOffsets();
-    SmallVector<Value> offsetVals(prevOffsets.end() - destTy.getRank(),
+    SmallVector<Value> offsetVals(prevOffsets.end() - dstTy.getRank(),
                                   prevOffsets.end());
     // Advance the pointer and keep the opOffsets as the new shape
     smemObj = SharedMemoryObject(b.gep(elemPtrTy, llvmElemTy, base, offset),

lib/Dialect/TritonGPU/IR/Ops.cpp

@@ -690,19 +690,17 @@ LogicalResult MemDescIndexOp::verify() {
   if (srcTy.getElementType() != dstTy.getElementType()) {
     return emitError("result element type must match desc element type");
   }
-  bool is1D =
-      srcTy.getRank() == 1 && dstTy.getRank() == 1 && dstTy.getDimSize(0) == 1;
-  bool correctRank = srcTy.getRank() == dstTy.getRank() + 1 || is1D;
+  // memdesc_index reduces rank by 1 and preserves the trailing shape.
+  bool correctRank = srcTy.getRank() == dstTy.getRank() + 1;
   if (!correctRank) {
-    return emitError(
-        "result rank must be less than or equal to input rank or 1D -> 1D");
+    return emitError("result rank must be input rank - 1");
   }
   if (srcTy.getAllocShape().size() != srcTy.getRank()) {
     return emitError("We don't allow taking memdesc_index of a memdesc_index");
   }
 
-  if (!is1D && ArrayRef(srcTy.getShape()).take_back(dstTy.getRank()) !=
-                   dstTy.getShape()) {
+  if (ArrayRef(srcTy.getShape()).take_back(dstTy.getRank()) !=
+      dstTy.getShape()) {
     return emitError("result shape must equal to srcShape[1:]");
   }
 

lib/Dialect/TritonGPU/Transforms/Pipeliner/PipeliningUtility.cpp

@@ -459,7 +459,7 @@ Value mlir::triton::createScalarAlloc(ImplicitLocOpBuilder &rewriter, Type type,
   auto barrierEncoding =
       ttg::SwizzledSharedEncodingAttr::get(ctx, 1, 1, 1, {0}, barrierCTALayout);
   ttg::MemDescType memDescType = ttg::MemDescType::get(
-      {numBuffers}, type, barrierEncoding, sharedMemorySpace,
+      {numBuffers, 1}, type, barrierEncoding, sharedMemorySpace,
       /*mutableMemory=*/true);
   return rewriter.create<ttg::LocalAllocOp>(memDescType, Value());
 }
@@ -653,12 +653,10 @@ triton::createSingleBufferView(OpBuilder &builder, Value alloc, Value idx) {
   assert(isa<ttg::MemDescType>(alloc.getType()) && "Expected MemDescType");
   auto allocDescType = cast<ttg::MemDescType>(alloc.getType());
   SmallVector<int64_t> shape;
-  if (allocDescType.getShape().size() > 1) {
-    shape.insert(shape.end(), allocDescType.getShape().begin() + 1,
-                 allocDescType.getShape().end());
-  } else {
-    shape.push_back(1);
-  }
+  assert(allocDescType.getShape().size() > 1 &&
+         "Expected multi-dimensional memdesc (e.g., Nx...) for subview");
+  shape.insert(shape.end(), allocDescType.getShape().begin() + 1,
+               allocDescType.getShape().end());
   auto viewDescType = ttg::MemDescType::get(
       shape, allocDescType.getElementType(), allocDescType.getEncoding(),
       allocDescType.getMemorySpace(), allocDescType.getMutableMemory(),

lib/Dialect/TritonGPU/Transforms/WarpSpecialization/LoadMMASpecialization.cpp

@@ -785,25 +785,32 @@ static LogicalResult pipelineMMA(scf::ForOp &loop, PipelinedMMA &mma,
     Value emptyBar = createBarrierAlloc(loop, /*numBarriers=*/1);
     Value readyBar = createBarrierAlloc(loop, /*numBarriers=*/1);
     PartitionBuilder b(defs.front()->getLoc(), loop);
-    b.create<ttng::ArriveBarrierOp>(emptyBar, /*arriveCount=*/1);
+    // For Nx1 barrier allocations, pass a 1D view into barrier ops.
+    Value emptyView0 = createSingleBufferView(b, emptyBar, b.intCst(0));
+    b.create<ttng::ArriveBarrierOp>(emptyView0, /*arriveCount=*/1);
 
     Operation *domOp = findNearestCommonDominator(defs, domInfo);
     Operation *lastOp = findNearestCommonPostDominator(defs, postDomInfo);
 
     auto [index, phase] = addIndexAndPhase(b, loop, /*numStages=*/1);
     StageCluster srcStageCluster = getStageCluster(domOp);
     b.setInsertionPoint(domOp);
-    b.createInto<ttng::WaitBarrierOp>(*partition, srcStageCluster, emptyBar,
+    Value emptyView = createSingleBufferView(b, emptyBar, index);
+    b.createInto<ttng::WaitBarrierOp>(*partition, srcStageCluster, emptyView,
                                       phase);
 
     b.setInsertionPointAfter(lastOp);
-    b.createInto<ttng::ArriveBarrierOp>(*partition, srcStageCluster, readyBar,
+    Value readyView = createSingleBufferView(b, readyBar, index);
+    b.createInto<ttng::ArriveBarrierOp>(*partition, srcStageCluster, readyView,
                                         1);
 
     b.setInsertionPoint(mmaOp);
+    Value readyView2 = createSingleBufferView(b, readyBar, index);
     b.createInto<ttng::WaitBarrierOp>(*schedule.getPartition(mmaOp),
-                                      getStageCluster(mmaOp), readyBar, phase);
-    mmaOp.addCompletionBarrier(emptyBar, b.boolCst(true));
+                                      getStageCluster(mmaOp), readyView2,
+                                      phase);
+    Value emptyView2 = createSingleBufferView(b, emptyBar, index);
+    mmaOp.addCompletionBarrier(emptyView2, b.boolCst(true));
     mmaOp.setIsAsync(true);
   }
 

test/Analysis/test-allocation.mlir

@@ -832,9 +832,9 @@ tt.func @aliasing_in_partition() {
   }
   partition0() num_warps(4) {
     // expected-remark @below {{offset = 0, size = 16}}
-    %0 = ttg.local_alloc : () -> !ttg.memdesc<2xi64, #A_SHARED, #smem, mutable>
+    %0 = ttg.local_alloc : () -> !ttg.memdesc<2x1xi64, #A_SHARED, #smem, mutable>
     %c0_i32 = arith.constant 0 : i32
-    %1 = ttg.memdesc_index %0, %c0_i32 : !ttg.memdesc<2xi64, #A_SHARED, #smem, mutable> -> !ttg.memdesc<1xi64, #A_SHARED, #smem, mutable>
+    %1 = ttg.memdesc_index %0, %c0_i32 : !ttg.memdesc<2x1xi64, #A_SHARED, #smem, mutable> -> !ttg.memdesc<1xi64, #A_SHARED, #smem, mutable>
     // expected-remark @below {{offset = 16, size = 16}}
     %2 = ttg.local_alloc : () -> !ttg.memdesc<2xi64, #A_SHARED, #smem, mutable>
     "use"(%1) : (!ttg.memdesc<1xi64, #A_SHARED, #smem, mutable>) -> ()

test/Conversion/tritongpu_to_llvm.mlir

@@ -552,14 +552,12 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
   // CHECK-LABEL: rank_reducing_subview
   tt.func @rank_reducing_subview() {
     // CHECK: llvm.mlir.addressof @global_smem
-    // CHECK: llvm.extractvalue
+    // CHECK: llvm.mlir.constant(512 : i32) : i32
+    // CHECK-NEXT: llvm.mul
+    // CHECK-NEXT: llvm.extractvalue
     // CHECK-NEXT: llvm.extractvalue
     // CHECK-NEXT: llvm.extractvalue
     // CHECK-NEXT: llvm.extractvalue
-    // CHECK-NEXT: llvm.mlir.constant(1 : i32) : i32
-    // CHECK-NEXT: llvm.mlir.constant(32 : i32) : i32
-    // CHECK-NEXT: llvm.mlir.constant(512 : i32) : i32
-    // CHECK-NEXT: llvm.mul
     // CHECK-NEXT: llvm.getelementptr
     %index = arith.constant 1 : i32
     %zero = arith.constant 0 : i32
@@ -2111,8 +2109,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.targ
   // CHECK: llvm.store
   tt.func public @test_local_store_subview(%arg0: tensor<1xf32, #blocked>) {
     %c0_i32 = arith.constant 0 : i32
-    %0 = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<1xf32, #shared, #smem, mutable>
-    %sv = ttg.memdesc_index %0, %c0_i32 : !ttg.memdesc<1xf32, #shared, #smem, mutable> -> !ttg.memdesc<1xf32, #shared, #smem, mutable>
+    %0 = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<1x1xf32, #shared, #smem, mutable>
+    %sv = ttg.memdesc_index %0, %c0_i32 : !ttg.memdesc<1x1xf32, #shared, #smem, mutable> -> !ttg.memdesc<1xf32, #shared, #smem, mutable>
     ttg.local_store %arg0, %sv : tensor<1xf32, #blocked> -> !ttg.memdesc<1xf32, #shared, #smem, mutable>
     tt.return
   }

test/TritonGPU/invalid.mlir

@@ -79,7 +79,7 @@ tt.func public @result_rank_too_large(%arg0: !ttg.memdesc<3x8x16xf32, #shared, #
 #smem = #ttg.shared_memory
 tt.func public @result_1d_to_1d(%arg0: !ttg.memdesc<8xf32, #shared, #smem>) {
     %zero = arith.constant 0 : i32
-    // expected-error @+1 {{1D -> 1D}}
+    // expected-error @+1 {{result rank}}
     %a = ttg.memdesc_index %arg0, %zero : !ttg.memdesc<8xf32, #shared, #smem> -> !ttg.memdesc<2xf32, #shared, #smem>
     tt.return
 }