Improve verification

Author: akroviakov

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td

@@ -228,12 +228,12 @@ def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> {
                     "FailureOr<SmallVector<Value>>",
                     "delinearizeId",
                     (ins "OpBuilder &": $builder, "Location":$loc, "Value":$linearId)>,
-    InterfaceMethod<[{Generates instructions to compute multidimensional offsets for dist units
+    InterfaceMethod<[{Generates instructions to compute multidimensional coordinates for dist units
                       assigned to a level identified by linearId. The shape parameter
                       represents the higher-level problem size. Each level may access
                       multiple blocks according to round-robin distribution rules.}],
                     "FailureOr<SmallVector<SmallVector<Value>>>",
-                    "computeDistributedOffsets",
+                    "computeDistributedCoords",
                     (ins "OpBuilder &": $builder, "Location":$loc, "Value":$linearId, "ArrayRef<int64_t>":$shape)>,
     InterfaceMethod</*desc=*/[{Check if this layout can be achieved by applying a transpose
                      to some other layout according to given permutation of (0...n-1).}],
@@ -481,12 +481,12 @@ def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout", [DistributeLayoutAttr]> {
     FailureOr<SmallVector<Value>>
     delinearizeId(OpBuilder &builder, Location loc, Value linearId);
 
-    /// Generates instructions to compute multidimensional offsets for dist units
+    /// Generates instructions to compute multidimensional coordinates for dist units
     /// assigned to a level identified by linearId. The shape parameter
     /// represents the higher-level problem size. Each `level` may access
     /// multiple blocks according to round-robin distribution rules.
     FailureOr<SmallVector<SmallVector<Value>>>
-    computeDistributedOffsets(OpBuilder &builder, Location loc, Value linearId, ArrayRef<int64_t> shape);
+    computeDistributedCoords(OpBuilder &builder, Location loc, Value linearId, ArrayRef<int64_t> shape);
 
     /// Check if this is slice of some other layout.
     bool isSliceOf(const xegpu::DistributeLayoutAttr &other) { return false; }
@@ -645,13 +645,13 @@ def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [DistributeLayoutAttr]> {
     FailureOr<SmallVector<Value>>
     delinearizeId(OpBuilder &builder, Location loc, Value linearId);
 
-    /// Generates instructions to compute multidimensional offsets for blocks
+    /// Generates instructions to compute multidimensional coordinates for blocks
     /// assigned to a subgroup identified by linearId. The shape parameter
     /// represents the workgroup-level problem size. Each subgroup may access
     /// multiple blocks according to round-robin distribution rules.
 
     FailureOr<SmallVector<SmallVector<Value>>>
-    computeDistributedOffsets(OpBuilder &builder, Location loc, Value linearId,ArrayRef<int64_t> shape);
+    computeDistributedCoords(OpBuilder &builder, Location loc, Value linearId,ArrayRef<int64_t> shape);
 
     /// Check if this is slice of some other layout.
     bool isSliceOf(const xegpu::DistributeLayoutAttr &other);

mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp

@@ -47,10 +47,11 @@ void XeGPUDialect::initialize() {
 // `delinearizedId` is used to identify a 16x32 of a subgroup in each
 // distribution unit.
 static SmallVector<SmallVector<Value>>
-genOffsets(OpBuilder &builder, Location loc, SmallVector<Value> delinearizedId,
-           ArrayRef<int64_t> subShapesLayout, ArrayRef<int64_t> subShape,
-           ArrayRef<int64_t> srcShape) {
-  SmallVector<SmallVector<Value>> offsets;
+genCoordinates(OpBuilder &builder, Location loc,
+               SmallVector<Value> delinearizedId,
+               ArrayRef<int64_t> subShapesLayout, ArrayRef<int64_t> subShape,
+               ArrayRef<int64_t> srcShape) {
+  SmallVector<SmallVector<Value>> coordinates;
 
   // A distribution unit must be less than or equal to `srcShape`
   SmallVector<int64_t> distUnitShape = llvm::map_to_vector(
@@ -89,9 +90,9 @@ genOffsets(OpBuilder &builder, Location loc, SmallVector<Value> delinearizedId,
               arith::ConstantIndexOp::create(builder, loc, std::get<1>(t)));
         });
 
-    offsets.push_back(mods);
+    coordinates.push_back(mods);
   }
-  return offsets;
+  return coordinates;
 }
 
 // Checks if the given shape can be evenly distributed based on the layout
@@ -298,12 +299,12 @@ LayoutAttr::delinearizeId(OpBuilder &builder, Location loc, Value linearId) {
   return affine::delinearizeIndex(builder, loc, linearId, dims);
 }
 
-/// Implements DistributeLayoutAttr::computeDistributedOffsets to generate
+/// Implements DistributeLayoutAttr::computeDistributedCoords to generate
 /// instructions for computing multi-dimensional offsets when distributed by
 /// LayoutAttr.
 FailureOr<SmallVector<SmallVector<Value>>>
-LayoutAttr::computeDistributedOffsets(OpBuilder &builder, Location loc,
-                                      Value linearId, ArrayRef<int64_t> shape) {
+LayoutAttr::computeDistributedCoords(OpBuilder &builder, Location loc,
+                                     Value linearId, ArrayRef<int64_t> shape) {
   SmallVector<int64_t> layout;
   SmallVector<int64_t> subShape;
   if (isForWorkgroup()) {
@@ -328,7 +329,7 @@ LayoutAttr::computeDistributedOffsets(OpBuilder &builder, Location loc,
     return failure();
   SmallVector<Value> ids = *maybeIds;
 
-  return genOffsets(builder, loc, ids, layout, subShape, shape);
+  return genCoordinates(builder, loc, ids, layout, subShape, shape);
 }
 
 //===----------------------------------------------------------------------===//
@@ -388,12 +389,12 @@ SliceAttr::delinearizeId(OpBuilder &builder, Location loc, Value linearId) {
   return parent.delinearizeId(builder, loc, linearId);
 }
 
-// Implements DistributeLayoutAttr::computeDistributedOffsets to generate
+// Implements DistributeLayoutAttr::computeDistributedCoords to generate
 // instructions for computing multi-dimensional offsets when distributed by
 // LayoutAttr.
 FailureOr<SmallVector<SmallVector<Value>>>
-SliceAttr::computeDistributedOffsets(OpBuilder &builder, Location loc,
-                                     Value linearId, ArrayRef<int64_t> shape) {
+SliceAttr::computeDistributedCoords(OpBuilder &builder, Location loc,
+                                    Value linearId, ArrayRef<int64_t> shape) {
   assert(getRank() == static_cast<int64_t>(shape.size()) && "invalid shape.");
   if (!isForWorkgroup())
     return failure();
@@ -428,7 +429,7 @@ SliceAttr::computeDistributedOffsets(OpBuilder &builder, Location loc,
   SmallVector<Value> sgIds =
       XeGPUDialect::slice(ArrayRef<Value>(*maybeIds), dims);
 
-  return genOffsets(builder, loc, sgIds, layout, subShape, shape);
+  return genCoordinates(builder, loc, sgIds, layout, subShape, shape);
 }
 
 bool SliceAttr::isSliceOf(const xegpu::DistributeLayoutAttr &other) {

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

@@ -175,7 +175,7 @@ isValidGatherScatterBufferParams(Type offsetsTy, Type maskTy,
 
 LogicalResult
 IsValidMatrixOpParams(VectorType dataTy, MemDescType mdescTy,
-                      UnitAttr subgroup_block_io,
+                      UnitAttr subgroup_block_io, DistributeLayoutAttr layout,
                       function_ref<InFlightDiagnostic()> emitError) {
 
   if (!dataTy) {
@@ -191,7 +191,20 @@ IsValidMatrixOpParams(VectorType dataTy, MemDescType mdescTy,
 
   ArrayRef<int64_t> dataShape = dataTy.getShape();
   ArrayRef<int64_t> mdescShape = mdescTy.getShape();
-
+  if (subgroup_block_io && layout) {
+    auto laneData = layout.getEffectiveLaneDataAsInt();
+    if (!laneData.empty()) {
+      bool isLaneDataLinear =
+          std::all_of(laneData.begin(), std::prev(laneData.end()),
+                      [](int x) { return x == 1; });
+      if (!isLaneDataLinear)
+        return emitError()
+               << "With subgroup_block_io, lane data must be linear.";
+      if (isLaneDataLinear && laneData.back() != 1)
+        return emitError()
+               << "With subgroup_block_io, lane data must be coalesced.";
+    }
+  }
   if (dataShape.size() == 2) {
     if (llvm::any_of(llvm::zip_equal(dataShape, mdescShape),
                      [](auto p) { return std::get<0>(p) > std::get<1>(p); }))
@@ -1102,7 +1115,7 @@ LogicalResult LoadMatrixOp::verify() {
   MemDescType mdescTy = getMemDesc().getType();
 
   return IsValidMatrixOpParams(resTy, mdescTy, subgroup_block_io,
-                               [&]() { return emitError(); });
+                               getLayoutAttr(), [&]() { return emitError(); });
 }
 
 //===----------------------------------------------------------------------===//
@@ -1126,7 +1139,7 @@ LogicalResult StoreMatrixOp::verify() {
   UnitAttr subgroup_block_io = getSubgroupBlockIoAttr();
   MemDescType mdescTy = getMemDesc().getType();
   return IsValidMatrixOpParams(dataTy, mdescTy, subgroup_block_io,
-                               [&]() { return emitError(); });
+                               getLayoutAttr(), [&]() { return emitError(); });
 }
 
 namespace mlir {

mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp

@@ -907,22 +907,22 @@ struct StoreDistribution final : public gpu::WarpDistributionPattern {
   }
 };
 
-static SmallVector<Value> computeDistributedOffsetsForMatrixOp(
+static SmallVector<Value> computeDistributedCoordinatesForMatrixOp(
     PatternRewriter &rewriter, Location loc, xegpu::DistributeLayoutAttr layout,
     Value laneId, ArrayRef<int64_t> payloadShape, ValueRange origOffsets) {
-  SmallVector<Value> newOffsets;
-  auto maybeDescOffsets =
-      layout.computeDistributedOffsets(rewriter, loc, laneId, payloadShape);
-  if (failed(maybeDescOffsets))
+  SmallVector<Value> newCoods;
+  auto maybeCoords =
+      layout.computeDistributedCoords(rewriter, loc, laneId, payloadShape);
+  if (failed(maybeCoords))
     return {};
-  assert(maybeDescOffsets.value().size() == 1 &&
+  assert(maybeCoords.value().size() == 1 &&
          "Expected one set of distributed offsets");
   SmallVector<OpFoldResult> ofrVec = xegpu::addWithRightAligned(
-      rewriter, loc, getAsOpFoldResult(maybeDescOffsets.value()[0]),
+      rewriter, loc, getAsOpFoldResult(maybeCoords.value()[0]),
       getAsOpFoldResult(origOffsets));
-  newOffsets = llvm::to_vector(llvm::map_range(
+  newCoods = llvm::to_vector(llvm::map_range(
       ofrVec, [&](OpFoldResult ofr) -> Value { return cast<Value>(ofr); }));
-  return newOffsets;
+  return newCoods;
 }
 
 /// Pattern for distributing xegpu::LoadMatrixOp.
@@ -969,7 +969,7 @@ struct LoadMatrixDistribution final : public gpu::WarpDistributionPattern {
         getDistVecTypeBasedOnLaneLayout(layout, sgPayloadTy);
     if (failed(distPayloadByWarpOpOrFailure))
       return rewriter.notifyMatchFailure(
-          matrixOp, "The matrix op payload has no layout.");
+          matrixOp, "Failed to distribute matrix op payload based on layout.");
 
     SmallVector<Value> operands = {matrixOp.getMemDesc()};
     const unsigned offsetsStartIdx = operands.size();
@@ -992,17 +992,17 @@ struct LoadMatrixDistribution final : public gpu::WarpDistributionPattern {
     ValueRange currentOffsets =
         ValueRange(newOperands).drop_front(offsetsStartIdx);
 
-    SmallVector<Value> newOffsets = currentOffsets;
+    SmallVector<Value> newCoords = currentOffsets;
     rewriter.setInsertionPointAfter(newWarpOp);
 
     if (!matrixOp.getSubgroupBlockIoAttr()) {
-      newOffsets = computeDistributedOffsetsForMatrixOp(
+      newCoords = computeDistributedCoordinatesForMatrixOp(
           rewriter, loc, layout, newWarpOp.getLaneid(), sgPayloadTy.getShape(),
           currentOffsets);
     }
     xegpu::LoadMatrixOp newOp = xegpu::LoadMatrixOp::create(
         rewriter, newWarpOp.getLoc(), *distPayloadByWarpOpOrFailure,
-        newOperands[0], ValueRange(newOffsets), newConstOffsetsAttr,
+        newOperands[0], ValueRange(newCoords), newConstOffsetsAttr,
         matrixOp.getSubgroupBlockIoAttr(), xegpu::DistributeLayoutAttr{});
     // Resolve the output type and replace all uses.
     rewriter.replaceAllUsesWith(
@@ -1045,7 +1045,7 @@ struct StoreMatrixDistribution final : public gpu::WarpDistributionPattern {
         getDistVecTypeBasedOnLaneLayout(layout, sgPayloadTy);
     if (failed(distPayloadByWarpOpOrFailure))
       return rewriter.notifyMatchFailure(
-          matrixOp, "The matrix op payload has no layout.");
+          matrixOp, "Failed to distribute matrix op payload based on layout.");
 
     SmallVector<Value> operands = {matrixOp.getData(), matrixOp.getMemDesc()};
     const unsigned offsetsStartIdx = operands.size();
@@ -1069,18 +1069,18 @@ struct StoreMatrixDistribution final : public gpu::WarpDistributionPattern {
     ValueRange currentOffsets =
         ValueRange(newOperands).drop_front(offsetsStartIdx);
 
-    SmallVector<Value> newOffsets = currentOffsets;
+    SmallVector<Value> newCoords = currentOffsets;
     rewriter.setInsertionPointAfter(newWarpOp);
 
     if (!matrixOp.getSubgroupBlockIoAttr()) {
-      newOffsets = computeDistributedOffsetsForMatrixOp(
+      newCoords = computeDistributedCoordinatesForMatrixOp(
           rewriter, loc, layout, newWarpOp.getLaneid(), sgPayloadTy.getShape(),
           currentOffsets);
     }
 
     xegpu::StoreMatrixOp::create(
         rewriter, loc, TypeRange{}, newOperands[0], newOperands[1],
-        ValueRange(newOffsets), newConstOffsetsAttr,
+        ValueRange(newCoords), newConstOffsetsAttr,
         matrixOp.getSubgroupBlockIoAttr(), xegpu::DistributeLayoutAttr{});
     rewriter.eraseOp(matrixOp);
     return success();

mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp

@@ -115,7 +115,7 @@ genOffsetsList(ConversionPatternRewriter &rewriter, OpType op,
   // descriptors to be accessed, based on the layout information.
   ArrayRef<int64_t> wgShape = op.getDataShape();
   auto maybeDescOffsets =
-      layout.computeDistributedOffsets(rewriter, loc, sgId, wgShape);
+      layout.computeDistributedCoords(rewriter, loc, sgId, wgShape);
   if (failed(maybeDescOffsets))
     return failure();
 
@@ -832,7 +832,7 @@ struct WgToSgArithConstantOp : public OpConversionPattern<arith::ConstantOp> {
       Value sgId =
           gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);
       auto sgOffsets =
-          layout.computeDistributedOffsets(rewriter, loc, sgId, wgShape);
+          layout.computeDistributedCoords(rewriter, loc, sgId, wgShape);
       if (failed(sgOffsets))
         return failure();
 
@@ -1054,7 +1054,7 @@ struct WgToSgVectorStepOp : public OpConversionPattern<vector::StepOp> {
     Value sgId =
         gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);
     auto sgOffsets =
-        layout.computeDistributedOffsets(rewriter, loc, sgId, wgShape);
+        layout.computeDistributedCoords(rewriter, loc, sgId, wgShape);
     if (failed(sgOffsets))
       return failure();
 

mlir/test/Dialect/XeGPU/invalid.mlir

@@ -890,3 +890,19 @@ func.func @store_mem_desc_invalid_rank(%arg0: !xegpu.mem_desc<64xf16>, %arg1: ve
   xegpu.store_matrix %arg1, %arg0[32] : vector<32xf16>, !xegpu.mem_desc<64xf16>
   return
 }
+
+// -----
+func.func @simt_store_matrix_vector_nonlinear(%arg0: !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32]>>, %arg1: vector<2x16xf32>) {
+  // expected-error@+1 {{With subgroup_block_io, lane data must be linear}}
+  xegpu.store_matrix %arg1, %arg0[0, 0] {subgroup_block_io, layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>} :
+        vector<2x16xf32>, !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32]>>
+  return
+}
+
+// -----
+func.func @simt_store_matrix_vector_noncoalesced(%arg0: !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32]>>, %arg1: vector<16x2xf32>) {
+  // expected-error@+1 {{With subgroup_block_io, lane data must be coalesced}}
+  xegpu.store_matrix %arg1, %arg0[0, 0] {subgroup_block_io, layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>} :
+        vector<16x2xf32>, !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32]>>
+  return
+}

mlir/test/Dialect/XeGPU/subgroup-distribute.mlir

@@ -321,9 +321,9 @@ gpu.module @xevm_module{
   gpu.func @load_store_matrix_3(%arg0: !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>) {
     %c0 = arith.constant 0 : index
     %c1 = arith.constant 1 : index
-    %1 = xegpu.load_matrix %arg0[%c0, %c1] {subgroup_block_io, layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>} :
+    %1 = xegpu.load_matrix %arg0[%c0, %c1] {subgroup_block_io, layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
       !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>, index, index -> vector<2x16xf32>
-    xegpu.store_matrix %1, %arg0[%c0, %c1] {subgroup_block_io, layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>} :
+    xegpu.store_matrix %1, %arg0[%c0, %c1] {subgroup_block_io, layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
       vector<2x16xf32>, !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>, index, index
     gpu.return
   }

mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp

@@ -201,7 +201,7 @@ class TestStepOpPattern : public OpConversionPattern<vector::StepOp> {
     Value sgId =
         gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);
     auto maybeOffsets =
-        sliceAttr.computeDistributedOffsets(rewriter, loc, sgId, wgShape);
+        sliceAttr.computeDistributedCoords(rewriter, loc, sgId, wgShape);
     if (failed(maybeOffsets))
       return failure();