[Mlir-commits] [mlir] [MLIR][XeGPU] Optimize create_mask sg-to-lane lowering (PR #221198)
Artem Kroviakov
llvmlistbot at llvm.org
Mon Sep 21 06:38:10 PDT 2026
https://github.com/akroviakov updated https://github.com/llvm/llvm-project/pull/221198
>From 33ad148503078807564b3bf98b1ebaee57f12238 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Fri, 4 Sep 2026 11:15:02 +0000
Subject: [PATCH 1/3] [MLIR][XeGPU] Optimize create_mask sg-to-lane lowering
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 128 +++++++++-----
.../XeGPU/sg-to-lane-distribute-unit.mlir | 162 ++++++++++++++----
2 files changed, 208 insertions(+), 82 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 812233047fb88..60ad08870a539 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -864,28 +864,8 @@ 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>
-///
-/// 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>
+/// to lane-level. An element is in-bounds when every dim satisfies
+/// `coord[d] < bound[d]`.
template <typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
OpType, vector::CreateMaskOp, vector::ConstantMaskOp>::value>>
@@ -933,36 +913,90 @@ 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> resultStrides = 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>> staticLaneDataOffsetsOrigShape =
+ layout.computeStaticDistributedCoords(/*linearId=*/0, origShape);
+ if (staticLaneDataOffsetsOrigShape.empty() ||
+ staticLaneDataOffsetsOrigShape.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.
+ SmallVector<SmallVector<int64_t>> staticElemOffsetInSrcDims(
+ rank, SmallVector<int64_t>(numElements));
+ // For each lane_data block in the overall lane's shape
+ for (SmallVector<int64_t> laneDataOffsetInResult :
+ StaticTileOffsetRange(distShape, laneData)) {
+ ArrayRef<int64_t> staticLaneDataOffsetInSource =
+ staticLaneDataOffsetsOrigShape[linearLaneDataIdx++];
+ // For each element in the lane_data block
+ for (SmallVector<int64_t> elementOffsetInLaneData :
+ StaticTileOffsetRange(laneData, unitTile)) {
+ // For each dim of indexing space
+ SmallVector<int64_t> elementOffsetInResult(rank);
+ for (int64_t d = 0; d < rank; d++)
+ elementOffsetInResult[d] =
+ laneDataOffsetInResult[d] + elementOffsetInLaneData[d];
+ int64_t elemLinearizedIdx =
+ linearize(elementOffsetInResult, resultStrides);
+ for (int64_t d = 0; d < rank; d++)
+ staticElemOffsetInSrcDims[d][elemLinearizedIdx] =
+ staticLaneDataOffsetInSource[d] + elementOffsetInLaneData[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++) {
+ Value materializedStaticOffset = arith::ConstantOp::create(
+ rewriter, loc,
+ rewriter.getIndexVectorAttr(staticElemOffsetInSrcDims[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);
+ 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 415447351ca45..00f194110a789 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -601,11 +601,13 @@ 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.from_elements
+// CHECK-NOT: vector.shape_cast
// CHECK: gpu.return
gpu.func @create_mask_1d(%m0: index) {
%mask = vector.create_mask %m0
@@ -618,12 +620,13 @@ 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-NOT: vector.from_elements
// CHECK: gpu.return
gpu.func @constant_mask_1d() {
%mask = vector.constant_mask [4]
@@ -635,19 +638,55 @@ 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-NOT: vector.from_elements
+// 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-NOT: vector.from_elements
+// 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-NOT: vector.from_elements
// CHECK: gpu.return
gpu.func @create_mask_2d(%m0: index, %m1: index) {
%mask = vector.create_mask %m0, %m1
@@ -660,19 +699,19 @@ 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-NOT: vector.from_elements
// CHECK: gpu.return
gpu.func @constant_mask_2d() {
%mask = vector.constant_mask [2, 3]
@@ -684,6 +723,59 @@ 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-NOT: vector.from_elements
+// 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-NOT: vector.from_elements
+// 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 @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>
>From 43d365e9766762adc477d012166136134b7db864 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Mon, 21 Sep 2026 13:34:08 +0000
Subject: [PATCH 2/3] Add description, optimize bounds checks
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 93 ++++++++++++++-----
.../XeGPU/sg-to-lane-distribute-unit.mlir | 67 +++++++++++--
2 files changed, 130 insertions(+), 30 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 60ad08870a539..8f78edda68db6 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -864,8 +864,44 @@ struct SgToLaneVectorBitcast : public OpConversionPattern<vector::BitCastOp> {
};
/// Distributes a subgroup-level vector.create_mask or vector.constant_mask op
-/// to lane-level. An element is in-bounds when every dim satisfies
-/// `coord[d] < bound[d]`.
+/// to lane-level.
+/// The pattern performs the following in-bounds condition:
+/// ```
+/// 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,
+/// broadcast(%baseDistanceToBound_0) %mask_1 = arith.cmpi slt,
+/// %staticOffset_1, broadcast(%baseDistanceToBound_1) %mask =
+/// shape_cast(arith.andi %mask_0, %mask_1) : vector<16xi1> to vector<8x2xi1>
+///
template <typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
OpType, vector::CreateMaskOp, vector::ConstantMaskOp>::value>>
@@ -924,17 +960,17 @@ struct SgToLaneCreateMask : public OpConversionPattern<OpType> {
return rewriter.notifyMatchFailure(
op, "lane_data does not tile the distributed vector");
- SmallVector<int64_t> resultStrides = computeStrides(distShape);
+ 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>> staticLaneDataOffsetsOrigShape =
+ SmallVector<SmallVector<int64_t>> staticLaneDataOffsetsOrig =
layout.computeStaticDistributedCoords(/*linearId=*/0, origShape);
- if (staticLaneDataOffsetsOrigShape.empty() ||
- staticLaneDataOffsetsOrigShape.size() != laneDataCoords.size())
+ if (staticLaneDataOffsetsOrig.empty() ||
+ staticLaneDataOffsetsOrig.size() != laneDataCoords.size())
return rewriter.notifyMatchFailure(
op, "static and dynamic coordinates disagree on the distribution "
"unit count");
@@ -942,27 +978,29 @@ struct SgToLaneCreateMask : public OpConversionPattern<OpType> {
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.
- SmallVector<SmallVector<int64_t>> staticElemOffsetInSrcDims(
+ // 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 overall lane's shape
- for (SmallVector<int64_t> laneDataOffsetInResult :
+ // For each lane_data block in the lane's dist shape
+ for (SmallVector<int64_t> laneDataOffsetDist :
StaticTileOffsetRange(distShape, laneData)) {
- ArrayRef<int64_t> staticLaneDataOffsetInSource =
- staticLaneDataOffsetsOrigShape[linearLaneDataIdx++];
+ ArrayRef<int64_t> staticLaneDataOffsetOrig =
+ staticLaneDataOffsetsOrig[linearLaneDataIdx++];
// For each element in the lane_data block
- for (SmallVector<int64_t> elementOffsetInLaneData :
+ for (SmallVector<int64_t> elemOffsetInLaneData :
StaticTileOffsetRange(laneData, unitTile)) {
// For each dim of indexing space
- SmallVector<int64_t> elementOffsetInResult(rank);
+ SmallVector<int64_t> elementOffsetDist(rank);
for (int64_t d = 0; d < rank; d++)
- elementOffsetInResult[d] =
- laneDataOffsetInResult[d] + elementOffsetInLaneData[d];
- int64_t elemLinearizedIdx =
- linearize(elementOffsetInResult, resultStrides);
+ elementOffsetDist[d] =
+ laneDataOffsetDist[d] + elemOffsetInLaneData[d];
+ int64_t elemLinearizedIdxDist =
+ linearize(elementOffsetDist, distStrides);
for (int64_t d = 0; d < rank; d++)
- staticElemOffsetInSrcDims[d][elemLinearizedIdx] =
- staticLaneDataOffsetInSource[d] + elementOffsetInLaneData[d];
+ staticElemOffset[d][elemLinearizedIdxDist] =
+ staticLaneDataOffsetOrig[d] + elemOffsetInLaneData[d];
}
}
@@ -974,9 +1012,12 @@ struct SgToLaneCreateMask : public OpConversionPattern<OpType> {
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(staticElemOffsetInSrcDims[d]));
+ 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],
@@ -994,6 +1035,14 @@ struct SgToLaneCreateMask : public OpConversionPattern<OpType> {
.getResult()
: elementMaskInDim;
}
+ // 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);
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 00f194110a789..3f1cbf952df44 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -606,7 +606,6 @@ gpu.func @vector_bitcast() {
// 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.from_elements
// CHECK-NOT: vector.shape_cast
// CHECK: gpu.return
gpu.func @create_mask_1d(%m0: index) {
@@ -626,7 +625,6 @@ gpu.func @create_mask_1d(%m0: index) {
// 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-NOT: vector.from_elements
// CHECK: gpu.return
gpu.func @constant_mask_1d() {
%mask = vector.constant_mask [4]
@@ -644,7 +642,6 @@ gpu.func @constant_mask_1d() {
// 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-NOT: vector.from_elements
// CHECK: gpu.return
gpu.func @create_mask_1d_multi_unit(%m0: index) {
%mask = vector.create_mask %m0
@@ -662,7 +659,6 @@ gpu.func @create_mask_1d_multi_unit(%m0: index) {
// 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-NOT: vector.from_elements
// CHECK: gpu.return
gpu.func @create_mask_1d_lane_data(%m0: index) {
%mask = vector.create_mask %m0
@@ -686,7 +682,6 @@ gpu.func @create_mask_1d_lane_data(%m0: index) {
// 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-NOT: vector.from_elements
// CHECK: gpu.return
gpu.func @create_mask_2d(%m0: index, %m1: index) {
%mask = vector.create_mask %m0, %m1
@@ -711,7 +706,6 @@ gpu.func @create_mask_2d(%m0: index, %m1: index) {
// 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-NOT: vector.from_elements
// CHECK: gpu.return
gpu.func @constant_mask_2d() {
%mask = vector.constant_mask [2, 3]
@@ -735,7 +729,6 @@ gpu.func @constant_mask_2d() {
// 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-NOT: vector.from_elements
// CHECK: gpu.return
gpu.func @create_mask_2d_lane_data(%m0: index, %m1: index) {
%mask = vector.create_mask %m0, %m1
@@ -764,7 +757,6 @@ gpu.func @create_mask_2d_lane_data(%m0: index, %m1: index) {
// 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-NOT: vector.from_elements
// CHECK: gpu.return
gpu.func @create_mask_3d(%m0: index, %m1: index, %m2: index) {
%mask = vector.create_mask %m0, %m1, %m2
@@ -776,6 +768,65 @@ gpu.func @create_mask_3d(%m0: index, %m1: index, %m2: index) {
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>
>From 9e691408e4812f23447a712b3a47b6184b46c523 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Mon, 21 Sep 2026 13:37:52 +0000
Subject: [PATCH 3/3] Improve readability
---
.../Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 9 ++++-----
1 file changed, 4 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 8f78edda68db6..50baa3d5d756e 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -865,7 +865,7 @@ struct SgToLaneVectorBitcast : public OpConversionPattern<vector::BitCastOp> {
/// Distributes a subgroup-level vector.create_mask or vector.constant_mask op
/// to lane-level.
-/// The pattern performs the following in-bounds condition:
+/// The pattern constructs a mask based on the following bounds check:
/// ```
/// for d in [0, ..., maskRank):
/// mask &= (staticOffset[d] < (bound[d] - base[d]))
@@ -897,10 +897,9 @@ struct SgToLaneVectorBitcast : public OpConversionPattern<vector::BitCastOp> {
/// %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,
-/// broadcast(%baseDistanceToBound_0) %mask_1 = arith.cmpi slt,
-/// %staticOffset_1, broadcast(%baseDistanceToBound_1) %mask =
-/// shape_cast(arith.andi %mask_0, %mask_1) : vector<16xi1> to vector<8x2xi1>
+/// %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>
///
template <typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
More information about the Mlir-commits
mailing list