[Mlir-commits] [mlir] 93aede0 - [MLIR][XeGPU] Use all ones as the neutral value of an AND reduction (#232828)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 26 08:10:12 PDT 2026
Author: Md Abdullah Shahneous Bari
Date: 2026-10-26T15:10:07Z
New Revision: 93aede0caccca0caad0df5cd3d3a3a767cf3cf17
URL: https://github.com/llvm/llvm-project/commit/93aede0caccca0caad0df5cd3d3a3a767cf3cf17
DIFF: https://github.com/llvm/llvm-project/commit/93aede0caccca0caad0df5cd3d3a3a767cf3cf17.diff
LOG: [MLIR][XeGPU] Use all ones as the neutral value of an AND reduction (#232828)
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>
Added:
Modified:
mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
Removed:
################################################################################
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..0c4bf0428cc5e9 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -963,6 +963,57 @@ gpu.module @test_distribution {
gpu.return
}
+ // 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>
+ %r = vector.multi_reduction <and>, %load, %cst [1]
+ : vector<256x64xi32> to vector<256xi32>
+ %a = xegpu.convert_layout %r
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xi32>
+ 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
+ }
+
// 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]>>
More information about the Mlir-commits
mailing list