[Mlir-commits] [mlir] [MLIR][XeGPU] Add non-splat constant distribution in SgToLane pass (PR #205575)
Nishant Patel
llvmlistbot at llvm.org
Tue Jul 7 17:10:32 PDT 2026
https://github.com/nbpatel updated https://github.com/llvm/llvm-project/pull/205575
>From 0b3f5844be4ec0ca4d1d0fac641c4b299536d1ed Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Mon, 22 Jun 2026 21:15:54 +0000
Subject: [PATCH 1/3] Add non-splat constant distribution for SgToLane
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 68 ++++++++++++++++---
.../XeGPU/sg-to-lane-distribute-unit.mlir | 20 ++++++
2 files changed, 79 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 75a87f84b3da8..25fcd6054e569 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -355,6 +355,13 @@ struct SgToLaneElementWise : public ConversionPattern {
/// Distributes a subgroup-level arith ConstantOp to lane-level arith
/// ConstantOp.
+///
+/// Splat constants are distributed by simply rebuilding the splat with the
+/// lane-local vector type. Non-splat dense constants are distributed as :
+/// `computeDistributedCoords` yields the coordinates each lane owns, each
+/// element is extracted from the full (subgroup-level) constant, and the
+/// per-lane elements are assembled into the distributed vector with
+/// `vector.from_elements` (or `vector.broadcast` for a single element).
struct SgToLaneArithConstant : public OpConversionPattern<arith::ConstantOp> {
using OpConversionPattern<arith::ConstantOp>::OpConversionPattern;
@@ -365,11 +372,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()));
@@ -385,12 +392,55 @@ 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();
+ int64_t numElements = newResultType.getNumElements();
+ assert(static_cast<int64_t>(coordsVec.size()) == numElements &&
+ "number of coordinate sets must match number of distributed "
+ "elements");
+
+ SmallVector<Value> elements;
+ for (auto &coords : coordsVec) {
+ SmallVector<OpFoldResult> mixedPos = getAsOpFoldResult(coords);
+ elements.push_back(vector::ExtractOp::create(
+ rewriter, loc, fullConst.getResult(), mixedPos));
+ }
+
+ // Assemble the distributed vector.
+ Value result;
+ if (numElements == 1) {
+ result = vector::BroadcastOp::create(rewriter, loc, newResultType,
+ elements[0]);
+ } else {
+ result = vector::FromElementsOp::create(rewriter, loc, newResultType,
+ elements);
+ }
+ 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..61f4cc5ceb675 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,26 @@ gpu.func @arith_constant() {
gpu.return
}
+// A non-splat constant is distributed by extracting the element each lane owns
+// from the full constant (using the layout's distributed coordinates).
+// 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: %[[ELEM:.*]] = vector.extract %[[CST]][%{{.*}}] : index from vector<16xindex>
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ELEM]] : index to 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
+}
+
// CHECK-LABEL: gpu.func @prefetch_nd
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: xegpu.prefetch_nd %{{.*}}[%[[C0]], %[[C0]]] : !xegpu.tensor_desc<16x16xf16>
>From 61cfd91883db7075b92c45d08c66e3e1fdf28e7b Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Tue, 30 Jun 2026 14:47:29 +0000
Subject: [PATCH 2/3] Address feedback
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 42 +++++++++++++++----
.../XeGPU/sg-to-lane-distribute-unit.mlir | 25 ++++++++++-
2 files changed, 57 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 25fcd6054e569..4cd242b5ddca5 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"
@@ -356,12 +357,10 @@ struct SgToLaneElementWise : public ConversionPattern {
/// Distributes a subgroup-level arith ConstantOp to lane-level arith
/// ConstantOp.
///
-/// Splat constants are distributed by simply rebuilding the splat with the
-/// lane-local vector type. Non-splat dense constants are distributed as :
-/// `computeDistributedCoords` yields the coordinates each lane owns, each
-/// element is extracted from the full (subgroup-level) constant, and the
-/// per-lane elements are assembled into the distributed vector with
-/// `vector.from_elements` (or `vector.broadcast` for a single element).
+/// Splat constants are rebuilt with the lane-local vector type. Non-splat
+/// constants are distributed by extracting the elements each lane owns from
+/// the full constant and assembling them with vector.from_elements (or
+/// vector.broadcast for a single element).
struct SgToLaneArithConstant : public OpConversionPattern<arith::ConstantOp> {
using OpConversionPattern<arith::ConstantOp>::OpConversionPattern;
@@ -420,12 +419,39 @@ struct SgToLaneArithConstant : public OpConversionPattern<arith::ConstantOp> {
SmallVector<SmallVector<Value>> coordsVec = maybeCoordsVec.value();
int64_t numElements = newResultType.getNumElements();
- assert(static_cast<int64_t>(coordsVec.size()) == numElements &&
+
+ // computeDistributedCoords returns the block start each lane owns. With
+ // all-ones lane_data the block is a single element. Otherwise expand each
+ // block start into its lane_data-sized element coordinates (row-major, to
+ // match vector.from_elements below).
+ SmallVector<int64_t> laneData = layout.getEffectiveLaneDataAsInt();
+ bool unitLaneData =
+ llvm::all_of(laneData, [](int64_t d) { return d == 1; });
+
+ SmallVector<SmallVector<Value>> elementCoords;
+ if (unitLaneData) {
+ elementCoords = std::move(coordsVec);
+ } else {
+ SmallVector<int64_t> unitTile(laneData.size(), 1);
+ for (const SmallVector<Value> &start : coordsVec) {
+ for (SmallVector<int64_t> off :
+ StaticTileOffsetRange(laneData, unitTile)) {
+ SmallVector<Value> coord(start.size());
+ for (size_t i = 0; i < start.size(); ++i)
+ coord[i] = arith::AddIOp::create(
+ rewriter, loc, start[i],
+ arith::ConstantIndexOp::create(rewriter, loc, off[i]));
+ elementCoords.push_back(std::move(coord));
+ }
+ }
+ }
+
+ assert(static_cast<int64_t>(elementCoords.size()) == numElements &&
"number of coordinate sets must match number of distributed "
"elements");
SmallVector<Value> elements;
- for (auto &coords : coordsVec) {
+ for (auto &coords : elementCoords) {
SmallVector<OpFoldResult> mixedPos = getAsOpFoldResult(coords);
elements.push_back(vector::ExtractOp::create(
rewriter, loc, fullConst.getResult(), mixedPos));
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 61f4cc5ceb675..fd16dc3cb6b4c 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -160,8 +160,8 @@ gpu.func @arith_constant() {
gpu.return
}
-// A non-splat constant is distributed by extracting the element each lane owns
-// from the full constant (using the layout's distributed coordinates).
+// Non-splat constant: each lane extracts the element it owns from the full
+// constant.
// 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
@@ -180,6 +180,27 @@ gpu.func @arith_constant_non_splat() {
gpu.return
}
+// With lane_data > 1 each lane owns a lane_data-sized block, so multiple
+// elements are extracted and assembled with vector.from_elements.
+// 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: %[[ELEM0:.*]] = vector.extract %[[CST]][%{{.*}}] : index from vector<32xindex>
+// CHECK: %[[ELEM1:.*]] = vector.extract %[[CST]][%{{.*}}] : index from vector<32xindex>
+// CHECK: %[[RES:.*]] = vector.from_elements %[[ELEM0]], %[[ELEM1]] : 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
+}
+
// CHECK-LABEL: gpu.func @prefetch_nd
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: xegpu.prefetch_nd %{{.*}}[%[[C0]], %[[C0]]] : !xegpu.tensor_desc<16x16xf16>
>From 678f39ce0b51a9d996ea15f3bbd02dd459c934ec Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Mon, 6 Jul 2026 20:54:53 +0000
Subject: [PATCH 3/3] Address feedback
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 93 ++++++++++---------
.../XeGPU/sg-to-lane-distribute-unit.mlir | 52 ++++++++++-
2 files changed, 95 insertions(+), 50 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 4cd242b5ddca5..0f06f16e2f8e7 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -358,9 +358,9 @@ struct SgToLaneElementWise : public ConversionPattern {
/// ConstantOp.
///
/// Splat constants are rebuilt with the lane-local vector type. Non-splat
-/// constants are distributed by extracting the elements each lane owns from
-/// the full constant and assembling them with vector.from_elements (or
-/// vector.broadcast for a single element).
+/// 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;
@@ -418,54 +418,55 @@ struct SgToLaneArithConstant : public OpConversionPattern<arith::ConstantOp> {
op, "failed to compute distributed coordinates from layout");
SmallVector<SmallVector<Value>> coordsVec = maybeCoordsVec.value();
- int64_t numElements = newResultType.getNumElements();
-
- // computeDistributedCoords returns the block start each lane owns. With
- // all-ones lane_data the block is a single element. Otherwise expand each
- // block start into its lane_data-sized element coordinates (row-major, to
- // match vector.from_elements below).
SmallVector<int64_t> laneData = layout.getEffectiveLaneDataAsInt();
- bool unitLaneData =
- llvm::all_of(laneData, [](int64_t d) { return d == 1; });
-
- SmallVector<SmallVector<Value>> elementCoords;
- if (unitLaneData) {
- elementCoords = std::move(coordsVec);
- } else {
- SmallVector<int64_t> unitTile(laneData.size(), 1);
- for (const SmallVector<Value> &start : coordsVec) {
- for (SmallVector<int64_t> off :
- StaticTileOffsetRange(laneData, unitTile)) {
- SmallVector<Value> coord(start.size());
- for (size_t i = 0; i < start.size(); ++i)
- coord[i] = arith::AddIOp::create(
- rewriter, loc, start[i],
- arith::ConstantIndexOp::create(rewriter, loc, off[i]));
- elementCoords.push_back(std::move(coord));
- }
+ 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));
}
- }
-
- assert(static_cast<int64_t>(elementCoords.size()) == numElements &&
- "number of coordinate sets must match number of distributed "
- "elements");
- SmallVector<Value> elements;
- for (auto &coords : elementCoords) {
- SmallVector<OpFoldResult> mixedPos = getAsOpFoldResult(coords);
- elements.push_back(vector::ExtractOp::create(
- rewriter, loc, fullConst.getResult(), mixedPos));
+ // 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);
}
- // Assemble the distributed vector.
- Value result;
- if (numElements == 1) {
- result = vector::BroadcastOp::create(rewriter, loc, newResultType,
- elements[0]);
- } else {
- result = vector::FromElementsOp::create(rewriter, loc, newResultType,
- elements);
- }
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 fd16dc3cb6b4c..fe356e6af35c1 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -161,12 +161,14 @@ gpu.func @arith_constant() {
}
// Non-splat constant: each lane extracts the element it owns from the full
-// constant.
+// 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: %[[BCAST:.*]] = vector.broadcast %[[ELEM]] : index to vector<1xindex>
+// 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
@@ -181,13 +183,16 @@ gpu.func @arith_constant_non_splat() {
}
// With lane_data > 1 each lane owns a lane_data-sized block, so multiple
-// elements are extracted and assembled with vector.from_elements.
+// 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: %[[RES:.*]] = vector.from_elements %[[ELEM0]], %[[ELEM1]] : vector<2xindex>
+// 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
@@ -201,6 +206,45 @@ gpu.func @arith_constant_non_splat_lane_data() {
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