[Mlir-commits] [mlir] [MLIR][XeGPU] Use all ones as the neutral value of an AND reduction (PR #232828)
Md Abdullah Shahneous Bari
llvmlistbot at llvm.org
Mon Oct 19 13:53:27 PDT 2026
https://github.com/mshahneo updated https://github.com/llvm/llvm-project/pull/232828
>From 1a9e57f871506e7433fc8bd72aab4fdb4931974f Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Fri, 16 Oct 2026 16:13:39 +0000
Subject: [PATCH 1/2] [MLIR][XeGPU] Use all ones as the neutral value of an AND
reduction
createReductionNeutralValue returns 1 for an AND reduction. The neutral
value of AND is all ones. With 1, the reduction keeps only bit 0 of the
result.
The WG-to-SG distribution of vector.multi_reduction starts its local
reductions and the final cross-subgroup reduction from this value. The
peephole split of a 2-D reduction does the same. So an `<and>`
reduction on an integer type wider than i1 gives a wrong result. For
i1, 1 is already all ones.
Return all ones for integer and index types.
Co-Authored-By: Claude Opus 5.5 <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 11 +++---
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 9 ++++-
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir | 34 +++++++++++++++++++
3 files changed, 48 insertions(+), 6 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index ab526793f2aa8..f95478dae74cf 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -147,11 +147,12 @@ Value lowerToVectorReductions(TypedValue<VectorType> src,
Location loc, PatternRewriter &rewriter);
/// Creates a constant filled with the neutral (identity) value for the
-/// given reduction kind. For example: 0 for ADD/OR/XOR, 1 for MUL/AND,
-/// max/min signed/unsigned int for MINSI/MINUI/MAXSI/MAXUI, and +/-infinity
-/// for float min/max operations. If \p type is a VectorType, returns a splat
-/// vector constant; otherwise returns a scalar constant. Returns nullptr if
-/// the element type is incompatible with the requested reduction kind.
+/// given reduction kind. For example: 0 for ADD/OR/XOR, 1 for MUL, all ones
+/// for AND, max/min signed/unsigned int for MINSI/MINUI/MAXSI/MAXUI, and
+/// +/-infinity for float min/max operations. If \p type is a VectorType,
+/// returns a splat vector constant; otherwise returns a scalar constant.
+/// Returns nullptr if the element type is incompatible with the requested
+/// reduction kind.
Value createReductionNeutralValue(OpBuilder &builder, Location loc, Type type,
vector::CombiningKind kind);
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index e22364df9e590..b83c41615019e 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -747,9 +747,16 @@ Value xegpu::createReductionNeutralValue(OpBuilder &builder, Location loc,
return makeConst(builder.getZeroAttr(elemTy));
case vector::CombiningKind::MUL:
- case vector::CombiningKind::AND:
return makeConst(builder.getOneAttr(elemTy));
+ case vector::CombiningKind::AND:
+ if (auto intTy = dyn_cast<IntegerType>(elemTy))
+ return makeConst(
+ builder.getIntegerAttr(elemTy, APInt::getAllOnes(intTy.getWidth())));
+ if (elemTy.isIndex())
+ return makeConst(builder.getIndexAttr(-1));
+ return nullptr;
+
case vector::CombiningKind::MINSI:
if (auto intTy = dyn_cast<IntegerType>(elemTy))
return makeConst(builder.getIntegerAttr(
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index 8e53c40b2078c..2acad3cde0e3b 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -963,6 +963,40 @@ gpu.module @test_distribution {
gpu.return
}
+ // AND has all ones as neutral element, not 1.
+ // CHECK-LABEL: gpu.func @vector_reduce_and_neutral
+ gpu.func @vector_reduce_and_neutral(%src: memref<256x64xi32>, %src8: memref<256x64xi8>) {
+ // CHECK: %[[LOAD:.*]] = xegpu.load_nd {{.*}} -> vector<16x64xi32>
+ // CHECK: %[[N0:.*]] = arith.constant dense<-1> : vector<16xi32>
+ // CHECK: vector.multi_reduction <and>, %[[LOAD]], %[[N0]] [1] : vector<16x64xi32> to vector<16xi32>
+ // CHECK: %[[LOAD8:.*]] = xegpu.load_nd {{.*}} -> vector<16x64xi8>
+ // CHECK: %[[N1:.*]] = arith.constant dense<-1> : vector<16xi8>
+ // CHECK: vector.multi_reduction <and>, %[[LOAD8]], %[[N1]] [1] : vector<16x64xi8> to vector<16xi8>
+ %cst = arith.constant dense<-1> : vector<256xi32>
+ %tdesc = xegpu.create_nd_tdesc %src : memref<256x64xi32>
+ -> !xegpu.tensor_desc<256x64xi32>
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>}>
+ : !xegpu.tensor_desc<256x64xi32>
+ -> vector<256x64xi32>
+ %r0 = vector.multi_reduction <and>, %load, %cst [1]
+ : vector<256x64xi32> to vector<256xi32>
+ %a0 = xegpu.convert_layout %r0
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xi32>
+ %cst8 = arith.constant dense<-1> : vector<256xi8>
+ %tdesc8 = xegpu.create_nd_tdesc %src8 : memref<256x64xi8>
+ -> !xegpu.tensor_desc<256x64xi8>
+ %load8 = xegpu.load_nd %tdesc8[0, 0] <{layout = #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>}>
+ : !xegpu.tensor_desc<256x64xi8>
+ -> vector<256x64xi8>
+ %r1 = vector.multi_reduction <and>, %load8, %cst8 [1]
+ : vector<256x64xi8> to vector<256xi8>
+ %a1 = xegpu.convert_layout %r1
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xi8>
+ gpu.return
+ }
+
// CHECK-LABEL: load_nd_tdesc_with_anchor_layout
gpu.func @load_nd_tdesc_with_anchor_layout(%src: memref<256x128xf32>) {
//CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<256x128xf32> -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
>From 324804a719a35564c7c9f579e3d10e4b381dd8a0 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Mon, 19 Oct 2026 20:42:41 +0000
Subject: [PATCH 2/2] [MLIR][XeGPU] Make the AND neutral test tell the neutral
from the accumulator
The test used an accumulator of all ones, the same value as the neutral
element. So the CHECKs could not show which constant the reduction
starts from, and `andi %x, -1` folded the accumulator combine away.
Use an accumulator of 7. The CHECKs now show the local reduction
starting from -1 and the accumulator combined afterwards. Drop the i8
case, which takes the same path as i32. Add a cross-subgroup case: the
final reduction after the SLM round trip also starts from the neutral
element.
Co-Authored-By: Claude Opus 5.5 <noreply at anthropic.com>
---
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir | 63 +++++++++++++--------
1 file changed, 40 insertions(+), 23 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index 2acad3cde0e3b..0c4bf0428cc5e 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -963,37 +963,54 @@ gpu.module @test_distribution {
gpu.return
}
- // AND has all ones as neutral element, not 1.
- // CHECK-LABEL: gpu.func @vector_reduce_and_neutral
- gpu.func @vector_reduce_and_neutral(%src: memref<256x64xi32>, %src8: memref<256x64xi8>) {
- // CHECK: %[[LOAD:.*]] = xegpu.load_nd {{.*}} -> vector<16x64xi32>
- // CHECK: %[[N0:.*]] = arith.constant dense<-1> : vector<16xi32>
- // CHECK: vector.multi_reduction <and>, %[[LOAD]], %[[N0]] [1] : vector<16x64xi32> to vector<16xi32>
- // CHECK: %[[LOAD8:.*]] = xegpu.load_nd {{.*}} -> vector<16x64xi8>
- // CHECK: %[[N1:.*]] = arith.constant dense<-1> : vector<16xi8>
- // CHECK: vector.multi_reduction <and>, %[[LOAD8]], %[[N1]] [1] : vector<16x64xi8> to vector<16xi8>
- %cst = arith.constant dense<-1> : vector<256xi32>
+ // AND has all ones as neutral element, not 1. Each subgroup's local
+ // reduction starts from the neutral element; the original accumulator is
+ // combined afterwards.
+ // CHECK-LABEL: gpu.func @vector_reduce_and_neutral(
+ // CHECK: %[[ACC:.*]] = arith.constant dense<7> : vector<16xi32>
+ // CHECK: %[[LOAD:.*]] = xegpu.load_nd {{.*}} -> vector<16x64xi32>
+ // CHECK: %[[N:.*]] = arith.constant dense<-1> : vector<16xi32>
+ // CHECK: %[[RED:.*]] = vector.multi_reduction <and>, %[[LOAD]], %[[N]] [1] : vector<16x64xi32> to vector<16xi32>
+ // CHECK: arith.andi %[[RED]], %[[ACC]] : vector<16xi32>
+ gpu.func @vector_reduce_and_neutral(%src: memref<256x64xi32>) {
+ %cst = arith.constant dense<7> : vector<256xi32>
%tdesc = xegpu.create_nd_tdesc %src : memref<256x64xi32>
-> !xegpu.tensor_desc<256x64xi32>
%load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>}>
: !xegpu.tensor_desc<256x64xi32>
-> vector<256x64xi32>
- %r0 = vector.multi_reduction <and>, %load, %cst [1]
+ %r = vector.multi_reduction <and>, %load, %cst [1]
: vector<256x64xi32> to vector<256xi32>
- %a0 = xegpu.convert_layout %r0
+ %a = xegpu.convert_layout %r
<{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
: vector<256xi32>
- %cst8 = arith.constant dense<-1> : vector<256xi8>
- %tdesc8 = xegpu.create_nd_tdesc %src8 : memref<256x64xi8>
- -> !xegpu.tensor_desc<256x64xi8>
- %load8 = xegpu.load_nd %tdesc8[0, 0] <{layout = #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>}>
- : !xegpu.tensor_desc<256x64xi8>
- -> vector<256x64xi8>
- %r1 = vector.multi_reduction <and>, %load8, %cst8 [1]
- : vector<256x64xi8> to vector<256xi8>
- %a1 = xegpu.convert_layout %r1
- <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
- : vector<256xi8>
+ gpu.return
+ }
+
+ // The final cross-subgroup reduction also starts from the neutral element.
+ // CHECK-LABEL: gpu.func @vector_reduce_and_neutral_cross_sg(
+ // CHECK: %[[ACC:.*]] = arith.constant 7 : i32
+ // CHECK: %[[LOAD:.*]] = xegpu.load_nd {{.*}} -> vector<8x8xi32>
+ // CHECK: %[[N0:.*]] = arith.constant -1 : i32
+ // CHECK: vector.multi_reduction <and>, %[[LOAD]], %[[N0]] [0, 1] : vector<8x8xi32> to i32
+ // CHECK: gpu.barrier
+ // CHECK: %[[LOAD_SLM:.*]] = xegpu.load_matrix {{.*}} -> vector<4x4xi32>
+ // CHECK: %[[N1:.*]] = arith.constant -1 : i32
+ // CHECK: %[[FINAL:.*]] = vector.multi_reduction <and>, %[[LOAD_SLM]], %[[N1]] [0, 1] : vector<4x4xi32> to i32
+ // CHECK: arith.andi %[[FINAL]], %[[ACC]] : i32
+ gpu.func @vector_reduce_and_neutral_cross_sg(%src: memref<32x32xi32>) {
+ %cst = arith.constant 7 : i32
+ %tdesc = xegpu.create_nd_tdesc %src : memref<32x32xi32>
+ -> !xegpu.tensor_desc<32x32xi32>
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [8, 8]>}>
+ : !xegpu.tensor_desc<32x32xi32>
+ -> vector<32x32xi32>
+ %reduce = vector.multi_reduction <and>, %load, %cst [0, 1]
+ : vector<32x32xi32> to i32
+ %anchor = xegpu.convert_layout %reduce
+ <{
+ target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [4, 4], sg_data = [8, 8]>, dims = [0, 1]>
+ }> : i32
gpu.return
}
More information about the Mlir-commits
mailing list