[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 26 07:55:36 PDT 2026
https://github.com/mshahneo updated https://github.com/llvm/llvm-project/pull/232828
>From af06904355e92fcb0d64c8964a0ac663b46e4eaa 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 ab526793f2aa82..f95478dae74cf2 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 406712b3f5b1b5..a12588945ded36 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 8e53c40b2078cf..2acad3cde0e3b0 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 e1fe990b481fc4cf38cd3b392befc92942e56571 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 2acad3cde0e3b0..0c4bf0428cc5e9 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