[Mlir-commits] [mlir] b6a1653 - [mlir][xegpu] Distribute extract_strided_slice over multiple dims (#227478)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 1 10:50:00 PDT 2026
Author: Jianhui Li
Date: 2026-10-01T10:49:52-07:00
New Revision: b6a1653f64c5250418b3e780e5f5b96d5cc85094
URL: https://github.com/llvm/llvm-project/commit/b6a1653f64c5250418b3e780e5f5b96d5cc85094
DIFF: https://github.com/llvm/llvm-project/commit/b6a1653f64c5250418b3e780e5f5b96d5cc85094.diff
LOG: [mlir][xegpu] Distribute extract_strided_slice over multiple dims (#227478)
SgToLaneVectorExtractStridedSlice only handled a single distributed
dimension, and this PR handles distributing multiple dimensions. It
scales each distributed dim by its own lane_layout entry. Both the size
and the offset along a dim shrink by the number of lanes that split it.
A dim the lanes do not split keeps its offset untouched, and may still
carry non-unit lane_data as before.
assisted-by-claude
---------
Co-authored-by: Claude Opus 5 (1M context) <noreply at anthropic.com>
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 06f769cefb4fd..424d6da5ec33a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1359,52 +1359,34 @@ struct SgToLaneVectorExtractStridedSlice
updatedStrides.push_back(rewriter.getI64IntegerAttr(1));
}
- // If the result is distributed, adjust offsets and sizes in the
- // distributed dimension.
+ // Each distributed dim shrinks by its own lane count, so its size and
+ // offset are rescaled by that count.
if (!distributedDims.empty()) {
- if (distributedDims.size() != 1)
- return rewriter.notifyMatchFailure(
- op, "only single dimension distribution is supported");
- int64_t distDim = distributedDims[0];
- const auto *uArch =
- xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
- if (!uArch)
- return rewriter.notifyMatchFailure(
- op, "target attribute required to determine subgroup size");
- int subgroupSize = uArch->getSubgroupSize();
auto sourceLayout = xegpu::getTemporaryLayout(op->getOpOperand(0));
if (!sourceLayout || sourceLayout.getEffectiveLaneLayoutAsInt().empty())
return rewriter.notifyMatchFailure(
op, "source of extract_strided_slice lacks distribution layout");
- int sourceDistrDimSize = op.getSourceVectorType().getShape()[distDim];
- auto laneLayout = sourceLayout.getEffectiveLaneLayoutAsInt();
- // Effective subgroup size needs to be adjusted if laneLayout along
- // the distributed dimension is smaller than subgroup size.
- if (laneLayout[distDim] < subgroupSize &&
- subgroupSize % laneLayout[distDim] == 0)
- subgroupSize = laneLayout[distDim];
- if (sourceDistrDimSize % subgroupSize != 0)
- return rewriter.notifyMatchFailure(
- op, "source size along distributed dim is not a multiple of "
- "subgroup size");
- auto sourceLaneData = sourceLayout.getEffectiveLaneDataAsInt();
- // Only check lane_data for the distributed dimension. Non-distributed
- // dimensions may have non-unit lane_data (e.g., packed layouts).
- if (distDim < static_cast<int64_t>(sourceLaneData.size()) &&
- sourceLaneData[distDim] != 1)
- return rewriter.notifyMatchFailure(
- op, "expecting unit lane data along the distributed dimension");
- int64_t distrDimOffset =
- cast<IntegerAttr>(updatedOffsets[distDim]).getInt();
- if (distrDimOffset % subgroupSize != 0)
- return rewriter.notifyMatchFailure(
- op, "offset along distributed dim is not a multiple of "
- "subgroup size");
- // Adjust sizes and offsets for the distributed dimension.
- updatedSizes[distDim] =
- rewriter.getI64IntegerAttr(distResultTy.getDimSize(distDim));
- updatedOffsets[distDim] =
- rewriter.getI64IntegerAttr(distrDimOffset / subgroupSize);
+ SmallVector<int64_t> laneLayout =
+ sourceLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> laneData = sourceLayout.getEffectiveLaneDataAsInt();
+ ArrayRef<int64_t> sourceShape = op.getSourceVectorType().getShape();
+ for (int64_t distDim : distributedDims) {
+ int64_t lanes = laneLayout[distDim];
+ if (lanes == 0 || sourceShape[distDim] % lanes != 0)
+ return rewriter.notifyMatchFailure(
+ op, "source size along a distributed dim is not a multiple of "
+ "its lane count");
+ int64_t distrDimOffset =
+ cast<IntegerAttr>(updatedOffsets[distDim]).getInt();
+ if (distrDimOffset % (lanes * laneData[distDim]) != 0)
+ return rewriter.notifyMatchFailure(
+ op, "offset along a distributed dim is not a multiple of its "
+ "lane tile");
+ updatedSizes[distDim] =
+ rewriter.getI64IntegerAttr(distResultTy.getDimSize(distDim));
+ updatedOffsets[distDim] =
+ rewriter.getI64IntegerAttr(distrDimOffset / lanes);
+ }
}
auto newOp = vector::ExtractStridedSliceOp::create(
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 e69fabbce0dc3..2431ac9492126 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -998,6 +998,22 @@ gpu.func @vector_extract_strided_slice_partial_offsets() {
gpu.return
}
+// Only 8 of the subgroup's lanes split dim 1, so that dim's size and offset are
+// rescaled by 8 rather than by the subgroup size.
+// CHECK-LABEL: gpu.func @vector_extract_strided_slice_lane_count_below_subgroup_size
+// CHECK: %[[ESS:.*]] = vector.extract_strided_slice %{{.*}} offsets = [2, 2], sizes = [4, 2], strides = [1, 1] : vector<8x4xf32> to vector<4x2xf32>
+gpu.func @vector_extract_strided_slice_lane_count_below_subgroup_size() {
+ %0 = "test.some_op"()
+ : () -> vector<8x32xf32>
+ %1 = vector.extract_strided_slice %0 offsets = [2, 16], sizes = [4, 16], strides = [1, 1]
+ : vector<8x32xf32> to vector<4x16xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ target_layout = #xegpu.layout<lane_layout = [1, 8], lane_data = [1, 1]>
+ }> : vector<4x16xf32>
+ gpu.return
+}
+
// A convert_layout that only repacks lane_data along the non-distributed outer
// dimension (from [4, 1] to [1, 1] with order = [1, 0]), keeping lane_layout
// unchanged, folds to its source when consumed by exactly 4 (== outer
@@ -1736,3 +1752,48 @@ gpu.func @convert_layout_partial_subgroup() {
gpu.return
}
}
+
+// -----
+// Two distributed dims: lane_layout [1, 2, 8] splits dims 1 and 2, so each
+// carries its own lane count and both the size and the offset along it are
+// rescaled by that count. Dim 0 is not split and keeps its offset.
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @extract_strided_slice_two_distributed_dims
+// CHECK: %[[SRC:.*]] = "test.some_op"()
+// CHECK: %[[DIST:.*]] = builtin.unrealized_conversion_cast %[[SRC]] : vector<16x32x32xbf16> to vector<16x16x4xbf16>
+// CHECK: vector.extract_strided_slice %[[DIST]] offsets = [3, 2, 2], sizes = [1, 2, 2], strides = [1, 1, 1] : vector<16x16x4xbf16> to vector<1x2x2xbf16>
+gpu.func @extract_strided_slice_two_distributed_dims() {
+ %src = "test.some_op"() : () -> vector<16x32x32xbf16>
+ %0 = vector.extract_strided_slice %src offsets = [3, 4, 16], sizes = [1, 4, 16], strides = [1, 1, 1]
+ : vector<16x32x32xbf16> to vector<1x4x16xbf16>
+ // Anchors the layout on %0; recoverTemporaryLayouts derives the operand's
+ // layout from it.
+ %1 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 1]>
+ }> : vector<1x4x16xbf16>
+ gpu.return
+}
+}
+
+// -----
+// A distributed dim that is also packed: lane_data 4 on dim 1, so a lane owns 4
+// contiguous columns. The slice covers the whole dim, so only its size rescales.
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @extract_strided_slice_packed_distributed_dim
+// CHECK: %[[SRC:.*]] = "test.some_op"()
+// CHECK: %[[DIST:.*]] = builtin.unrealized_conversion_cast %[[SRC]] : vector<8x64xbf16> to vector<8x4xbf16>
+// CHECK: vector.extract_strided_slice %[[DIST]] offsets = [0, 0], sizes = [1, 4], strides = [1, 1] : vector<8x4xbf16> to vector<1x4xbf16>
+gpu.func @extract_strided_slice_packed_distributed_dim() {
+ %src = "test.some_op"() : () -> vector<8x64xbf16>
+ %0 = vector.extract_strided_slice %src offsets = [0, 0], sizes = [1, 64], strides = [1, 1]
+ : vector<8x64xbf16> to vector<1x64xbf16>
+ %1 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>
+ }> : vector<1x64xbf16>
+ gpu.return
+}
+}
More information about the Mlir-commits
mailing list