[Mlir-commits] [mlir] [MLIR][XeGPU] Add 2D `vector.multi_reduction` optimization (PR #171154)
Artem Kroviakov
llvmlistbot at llvm.org
Thu Dec 18 02:53:15 PST 2025
https://github.com/akroviakov updated https://github.com/llvm/llvm-project/pull/171154
>From 98ad0818366dca893ef4b096e17934386338cc9a Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Fri, 5 Dec 2025 16:08:04 +0000
Subject: [PATCH 1/9] [MLIR][XeGPU] Add sg layout propagation
---
.../mlir/Dialect/XeGPU/Transforms/Passes.td | 6 +-
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 59 +++++++++++++++++--
.../Dialect/XeGPU/propagate-layout-sg.mlir | 53 +++++++++++++++++
3 files changed, 112 insertions(+), 6 deletions(-)
create mode 100644 mlir/test/Dialect/XeGPU/propagate-layout-sg.mlir
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
index 0ca58426ecfcb..c682e6fdad1df 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
@@ -50,6 +50,10 @@ def XeGPUPropagateLayout : Pass<"xegpu-propagate-layout"> {
- `lane`
Propagate the `lane_layout` and `lane_data` fields of the layout attribute.
Default values are selected to align with hardware.
+
+ - `sg`
+ Propagate the `sg_layout` and `sg_data` fields of the layout attribute.
+ Default values are selected to align with hardware.
}];
let dependentDialects = ["memref::MemRefDialect", "xegpu::XeGPUDialect",
"vector::VectorDialect"];
@@ -60,7 +64,7 @@ def XeGPUPropagateLayout : Pass<"xegpu-propagate-layout"> {
Option<
"layoutKind", "layout-kind", "std::string",
/*default=*/"\"lane\"",
- "Propagate `inst` / `lane` level of xegpu layouts.">
+ "Propagate `sg` / `inst` / `lane` level of xegpu layouts.">
];
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index dc9eb96c169b4..a3d057981024d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -53,7 +53,7 @@ using namespace mlir::dataflow;
namespace {
-enum class LayoutKind { Lane, InstData };
+enum class LayoutKind { Lane, InstData, Subgroup };
//===----------------------------------------------------------------------===//
// LayoutInfo
@@ -109,6 +109,12 @@ struct LayoutInfo {
SmallVector<int> getInstData() const;
+ SmallVector<int> getSgLayout() const;
+
+ SmallVector<int> getSgData() const;
+
+ SmallVector<int> getOrder() const;
+
bool isSliceLayout() const {
if (!isAssigned())
return false;
@@ -127,8 +133,6 @@ struct LayoutInfo {
SmallVector<int> LayoutInfo::getLaneLayout() const {
if (!isAssigned())
return {};
- assert(storage.getEffectiveLaneLayoutAsInt().size() &&
- "Expected lane layout to be assigned");
return llvm::map_to_vector(storage.getEffectiveLaneLayoutAsInt(),
[](int64_t val) { return static_cast<int>(val); });
}
@@ -136,8 +140,6 @@ SmallVector<int> LayoutInfo::getLaneLayout() const {
SmallVector<int> LayoutInfo::getLaneData() const {
if (!isAssigned())
return {};
- assert(storage.getEffectiveLaneDataAsInt().size() &&
- "Expected lane data to be assigned");
return llvm::map_to_vector(storage.getEffectiveLaneDataAsInt(),
[](int64_t val) { return static_cast<int>(val); });
}
@@ -149,6 +151,27 @@ SmallVector<int> LayoutInfo::getInstData() const {
[](int64_t val) { return static_cast<int>(val); });
}
+SmallVector<int> LayoutInfo::getSgLayout() const {
+ if (!isAssigned())
+ return {};
+ return llvm::map_to_vector(storage.getEffectiveSgLayoutAsInt(),
+ [](int64_t val) { return static_cast<int>(val); });
+}
+
+SmallVector<int> LayoutInfo::getSgData() const {
+ if (!isAssigned())
+ return {};
+ return llvm::map_to_vector(storage.getEffectiveSgDataAsInt(),
+ [](int64_t val) { return static_cast<int>(val); });
+}
+
+SmallVector<int> LayoutInfo::getOrder() const {
+ if (!isAssigned() || !storage.getOrder())
+ return {};
+ return llvm::map_to_vector(storage.getOrder().asArrayRef(),
+ [](int64_t val) { return static_cast<int>(val); });
+}
+
void LayoutInfo::print(raw_ostream &os) const {
if (isAssigned()) {
os << storage;
@@ -188,6 +211,10 @@ LayoutInfo LayoutInfo::transpose(ArrayRef<int64_t> permutation) const {
SmallVector<int32_t> laneLayout;
SmallVector<int32_t> laneData;
SmallVector<int32_t> instData;
+ SmallVector<int32_t> sgLayout;
+ SmallVector<int32_t> sgData;
+ SmallVector<int32_t> order;
+
for (int64_t idx : permutation) {
if (getLaneLayout().size()) {
laneLayout.push_back(static_cast<int32_t>(getLaneLayout()[idx]));
@@ -195,13 +222,30 @@ LayoutInfo LayoutInfo::transpose(ArrayRef<int64_t> permutation) const {
}
if (getInstData().size())
instData.push_back(static_cast<int32_t>(getInstData()[idx]));
+ if (getSgData().size()) {
+ sgLayout.push_back(static_cast<int32_t>(getSgLayout()[idx]));
+ sgData.push_back(static_cast<int32_t>(getSgData()[idx]));
+ }
+ if (getOrder().size()) {
+ order.push_back(static_cast<int32_t>(getOrder()[idx]));
+ }
}
+ auto orderAttr = order.size()
+ ? DenseI32ArrayAttr::get(storage.getContext(), order)
+ : nullptr;
xegpu::LayoutAttr layoutAttr;
if (getLaneLayout().size())
layoutAttr =
xegpu::LayoutAttr::get(storage.getContext(), laneLayout, laneData);
if (getInstData().size())
layoutAttr = xegpu::LayoutAttr::get(storage.getContext(), instData);
+ if (getSgData().size())
+ layoutAttr = xegpu::LayoutAttr::get(
+ storage.getContext(),
+ DenseI32ArrayAttr::get(storage.getContext(), sgLayout),
+ DenseI32ArrayAttr::get(storage.getContext(), sgData),
+ /*inst_data =*/nullptr, /*lane_layout =*/nullptr,
+ /*lane_data =*/nullptr, orderAttr);
return LayoutInfo(layoutAttr);
}
@@ -487,6 +531,9 @@ bool LayoutInfoPropagation::hasParamsOfLayoutKind(
} else if (layoutKind == LayoutKind::Lane) {
return !(anchorLayout.getEffectiveLaneLayoutAsInt().empty() ||
anchorLayout.getEffectiveLaneDataAsInt().empty());
+ } else if (layoutKind == LayoutKind::Subgroup) {
+ return !(anchorLayout.getEffectiveSgLayoutAsInt().empty() ||
+ anchorLayout.getEffectiveSgDataAsInt().empty());
}
return false;
}
@@ -1311,6 +1358,8 @@ void XeGPUPropagateLayoutPass::runOnOperation() {
layoutKind = LayoutKind::Lane;
} else if (this->layoutKind == "inst") {
layoutKind = LayoutKind::InstData;
+ } else if (this->layoutKind == "sg") {
+ layoutKind = LayoutKind::Subgroup;
} else {
getOperation()->emitError("Unsupported layout kind option: " +
this->layoutKind);
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-sg.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-sg.mlir
new file mode 100644
index 0000000000000..5659e9995b22a
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-sg.mlir
@@ -0,0 +1,53 @@
+// RUN: mlir-opt -xevm-attach-target='chip=pvc' -xegpu-propagate-layout="layout-kind=sg" -split-input-file %s | FileCheck %s
+
+gpu.module @test {
+ // CHECK-LABEL: store_nd
+ // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>
+ func.func @store_nd(%src: memref<256x128xf32>) {
+ // CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG_0]] : memref<256x128xf32>
+ // CHECK-SAME: -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>>
+ // CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}>
+ // CHECK-SAME: {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}
+ // CHECK-SAME: : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>>
+ // CHECK-SAME: -> vector<256x128xf32>
+ // CHECK: xegpu.store_nd %[[LOAD]], %[[TDESC]] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}>
+ // CHECK-SAME: : vector<256x128xf32>, !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>>
+ %tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32> -> !xegpu.tensor_desc<256x128xf32>
+ %load = xegpu.load_nd %tdesc : !xegpu.tensor_desc<256x128xf32> -> vector<256x128xf32>
+ xegpu.store_nd %load, %tdesc {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}
+ : vector<256x128xf32>, !xegpu.tensor_desc<256x128xf32>
+ return
+ }
+}
+
+// -----
+
+gpu.module @test {
+ // CHECK-LABEL: vector_transpose
+ // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>
+ // CHECK-SAME: %[[ARG_1:.*]]: memref<128x256xf32>
+ func.func @vector_transpose(%src: memref<256x128xf32>, %src1: memref<128x256xf32>) {
+ // CHECK: %[[TDESC_LD:.*]] = xegpu.create_nd_tdesc %[[ARG_0]] : memref<256x128xf32> ->
+ // CHECK-SAME: !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [4, 8], sg_data = [64, 32], order = [0, 1]>>
+ // CHECK: %[[TDESC_ST:.*]] = xegpu.create_nd_tdesc %[[ARG_1]] : memref<128x256xf32> ->
+ // CHECK-SAME: !xegpu.tensor_desc<128x256xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], order = [1, 0]>>
+
+ // CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC_LD]][0, 0] <{layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [64, 32], order = [0, 1]>}>
+ // CHECK-SAME: {layout_result_0 = #xegpu.layout<sg_layout = [4, 8], sg_data = [64, 32], order = [0, 1]>} :
+ // CHECK-SAME: !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [4, 8], sg_data = [64, 32], order = [0, 1]>> -> vector<256x128xf32>
+
+ // CHECK: %[[TRANSPOSED:.*]] = vector.transpose %2, [1, 0]
+ // CHECK-SAME {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], order = [1, 0]>} : vector<256x128xf32> to vector<128x256xf32>
+
+ // CHECK: xegpu.store_nd %[[TRANSPOSED]], %[[TDESC_ST]][0, 0]
+ // CHECK-SAME: <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], order = [1, 0]>}> : vector<128x256xf32>,
+ // CHECK-SAME: !xegpu.tensor_desc<128x256xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], order = [1, 0]>>
+ %tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32> -> !xegpu.tensor_desc<256x128xf32>
+ %tdesc1 = xegpu.create_nd_tdesc %src1 : memref<128x256xf32> -> !xegpu.tensor_desc<128x256xf32>
+ %load = xegpu.load_nd %tdesc[0, 0] : !xegpu.tensor_desc<256x128xf32> -> vector<256x128xf32>
+ %trans = vector.transpose %load, [1, 0] : vector<256x128xf32> to vector<128x256xf32>
+ xegpu.store_nd %trans, %tdesc1[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], order = [1, 0]>}
+ : vector<128x256xf32>, !xegpu.tensor_desc<128x256xf32>
+ return
+ }
+}
>From 353c945276f683a2c3f63d64dd1373f8bbc5004f Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Tue, 16 Dec 2025 16:15:40 +0000
Subject: [PATCH 2/9] Rename subgroup prop option
---
mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td | 4 ++--
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp | 2 +-
...pagate-layout-sg.mlir => propagate-layout-subgroup.mlir} | 6 +++---
3 files changed, 6 insertions(+), 6 deletions(-)
rename mlir/test/Dialect/XeGPU/{propagate-layout-sg.mlir => propagate-layout-subgroup.mlir} (92%)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
index c682e6fdad1df..3ff7805263f0e 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
@@ -51,7 +51,7 @@ def XeGPUPropagateLayout : Pass<"xegpu-propagate-layout"> {
Propagate the `lane_layout` and `lane_data` fields of the layout attribute.
Default values are selected to align with hardware.
- - `sg`
+ - `subgroup`
Propagate the `sg_layout` and `sg_data` fields of the layout attribute.
Default values are selected to align with hardware.
}];
@@ -64,7 +64,7 @@ def XeGPUPropagateLayout : Pass<"xegpu-propagate-layout"> {
Option<
"layoutKind", "layout-kind", "std::string",
/*default=*/"\"lane\"",
- "Propagate `sg` / `inst` / `lane` level of xegpu layouts.">
+ "Propagate `subgroup` / `inst` / `lane` level of xegpu layouts.">
];
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index a3d057981024d..cbd91154ce0aa 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1358,7 +1358,7 @@ void XeGPUPropagateLayoutPass::runOnOperation() {
layoutKind = LayoutKind::Lane;
} else if (this->layoutKind == "inst") {
layoutKind = LayoutKind::InstData;
- } else if (this->layoutKind == "sg") {
+ } else if (this->layoutKind == "subgroup") {
layoutKind = LayoutKind::Subgroup;
} else {
getOperation()->emitError("Unsupported layout kind option: " +
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-sg.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
similarity index 92%
rename from mlir/test/Dialect/XeGPU/propagate-layout-sg.mlir
rename to mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index 5659e9995b22a..c7dfc9fb7b1f1 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt -xevm-attach-target='chip=pvc' -xegpu-propagate-layout="layout-kind=sg" -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -xevm-attach-target='chip=pvc' -xegpu-propagate-layout="layout-kind=subgroup" -split-input-file %s | FileCheck %s
gpu.module @test {
// CHECK-LABEL: store_nd
@@ -14,7 +14,7 @@ gpu.module @test {
// CHECK-SAME: : vector<256x128xf32>, !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>>
%tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32> -> !xegpu.tensor_desc<256x128xf32>
%load = xegpu.load_nd %tdesc : !xegpu.tensor_desc<256x128xf32> -> vector<256x128xf32>
- xegpu.store_nd %load, %tdesc {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}
+ xegpu.store_nd %load, %tdesc <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}>
: vector<256x128xf32>, !xegpu.tensor_desc<256x128xf32>
return
}
@@ -46,7 +46,7 @@ gpu.module @test {
%tdesc1 = xegpu.create_nd_tdesc %src1 : memref<128x256xf32> -> !xegpu.tensor_desc<128x256xf32>
%load = xegpu.load_nd %tdesc[0, 0] : !xegpu.tensor_desc<256x128xf32> -> vector<256x128xf32>
%trans = vector.transpose %load, [1, 0] : vector<256x128xf32> to vector<128x256xf32>
- xegpu.store_nd %trans, %tdesc1[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], order = [1, 0]>}
+ xegpu.store_nd %trans, %tdesc1[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64], order = [1, 0]>}>
: vector<128x256xf32>, !xegpu.tensor_desc<128x256xf32>
return
}
>From da550e877c7ba619b8a790cdd7cfb22413d7f810 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 3/9] [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 4eccaef4edca5c828e30326d5a42a551a4c75266 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 4/9] 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 dbd5a50489173..18c15596b33be 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 b40325ac863f004e4b042dead7ce8d7df2f8d3fe 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 5/9] 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
>From c797046604d07f55b8f0dcefa73242bf1a51f14b Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Tue, 16 Dec 2025 10:25:41 +0000
Subject: [PATCH 6/9] Add peephole to master pass
---
mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp | 1 +
1 file changed, 1 insertion(+)
diff --git a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
index 38313dc3c01d5..269acc4382eb4 100644
--- a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
+++ b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
@@ -77,6 +77,7 @@ void buildGPUPassPipeline(OpPassManager &pm,
layoutOptions.layoutKind = "lane";
pm.addNestedPass<gpu::GPUModuleOp>(
xegpu::createXeGPUPropagateLayout(layoutOptions));
+ pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUOptimizeBlockLoads());
pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUSubgroupDistribute());
pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());
pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
>From 4763a688dd824ee06eb99f74e12216dfd20e122e Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Tue, 16 Dec 2025 14:46:43 +0000
Subject: [PATCH 7/9] Cleanup
---
.../Transforms/XeGPUOptimizeBlockLoads.cpp | 33 ++++++-------------
.../Dialect/XeGPU/optimize-2d-reduction.mlir | 24 +++++++-------
.../Dialect/XeGPU/subgroup-distribute.mlir | 30 ++++++++---------
3 files changed, 35 insertions(+), 52 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
index 4a396909ff539..bc2703e7ffbda 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
@@ -421,15 +421,12 @@ class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
LogicalResult
matchAndRewrite(vector::MultiDimReductionOp reductionOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
- if (reductionOp.getReductionDims().size() != 2)
- return rewriter.notifyMatchFailure(reductionOp,
- "Expected 2D multi reduction");
- // Retrieve layouts.
+ auto sourceVecType = reductionOp.getSourceVectorType();
+ if (reductionOp.getReductionDims().size() != 2 ||
+ sourceVecType.getRank() != 2)
+ return rewriter.notifyMatchFailure(
+ reductionOp, "Expected 2D multi reduction of a 2D source");
auto resLayout = xegpu::getDistributeLayoutAttr(reductionOp.getResult());
- auto srcLayout = xegpu::getDistributeLayoutAttr(reductionOp.getSource());
- assert(isa<xegpu::LayoutAttr>(srcLayout) &&
- "Currently we do not support sliced inputs");
-
// Retrieve and order dims for 1D decomposition (prefer intra-lane first).
auto dims = llvm::to_vector(reductionOp.getReductionDims());
auto [intraLaneDim, crossLaneDim] = getReductionDimOrder(dims, resLayout);
@@ -439,14 +436,7 @@ class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
crossLaneDim = dims[1];
}
auto loc = reductionOp.getLoc();
- 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());
- // TODO: remove scalar acc assumption (need more complex layout adjustments
- // for sliced inputs).
- assert(scalarAcc && "Expected scalar acc");
// The first reduction's dist attribute does not have the cross lane dim.
auto resSliceLayoutAttr = cast<xegpu::SliceAttr>(resLayout);
@@ -459,22 +449,19 @@ class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
reductionOp.getContext(), resSliceLayoutAttr.getParent(),
DenseI64ArrayAttr::get(getContext(), sliceDims));
- // 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
- Value firstRedAcc = acc;
- if (firstRedAcc) {
- firstRedAcc = vector::BroadcastOp::create(
+ if (acc) {
+ acc = vector::BroadcastOp::create(
rewriter, loc,
VectorType::get(accShape, sourceVecType.getElementType()), acc);
xegpu::setDistributeLayoutAttr(
- llvm::dyn_cast<OpResult>(firstRedAcc),
+ llvm::dyn_cast<OpResult>(acc),
cast<xegpu::DistributeLayoutAttr>(intraLaneRedResLayout));
}
Value intraLaneReduced = vector::MultiDimReductionOp::create(
- rewriter, loc, reductionOp.getKind(), reductionOp.getSource(),
- firstRedAcc, ArrayRef<int64_t>(intraLaneDim));
+ rewriter, loc, reductionOp.getKind(), reductionOp.getSource(), acc,
+ ArrayRef<int64_t>(intraLaneDim));
xegpu::setDistributeLayoutAttr(
llvm::dyn_cast<OpResult>(intraLaneReduced),
cast<xegpu::DistributeLayoutAttr>(intraLaneRedResLayout));
diff --git a/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir b/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir
index dcd9ff1154dcc..d3a1c33ced0e9 100644
--- a/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir
+++ b/mlir/test/Dialect/XeGPU/optimize-2d-reduction.mlir
@@ -3,33 +3,33 @@
// CHECK-LABEL: gpu.func @vector_reduce_2d(
// 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.broadcast %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, dims = [0]>} : f32 to vector<16xf32>
+// CHECK: %[[ACC:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 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 = [1, 1]>>
+// CHECK: %[[LOADED:.*]] = xegpu.load_nd %[[TDESC]][0, 0] : !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<4x16xf32>
+// CHECK: %[[ACC_VEC:.*]] = vector.broadcast %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 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-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} [0] : vector<4x16xf32> to vector<16xf32>
// 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
+// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 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
+ %cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 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]>>
+ -> !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
%load = xegpu.load_nd %tdesc[0, 0]
- : !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>>
+ : !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 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]>}
+ {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 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]>}
+ {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>}
: f32 to vector<16xf32>
%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>
- 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>
+ xegpu.store %reduce_bcast, %dst[%offset], %mask {layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 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 dadefa74659e3..c8dc1b7a1c6d6 100644
--- a/mlir/test/Dialect/XeGPU/subgroup-distribute.mlir
+++ b/mlir/test/Dialect/XeGPU/subgroup-distribute.mlir
@@ -393,8 +393,6 @@ 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
@@ -402,8 +400,10 @@ gpu.module @xevm_test {
// 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-DAG: %[[CST_0:.*]] = arith.constant dense<true> : vector<1xi1>
+ // CHECK-DAG: %[[CST:.*]] = arith.constant dense<0> : vector<1xindex>
// 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:.*]] = xegpu.load_nd %[[TDESC]][0, 0] : !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
@@ -413,24 +413,20 @@ gpu.module @xevm_test {
// 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: %[[VEC_RED:.*]] = vector.broadcast %[[VEC_RED_3]] : 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>
+ // CHECK-SAME: <{layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 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 = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 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 = [1, 1]>>
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<4x16xf32>
+ %2 = vector.broadcast %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} : f32 to vector<16xf32>
+ %3 = vector.multi_reduction <add>, %1, %2 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} [0] : vector<4x16xf32> to vector<16xf32>
+ %4 = vector.reduction <add>, %3 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0, 1]>} : vector<16xf32> into f32
+ %5 = vector.broadcast %4 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 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>
+ xegpu.store %5, %arg1[%cst_0], %cst_1 <{layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>}> : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
gpu.return
}
}
>From 711bea1660bdba304cdc86996fa7438a332352cb Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Wed, 17 Dec 2025 16:24:53 +0000
Subject: [PATCH 8/9] Add propagation
---
mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp | 2 ++
1 file changed, 2 insertions(+)
diff --git a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
index 269acc4382eb4..54c7b05eb4f7a 100644
--- a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
+++ b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
@@ -78,6 +78,8 @@ void buildGPUPassPipeline(OpPassManager &pm,
pm.addNestedPass<gpu::GPUModuleOp>(
xegpu::createXeGPUPropagateLayout(layoutOptions));
pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUOptimizeBlockLoads());
+ pm.addNestedPass<gpu::GPUModuleOp>(
+ xegpu::createXeGPUPropagateLayout(layoutOptions));
pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUSubgroupDistribute());
pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());
pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
>From bbaf4f7b576421dd914540a57bee9c310dba7310 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Thu, 18 Dec 2025 10:50:03 +0000
Subject: [PATCH 9/9] Feedback
---
.../mlir/Dialect/XeGPU/IR/XeGPUAttrs.td | 21 +++++++++++----
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 20 ++++++++++++++
.../Transforms/XeGPUOptimizeBlockLoads.cpp | 26 +++++++------------
3 files changed, 46 insertions(+), 21 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
index eae0bd4e68a84..c47db4c4ba688 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
@@ -299,7 +299,12 @@ def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> {
InterfaceMethod</*desc=*/[{Check if this layout is identical to another layout.}],
/*retTy=*/"bool",
/*methodName=*/"isEqualTo",
- /*args=*/(ins "const xegpu::DistributeLayoutAttr&": $other)>
+ /*args=*/(ins "const xegpu::DistributeLayoutAttr&": $other)>,
+
+ InterfaceMethod</*desc=*/[{Drops specified dimensions from slice.}],
+ /*retTy=*/"xegpu::DistributeLayoutAttr",
+ /*methodName=*/"dropSliceDims",
+ /*args=*/(ins "ArrayRef<int64_t>": $dimsToDrop)>
];
}
@@ -520,9 +525,12 @@ def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout", [DistributeLayoutAttr]> {
/// Check if this is slice of some other layout.
bool isSliceOf(const xegpu::DistributeLayoutAttr &other) { return false; }
-
+
+ /// Drop the slice dims to get the original layout.
+ xegpu::DistributeLayoutAttr dropSliceDims(ArrayRef<int64_t> dimsToDrop) { return *this; }
+
/// Check if this is identical to some other layout.
- bool isEqualTo(const xegpu::DistributeLayoutAttr &other);
+ bool isEqualTo(const xegpu::DistributeLayoutAttr &other);
}];
@@ -698,9 +706,12 @@ def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [DistributeLayoutAttr]> {
/// Check if this is slice of some other layout.
bool isSliceOf(const xegpu::DistributeLayoutAttr &other);
-
+
+ /// Drop the slice dims to get the original layout.
+ xegpu::DistributeLayoutAttr dropSliceDims(ArrayRef<int64_t> dimsToDrop);
+
/// Check if this is identical to some other layout.
- bool isEqualTo(const xegpu::DistributeLayoutAttr &other);
+ bool isEqualTo(const xegpu::DistributeLayoutAttr &other);
}];
let assemblyFormat = "`<` qualified($parent) `,` `dims` `=` $dims `>`";
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 1a19ab5fd970b..c8a95f6d80fb3 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -590,6 +590,26 @@ bool SliceAttr::isSliceOf(const xegpu::DistributeLayoutAttr &other) {
[&](int64_t dim) { return thisDims.contains(dim); });
}
+xegpu::DistributeLayoutAttr
+SliceAttr::dropSliceDims(ArrayRef<int64_t> dimsToDrop) {
+ auto flattenedThis = flatten();
+ if (dimsToDrop.empty())
+ return flattenedThis;
+ SmallVector<int64_t> sliceDims{flattenedThis.getDims().asArrayRef()};
+ for (auto dim : dimsToDrop) {
+ auto foundIt = std::find(sliceDims.begin(), sliceDims.end(), dim);
+ assert(foundIt != sliceDims.end() &&
+ "Expected to find the specified reduction dim in slice dims");
+ sliceDims.erase(foundIt);
+ }
+
+ auto sliceWithoutDims = xegpu::SliceAttr::get(
+ this->getContext(), flattenedThis.getParent(),
+ DenseI64ArrayAttr::get(this->getContext(), sliceDims));
+
+ return sliceWithoutDims;
+}
+
bool SliceAttr::isEqualTo(const xegpu::DistributeLayoutAttr &other) {
if (dyn_cast<xegpu::LayoutAttr>(other))
return false;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
index bc2703e7ffbda..72e12eb9123bb 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUOptimizeBlockLoads.cpp
@@ -416,7 +416,8 @@ class VectorExtractOpPattern final
}
};
-class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
+class MultiRed2dOpPattern
+ : public OpConversionPattern<vector::MultiDimReductionOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(vector::MultiDimReductionOp reductionOp, OpAdaptor adaptor,
@@ -440,14 +441,9 @@ class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
// 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 intraLaneRedResLayout = xegpu::SliceAttr::get(
- reductionOp.getContext(), resSliceLayoutAttr.getParent(),
- DenseI64ArrayAttr::get(getContext(), sliceDims));
+ SmallVector<int64_t> dropDims{crossLaneDim};
+ auto intraLaneRedResLayout =
+ cast<xegpu::SliceAttr>(resSliceLayoutAttr.dropSliceDims(dropDims));
SmallVector<int64_t> accShape(sourceVecType.getShape());
accShape.erase(accShape.begin() + intraLaneDim);
@@ -485,12 +481,9 @@ class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
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());
- }
+ if (auto layoutSliceAttr = dyn_cast<xegpu::SliceAttr>(layout))
+ layoutAttr =
+ dyn_cast<xegpu::LayoutAttr>(layoutSliceAttr.flatten().getParent());
assert(layoutAttr);
SmallVector<int64_t> laneLayout = layoutAttr.getEffectiveLaneLayoutAsInt();
@@ -511,7 +504,8 @@ class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
void xegpu::populateXeGPUOptimizeBlockLoadsPatterns(
RewritePatternSet &patterns) {
patterns.add<XeGPUCreateNdDescOpPattern, XeGPULoadNdDescOpPattern,
- VectorExtractOpPattern, MultiRed2dOp>(patterns.getContext());
+ VectorExtractOpPattern, MultiRed2dOpPattern>(
+ patterns.getContext());
}
namespace {
More information about the Mlir-commits
mailing list