[Mlir-commits] [mlir] [MLIR][XeGPU] Add 2D `vector.multi_reduction` optimization (PR #171154)
Artem Kroviakov
llvmlistbot at llvm.org
Mon Dec 15 09:01:01 PST 2025
https://github.com/akroviakov updated https://github.com/llvm/llvm-project/pull/171154
>From a83babef75d79905b41b479b97165206ec93da6a Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Mon, 8 Dec 2025 16:33:23 +0000
Subject: [PATCH 1/3] [MLIR][XeGPU] Add 2D `vector.multi_reduction`
optimization
---
.../Transforms/XeGPUOptimizeBlockLoads.cpp | 132 +++++++++++++++++-
.../Dialect/XeGPU/optimize-2d-reduction.mlir | 85 +++++++++++
2 files changed, 216 insertions(+), 1 deletion(-)
create mode 100644 mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
index ab41fe4298d99..238599e21f65a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
@@ -416,12 +416,131 @@ class VectorExtractOpPattern final
}
};
+class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(vector::MultiDimReductionOp reductionOp, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ if (reductionOp.getReductionDims().size() != 2)
+ return rewriter.notifyMatchFailure(reductionOp,
+ "Expected 2D multi reduction");
+
+ auto layout = xegpu::getDistributeLayoutAttr(reductionOp.getResult());
+
+ auto dims = llvm::to_vector(reductionOp.getReductionDims());
+ auto [intraLaneDim, crossLaneDim] = getReductionDimOrder(dims, layout);
+ // Order does not matter
+ if (intraLaneDim == -1 || crossLaneDim == -1) {
+ intraLaneDim = dims[0];
+ crossLaneDim = dims[1];
+ }
+ auto loc = reductionOp.getLoc();
+ // XeGPU transforms expect vector types
+ auto sourceVecType = reductionOp.getSourceVectorType();
+ auto acc = reductionOp.getAcc();
+ bool scalarAcc = !isa<VectorType>(acc.getType());
+ if (scalarAcc)
+ acc = vector::FromElementsOp::create(
+ rewriter, loc, VectorType::get({1}, sourceVecType.getElementType()),
+ acc);
+
+ // Preserve layout in the intermediate reduction (apart from the reduced
+ // dim)
+ auto sourceSliceLayoutAttr = cast<xegpu::SliceAttr>(layout);
+ SmallVector<int64_t> sliceDims{
+ sourceSliceLayoutAttr.getDims().asArrayRef()};
+ auto foundIt = std::find(sliceDims.begin(), sliceDims.end(), crossLaneDim);
+ assert(foundIt != sliceDims.end() &&
+ "Expected to find reduction dim in slice dims");
+ sliceDims.erase(foundIt);
+ auto intraLaneLayout = xegpu::SliceAttr::get(
+ reductionOp.getContext(), sourceSliceLayoutAttr.getParent(),
+ DenseI64ArrayAttr::get(getContext(), sliceDims));
+
+ // First we do intra-lane reduction
+ SmallVector<int64_t> accShape(sourceVecType.getShape());
+ accShape.erase(accShape.begin() + intraLaneDim);
+ // Add a dim to the lower-dim user-supplied acc
+ Value firstRedAcc = acc;
+ if (firstRedAcc) {
+ firstRedAcc = vector::BroadcastOp::create(
+ rewriter, loc,
+ VectorType::get(accShape, sourceVecType.getElementType()), acc);
+ xegpu::setDistributeLayoutAttr(
+ llvm::dyn_cast<OpResult>(firstRedAcc),
+ cast<xegpu::DistributeLayoutAttr>(intraLaneLayout));
+ }
+ Value intraLaneReduced = vector::MultiDimReductionOp::create(
+ rewriter, loc, reductionOp.getKind(), reductionOp.getSource(),
+ firstRedAcc, ArrayRef<int64_t>(intraLaneDim));
+ xegpu::setDistributeLayoutAttr(
+ llvm::dyn_cast<OpResult>(intraLaneReduced),
+ cast<xegpu::DistributeLayoutAttr>(intraLaneLayout));
+
+ // For scalar results, add a unit dim where intra lane dim was
+ if (scalarAcc) {
+ SmallVector<int64_t> vecTypeWithUnitDim{sourceVecType.getShape()};
+ vecTypeWithUnitDim[intraLaneDim] = 1;
+ intraLaneReduced = vector::ShapeCastOp::create(
+ rewriter, loc,
+ VectorType::get(vecTypeWithUnitDim, sourceVecType.getElementType()),
+ intraLaneReduced);
+ // Layout matches last reduction
+ xegpu::setDistributeLayoutAttr(llvm::dyn_cast<OpResult>(intraLaneReduced),
+ layout);
+ } else
+ crossLaneDim -= static_cast<int64_t>(intraLaneDim < crossLaneDim);
+ // Do cross-lane reduction
+ Value crossLaneReduced = vector::MultiDimReductionOp::create(
+ rewriter, loc, reductionOp.getKind(), intraLaneReduced, acc,
+ ArrayRef<int64_t>(crossLaneDim));
+ xegpu::setDistributeLayoutAttr(llvm::dyn_cast<OpResult>(crossLaneReduced),
+ layout);
+
+ if (scalarAcc)
+ crossLaneReduced =
+ vector::ExtractOp::create(rewriter, loc, crossLaneReduced, 0);
+ assert(crossLaneReduced.getType() == reductionOp.getResult().getType() &&
+ "Type mismatch");
+ rewriter.replaceOp(reductionOp, crossLaneReduced);
+ return success();
+ }
+
+private:
+ std::pair<int64_t, int64_t>
+ getReductionDimOrder(ArrayRef<int64_t> reductionDims,
+ xegpu::DistributeLayoutAttr layout) const {
+ assert(layout.isForSubgroup() && "Must know the lane layout");
+ assert(reductionDims.size() == 2 && "Expected 2D reduction");
+ int64_t intra, cross = -1;
+ xegpu::LayoutAttr layoutAttr = dyn_cast<xegpu::LayoutAttr>(layout);
+ if (auto layoutSliceAttr = dyn_cast<xegpu::SliceAttr>(layout)) {
+ while (dyn_cast<xegpu::SliceAttr>(layoutSliceAttr.getParent()))
+ layoutSliceAttr =
+ dyn_cast<xegpu::SliceAttr>(layoutSliceAttr.getParent());
+ layoutAttr = dyn_cast<xegpu::LayoutAttr>(layoutSliceAttr.getParent());
+ }
+ assert(layoutAttr);
+ SmallVector<int64_t> laneLayout = layoutAttr.getEffectiveLaneLayoutAsInt();
+
+ assert(laneLayout.size() && "Expected a non-empty layout");
+ // try to pick a dim that does not communicate
+ for (auto dim : reductionDims) {
+ if (laneLayout[dim] == 1)
+ intra = dim;
+ else
+ cross = dim;
+ }
+ return {intra, cross};
+ }
+};
+
} // namespace
void xegpu::populateXeGPUOptimizeBlockLoadsPatterns(
RewritePatternSet &patterns) {
patterns.add<XeGPUCreateNdDescOpPattern, XeGPULoadNdDescOpPattern,
- VectorExtractOpPattern>(patterns.getContext());
+ VectorExtractOpPattern, MultiRed2dOp>(patterns.getContext());
}
namespace {
@@ -472,6 +591,17 @@ struct XeGPUOptimizeBlockLoadsPass final
auto laneData = layout.getEffectiveLaneDataAsInt();
return !canBeOptimizedForTranspose(laneLayout, laneData);
});
+
+ target.addDynamicallyLegalOp<vector::MultiDimReductionOp>(
+ [=](Operation *op) -> bool {
+ auto layout = xegpu::getDistributeLayoutAttr(op->getResult(0));
+ if (!layout || !layout.isForSubgroup())
+ return true;
+ if (auto reductionOp = dyn_cast<vector::MultiDimReductionOp>(op))
+ return reductionOp.getReductionDims().size() != 2;
+ return true;
+ });
+
converter.addConversion([](Type type) { return type; });
target.addLegalDialect<arith::ArithDialect, memref::MemRefDialect,
diff --git a/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir b/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir
new file mode 100644
index 0000000000000..754825193a10f
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir
@@ -0,0 +1,85 @@
+// RUN: mlir-opt --xevm-attach-target='module=xevm_* chip=pvc' \
+// RUN: --xegpu-optimize-block-loads --split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: gpu.func @vector_reduce_2d(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<4x16xf32>) {
+// CHECK: %[[ACC:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} 1.000000e+00 : f32
+// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<4x16xf32> -> !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>>
+// CHECK: %[[LOADED:.*]] = xegpu.load_nd %[[TDESC]][0, 0] : !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>> -> vector<4x16xf32>
+// CHECK: %[[ACC_VEC:.*]] = vector.from_elements %[[ACC]] : vector<1xf32>
+// CHECK: %[[ACC_VEC_FOR_INTRA:.*]] = vector.broadcast %[[ACC_VEC]]
+// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>} : vector<1xf32> to vector<16xf32>
+// CHECK: %[[LOADED_REDUCED:.*]] = vector.multi_reduction <add>, %[[LOADED]], %[[ACC_VEC_FOR_INTRA]]
+// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>} [0] : vector<4x16xf32> to vector<16xf32>
+// CHECK: %[[LOADED_REDUCED_FOR_CROSS:.*]] = vector.shape_cast %[[LOADED_REDUCED]]
+// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} : vector<16xf32> to vector<1x16xf32>
+// CHECK: %[[LOADED_REDUCED_2D:.*]] = vector.multi_reduction <add>, %[[LOADED_REDUCED_FOR_CROSS]], %[[ACC_VEC]]
+// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} [1] : vector<1x16xf32> to vector<1xf32>
+// CHECK: %[[SCALAR_RES:.*]] = vector.extract %[[LOADED_REDUCED_2D]][0] : f32 from vector<1xf32>
+gpu.module @xevm_test {
+ gpu.func @vector_reduce_2d(%src: memref<4x16xf32>) {
+ %cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} 1.0 : f32
+ %tdesc = xegpu.create_nd_tdesc %src : memref<4x16xf32>
+ -> !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>>
+ %load = xegpu.load_nd %tdesc[0, 0]
+ : !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>>
+ -> vector<4x16xf32>
+ %reduce = vector.multi_reduction <add>, %load, %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} [0, 1]
+ : vector<4x16xf32> to f32
+ gpu.return
+ }
+}
+
+// -----
+// CHECK-LABEL: gpu.func @vector_reduce_2d(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<4x16xf32>) {
+// CHECK: %[[ACC:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>, dims = [0, 1]>} dense<1.000000e+00> : vector<1xf32>
+// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<4x16xf32> -> !xegpu.tensor_desc<4x16xf32>
+// CHECK: %[[LOADED:.*]] = xegpu.load_nd %[[TDESC]][0, 0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>} : !xegpu.tensor_desc<4x16xf32> -> vector<4x16xf32>
+// CHECK: %[[LOADED_LEADING_UNIT:.*]] = vector.shape_cast %[[LOADED]]
+// CHECK-SAME: {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>} : vector<4x16xf32> to vector<1x4x16xf32
+// CHECK: %[[ACC_VEC_FOR_INTRA:.*]] = vector.broadcast %[[ACC]]
+// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>, dims = [1]>} : vector<1xf32> to vector<1x16xf32>
+// CHECK: %[[LOADED_REDUCED:.*]] = vector.multi_reduction <add>, %[[LOADED_LEADING_UNIT]], %[[ACC_VEC_FOR_INTRA]]
+// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>, dims = [1]>} [1] : vector<1x4x16xf32> to vector<1x16xf32>
+// CHECK: %[[LOADED_REDUCED_2D:.*]] = vector.multi_reduction <add>, %[[LOADED_REDUCED]], %[[ACC]]
+// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>, dims = [1, 2]>} [1] : vector<1x16xf32> to vector<1xf32>
+gpu.module @xevm_test {
+ gpu.func @vector_reduce_2d(%src: memref<4x16xf32>) {
+ %cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>, dims = [0, 1]>} dense<1.0> : vector<1xf32>
+ %tdesc = xegpu.create_nd_tdesc %src : memref<4x16xf32>
+ -> !xegpu.tensor_desc<4x16xf32>
+ %load = xegpu.load_nd %tdesc[0, 0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>} : !xegpu.tensor_desc<4x16xf32> -> vector<4x16xf32>
+ %load_with_dim = vector.shape_cast %load {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>} : vector<4x16xf32> to vector<1x4x16xf32>
+ %reduce = vector.multi_reduction <add>, %load_with_dim, %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1,1, 16], lane_data = [1, 4, 1]>, dims = [1, 2]>} [1, 2]
+ : vector<1x4x16xf32> to vector<1xf32>
+ gpu.return
+ }
+}
+
+// -----
+// CHECK-LABEL: gpu.func @vector_reduce_2d(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<4x64xf32>) {
+// CHECK: %[[ACC:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>, dims = [0, 1]>} dense<1.000000e+00> : vector<1xf32>
+// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<4x64xf32> -> !xegpu.tensor_desc<1x64xf32>
+// CHECK: %[[LOADED:.*]] = xegpu.load_nd %[[TDESC]][0, 0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>} : !xegpu.tensor_desc<1x64xf32> -> vector<1x64xf32>
+// CHECK: %[[LOADED_LEADING_UNIT:.*]] = vector.shape_cast %[[LOADED]]
+// CHECK-SAME: {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>} : vector<1x64xf32> to vector<1x1x64xf32>
+// CHECK: %[[ACC_VEC_FOR_INTRA:.*]] = vector.broadcast %[[ACC]]
+// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>, dims = [1]>} : vector<1xf32> to vector<1x64xf32>
+// CHECK: %[[LOADED_REDUCED:.*]] = vector.multi_reduction <add>, %[[LOADED_LEADING_UNIT]], %[[ACC_VEC_FOR_INTRA]]
+// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>, dims = [1]>} [1] : vector<1x1x64xf32> to vector<1x64xf32>
+// CHECK: %[[LOADED_REDUCED_2D:.*]] = vector.multi_reduction <add>, %[[LOADED_REDUCED]], %[[ACC]]
+// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>, dims = [1, 2]>} [1] : vector<1x64xf32> to vector<1xf32>
+gpu.module @xevm_test {
+ gpu.func @vector_reduce_2d(%src: memref<4x64xf32>) {
+ %cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>, dims = [0, 1]>} dense<1.0> : vector<1xf32>
+ %tdesc = xegpu.create_nd_tdesc %src : memref<4x64xf32>
+ -> !xegpu.tensor_desc<1x64xf32>
+ %load = xegpu.load_nd %tdesc[0, 0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>} : !xegpu.tensor_desc<1x64xf32> -> vector<1x64xf32>
+ %load_with_dim = vector.shape_cast %load {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>} : vector<1x64xf32> to vector<1x1x64xf32>
+ %reduce = vector.multi_reduction <add>, %load_with_dim, %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>, dims = [1, 2]>} [1, 2]
+ : vector<1x1x64xf32> to vector<1xf32>
+ gpu.return
+ }
+}
>From 7bfb48d031587f7cf060239e753a1881fb679601 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Thu, 11 Dec 2025 09:06:39 +0000
Subject: [PATCH 2/3] Add layout handling, disallow leading unit dims/slices
for now
---
.../Transforms/XeGPUOptimizeBlockLoads.cpp | 64 +++++++++++-----
.../Transforms/XeGPUSubgroupDistribute.cpp | 15 ++--
.../Dialect/XeGPU/optimize-2d-reduction.mlir | 73 ++++---------------
.../Dialect/XeGPU/subgroup-distribute.mlir | 49 ++++++++++++-
4 files changed, 114 insertions(+), 87 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
index 238599e21f65a..3c6513ff34fed 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
@@ -424,40 +424,50 @@ class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
if (reductionOp.getReductionDims().size() != 2)
return rewriter.notifyMatchFailure(reductionOp,
"Expected 2D multi reduction");
+ // Retrieve layouts.
+ auto resLayout = xegpu::getDistributeLayoutAttr(reductionOp.getResult());
+ auto srcLayout = xegpu::getDistributeLayoutAttr(reductionOp.getSource());
+ assert(isa<xegpu::LayoutAttr>(srcLayout) &&
+ "Currently we do not support sliced inputs");
- auto layout = xegpu::getDistributeLayoutAttr(reductionOp.getResult());
-
+ // Retrieve and order dims for 1D decomposition (prefer intra-lane first).
auto dims = llvm::to_vector(reductionOp.getReductionDims());
- auto [intraLaneDim, crossLaneDim] = getReductionDimOrder(dims, layout);
+ auto [intraLaneDim, crossLaneDim] = getReductionDimOrder(dims, resLayout);
// Order does not matter
if (intraLaneDim == -1 || crossLaneDim == -1) {
intraLaneDim = dims[0];
crossLaneDim = dims[1];
}
auto loc = reductionOp.getLoc();
- // XeGPU transforms expect vector types
auto sourceVecType = reductionOp.getSourceVectorType();
auto acc = reductionOp.getAcc();
+ // If the accumulator is scalar, convert to 1-element vector and assign the
+ // result layout
bool scalarAcc = !isa<VectorType>(acc.getType());
- if (scalarAcc)
+ // TODO: remove scalar acc assumption (need more complex layout adjustments
+ // for sliced inputs).
+ assert(scalarAcc && "Expected scalar acc");
+ if (scalarAcc) {
acc = vector::FromElementsOp::create(
rewriter, loc, VectorType::get({1}, sourceVecType.getElementType()),
acc);
+ xegpu::setDistributeLayoutAttr(
+ llvm::dyn_cast<OpResult>(acc),
+ cast<xegpu::DistributeLayoutAttr>(resLayout));
+ }
- // Preserve layout in the intermediate reduction (apart from the reduced
- // dim)
- auto sourceSliceLayoutAttr = cast<xegpu::SliceAttr>(layout);
- SmallVector<int64_t> sliceDims{
- sourceSliceLayoutAttr.getDims().asArrayRef()};
+ // The first reduction's dist attribute does not have the cross lane dim.
+ auto resSliceLayoutAttr = cast<xegpu::SliceAttr>(resLayout);
+ SmallVector<int64_t> sliceDims{resSliceLayoutAttr.getDims().asArrayRef()};
auto foundIt = std::find(sliceDims.begin(), sliceDims.end(), crossLaneDim);
assert(foundIt != sliceDims.end() &&
"Expected to find reduction dim in slice dims");
sliceDims.erase(foundIt);
- auto intraLaneLayout = xegpu::SliceAttr::get(
- reductionOp.getContext(), sourceSliceLayoutAttr.getParent(),
+ auto intraLaneRedResLayout = xegpu::SliceAttr::get(
+ reductionOp.getContext(), resSliceLayoutAttr.getParent(),
DenseI64ArrayAttr::get(getContext(), sliceDims));
- // First we do intra-lane reduction
+ // We reduce only one dim first, adjsut accumulator.
SmallVector<int64_t> accShape(sourceVecType.getShape());
accShape.erase(accShape.begin() + intraLaneDim);
// Add a dim to the lower-dim user-supplied acc
@@ -468,34 +478,50 @@ class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
VectorType::get(accShape, sourceVecType.getElementType()), acc);
xegpu::setDistributeLayoutAttr(
llvm::dyn_cast<OpResult>(firstRedAcc),
- cast<xegpu::DistributeLayoutAttr>(intraLaneLayout));
+ cast<xegpu::DistributeLayoutAttr>(intraLaneRedResLayout));
}
Value intraLaneReduced = vector::MultiDimReductionOp::create(
rewriter, loc, reductionOp.getKind(), reductionOp.getSource(),
firstRedAcc, ArrayRef<int64_t>(intraLaneDim));
xegpu::setDistributeLayoutAttr(
llvm::dyn_cast<OpResult>(intraLaneReduced),
- cast<xegpu::DistributeLayoutAttr>(intraLaneLayout));
+ cast<xegpu::DistributeLayoutAttr>(intraLaneRedResLayout));
- // For scalar results, add a unit dim where intra lane dim was
+ xegpu::DistributeLayoutAttr nextMultiRedLayout = intraLaneRedResLayout;
+ // Example: vector<2x4> got reduced to vector<2>, next reduction returns a
+ // scalar, distribution passes do not support this result type. Expand to
+ // vector<2x1>, so that the second reduction result is vector<1>. Restore
+ // this dim in layout, but lane data is 1.
if (scalarAcc) {
+ SmallVector<int> srcLaneData(srcLayout.getRank(), 1);
+ auto laneLayoutSrc = srcLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int> srcLaneLayout(laneLayoutSrc.begin(),
+ laneLayoutSrc.end());
+ nextMultiRedLayout = xegpu::LayoutAttr::get(
+ reductionOp.getContext(),
+ DenseI32ArrayAttr::get(reductionOp.getContext(), srcLaneLayout),
+ DenseI32ArrayAttr::get(reductionOp.getContext(), srcLaneData));
+
SmallVector<int64_t> vecTypeWithUnitDim{sourceVecType.getShape()};
vecTypeWithUnitDim[intraLaneDim] = 1;
intraLaneReduced = vector::ShapeCastOp::create(
rewriter, loc,
VectorType::get(vecTypeWithUnitDim, sourceVecType.getElementType()),
intraLaneReduced);
- // Layout matches last reduction
xegpu::setDistributeLayoutAttr(llvm::dyn_cast<OpResult>(intraLaneReduced),
- layout);
+ nextMultiRedLayout);
} else
crossLaneDim -= static_cast<int64_t>(intraLaneDim < crossLaneDim);
// Do cross-lane reduction
+ // TODO: why use accumulator again?
Value crossLaneReduced = vector::MultiDimReductionOp::create(
rewriter, loc, reductionOp.getKind(), intraLaneReduced, acc,
ArrayRef<int64_t>(crossLaneDim));
+ auto crossLaneLayout = xegpu::SliceAttr::get(
+ reductionOp.getContext(), nextMultiRedLayout,
+ DenseI64ArrayAttr::get(getContext(), crossLaneDim));
xegpu::setDistributeLayoutAttr(llvm::dyn_cast<OpResult>(crossLaneReduced),
- layout);
+ crossLaneLayout);
if (scalarAcc)
crossLaneReduced =
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
index ca81c3cd7be42..3dd9bf4556bb0 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
@@ -875,12 +875,15 @@ struct StoreDistribution final : public gpu::WarpDistributionPattern {
std::string layoutMaskName =
xegpu::getLayoutName(storeScatterOp->getOpOperand(3));
- xegpu::LayoutAttr layoutPayload =
- storeScatterOp->getAttrOfType<xegpu::LayoutAttr>(layoutPayloadName);
- xegpu::LayoutAttr layoutOffsets =
- storeScatterOp->getAttrOfType<xegpu::LayoutAttr>(layoutOffsetsName);
- xegpu::LayoutAttr layoutMask =
- storeScatterOp->getAttrOfType<xegpu::LayoutAttr>(layoutMaskName);
+ xegpu::DistributeLayoutAttr layoutPayload =
+ storeScatterOp->getAttrOfType<xegpu::DistributeLayoutAttr>(
+ layoutPayloadName);
+ xegpu::DistributeLayoutAttr layoutOffsets =
+ storeScatterOp->getAttrOfType<xegpu::DistributeLayoutAttr>(
+ layoutOffsetsName);
+ xegpu::DistributeLayoutAttr layoutMask =
+ storeScatterOp->getAttrOfType<xegpu::DistributeLayoutAttr>(
+ layoutMaskName);
FailureOr<VectorType> distStoreVecByWarpOpOrFailure =
getDistVecTypeBasedOnLaneLayout(layoutPayload, storeVecTy);
diff --git a/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir b/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir
index 754825193a10f..6d69b351bb5d9 100644
--- a/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir
+++ b/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir
@@ -2,84 +2,39 @@
// RUN: --xegpu-optimize-block-loads --split-input-file %s | FileCheck %s
// CHECK-LABEL: gpu.func @vector_reduce_2d(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<4x16xf32>) {
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<4x16xf32>, %[[ARG2:[0-9a-zA-Z]+]]: memref<256xf32>) {
// CHECK: %[[ACC:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} 1.000000e+00 : f32
// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<4x16xf32> -> !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>>
// CHECK: %[[LOADED:.*]] = xegpu.load_nd %[[TDESC]][0, 0] : !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>> -> vector<4x16xf32>
-// CHECK: %[[ACC_VEC:.*]] = vector.from_elements %[[ACC]] : vector<1xf32>
+// CHECK: %[[ACC_VEC:.*]] = vector.from_elements %[[ACC]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} : vector<1xf32>
// CHECK: %[[ACC_VEC_FOR_INTRA:.*]] = vector.broadcast %[[ACC_VEC]]
// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>} : vector<1xf32> to vector<16xf32>
// CHECK: %[[LOADED_REDUCED:.*]] = vector.multi_reduction <add>, %[[LOADED]], %[[ACC_VEC_FOR_INTRA]]
// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>} [0] : vector<4x16xf32> to vector<16xf32>
// CHECK: %[[LOADED_REDUCED_FOR_CROSS:.*]] = vector.shape_cast %[[LOADED_REDUCED]]
-// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} : vector<16xf32> to vector<1x16xf32>
+// CHECK-SAME: {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16xf32> to vector<1x16xf32>
// CHECK: %[[LOADED_REDUCED_2D:.*]] = vector.multi_reduction <add>, %[[LOADED_REDUCED_FOR_CROSS]], %[[ACC_VEC]]
-// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} [1] : vector<1x16xf32> to vector<1xf32>
+// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>} [1] : vector<1x16xf32> to vector<1xf32>
// CHECK: %[[SCALAR_RES:.*]] = vector.extract %[[LOADED_REDUCED_2D]][0] : f32 from vector<1xf32>
gpu.module @xevm_test {
- gpu.func @vector_reduce_2d(%src: memref<4x16xf32>) {
+ gpu.func @vector_reduce_2d(%src: memref<4x16xf32>, %dst: memref<256xf32>) {
%cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} 1.0 : f32
%tdesc = xegpu.create_nd_tdesc %src : memref<4x16xf32>
-> !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>>
%load = xegpu.load_nd %tdesc[0, 0]
: !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>>
-> vector<4x16xf32>
- %reduce = vector.multi_reduction <add>, %load, %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} [0, 1]
- : vector<4x16xf32> to f32
- gpu.return
- }
-}
+ %reduce = vector.multi_reduction <add>, %load, %cst
+ {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>}
+ [0, 1] : vector<4x16xf32> to f32
+ %reduce_bcast = vector.broadcast %reduce
+ {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>}
+ : f32 to vector<16xf32>
-// -----
-// CHECK-LABEL: gpu.func @vector_reduce_2d(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<4x16xf32>) {
-// CHECK: %[[ACC:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>, dims = [0, 1]>} dense<1.000000e+00> : vector<1xf32>
-// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<4x16xf32> -> !xegpu.tensor_desc<4x16xf32>
-// CHECK: %[[LOADED:.*]] = xegpu.load_nd %[[TDESC]][0, 0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>} : !xegpu.tensor_desc<4x16xf32> -> vector<4x16xf32>
-// CHECK: %[[LOADED_LEADING_UNIT:.*]] = vector.shape_cast %[[LOADED]]
-// CHECK-SAME: {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>} : vector<4x16xf32> to vector<1x4x16xf32
-// CHECK: %[[ACC_VEC_FOR_INTRA:.*]] = vector.broadcast %[[ACC]]
-// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>, dims = [1]>} : vector<1xf32> to vector<1x16xf32>
-// CHECK: %[[LOADED_REDUCED:.*]] = vector.multi_reduction <add>, %[[LOADED_LEADING_UNIT]], %[[ACC_VEC_FOR_INTRA]]
-// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>, dims = [1]>} [1] : vector<1x4x16xf32> to vector<1x16xf32>
-// CHECK: %[[LOADED_REDUCED_2D:.*]] = vector.multi_reduction <add>, %[[LOADED_REDUCED]], %[[ACC]]
-// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>, dims = [1, 2]>} [1] : vector<1x16xf32> to vector<1xf32>
-gpu.module @xevm_test {
- gpu.func @vector_reduce_2d(%src: memref<4x16xf32>) {
- %cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>, dims = [0, 1]>} dense<1.0> : vector<1xf32>
- %tdesc = xegpu.create_nd_tdesc %src : memref<4x16xf32>
- -> !xegpu.tensor_desc<4x16xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>} : !xegpu.tensor_desc<4x16xf32> -> vector<4x16xf32>
- %load_with_dim = vector.shape_cast %load {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 4, 1]>} : vector<4x16xf32> to vector<1x4x16xf32>
- %reduce = vector.multi_reduction <add>, %load_with_dim, %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1,1, 16], lane_data = [1, 4, 1]>, dims = [1, 2]>} [1, 2]
- : vector<1x4x16xf32> to vector<1xf32>
- gpu.return
- }
-}
+ %offset = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<0> : vector<16xindex>
+ %mask = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<1> : vector<16xi1>
-// -----
-// CHECK-LABEL: gpu.func @vector_reduce_2d(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<4x64xf32>) {
-// CHECK: %[[ACC:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>, dims = [0, 1]>} dense<1.000000e+00> : vector<1xf32>
-// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<4x64xf32> -> !xegpu.tensor_desc<1x64xf32>
-// CHECK: %[[LOADED:.*]] = xegpu.load_nd %[[TDESC]][0, 0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>} : !xegpu.tensor_desc<1x64xf32> -> vector<1x64xf32>
-// CHECK: %[[LOADED_LEADING_UNIT:.*]] = vector.shape_cast %[[LOADED]]
-// CHECK-SAME: {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>} : vector<1x64xf32> to vector<1x1x64xf32>
-// CHECK: %[[ACC_VEC_FOR_INTRA:.*]] = vector.broadcast %[[ACC]]
-// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>, dims = [1]>} : vector<1xf32> to vector<1x64xf32>
-// CHECK: %[[LOADED_REDUCED:.*]] = vector.multi_reduction <add>, %[[LOADED_LEADING_UNIT]], %[[ACC_VEC_FOR_INTRA]]
-// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>, dims = [1]>} [1] : vector<1x1x64xf32> to vector<1x64xf32>
-// CHECK: %[[LOADED_REDUCED_2D:.*]] = vector.multi_reduction <add>, %[[LOADED_REDUCED]], %[[ACC]]
-// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>, dims = [1, 2]>} [1] : vector<1x64xf32> to vector<1xf32>
-gpu.module @xevm_test {
- gpu.func @vector_reduce_2d(%src: memref<4x64xf32>) {
- %cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>, dims = [0, 1]>} dense<1.0> : vector<1xf32>
- %tdesc = xegpu.create_nd_tdesc %src : memref<4x64xf32>
- -> !xegpu.tensor_desc<1x64xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>} : !xegpu.tensor_desc<1x64xf32> -> vector<1x64xf32>
- %load_with_dim = vector.shape_cast %load {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>} : vector<1x64xf32> to vector<1x1x64xf32>
- %reduce = vector.multi_reduction <add>, %load_with_dim, %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>, dims = [1, 2]>} [1, 2]
- : vector<1x1x64xf32> to vector<1xf32>
+ xegpu.store %reduce_bcast, %dst[%offset], %mask {layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>} : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
gpu.return
}
}
diff --git a/mlir/test/Dialect/XeGPU/subgroup-distribute.mlir b/mlir/test/Dialect/XeGPU/subgroup-distribute.mlir
index e5e3d2a1c1ad5..dadefa74659e3 100644
--- a/mlir/test/Dialect/XeGPU/subgroup-distribute.mlir
+++ b/mlir/test/Dialect/XeGPU/subgroup-distribute.mlir
@@ -357,12 +357,12 @@ gpu.module @xevm_module{
%c0 = arith.constant 0 : index
%mask = vector.constant_mask [16] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>}: vector<16xi1>
%1 = xegpu.load %arg0[%c0], %mask {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>}: memref<16xf16>, index, vector<16xi1> -> vector<16xf16>
-
+
%11 = vector.shape_cast %1 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16xf16> to vector<16x1xf16>
%2 = vector.broadcast %11 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x1xf16> to vector<16x16xf16>
// CHECK-NOT: vector.broadcast
// CHECK-NOT: vector.shape_cast
-
+
%tdesc1 = xegpu.create_nd_tdesc %arg1 : memref<16x16xf16>
-> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd {{.*}}, {{.*}}[{{.*}}, {{.*}}]
@@ -390,4 +390,47 @@ gpu.module @xevm_module{
}
}
-
+// -----
+gpu.module @xevm_test {
+ // CHECK-LABEL: gpu.func @vector_reduce_2d
+ // CHECK-DAG: %[[CST:.*]] = arith.constant dense<0> : vector<1xindex>
+ // CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<true> : vector<1xi1>
+ // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : i32
+ // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : i32
+ // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : i32
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32
+ // CHECK-DAG: %[[C16:.*]] = arith.constant 16 : i32
+ // CHECK-DAG: %[[CST_1:.*]] = arith.constant 1.000000e+00 : f32
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %arg0 : memref<4x16xf32> -> !xegpu.tensor_desc<4x16xf32>
+ // CHECK: %[[LOADED:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{packed}> : !xegpu.tensor_desc<4x16xf32> -> vector<4xf32>
+ // CHECK: %[[LOADED_REDUCED:.*]] = vector.reduction <add>, %[[LOADED]], %[[CST_1]] : vector<4xf32> into f32
+ // CHECK: %[[SHUFFLE_0:.*]], %{{.*}} = gpu.shuffle xor %[[LOADED_REDUCED]], %[[C1]], %[[C16]] : f32
+ // CHECK: %[[VEC_RED_0:.*]] = arith.addf %[[LOADED_REDUCED]], %[[SHUFFLE_0]] : f32
+ // CHECK: %[[SHUFFLE_1:.*]], %{{.*}} = gpu.shuffle xor %[[VEC_RED_0]], %[[C2]], %[[C16]] : f32
+ // CHECK: %[[VEC_RED_1:.*]] = arith.addf %[[VEC_RED_0]], %[[SHUFFLE_1]] : f32
+ // CHECK: %[[SHUFFLE_2:.*]], %{{.*}} = gpu.shuffle xor %[[VEC_RED_1]], %[[C4]], %[[C16]] : f32
+ // CHECK: %[[VEC_RED_2:.*]] = arith.addf %[[VEC_RED_1]], %[[SHUFFLE_2]] : f32
+ // CHECK: %[[SHUFFLE_3:.*]], %{{.*}} = gpu.shuffle xor %[[VEC_RED_2]], %[[C8]], %[[C16]] : f32
+ // CHECK: %[[VEC_RED_3:.*]] = arith.addf %[[VEC_RED_2]], %[[SHUFFLE_3]] : f32
+ // CHECK: %[[VEC_RED_4:.*]] = arith.addf %[[VEC_RED_3]], %[[CST_1]] : f32
+ // CHECK: %[[VEC_RED:.*]] = vector.broadcast %[[VEC_RED_4]] : f32 to vector<1xf32>
+ // CHECK: xegpu.store %[[VEC_RED]], %arg1[%[[CST]]], %[[CST_0]]
+ // CHECK-SAME: <{layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>}> : vector<1xf32>, memref<256xf32>, vector<1xindex>, vector<1xi1>
+ gpu.func @vector_reduce_2d(%arg0: memref<4x16xf32>, %arg1: memref<256xf32>) {
+ %cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} 1.000000e+00 : f32
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<4x16xf32> -> !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>>
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>> -> vector<4x16xf32>
+ %2 = vector.from_elements %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} : vector<1xf32>
+ %3 = vector.broadcast %2 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>} : vector<1xf32> to vector<16xf32>
+ %4 = vector.multi_reduction <add>, %1, %3 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>} [0] : vector<4x16xf32> to vector<16xf32>
+ %5 = vector.shape_cast %4 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16xf32> to vector<1x16xf32>
+ %6 = vector.multi_reduction <add>, %5, %2 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>} [1] : vector<1x16xf32> to vector<1xf32>
+ %7 = vector.extract %6[0] : f32 from vector<1xf32>
+ %8 = vector.broadcast %7 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>} : f32 to vector<16xf32>
+ %cst_0 = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<0> : vector<16xindex>
+ %cst_1 = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
+ xegpu.store %8, %arg1[%cst_0], %cst_1 <{layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>}> : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
+ gpu.return
+ }
+}
>From aaed7d3d5bc8f0dfff16c14068663e8435924dc1 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Mon, 15 Dec 2025 17:00:44 +0000
Subject: [PATCH 3/3] Simplify 2d optimization
---
.../Transforms/XeGPUOptimizeBlockLoads.cpp | 54 +++----------------
.../Dialect/XeGPU/optimize-2d-reduction.mlir | 13 ++---
2 files changed, 10 insertions(+), 57 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
index 3c6513ff34fed..4a396909ff539 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
@@ -447,14 +447,6 @@ class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
// TODO: remove scalar acc assumption (need more complex layout adjustments
// for sliced inputs).
assert(scalarAcc && "Expected scalar acc");
- if (scalarAcc) {
- acc = vector::FromElementsOp::create(
- rewriter, loc, VectorType::get({1}, sourceVecType.getElementType()),
- acc);
- xegpu::setDistributeLayoutAttr(
- llvm::dyn_cast<OpResult>(acc),
- cast<xegpu::DistributeLayoutAttr>(resLayout));
- }
// The first reduction's dist attribute does not have the cross lane dim.
auto resSliceLayoutAttr = cast<xegpu::SliceAttr>(resLayout);
@@ -467,7 +459,7 @@ class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
reductionOp.getContext(), resSliceLayoutAttr.getParent(),
DenseI64ArrayAttr::get(getContext(), sliceDims));
- // We reduce only one dim first, adjsut accumulator.
+ // We reduce intra-lane, acc is source without intra lane.
SmallVector<int64_t> accShape(sourceVecType.getShape());
accShape.erase(accShape.begin() + intraLaneDim);
// Add a dim to the lower-dim user-supplied acc
@@ -487,45 +479,11 @@ class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
llvm::dyn_cast<OpResult>(intraLaneReduced),
cast<xegpu::DistributeLayoutAttr>(intraLaneRedResLayout));
- xegpu::DistributeLayoutAttr nextMultiRedLayout = intraLaneRedResLayout;
- // Example: vector<2x4> got reduced to vector<2>, next reduction returns a
- // scalar, distribution passes do not support this result type. Expand to
- // vector<2x1>, so that the second reduction result is vector<1>. Restore
- // this dim in layout, but lane data is 1.
- if (scalarAcc) {
- SmallVector<int> srcLaneData(srcLayout.getRank(), 1);
- auto laneLayoutSrc = srcLayout.getEffectiveLaneLayoutAsInt();
- SmallVector<int> srcLaneLayout(laneLayoutSrc.begin(),
- laneLayoutSrc.end());
- nextMultiRedLayout = xegpu::LayoutAttr::get(
- reductionOp.getContext(),
- DenseI32ArrayAttr::get(reductionOp.getContext(), srcLaneLayout),
- DenseI32ArrayAttr::get(reductionOp.getContext(), srcLaneData));
-
- SmallVector<int64_t> vecTypeWithUnitDim{sourceVecType.getShape()};
- vecTypeWithUnitDim[intraLaneDim] = 1;
- intraLaneReduced = vector::ShapeCastOp::create(
- rewriter, loc,
- VectorType::get(vecTypeWithUnitDim, sourceVecType.getElementType()),
- intraLaneReduced);
- xegpu::setDistributeLayoutAttr(llvm::dyn_cast<OpResult>(intraLaneReduced),
- nextMultiRedLayout);
- } else
- crossLaneDim -= static_cast<int64_t>(intraLaneDim < crossLaneDim);
- // Do cross-lane reduction
- // TODO: why use accumulator again?
- Value crossLaneReduced = vector::MultiDimReductionOp::create(
- rewriter, loc, reductionOp.getKind(), intraLaneReduced, acc,
- ArrayRef<int64_t>(crossLaneDim));
- auto crossLaneLayout = xegpu::SliceAttr::get(
- reductionOp.getContext(), nextMultiRedLayout,
- DenseI64ArrayAttr::get(getContext(), crossLaneDim));
- xegpu::setDistributeLayoutAttr(llvm::dyn_cast<OpResult>(crossLaneReduced),
- crossLaneLayout);
-
- if (scalarAcc)
- crossLaneReduced =
- vector::ExtractOp::create(rewriter, loc, crossLaneReduced, 0);
+ Value crossLaneReduced = vector::ReductionOp::create(
+ rewriter, loc, reductionOp.getKind(), intraLaneReduced, nullptr);
+ xegpu::setDistributeLayoutAttr(
+ llvm::dyn_cast<OpResult>(crossLaneReduced),
+ cast<xegpu::DistributeLayoutAttr>(resLayout));
assert(crossLaneReduced.getType() == reductionOp.getResult().getType() &&
"Type mismatch");
rewriter.replaceOp(reductionOp, crossLaneReduced);
diff --git a/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir b/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir
index 6d69b351bb5d9..dcd9ff1154dcc 100644
--- a/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir
+++ b/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir
@@ -6,16 +6,11 @@
// CHECK: %[[ACC:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} 1.000000e+00 : f32
// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<4x16xf32> -> !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>>
// CHECK: %[[LOADED:.*]] = xegpu.load_nd %[[TDESC]][0, 0] : !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>> -> vector<4x16xf32>
-// CHECK: %[[ACC_VEC:.*]] = vector.from_elements %[[ACC]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} : vector<1xf32>
-// CHECK: %[[ACC_VEC_FOR_INTRA:.*]] = vector.broadcast %[[ACC_VEC]]
-// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>} : vector<1xf32> to vector<16xf32>
-// CHECK: %[[LOADED_REDUCED:.*]] = vector.multi_reduction <add>, %[[LOADED]], %[[ACC_VEC_FOR_INTRA]]
+// CHECK: %[[ACC_VEC:.*]] = vector.broadcast %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>} : f32 to vector<16xf32>
+// CHECK: %[[LOADED_REDUCED:.*]] = vector.multi_reduction <add>, %[[LOADED]], %[[ACC_VEC]]
// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>} [0] : vector<4x16xf32> to vector<16xf32>
-// CHECK: %[[LOADED_REDUCED_FOR_CROSS:.*]] = vector.shape_cast %[[LOADED_REDUCED]]
-// CHECK-SAME: {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16xf32> to vector<1x16xf32>
-// CHECK: %[[LOADED_REDUCED_2D:.*]] = vector.multi_reduction <add>, %[[LOADED_REDUCED_FOR_CROSS]], %[[ACC_VEC]]
-// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>} [1] : vector<1x16xf32> to vector<1xf32>
-// CHECK: %[[SCALAR_RES:.*]] = vector.extract %[[LOADED_REDUCED_2D]][0] : f32 from vector<1xf32>
+// CHECK: %[[LOADED_REDUCED_FOR_CROSS:.*]] = vector.reduction <add>, %[[LOADED_REDUCED]]
+// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} : vector<16xf32> into f32
gpu.module @xevm_test {
gpu.func @vector_reduce_2d(%src: memref<4x16xf32>, %dst: memref<256xf32>) {
%cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0, 1]>} 1.0 : f32
More information about the Mlir-commits
mailing list