[Mlir-commits] [mlir] [MLIR][XeGPU] Distribute `vector.step` with sliced layout (PR #182010)
Artem Kroviakov
llvmlistbot at llvm.org
Thu Feb 19 10:50:52 PST 2026
https://github.com/akroviakov updated https://github.com/llvm/llvm-project/pull/182010
>From b73fdec5ffd4814c2b62c750569cf704d8a0f048 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Wed, 18 Feb 2026 12:37:40 +0000
Subject: [PATCH 1/2] [MLIR][XeGPU] Distribute `vector.step` with sliced layout
---
.../Transforms/XeGPUSubgroupDistribute.cpp | 70 ++++++++++++++++++-
.../XeGPU/subgroup-distribute-unit.mlir | 27 +++++++
2 files changed, 95 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
index 99c2da386fab6..7e34bec86cf4f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
@@ -5,6 +5,7 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
+#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/GPU/Utils/DistributionUtils.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
@@ -1988,6 +1989,70 @@ struct VectorTransposeDistribution final : public gpu::WarpDistributionPattern {
}
};
+/// Distribute a vector::StepOp with the sliced result layout.
+/// The sliced layout must have exactly 1 effective lane dimension.
+/// The lane id is delinearized into the parent layout coordinate and only
+/// the effective dim coordinate is used.
+/// Example:
+/// ```
+/// %0 = gpu.warp_execute_on_lane_0(%arg0)[16] -> (vector<1xindex>) {
+/// %5 = vector.step {layout_result_0 =
+/// #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1,
+/// 1, 1, 1]>, dims = [0, 1, 3]>} : vector<1xindex> gpu.yield %5 :
+/// vector<1xindex>
+/// }
+/// ```
+/// is distributed to `arith.constant dense<0> : vector<1xindex>`
+/// because the effective lane dimension is dim 2 and the lane id is
+/// delinearized into 4D coordinate (0, 0, 0, laneid).
+struct VectorStepSliceDistribution final : public gpu::WarpDistributionPattern {
+ using gpu::WarpDistributionPattern::WarpDistributionPattern;
+ LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,
+ PatternRewriter &rewriter) const override {
+ OpOperand *operand = getWarpResult(warpOp, llvm::IsaPred<vector::StepOp>);
+ if (!operand)
+ return rewriter.notifyMatchFailure(
+ warpOp, "warp result is not a vector::StepOp op");
+ auto stepOp = operand->get().getDefiningOp<vector::StepOp>();
+ unsigned operandIdx = operand->getOperandNumber();
+ xegpu::DistributeLayoutAttr resultLayout =
+ xegpu::getTemporaryLayout(stepOp->getOpResult(0));
+ if (!resultLayout)
+ return rewriter.notifyMatchFailure(
+ stepOp, "the result vector of the step op lacks layout "
+ "attribute");
+ auto sliceLayout = dyn_cast<xegpu::SliceAttr>(resultLayout);
+ if (!sliceLayout)
+ return rewriter.notifyMatchFailure(
+ stepOp, "the result layout must be a slice layout");
+ if (sliceLayout.getEffectiveLaneLayoutAsInt().size() != 1)
+ return rewriter.notifyMatchFailure(
+ stepOp, "expecting 1 dim in the effective result layout");
+
+ rewriter.setInsertionPointAfter(warpOp);
+ auto parentLayout = cast<xegpu::LayoutAttr>(sliceLayout.getParent());
+ auto loc = stepOp.getLoc();
+ auto laneLayout = parentLayout.getEffectiveLaneLayoutAsInt();
+ auto laneLayoutValues = llvm::map_to_vector(laneLayout, [&](int64_t dim) {
+ return arith::ConstantIndexOp::create(rewriter, loc, dim).getResult();
+ });
+ auto laneIdsResult = affine::delinearizeIndex(
+ rewriter, loc, warpOp.getLaneid(), laneLayoutValues);
+ assert(!failed(laneIdsResult));
+ int expectedDimIdxSum = (laneLayout.size() * (laneLayout.size() - 1)) / 2;
+ auto sliceDims = sliceLayout.getDims().asArrayRef();
+ int actualSum = std::accumulate(sliceDims.begin(), sliceDims.end(), 0);
+ int missingDimIdx = expectedDimIdxSum - actualSum;
+ Value distributedVal = warpOp.getResult(operandIdx);
+ VectorType newVecTy = cast<VectorType>(distributedVal.getType());
+ Value laneIdVec =
+ vector::BroadcastOp::create(rewriter, warpOp.getLoc(), newVecTy,
+ laneIdsResult.value()[missingDimIdx]);
+ rewriter.replaceAllUsesWith(distributedVal, laneIdVec);
+ return success();
+ }
+};
+
} // namespace
namespace {
@@ -2014,8 +2079,9 @@ void xegpu::populateXeGPUSubgroupDistributePatterns(
patterns
.add<VectorShapeCastDistribution, VectorExtractStridedSliceDistribution,
VectorInsertStridedSliceDistribution, VectorBroadcastDistribution,
- SinkUniformOps>(patterns.getContext(),
- /*pattern benefit=*/PatternHierarchy::AboveRegular);
+ VectorStepSliceDistribution, SinkUniformOps>(
+ patterns.getContext(),
+ /*pattern benefit=*/PatternHierarchy::AboveRegular);
}
void xegpu::populateXeGPUMoveFuncBodyToWarpOpPatterns(
diff --git a/mlir/test/Dialect/XeGPU/subgroup-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/subgroup-distribute-unit.mlir
index fb23f38b44b46..b1bd082900207 100644
--- a/mlir/test/Dialect/XeGPU/subgroup-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/subgroup-distribute-unit.mlir
@@ -1132,4 +1132,31 @@ gpu.func
gpu.return
}
+// CHECK-LABEL: gpu.func @vector_step_slice
+// CHECK: (%[[SG_ID:[0-9a-zA-Z]+]]: index) {
+// CHECK-NEXT: %[[SG_ID_IN_SLICED_DIM:.*]] = affine.apply #map()[%[[SG_ID]]]
+// CHECK-NEXT: %[[SG_ID_IN_SLICED_DIM_VEC:.*]] = vector.broadcast %[[SG_ID_IN_SLICED_DIM]] : index to vector<1xindex>
+// CHECK-NEXT: "some_use"(%[[SG_ID_IN_SLICED_DIM_VEC]]) : (vector<1xindex>) -> ()
+ gpu.func @vector_step_slice(%arg0: index) {
+ %0 = gpu.warp_execute_on_lane_0(%arg0)[16] -> (vector<1xindex>) {
+ %5 = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1, 2]>} : vector<16xindex>
+ gpu.yield %5 : vector<16xindex>
+ }
+ "some_use"(%0) : (vector<1xindex>) -> ()
+ gpu.return
+ }
+
+ // CHECK-LABEL: gpu.func @vector_step_slice_unit
+ // CHECK: (%[[SG_ID:[0-9a-zA-Z]+]]: index) {
+ // CHECK-NEXT: %[[SG_ID_IN_SLICED_DIM_VEC:.*]] = arith.constant dense<0> : vector<1xindex>
+ // CHECK-NEXT: "some_use"(%[[SG_ID_IN_SLICED_DIM_VEC]]) : (vector<1xindex>) -> ()
+ gpu.func @vector_step_slice_unit(%arg0: index) {
+ %0 = gpu.warp_execute_on_lane_0(%arg0)[16] -> (vector<1xindex>) {
+ %5 = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1, 3]>} : vector<1xindex>
+ gpu.yield %5 : vector<1xindex>
+ }
+ "some_use"(%0) : (vector<1xindex>) -> ()
+ gpu.return
+ }
+
}
>From 3e06d830295eb6ae773a47450be8e066ec0b5214 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Thu, 19 Feb 2026 18:50:33 +0000
Subject: [PATCH 2/2] Enhance vector step sg distribution
---
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 2 -
.../Transforms/XeGPUSubgroupDistribute.cpp | 68 ++++++++++---------
.../XeGPU/subgroup-distribute-unit.mlir | 44 ++++++++++--
3 files changed, 74 insertions(+), 40 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 7ace00a746e21..8561b139dcedb 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -699,8 +699,6 @@ FailureOr<SmallVector<SmallVector<Value>>>
SliceAttr::computeDistributedCoords(OpBuilder &builder, Location loc,
Value linearId, ArrayRef<int64_t> shape) {
assert(getRank() == static_cast<int64_t>(shape.size()) && "invalid shape.");
- if (!isForWorkgroup())
- return failure();
SmallVector<int64_t> layout;
SmallVector<int64_t> subShape;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
index 7e34bec86cf4f..4fe619ffb0733 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
@@ -1991,20 +1991,7 @@ struct VectorTransposeDistribution final : public gpu::WarpDistributionPattern {
/// Distribute a vector::StepOp with the sliced result layout.
/// The sliced layout must have exactly 1 effective lane dimension.
-/// The lane id is delinearized into the parent layout coordinate and only
-/// the effective dim coordinate is used.
-/// Example:
-/// ```
-/// %0 = gpu.warp_execute_on_lane_0(%arg0)[16] -> (vector<1xindex>) {
-/// %5 = vector.step {layout_result_0 =
-/// #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1,
-/// 1, 1, 1]>, dims = [0, 1, 3]>} : vector<1xindex> gpu.yield %5 :
-/// vector<1xindex>
-/// }
-/// ```
-/// is distributed to `arith.constant dense<0> : vector<1xindex>`
-/// because the effective lane dimension is dim 2 and the lane id is
-/// delinearized into 4D coordinate (0, 0, 0, laneid).
+/// We completely resolve the vector::StepOp by computing the lane_data-sized subranges.
struct VectorStepSliceDistribution final : public gpu::WarpDistributionPattern {
using gpu::WarpDistributionPattern::WarpDistributionPattern;
LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,
@@ -2016,7 +2003,8 @@ struct VectorStepSliceDistribution final : public gpu::WarpDistributionPattern {
auto stepOp = operand->get().getDefiningOp<vector::StepOp>();
unsigned operandIdx = operand->getOperandNumber();
xegpu::DistributeLayoutAttr resultLayout =
- xegpu::getTemporaryLayout(stepOp->getOpResult(0));
+ xegpu::getTemporaryLayout(stepOp->getResult(0));
+ auto stepResultVecTy = stepOp.getResult().getType();
if (!resultLayout)
return rewriter.notifyMatchFailure(
stepOp, "the result vector of the step op lacks layout "
@@ -2030,25 +2018,43 @@ struct VectorStepSliceDistribution final : public gpu::WarpDistributionPattern {
stepOp, "expecting 1 dim in the effective result layout");
rewriter.setInsertionPointAfter(warpOp);
- auto parentLayout = cast<xegpu::LayoutAttr>(sliceLayout.getParent());
auto loc = stepOp.getLoc();
- auto laneLayout = parentLayout.getEffectiveLaneLayoutAsInt();
- auto laneLayoutValues = llvm::map_to_vector(laneLayout, [&](int64_t dim) {
- return arith::ConstantIndexOp::create(rewriter, loc, dim).getResult();
- });
- auto laneIdsResult = affine::delinearizeIndex(
- rewriter, loc, warpOp.getLaneid(), laneLayoutValues);
- assert(!failed(laneIdsResult));
- int expectedDimIdxSum = (laneLayout.size() * (laneLayout.size() - 1)) / 2;
- auto sliceDims = sliceLayout.getDims().asArrayRef();
- int actualSum = std::accumulate(sliceDims.begin(), sliceDims.end(), 0);
- int missingDimIdx = expectedDimIdxSum - actualSum;
Value distributedVal = warpOp.getResult(operandIdx);
VectorType newVecTy = cast<VectorType>(distributedVal.getType());
- Value laneIdVec =
- vector::BroadcastOp::create(rewriter, warpOp.getLoc(), newVecTy,
- laneIdsResult.value()[missingDimIdx]);
- rewriter.replaceAllUsesWith(distributedVal, laneIdVec);
+
+ auto laneDataBlockCoords = resultLayout.computeDistributedCoords(
+ rewriter, loc, warpOp.getLaneid(), stepResultVecTy.getShape());
+ if (failed(laneDataBlockCoords))
+ return rewriter.notifyMatchFailure(
+ stepOp, "failed to compute lane data block coordinates");
+ // No dist units at lane level
+ auto laneDataBlockCoordsVec = laneDataBlockCoords.value();
+ auto laneDataBlockLength = resultLayout.getEffectiveLaneDataAsInt()[0];
+ assert(laneDataBlockCoordsVec.size() ==
+ newVecTy.getNumElements() / laneDataBlockLength);
+ SmallVector<Value> stepOpVals;
+ // For the offset of each block of lane_data, get the "contiguous" slice
+ // from the sequence of vector.step. Example: vector.step
+ // {slice<layout<lane_layout=[2,4,2], lane_data=[1,2,1]>, slice[1,2]>} :
+ // vector<16xindex> Each logical lane holds 4 elements, with 2 blocks of 2
+ // elements each. The blocks are round-robin distributed, so logical lane id
+ // 0 will hold [0,1, 8,9] values.
+ for (int laneBlockIdx = 0; laneBlockIdx < laneDataBlockCoordsVec.size();
+ ++laneBlockIdx) {
+ auto laneDataBlockStartCoord = laneDataBlockCoordsVec[laneBlockIdx][0];
+ stepOpVals.push_back(laneDataBlockStartCoord);
+ for (int i = 1; i < laneDataBlockLength; ++i) {
+ auto offset = rewriter.create<arith::ConstantIndexOp>(loc, i);
+ stepOpVals.push_back(rewriter.create<arith::AddIOp>(
+ loc, laneDataBlockStartCoord, offset));
+ }
+ }
+ assert(stepOpVals.size() == newVecTy.getNumElements() &&
+ "Expecting the number of step op values to match the number of "
+ "elements in the vector");
+ auto stepOpVal =
+ vector::FromElementsOp::create(rewriter, loc, newVecTy, stepOpVals);
+ rewriter.replaceAllUsesWith(distributedVal, stepOpVal);
return success();
}
};
diff --git a/mlir/test/Dialect/XeGPU/subgroup-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/subgroup-distribute-unit.mlir
index b1bd082900207..31bb6704eece9 100644
--- a/mlir/test/Dialect/XeGPU/subgroup-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/subgroup-distribute-unit.mlir
@@ -1133,10 +1133,11 @@ gpu.func
}
// CHECK-LABEL: gpu.func @vector_step_slice
-// CHECK: (%[[SG_ID:[0-9a-zA-Z]+]]: index) {
-// CHECK-NEXT: %[[SG_ID_IN_SLICED_DIM:.*]] = affine.apply #map()[%[[SG_ID]]]
-// CHECK-NEXT: %[[SG_ID_IN_SLICED_DIM_VEC:.*]] = vector.broadcast %[[SG_ID_IN_SLICED_DIM]] : index to vector<1xindex>
-// CHECK-NEXT: "some_use"(%[[SG_ID_IN_SLICED_DIM_VEC]]) : (vector<1xindex>) -> ()
+// CHECK: (%[[LANE_ID:[0-9a-zA-Z]+]]: index) {
+// CHECK: %[[LANE_ID_IN_SLICED_DIM:.*]] = arith.remui %[[LANE_ID]], %c16 : index
+// CHECK-NEXT: %[[LANE_ID_IN_SLICED_DIM1:.*]] = arith.remui %[[LANE_ID_IN_SLICED_DIM]], %c16 : index
+// CHECK-NEXT: %[[LANE_ID_IN_SLICED_DIM_VEC:.*]] = vector.broadcast %[[LANE_ID_IN_SLICED_DIM1]] : index to vector<1xindex>
+// CHECK-NEXT: "some_use"(%[[LANE_ID_IN_SLICED_DIM_VEC]]) : (vector<1xindex>) -> ()
gpu.func @vector_step_slice(%arg0: index) {
%0 = gpu.warp_execute_on_lane_0(%arg0)[16] -> (vector<1xindex>) {
%5 = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1, 2]>} : vector<16xindex>
@@ -1147,9 +1148,9 @@ gpu.func
}
// CHECK-LABEL: gpu.func @vector_step_slice_unit
- // CHECK: (%[[SG_ID:[0-9a-zA-Z]+]]: index) {
- // CHECK-NEXT: %[[SG_ID_IN_SLICED_DIM_VEC:.*]] = arith.constant dense<0> : vector<1xindex>
- // CHECK-NEXT: "some_use"(%[[SG_ID_IN_SLICED_DIM_VEC]]) : (vector<1xindex>) -> ()
+ // CHECK: (%[[LANE_ID:[0-9a-zA-Z]+]]: index) {
+ // CHECK-NEXT: %[[LANE_ID_IN_SLICED_DIM_VEC:.*]] = arith.constant dense<0> : vector<1xindex>
+ // CHECK-NEXT: "some_use"(%[[LANE_ID_IN_SLICED_DIM_VEC]]) : (vector<1xindex>) -> ()
gpu.func @vector_step_slice_unit(%arg0: index) {
%0 = gpu.warp_execute_on_lane_0(%arg0)[16] -> (vector<1xindex>) {
%5 = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1, 3]>} : vector<1xindex>
@@ -1159,4 +1160,33 @@ gpu.func
gpu.return
}
+ // CHECK-LABEL: gpu.func @vector_step_slice_multi_dist_unit
+ // CHECK: (%[[LANE_ID:[0-9a-zA-Z]+]]: index) {
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[DIST_UNIT_SIZE:.*]] = arith.constant 8 : index
+ // CHECK-DAG: %[[SG_LEVEL_VECSIZE:.*]] = arith.constant 16 : index
+ // CHECK-DAG: %[[LANE_LAYOUT:.*]] = arith.constant 4 : index
+ // CHECK-DAG: %[[LANE_DATA:.*]] = arith.constant 2 : index
+ // CHECK-DAG: %[[LANE_DIST_UNIT_START_IDX:.*]] = arith.divui %[[LANE_ID]], %[[LANE_DATA]] : index
+ // CHECK-DAG: %[[DIST_UNIT_0_IDX:.*]] = arith.remui %[[LANE_DIST_UNIT_START_IDX]], %[[LANE_LAYOUT]] : index
+ // CHECK-DAG: %[[DIST_UNIT_0_OFFSET:.*]] = arith.muli %[[DIST_UNIT_0_IDX]], %[[LANE_DATA]] : index
+ // CHECK-DAG: %[[DIST_UNIT_0_SUBRANGE_START:.*]] = arith.remui %[[DIST_UNIT_0_OFFSET]], %[[SG_LEVEL_VECSIZE]] : index
+ // CHECK-DAG: %[[DIST_UNIT_1_OFFSET:.*]] = arith.addi %[[DIST_UNIT_0_OFFSET]], %[[DIST_UNIT_SIZE]] : index
+ // CHECK-DAG: %[[DIST_UNIT_1_SUBRANGE_START:.*]] = arith.remui %[[DIST_UNIT_1_OFFSET]], %[[SG_LEVEL_VECSIZE]] : index
+ // CHECK-DAG: %[[V6:.*]] = arith.addi %[[DIST_UNIT_0_SUBRANGE_START]], %[[C1]] : index
+ // CHECK-DAG: %[[V7:.*]] = arith.addi %[[DIST_UNIT_1_SUBRANGE_START]], %[[C1]] : index
+ // CHECK-DAG: %[[VEC:.*]] = vector.from_elements
+ // CHECK-SAME: %[[DIST_UNIT_0_SUBRANGE_START]], %[[V6]],
+ // CHECK-SAME: %[[DIST_UNIT_1_SUBRANGE_START]], %[[V7]]
+ // CHECK-SAME: : vector<4xindex>
+ // CHECK-NEXT: "some_use"(%[[VEC]]) : (vector<4xindex>) -> ()
+ gpu.func @vector_step_slice_multi_dist_unit(%arg0: index) {
+ %0 = gpu.warp_execute_on_lane_0(%arg0)[4] -> (vector<4xindex>) {
+ %5 = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [2, 4, 2], lane_data = [1,2,1]>, dims = [0, 2]>} : vector<16xindex>
+ gpu.yield %5 : vector<16xindex>
+ }
+ "some_use"(%0) : (vector<4xindex>) -> ()
+ gpu.return
+ }
+
}
More information about the Mlir-commits
mailing list