[Mlir-commits] [mlir] 1538aeb - [MLIR][XeGPU] Add non-splat constant distribution in SgToLane pass (#205575)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 8 07:59:40 PDT 2026
Author: Nishant Patel
Date: 2026-07-08T07:59:36-07:00
New Revision: 1538aebce77546372d39ebe25bfe75cee6ed39f8
URL: https://github.com/llvm/llvm-project/commit/1538aebce77546372d39ebe25bfe75cee6ed39f8
DIFF: https://github.com/llvm/llvm-project/commit/1538aebce77546372d39ebe25bfe75cee6ed39f8.diff
LOG: [MLIR][XeGPU] Add non-splat constant distribution in SgToLane pass (#205575)
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 f9a82bf2b1684..874487da10b30 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -10,6 +10,7 @@
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
@@ -357,6 +358,11 @@ struct SgToLaneElementWise : public ConversionPattern {
/// Distributes a subgroup-level arith ConstantOp to lane-level arith
/// ConstantOp.
+///
+/// Splat constants are rebuilt with the lane-local vector type. Non-splat
+/// constants are distributed by extracting each lane_data-sized block from
+/// the full constant and inserting it at the correct position in the
+/// distributed vector using insert_strided_slice.
struct SgToLaneArithConstant : public OpConversionPattern<arith::ConstantOp> {
using OpConversionPattern<arith::ConstantOp>::OpConversionPattern;
@@ -367,11 +373,11 @@ struct SgToLaneArithConstant : public OpConversionPattern<arith::ConstantOp> {
if (!resultType)
return failure();
- // Only handle dense vector constants
- auto dense = dyn_cast<SplatElementsAttr>(op.getValue());
- if (!dense)
+ // Only handle dense vector constants.
+ auto denseAttr = dyn_cast<DenseElementsAttr>(op.getValue());
+ if (!denseAttr)
return rewriter.notifyMatchFailure(
- op, "only dense splat vector constants are supported");
+ op, "only dense vector constants are supported");
xegpu::DistributeLayoutAttr layout =
xegpu::getTemporaryLayout(llvm::cast<OpResult>(op.getResult()));
@@ -387,12 +393,83 @@ struct SgToLaneArithConstant : public OpConversionPattern<arith::ConstantOp> {
op, "unable to compute lane vector type from the layout");
VectorType newResultType = laneShapeOrFailure.value();
- auto sclarValue = dense.getSplatValue<Attribute>();
- auto newDenseAttr = DenseElementsAttr::get(newResultType, sclarValue);
+ Location loc = op.getLoc();
- auto newOp = arith::ConstantOp::create(rewriter, op.getLoc(), newResultType,
- newDenseAttr);
- rewriter.replaceOp(op, newOp.getResult());
+ // Splat constants: every lane gets the same value, so just rebuild the
+ // splat with the distributed type.
+ if (denseAttr.isSplat()) {
+ auto scalarValue = denseAttr.getSplatValue<Attribute>();
+ auto newDenseAttr = DenseElementsAttr::get(newResultType, scalarValue);
+ auto newOp =
+ arith::ConstantOp::create(rewriter, loc, newResultType, newDenseAttr);
+ rewriter.replaceOp(op, newOp.getResult());
+ return success();
+ }
+
+ // Non-splat constants: each lane extracts the elements it owns from the
+ // full constant using the distributed coordinates from the layout.
+ auto fullConst =
+ arith::ConstantOp::create(rewriter, loc, resultType, denseAttr);
+
+ Value laneId = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
+ /*upperBound=*/mlir::IntegerAttr());
+ auto maybeCoordsVec = layout.computeDistributedCoords(
+ rewriter, loc, laneId, resultType.getShape());
+ if (failed(maybeCoordsVec))
+ return rewriter.notifyMatchFailure(
+ op, "failed to compute distributed coordinates from layout");
+
+ SmallVector<SmallVector<Value>> coordsVec = maybeCoordsVec.value();
+ SmallVector<int64_t> laneData = layout.getEffectiveLaneDataAsInt();
+ ArrayRef<int64_t> distShape = newResultType.getShape();
+ int64_t rank = newResultType.getRank();
+
+ // Each lane owns one lane_data-sized block per distribution unit.
+ // computeDistributedCoords returns those block starts in row-major order
+ // over the block grid (distShape / laneData).
+ SmallVector<int64_t> blockGridShape(rank);
+ for (int64_t d = 0; d < rank; d++)
+ blockGridShape[d] = distShape[d] / laneData[d];
+ SmallVector<int64_t> blockGridStrides = computeStrides(blockGridShape);
+
+ auto blockType = VectorType::get(laneData, newResultType.getElementType());
+ SmallVector<int64_t> unitTile(rank, 1);
+ SmallVector<int64_t> strides(rank, 1);
+
+ Value result = arith::ConstantOp::create(
+ rewriter, loc, newResultType, rewriter.getZeroAttr(newResultType));
+
+ for (auto [blockIdx, blockStart] : llvm::enumerate(coordsVec)) {
+ // Gather the block's elements from the full constant. The block start is
+ // lane-dynamic, so extract element-by-element (row-major over lane_data)
+ // instead.
+ SmallVector<Value> blockElems;
+ for (SmallVector<int64_t> off :
+ StaticTileOffsetRange(laneData, unitTile)) {
+ SmallVector<OpFoldResult> pos(rank);
+ for (int64_t d = 0; d < rank; d++)
+ pos[d] = getAsOpFoldResult(arith::AddIOp::create(
+ rewriter, loc, blockStart[d],
+ arith::ConstantIndexOp::create(rewriter, loc, off[d])));
+ blockElems.push_back(vector::ExtractOp::create(
+ rewriter, loc, fullConst.getResult(), pos));
+ }
+
+ // Rebuild the block keeping its lane_data shape, then place it with
+ // insert_strided_slice so the block keeps its orientation in the
+ // distributed vector (e.g. a [2, 1] block stays a vertical 2x1 slice).
+ Value block =
+ vector::FromElementsOp::create(rewriter, loc, blockType, blockElems);
+ SmallVector<int64_t> blockGridPos =
+ delinearize(blockIdx, blockGridStrides);
+ SmallVector<int64_t> offsets(rank);
+ for (int64_t d = 0; d < rank; d++)
+ offsets[d] = blockGridPos[d] * laneData[d];
+ result = vector::InsertStridedSliceOp::create(rewriter, loc, block,
+ result, offsets, strides);
+ }
+
+ rewriter.replaceOp(op, result);
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 fcc2da3d5005b..fe356e6af35c1 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -160,6 +160,91 @@ gpu.func @arith_constant() {
gpu.return
}
+// Non-splat constant: each lane extracts the element it owns from the full
+// constant and inserts its lane_data-sized block into the distributed vector.
+// CHECK-LABEL: gpu.func @arith_constant_non_splat
+// CHECK: %[[CST:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : vector<16xindex>
+// CHECK: %[[LANE:.*]] = gpu.lane_id
+// CHECK: %[[ZERO:.*]] = arith.constant dense<0> : vector<1xindex>
+// CHECK: %[[ELEM:.*]] = vector.extract %[[CST]][%{{.*}}] : index from vector<16xindex>
+// CHECK: %[[BLK:.*]] = vector.from_elements %[[ELEM]] : vector<1xindex>
+// CHECK: %[[RES:.*]] = vector.insert_strided_slice %[[BLK]], %[[ZERO]] {offsets = [0], strides = [1]} : vector<1xindex> into vector<1xindex>
+// CHECK: gpu.return
+gpu.func @arith_constant_non_splat() {
+ %0 = arith.constant
+ {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>}
+ dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : vector<16xindex>
+ %cl0 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>
+ }> : vector<16xindex>
+ gpu.return
+}
+
+// With lane_data > 1 each lane owns a lane_data-sized block, so multiple
+// elements are extracted, reassembled into the block, and placed with
+// insert_strided_slice.
+// CHECK-LABEL: gpu.func @arith_constant_non_splat_lane_data
+// CHECK: %[[CST:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]> : vector<32xindex>
+// CHECK: %[[LANE:.*]] = gpu.lane_id
+// CHECK: %[[ZERO:.*]] = arith.constant dense<0> : vector<2xindex>
+// CHECK: %[[ELEM0:.*]] = vector.extract %[[CST]][%{{.*}}] : index from vector<32xindex>
+// CHECK: %[[ELEM1:.*]] = vector.extract %[[CST]][%{{.*}}] : index from vector<32xindex>
+// CHECK: %[[BLK:.*]] = vector.from_elements %[[ELEM0]], %[[ELEM1]] : vector<2xindex>
+// CHECK: %[[RES:.*]] = vector.insert_strided_slice %[[BLK]], %[[ZERO]] {offsets = [0], strides = [1]} : vector<2xindex> into vector<2xindex>
+// CHECK: gpu.return
+gpu.func @arith_constant_non_splat_lane_data() {
+ %0 = arith.constant
+ {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>, dims = [0]>}
+ dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]> : vector<32xindex>
+ %cl0 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>, dims = [0]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>, dims = [0]>
+ }> : vector<32xindex>
+ gpu.return
+}
+
+// 2D non-splat constant with vertical lane_data [2, 1]. The distributed type
+// is vector<4x2xindex>. Each [2, 1] block is correctly placed vertically in
+// the result (same column, adjacent rows).
+// CHECK-LABEL: gpu.func @arith_constant_non_splat_2d_vertical_lanedata
+// CHECK: %[[CST:.*]] = arith.constant dense<{{.*}}> : vector<4x32xindex>
+// CHECK: %[[LANE:.*]] = gpu.lane_id
+// CHECK: %[[ZERO:.*]] = arith.constant dense<0> : vector<4x2xindex>
+// CHECK: %[[E0:.*]] = vector.extract %[[CST]][%{{.*}}, %{{.*}}] : index from vector<4x32xindex>
+// CHECK: %[[E1:.*]] = vector.extract %[[CST]][%{{.*}}, %{{.*}}] : index from vector<4x32xindex>
+// CHECK: %[[B0:.*]] = vector.from_elements %[[E0]], %[[E1]] : vector<2x1xindex>
+// CHECK: %[[I0:.*]] = vector.insert_strided_slice %[[B0]], %[[ZERO]] {offsets = [0, 0], strides = [1, 1]} : vector<2x1xindex> into vector<4x2xindex>
+// CHECK: %[[E2:.*]] = vector.extract %[[CST]][%{{.*}}, %{{.*}}] : index from vector<4x32xindex>
+// CHECK: %[[E3:.*]] = vector.extract %[[CST]][%{{.*}}, %{{.*}}] : index from vector<4x32xindex>
+// CHECK: %[[B1:.*]] = vector.from_elements %[[E2]], %[[E3]] : vector<2x1xindex>
+// CHECK: %[[I1:.*]] = vector.insert_strided_slice %[[B1]], %[[I0]] {offsets = [0, 1], strides = [1, 1]} : vector<2x1xindex> into vector<4x2xindex>
+// CHECK: %[[E4:.*]] = vector.extract %[[CST]][%{{.*}}, %{{.*}}] : index from vector<4x32xindex>
+// CHECK: %[[E5:.*]] = vector.extract %[[CST]][%{{.*}}, %{{.*}}] : index from vector<4x32xindex>
+// CHECK: %[[B2:.*]] = vector.from_elements %[[E4]], %[[E5]] : vector<2x1xindex>
+// CHECK: %[[I2:.*]] = vector.insert_strided_slice %[[B2]], %[[I1]] {offsets = [2, 0], strides = [1, 1]} : vector<2x1xindex> into vector<4x2xindex>
+// CHECK: %[[E6:.*]] = vector.extract %[[CST]][%{{.*}}, %{{.*}}] : index from vector<4x32xindex>
+// CHECK: %[[E7:.*]] = vector.extract %[[CST]][%{{.*}}, %{{.*}}] : index from vector<4x32xindex>
+// CHECK: %[[B3:.*]] = vector.from_elements %[[E6]], %[[E7]] : vector<2x1xindex>
+// CHECK: %[[I3:.*]] = vector.insert_strided_slice %[[B3]], %[[I2]] {offsets = [2, 1], strides = [1, 1]} : vector<2x1xindex> into vector<4x2xindex>
+// CHECK: gpu.return
+gpu.func @arith_constant_non_splat_2d_vertical_lanedata() {
+ %0 = arith.constant
+ {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ dense<[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31],
+ [32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63],
+ [64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95],
+ [96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127]]> : vector<4x32xindex>
+ %cl0 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>
+ }> : vector<4x32xindex>
+ gpu.return
+}
+
// CHECK-LABEL: gpu.func @prefetch_nd
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: xegpu.prefetch_nd %{{.*}}[%[[C0]], %[[C0]]] : !xegpu.tensor_desc<16x16xf16>
More information about the Mlir-commits
mailing list