[Mlir-commits] [mlir] cbfb9a6 - [MLIR][XeGPU] Optimize create_mask sg-to-lane lowering (#221198)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Sep 24 03:14:57 PDT 2026
Author: Artem Kroviakov
Date: 2026-09-24T12:14:49+02:00
New Revision: cbfb9a698f5c499c2f91a7bf37dd16de5a93f868
URL: https://github.com/llvm/llvm-project/commit/cbfb9a698f5c499c2f91a7bf37dd16de5a93f868
DIFF: https://github.com/llvm/llvm-project/commit/cbfb9a698f5c499c2f91a7bf37dd16de5a93f868.diff
LOG: [MLIR][XeGPU] Optimize create_mask sg-to-lane lowering (#221198)
Added:
Modified:
mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index b29f138e4ef21..8bc8a56bc3991 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -860,28 +860,43 @@ struct SgToLaneVectorBitcast : public OpConversionPattern<vector::BitCastOp> {
};
/// Distributes a subgroup-level vector.create_mask or vector.constant_mask op
-/// to lane-level. Uses `computeDistributedCoords()` to obtain the
-/// coordinates each lane owns, then compares each coordinate against the
-/// original mask bounds using `arith.cmpi slt`. The per-element boolean
-/// results are assembled into the distributed mask vector.
-///
-/// For multi-dimensional masks, the element is in-bounds when ALL dimensions
-/// satisfy `coord[i] < bound[i]`.
-///
-/// Example (1D):
-/// layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
-/// %mask = vector.create_mask %m0 : vector<16xi1>
-/// For lane k, computeDistributedCoords gives coord = [k], so:
-/// %in_bounds = arith.cmpi slt, %coord, %m0 → i1
-/// %mask = vector.broadcast %in_bounds : i1 to vector<1xi1>
+/// to lane-level.
+/// The pattern constructs a mask based on the following bounds check:
+/// ```
+/// for d in [0, ..., maskRank):
+/// mask &= (staticOffset[d] < (bound[d] - base[d]))
+/// ```
+/// where
+/// - `base` is the coordinate vector of the first distributed unit.
+/// - `staticOffset` is the offsets vector per element.
+/// - `bound` is the original mask bound for the corresponding dimension.
+/// The mask vector contains *all* elements (i.e., non-unit `lane_data` is
+/// flattened). For example,
+/// ```
+/// %mask = vector.create_mask %bound_0, %bound_1 {
+/// lane_layout = [1, 16], lane_data = [2, 1]
+/// }: vector<8x32xi1>
+/// ```
+/// Has 8 dist units (of shape [2, 1]) with offsets for lane 0:
+/// {
+/// [0, 0], [0, 16],
+/// [2, 0], [2, 16],
+/// [4, 0], [4, 16],
+/// [6, 0], [6, 16]
+/// }
+/// We check the distance to the bound from the first dist. unit:
+/// %base_0 = genCoords(gpu.lane_id, layout)[0][0]
+/// %base_1 = genCoords(gpu.lane_id, layout)[0][1]
+/// %baseDistanceToBound_0 = %bound_0 - %base_0
+/// %baseDistanceToBound_1 = %bound_1 - %base_1
+/// The pattern flattens the dimension d coordinates of each element:
+/// %staticOffset_0 = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7]
+/// %staticOffset_1 = [0, 16, 0, 16, 0, 16, 0, 16, 0, 16, 0, 16, 0, 16, 0, 16]
+/// and compares each coordinate against the corresponding mask bound dim:
+/// %mask_0 = arith.cmpi slt, %staticOffset_0, bcast(%baseDistanceToBound_0)
+/// %mask_1 = arith.cmpi slt, %staticOffset_1, bcast(%baseDistanceToBound_1)
+/// %mask = shape_cast(%mask_0 & %mask_1) : vector<16xi1> to vector<8x2xi1>
///
-/// Example (2D):
-/// layout = #xegpu.layout<lane_layout = [8, 2], lane_data = [1, 1]>
-/// %mask = vector.create_mask %m0, %m1 : vector<8x4xi1>
-/// Each WI owns a 1x2 slice. computeDistributedCoords returns 2 coords:
-/// [[r0, c0], [r0, c1]]
-/// For each coord: in_bounds = (r < m0) && (c < m1)
-/// %mask = vector.from_elements %bit0, %bit1 : vector<1x2xi1>
template <typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
OpType, vector::CreateMaskOp, vector::ConstantMaskOp>::value>>
@@ -929,36 +944,103 @@ struct SgToLaneCreateMask : public OpConversionPattern<OpType> {
return rewriter.notifyMatchFailure(
op, "failed to compute distributed coordinates from layout");
- SmallVector<SmallVector<Value>> coordsVec = maybeCoordsVec.value();
+ SmallVector<SmallVector<Value>> laneDataCoords = maybeCoordsVec.value();
+ SmallVector<int64_t> laneData = layout.getEffectiveLaneDataAsInt();
+ ArrayRef<int64_t> distShape = distType.getShape();
+ int64_t rank = distType.getRank();
int64_t numElements = distType.getNumElements();
- assert(static_cast<int64_t>(coordsVec.size()) == numElements &&
- "number of coordinate sets must match number of distributed "
- "elements");
-
- // For each element, compare all coordinates against bounds.
- Value trueVal =
- arith::ConstantIntOp::create(rewriter, loc, /*value=*/1, /*width=*/1);
- SmallVector<Value> maskBits;
- for (auto &coords : coordsVec) {
- Value inBounds = trueVal;
- for (size_t i = 0; i < coords.size(); ++i) {
- Value cmp = arith::CmpIOp::create(
- rewriter, loc, arith::CmpIPredicate::slt, coords[i], origBounds[i]);
- inBounds = arith::AndIOp::create(rewriter, loc, inBounds, cmp);
+
+ if (static_cast<int64_t>(laneData.size()) != rank ||
+ !computeShapeRatio(distShape, laneData))
+ return rewriter.notifyMatchFailure(
+ op, "lane_data does not tile the distributed vector");
+
+ SmallVector<int64_t> distStrides = computeStrides(distShape);
+ assert(static_cast<int64_t>(laneDataCoords.size()) *
+ computeProduct(laneData) ==
+ numElements &&
+ "number of coordinate sets must match number of lane_data blocks");
+
+ // Static offsets to be applied to a dynamic lane's dist units coordinates.
+ SmallVector<SmallVector<int64_t>> staticLaneDataOffsetsOrig =
+ layout.computeStaticDistributedCoords(/*linearId=*/0, origShape);
+ if (staticLaneDataOffsetsOrig.empty() ||
+ staticLaneDataOffsetsOrig.size() != laneDataCoords.size())
+ return rewriter.notifyMatchFailure(
+ op, "static and dynamic coordinates disagree on the distribution "
+ "unit count");
+
+ SmallVector<int64_t> unitTile(rank, 1);
+ int64_t linearLaneDataIdx = 0;
+ // Layout and shape information is static, compute static offset of lane's
+ // elements in the source dimensions. Each dim is a flat vec of element
+ // offsets. We store dim-major to have an easy materialization as one const
+ // vector.
+ SmallVector<SmallVector<int64_t>> staticElemOffset(
+ rank, SmallVector<int64_t>(numElements));
+ // For each lane_data block in the lane's dist shape
+ for (SmallVector<int64_t> laneDataOffsetDist :
+ StaticTileOffsetRange(distShape, laneData)) {
+ ArrayRef<int64_t> staticLaneDataOffsetOrig =
+ staticLaneDataOffsetsOrig[linearLaneDataIdx++];
+ // For each element in the lane_data block
+ for (SmallVector<int64_t> elemOffsetInLaneData :
+ StaticTileOffsetRange(laneData, unitTile)) {
+ // For each dim of indexing space
+ SmallVector<int64_t> elementOffsetDist(rank);
+ for (int64_t d = 0; d < rank; d++)
+ elementOffsetDist[d] =
+ laneDataOffsetDist[d] + elemOffsetInLaneData[d];
+ int64_t elemLinearizedIdxDist =
+ linearize(elementOffsetDist, distStrides);
+ for (int64_t d = 0; d < rank; d++)
+ staticElemOffset[d][elemLinearizedIdxDist] =
+ staticLaneDataOffsetOrig[d] + elemOffsetInLaneData[d];
}
- maskBits.push_back(inBounds);
}
- // Build the distributed mask vector.
- Value result;
- if (numElements == 1) {
- result =
- vector::BroadcastOp::create(rewriter, loc, distType, maskBits[0]);
- } else {
- result =
- vector::FromElementsOp::create(rewriter, loc, distType, maskBits);
+ // Check, whether ALL elements of the distributed mask are within the valid
+ // extent of EACH source dimension.
+ // Expressed as dyn_offset[d] + static_offset[d] < dyn_mask_bound[d],
+ // or equivalently,
+ // static_offset[d] < dyn_mask_bound[d] - dyn_offset[d]
+ auto flatIndexType = VectorType::get(numElements, rewriter.getIndexType());
+ Value inBounds;
+ for (int64_t d = 0; d < rank; d++) {
+ std::optional<int64_t> constBound = getConstantIntValue(origBounds[d]);
+ if (constBound && *constBound >= origShape[d])
+ continue;
+
+ Value materializedStaticOffset = arith::ConstantOp::create(
+ rewriter, loc, rewriter.getIndexVectorAttr(staticElemOffset[d]));
+ // Consider only the first unit, all others are compile-time multiple
+ // offsets of the first one and are already encoded in staticOffsets.
+ Value validDimExtent = arith::SubIOp::create(rewriter, loc, origBounds[d],
+ laneDataCoords[0][d]);
+ // Dim extent applies to all elements.
+ Value validDimExtentPerElement = vector::BroadcastOp::create(
+ rewriter, loc, flatIndexType, validDimExtent);
+ Value elementMaskInDim = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::slt, materializedStaticOffset,
+ validDimExtentPerElement);
+ // An element is valid iff it is within the valid extent of ALL
+ // dimensions.
+ inBounds = inBounds ? arith::AndIOp::create(rewriter, loc, inBounds,
+ elementMaskInDim)
+ .getResult()
+ : elementMaskInDim;
}
- rewriter.replaceOp(op, result);
+ // Every dim's bound saturates the mask extent, so no comparison was
+ // emitted at all: every element of every lane is in bounds.
+ if (!inBounds) {
+ rewriter.replaceOp(op, arith::ConstantOp::create(
+ rewriter, loc, distType,
+ DenseElementsAttr::get(distType, true)));
+ return success();
+ }
+ auto resMask =
+ rewriter.createOrFold<vector::ShapeCastOp>(loc, distType, inBounds);
+ rewriter.replaceOp(op, resMask);
return success();
}
};
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
index 1407acf9ee3dc..e69fabbce0dc3 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -565,11 +565,12 @@ gpu.func @vector_bitcast() {
// CHECK-LABEL: gpu.func @create_mask_1d
// CHECK-SAME: (%[[M0:.*]]: index)
-// CHECK-DAG: %[[LANE:.*]] = gpu.lane_id
-// CHECK-DAG: %[[TRUE:.*]] = arith.constant true
-// CHECK: %[[CMP:.*]] = arith.cmpi slt, %{{.*}}, %[[M0]] : index
-// CHECK: %[[AND:.*]] = arith.andi %[[TRUE]], %[[CMP]] : i1
-// CHECK: %[[MASK:.*]] = vector.broadcast %[[AND]] : i1 to vector<1xi1>
+// CHECK: %[[LANE:.*]] = gpu.lane_id
+// CHECK: %[[K0:.*]] = arith.constant dense<0> : vector<1xindex>
+// CHECK: %[[LEN0:.*]] = arith.subi %[[M0]], %{{.*}} : index
+// CHECK: %[[EXT0:.*]] = vector.broadcast %[[LEN0]] : index to vector<1xindex>
+// CHECK: %[[MASK:.*]] = arith.cmpi slt, %[[K0]], %[[EXT0]] : vector<1xindex>
+// CHECK-NOT: vector.shape_cast
// CHECK: gpu.return
gpu.func @create_mask_1d(%m0: index) {
%mask = vector.create_mask %m0
@@ -582,12 +583,12 @@ gpu.func @create_mask_1d(%m0: index) {
}
// CHECK-LABEL: gpu.func @constant_mask_1d
-// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
-// CHECK-DAG: %[[LANE:.*]] = gpu.lane_id
-// CHECK-DAG: %[[TRUE:.*]] = arith.constant true
-// CHECK: %[[CMP:.*]] = arith.cmpi slt, %{{.*}}, %[[C4]] : index
-// CHECK: %[[AND:.*]] = arith.andi %[[TRUE]], %[[CMP]] : i1
-// CHECK: %[[MASK:.*]] = vector.broadcast %[[AND]] : i1 to vector<1xi1>
+// CHECK: %[[C4:.*]] = arith.constant 4 : index
+// CHECK: %[[LANE:.*]] = gpu.lane_id
+// CHECK: %[[K0:.*]] = arith.constant dense<0> : vector<1xindex>
+// CHECK: %[[LEN0:.*]] = arith.subi %[[C4]], %{{.*}} : index
+// CHECK: %[[EXT0:.*]] = vector.broadcast %[[LEN0]] : index to vector<1xindex>
+// CHECK: %[[MASK:.*]] = arith.cmpi slt, %[[K0]], %[[EXT0]] : vector<1xindex>
// CHECK: gpu.return
gpu.func @constant_mask_1d() {
%mask = vector.constant_mask [4]
@@ -599,19 +600,52 @@ gpu.func @constant_mask_1d() {
gpu.return
}
+// CHECK-LABEL: gpu.func @create_mask_1d_multi_unit
+// CHECK-SAME: (%[[M0:.*]]: index)
+// CHECK: %[[K0:.*]] = arith.constant dense<[0, 4, 8, 12]> : vector<4xindex>
+// CHECK: %[[LEN0:.*]] = arith.subi %[[M0]], %{{.*}} : index
+// CHECK: %[[EXT0:.*]] = vector.broadcast %[[LEN0]] : index to vector<4xindex>
+// CHECK: %[[MASK:.*]] = arith.cmpi slt, %[[K0]], %[[EXT0]] : vector<4xindex>
+// CHECK: gpu.return
+gpu.func @create_mask_1d_multi_unit(%m0: index) {
+ %mask = vector.create_mask %m0
+ : vector<16xi1>
+ %mask_cl = xegpu.convert_layout %mask
+ <{
+ target_layout = #xegpu.layout<lane_layout = [4], lane_data = [1]>
+ }> : vector<16xi1>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @create_mask_1d_lane_data
+// CHECK-SAME: (%[[M0:.*]]: index)
+// CHECK: %[[K0:.*]] = arith.constant dense<[0, 1, 2, 3, 8, 9, 10, 11]> : vector<8xindex>
+// CHECK: %[[LEN0:.*]] = arith.subi %[[M0]], %{{.*}} : index
+// CHECK: %[[EXT0:.*]] = vector.broadcast %[[LEN0]] : index to vector<8xindex>
+// CHECK: %[[MASK:.*]] = arith.cmpi slt, %[[K0]], %[[EXT0]] : vector<8xindex>
+// CHECK: gpu.return
+gpu.func @create_mask_1d_lane_data(%m0: index) {
+ %mask = vector.create_mask %m0
+ : vector<16xi1>
+ %mask_cl = xegpu.convert_layout %mask
+ <{
+ target_layout = #xegpu.layout<lane_layout = [2], lane_data = [4]>
+ }> : vector<16xi1>
+ gpu.return
+}
+
// CHECK-LABEL: gpu.func @create_mask_2d
// CHECK-SAME: (%[[M0:.*]]: index, %[[M1:.*]]: index)
-// CHECK-DAG: %[[LANE:.*]] = gpu.lane_id
-// CHECK-DAG: %[[TRUE:.*]] = arith.constant true
-// CHECK: %[[CMP_R0:.*]] = arith.cmpi slt, %{{.*}}, %[[M0]] : index
-// CHECK: %[[AND0:.*]] = arith.andi %[[TRUE]], %[[CMP_R0]] : i1
-// CHECK: %[[CMP_C0:.*]] = arith.cmpi slt, %{{.*}}, %[[M1]] : index
-// CHECK: %[[BIT0:.*]] = arith.andi %[[AND0]], %[[CMP_C0]] : i1
-// CHECK: %[[CMP_R1:.*]] = arith.cmpi slt, %{{.*}}, %[[M0]] : index
-// CHECK: %[[AND1:.*]] = arith.andi %[[TRUE]], %[[CMP_R1]] : i1
-// CHECK: %[[CMP_C1:.*]] = arith.cmpi slt, %{{.*}}, %[[M1]] : index
-// CHECK: %[[BIT1:.*]] = arith.andi %[[AND1]], %[[CMP_C1]] : i1
-// CHECK: %[[MASK:.*]] = vector.from_elements %[[BIT0]], %[[BIT1]] : vector<1x2xi1>
+// CHECK: %[[K0:.*]] = arith.constant dense<0> : vector<2xindex>
+// CHECK: %[[LEN0:.*]] = arith.subi %[[M0]], %{{.*}} : index
+// CHECK: %[[EXT0:.*]] = vector.broadcast %[[LEN0]] : index to vector<2xindex>
+// CHECK: %[[P0:.*]] = arith.cmpi slt, %[[K0]], %[[EXT0]] : vector<2xindex>
+// CHECK: %[[K1:.*]] = arith.constant dense<[0, 2]> : vector<2xindex>
+// CHECK: %[[LEN1:.*]] = arith.subi %[[M1]], %{{.*}} : index
+// CHECK: %[[EXT1:.*]] = vector.broadcast %[[LEN1]] : index to vector<2xindex>
+// CHECK: %[[P1:.*]] = arith.cmpi slt, %[[K1]], %[[EXT1]] : vector<2xindex>
+// CHECK: %[[AND:.*]] = arith.andi %[[P0]], %[[P1]] : vector<2xi1>
+// CHECK: %[[MASK:.*]] = vector.shape_cast %[[AND]] : vector<2xi1> to vector<1x2xi1>
// CHECK: gpu.return
gpu.func @create_mask_2d(%m0: index, %m1: index) {
%mask = vector.create_mask %m0, %m1
@@ -624,19 +658,18 @@ gpu.func @create_mask_2d(%m0: index, %m1: index) {
}
// CHECK-LABEL: gpu.func @constant_mask_2d
-// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
-// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
-// CHECK-DAG: %[[LANE:.*]] = gpu.lane_id
-// CHECK-DAG: %[[TRUE:.*]] = arith.constant true
-// CHECK: %[[CMP_R0:.*]] = arith.cmpi slt, %{{.*}}, %[[C2]] : index
-// CHECK: %[[AND0:.*]] = arith.andi %[[TRUE]], %[[CMP_R0]] : i1
-// CHECK: %[[CMP_C0:.*]] = arith.cmpi slt, %{{.*}}, %[[C3]] : index
-// CHECK: %[[BIT0:.*]] = arith.andi %[[AND0]], %[[CMP_C0]] : i1
-// CHECK: %[[CMP_R1:.*]] = arith.cmpi slt, %{{.*}}, %[[C2]] : index
-// CHECK: %[[AND1:.*]] = arith.andi %[[TRUE]], %[[CMP_R1]] : i1
-// CHECK: %[[CMP_C1:.*]] = arith.cmpi slt, %{{.*}}, %[[C3]] : index
-// CHECK: %[[BIT1:.*]] = arith.andi %[[AND1]], %[[CMP_C1]] : i1
-// CHECK: %[[MASK:.*]] = vector.from_elements %[[BIT0]], %[[BIT1]] : vector<1x2xi1>
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[C3:.*]] = arith.constant 3 : index
+// CHECK: %[[K0:.*]] = arith.constant dense<0> : vector<2xindex>
+// CHECK: %[[LEN0:.*]] = arith.subi %[[C2]], %{{.*}} : index
+// CHECK: %[[EXT0:.*]] = vector.broadcast %[[LEN0]] : index to vector<2xindex>
+// CHECK: %[[P0:.*]] = arith.cmpi slt, %[[K0]], %[[EXT0]] : vector<2xindex>
+// CHECK: %[[K1:.*]] = arith.constant dense<[0, 2]> : vector<2xindex>
+// CHECK: %[[LEN1:.*]] = arith.subi %[[C3]], %{{.*}} : index
+// CHECK: %[[EXT1:.*]] = vector.broadcast %[[LEN1]] : index to vector<2xindex>
+// CHECK: %[[P1:.*]] = arith.cmpi slt, %[[K1]], %[[EXT1]] : vector<2xindex>
+// CHECK: %[[AND:.*]] = arith.andi %[[P0]], %[[P1]] : vector<2xi1>
+// CHECK: %[[MASK:.*]] = vector.shape_cast %[[AND]] : vector<2xi1> to vector<1x2xi1>
// CHECK: gpu.return
gpu.func @constant_mask_2d() {
%mask = vector.constant_mask [2, 3]
@@ -648,6 +681,116 @@ gpu.func @constant_mask_2d() {
gpu.return
}
+// CHECK-LABEL: gpu.func @create_mask_2d_lane_data
+// CHECK-SAME: (%[[M0:.*]]: index, %[[M1:.*]]: index)
+// CHECK: %[[K0:.*]] = arith.constant dense<[0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7]> : vector<16xindex>
+// CHECK: %[[LEN0:.*]] = arith.subi %[[M0]], %{{.*}} : index
+// CHECK: %[[EXT0:.*]] = vector.broadcast %[[LEN0]] : index to vector<16xindex>
+// CHECK: %[[P0:.*]] = arith.cmpi slt, %[[K0]], %[[EXT0]] : vector<16xindex>
+// CHECK: %[[K1:.*]] = arith.constant dense<[0, 16, 0, 16, 0, 16, 0, 16, 0, 16, 0, 16, 0, 16, 0, 16]> : vector<16xindex>
+// CHECK: %[[LEN1:.*]] = arith.subi %[[M1]], %{{.*}} : index
+// CHECK: %[[EXT1:.*]] = vector.broadcast %[[LEN1]] : index to vector<16xindex>
+// CHECK: %[[P1:.*]] = arith.cmpi slt, %[[K1]], %[[EXT1]] : vector<16xindex>
+// CHECK: %[[AND:.*]] = arith.andi %[[P0]], %[[P1]] : vector<16xi1>
+// CHECK: %[[MASK:.*]] = vector.shape_cast %[[AND]] : vector<16xi1> to vector<8x2xi1>
+// CHECK: gpu.return
+gpu.func @create_mask_2d_lane_data(%m0: index, %m1: index) {
+ %mask = vector.create_mask %m0, %m1
+ : vector<8x32xi1>
+ %mask_cl = xegpu.convert_layout %mask
+ <{
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>
+ }> : vector<8x32xi1>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @create_mask_3d
+// CHECK-SAME: (%[[M0:.*]]: index, %[[M1:.*]]: index, %[[M2:.*]]: index)
+// CHECK: %[[K0:.*]] = arith.constant dense<[0, 0, 0, 0, 1, 1, 1, 1]> : vector<8xindex>
+// CHECK: %[[LEN0:.*]] = arith.subi %[[M0]], %{{.*}} : index
+// CHECK: %[[EXT0:.*]] = vector.broadcast %[[LEN0]] : index to vector<8xindex>
+// CHECK: %[[P0:.*]] = arith.cmpi slt, %[[K0]], %[[EXT0]] : vector<8xindex>
+// CHECK: %[[K1:.*]] = arith.constant dense<[0, 2, 4, 6, 0, 2, 4, 6]> : vector<8xindex>
+// CHECK: %[[LEN1:.*]] = arith.subi %[[M1]], %{{.*}} : index
+// CHECK: %[[EXT1:.*]] = vector.broadcast %[[LEN1]] : index to vector<8xindex>
+// CHECK: %[[P1:.*]] = arith.cmpi slt, %[[K1]], %[[EXT1]] : vector<8xindex>
+// CHECK: %[[AND0:.*]] = arith.andi %[[P0]], %[[P1]] : vector<8xi1>
+// CHECK: %[[K2:.*]] = arith.constant dense<0> : vector<8xindex>
+// CHECK: %[[LEN2:.*]] = arith.subi %[[M2]], %{{.*}} : index
+// CHECK: %[[EXT2:.*]] = vector.broadcast %[[LEN2]] : index to vector<8xindex>
+// CHECK: %[[P2:.*]] = arith.cmpi slt, %[[K2]], %[[EXT2]] : vector<8xindex>
+// CHECK: %[[AND1:.*]] = arith.andi %[[AND0]], %[[P2]] : vector<8xi1>
+// CHECK: %[[MASK:.*]] = vector.shape_cast %[[AND1]] : vector<8xi1> to vector<2x4x1xi1>
+// CHECK: gpu.return
+gpu.func @create_mask_3d(%m0: index, %m1: index, %m2: index) {
+ %mask = vector.create_mask %m0, %m1, %m2
+ : vector<2x8x8xi1>
+ %mask_cl = xegpu.convert_layout %mask
+ <{
+ target_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 1]>
+ }> : vector<2x8x8xi1>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @constant_mask_2d_saturated_dim
+// CHECK: %[[C3:.*]] = arith.constant 3 : index
+// CHECK-NOT: arith.cmpi
+// CHECK: %[[K1:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex>
+// CHECK: %[[LEN1:.*]] = arith.subi %[[C3]], %{{.*}} : index
+// CHECK: %[[EXT1:.*]] = vector.broadcast %[[LEN1]] : index to vector<8xindex>
+// CHECK: %[[MASK:.*]] = arith.cmpi slt, %[[K1]], %[[EXT1]] : vector<8xindex>
+// CHECK-NOT: arith.andi
+// CHECK: vector.shape_cast %[[MASK]] : vector<8xi1> to vector<1x8xi1>
+// CHECK: gpu.return
+gpu.func @constant_mask_2d_saturated_dim() {
+ %mask = vector.constant_mask [16, 3]
+ : vector<16x8xi1>
+ %mask_cl = xegpu.convert_layout %mask
+ <{
+ target_layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
+ }> : vector<16x8xi1>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @create_mask_2d_saturated_distributed_dim
+// CHECK-SAME: (%[[M1:.*]]: index)
+// CHECK-NOT: arith.cmpi
+// CHECK: %[[K1:.*]] = arith.constant dense<[0, 16, 0, 16, 0, 16, 0, 16, 0, 16, 0, 16, 0, 16, 0, 16]> : vector<16xindex>
+// CHECK: %[[LEN1:.*]] = arith.subi %[[M1]], %{{.*}} : index
+// CHECK: %[[EXT1:.*]] = vector.broadcast %[[LEN1]] : index to vector<16xindex>
+// CHECK: %[[MASK:.*]] = arith.cmpi slt, %[[K1]], %[[EXT1]] : vector<16xindex>
+// CHECK-NOT: arith.andi
+// CHECK: vector.shape_cast %[[MASK]] : vector<16xi1> to vector<8x2xi1>
+// CHECK: gpu.return
+gpu.func @create_mask_2d_saturated_distributed_dim(%m1: index) {
+ %c8 = arith.constant 8 : index
+ %mask = vector.create_mask %c8, %m1
+ : vector<8x32xi1>
+ %mask_cl = xegpu.convert_layout %mask
+ <{
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>
+ }> : vector<8x32xi1>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @create_mask_2d_all_dims_saturated
+// CHECK-NOT: arith.cmpi
+// CHECK-NOT: vector.shape_cast
+// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<1x8xi1>
+// CHECK-NOT: arith.cmpi
+// CHECK: gpu.return
+gpu.func @create_mask_2d_all_dims_saturated() {
+ %c16 = arith.constant 16 : index
+ %c8 = arith.constant 8 : index
+ %mask = vector.create_mask %c16, %c8
+ : vector<16x8xi1>
+ %mask_cl = xegpu.convert_layout %mask
+ <{
+ target_layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
+ }> : vector<16x8xi1>
+ gpu.return
+}
+
// CHECK-LABEL: gpu.func @vector_multi_reduction_3d_leading_unit_dim_lane_local
// CHECK-DAG: %[[SRC:.*]] = arith.constant dense<0.000000e+00> : vector<1x16x2xf32>
// CHECK-DAG: %[[ACC:.*]] = arith.constant dense<0.000000e+00> : vector<1x2xf32>
More information about the Mlir-commits
mailing list