[Mlir-commits] [mlir] [MLIR][XeGPU] Preserve fast-math flags when lowering vector reductions (PR #233548)
Md Abdullah Shahneous Bari
llvmlistbot at llvm.org
Mon Nov 2 15:08:02 PST 2026
https://github.com/mshahneo updated https://github.com/llvm/llvm-project/pull/233548
>From 49f66df445b756b344b5b7ef2d0dbd3452fff31b Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Thu, 15 Oct 2026 15:53:26 +0000
Subject: [PATCH] [MLIR][XeGPU] Preserve fast-math flags when lowering vector
reductions
The XeGPU workgroup-to-subgroup (WG-to-SG), unroll, peephole and
subgroup-to-lane (SG-to-lane) passes rebuild vector.multi_reduction and
vector.reduction ops and drop their fast-math flags. Pass the flags to
the new reductions, and to the arith ops of the cross-lane and
accumulator combines. The XeGPU reduction helpers take an optional
`fastmath` argument.
Some of the new reductions start from a neutral value instead of the
original accumulator. They now run under the original flags, so the
neutral value must be one that both the element type can represent and
the flags allow. For this, createReductionNeutralValue now takes the
flags too:
- The kinds that ignore NaN operands use a quiet NaN if the type has one
and `nnan` is not set.
- Otherwise the reduction uses the infinity that never wins, unless
`ninf` is set or the type has no infinity.
- Otherwise it uses the lowest finite value for max and the highest
for min.
Without flags, the printed IR and the lowered code do not change.
Co-Authored-By: Claude Opus 5.5 <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 37 ++--
.../Transforms/XeGPUPeepHoleOptimizer.cpp | 8 +-
.../Transforms/XeGPUSgToLaneDistribute.cpp | 22 ++-
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 6 +-
.../Transforms/XeGPUWgToSgDistribute.cpp | 13 +-
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 68 ++++---
.../test/Dialect/XeGPU/peephole-optimize.mlir | 57 ++++++
.../XeGPU/sg-to-lane-distribute-unit.mlir | 45 +++++
.../Dialect/XeGPU/xegpu-unroll-patterns.mlir | 11 ++
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir | 180 ++++++++++++++++++
10 files changed, 388 insertions(+), 59 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index c996b59b73843e..62ebfc754dd2a8 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -133,8 +133,10 @@ SmallVector<OpFoldResult> addWithRightAligned(OpBuilder &builder, Location loc,
/// lanes given by `size`). This uses butterfly shuffles to perform the
/// reduction in a log2(size) number of steps.
/// NOTE: Implementation taken from TestVectorTransforms.cpp
+/// For float kinds, `fastmath` is set on every combining operation.
Value subgroupReduction(Location loc, OpBuilder &builder, Value input,
- vector::CombiningKind kind, uint32_t size);
+ vector::CombiningKind kind, uint32_t size,
+ arith::FastMathFlagsAttr fastmath = nullptr);
/// Given a `src` and an `acc` argumments from a vector::MultiDimReductionOp,
/// lower to a set of vector::ReductionOp ops over 1D slices extracted from
@@ -144,28 +146,35 @@ Value subgroupReduction(Location loc, OpBuilder &builder, Value input,
Value lowerToVectorReductions(TypedValue<VectorType> src,
TypedValue<VectorType> acc,
vector::CombiningKind kind, int64_t reductionDim,
- Location loc, PatternRewriter &rewriter);
+ Location loc, PatternRewriter &rewriter,
+ arith::FastMathFlagsAttr fastmath = nullptr);
/// Creates a constant filled with the neutral (identity) value for the
/// 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.
+/// for AND, and max/min signed/unsigned int for MINSI/MINUI/MAXSI/MAXUI.
+/// For float min/max kinds, the value must fit the element type and be
+/// allowed by \p fastmath:
+/// - A quiet NaN for the kinds that ignore NaN operands, if the type has NaN
+/// and `nnan` is not set.
+/// - Otherwise the infinity that never wins, if the type has it and `ninf`
+/// is not set.
+/// - Otherwise the lowest (for max) or highest (for min) finite value.
+/// 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);
+ vector::CombiningKind kind,
+ arith::FastMathFlagsAttr fastmath = nullptr);
/// Lowers cross-lane reductions to shuffle operations on a 2D vector.
/// Extracts slices along the reduction dimension, performs subgroup reductions
/// with shuffles across reductionSize work-items, and inserts the results back
/// into an accumulator vector.
-Value lowerCrossLaneReductionToShuffles(TypedValue<VectorType> src,
- TypedValue<VectorType> acc,
- vector::CombiningKind kind,
- int64_t reductionDim,
- int64_t reductionSize, Location loc,
- PatternRewriter &rewriter);
+Value lowerCrossLaneReductionToShuffles(
+ TypedValue<VectorType> src, TypedValue<VectorType> acc,
+ vector::CombiningKind kind, int64_t reductionDim, int64_t reductionSize,
+ Location loc, PatternRewriter &rewriter,
+ arith::FastMathFlagsAttr fastmath = nullptr);
/// Helper Function to find a proper instruction multiple for the user-supplied
/// sg-level data shape (diven by `dim`). `candidates` are uArch allowed shapes.
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
index 488a97d60f628d..24a1f13d78d21d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -522,18 +522,20 @@ class MultiRed2dOpPattern
accShape.erase(accShape.begin() + intraLaneDim);
Type eTy = sourceVecType.getElementType();
Value constNeutralVal = xegpu::createReductionNeutralValue(
- rewriter, loc, VectorType::get(accShape, eTy), reductionOp.getKind());
+ rewriter, loc, VectorType::get(accShape, eTy), reductionOp.getKind(),
+ reductionOp.getFastmathAttr());
Value intraLaneReduced = vector::MultiDimReductionOp::create(
rewriter, loc, reductionOp.getKind(), reductionOp.getSource(),
- constNeutralVal, ArrayRef<int64_t>(intraLaneDim));
+ constNeutralVal, ArrayRef<int64_t>(intraLaneDim),
+ reductionOp.getFastmath());
// Adjust crossLaneDim after the first reduction.
if (crossLaneDim > intraLaneDim)
crossLaneDim -= 1;
Value crossLaneReduced = vector::MultiDimReductionOp::create(
rewriter, loc, reductionOp.getKind(), intraLaneReduced, acc,
- ArrayRef<int64_t>(crossLaneDim));
+ ArrayRef<int64_t>(crossLaneDim), reductionOp.getFastmath());
assert(crossLaneReduced.getType() == reductionOp.getResult().getType() &&
"Type mismatch");
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 4f9a37851665ca..7462e87b5b08dc 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -643,13 +643,15 @@ struct SgToLaneVectorReduction
Value laneValVec = adaptor.getVector();
// Distribute and reduce across lanes in the subgroup.
- Value fullReduce = xegpu::subgroupReduction(
- op.getLoc(), rewriter, laneValVec, op.getKind(), sgSize);
+ Value fullReduce =
+ xegpu::subgroupReduction(op.getLoc(), rewriter, laneValVec,
+ op.getKind(), sgSize, op.getFastmathAttr());
// If there's an accumulator, combine it with the reduced value.
if (adaptor.getAcc())
fullReduce = vector::makeArithReduction(
- rewriter, op.getLoc(), op.getKind(), fullReduce, adaptor.getAcc());
+ rewriter, op.getLoc(), op.getKind(), fullReduce, adaptor.getAcc(),
+ op.getFastmathAttr());
rewriter.replaceOp(op, fullReduce);
return success();
@@ -690,13 +692,14 @@ struct SgToLaneMultiDimReduction
VectorType origSourceType = op.getSourceVectorType();
int64_t reductionDimSize = origSourceType.getShape()[reductionDim];
// Local reduction to scalar, then cross-lane butterfly shuffles.
- result =
- xegpu::subgroupReduction(op.getLoc(), rewriter, adaptor.getSource(),
- op.getKind(), reductionDimSize);
+ result = xegpu::subgroupReduction(op.getLoc(), rewriter,
+ adaptor.getSource(), op.getKind(),
+ reductionDimSize, op.getFastmathAttr());
// Combine with accumulator if present.
if (adaptor.getAcc())
result = vector::makeArithReduction(rewriter, op.getLoc(), op.getKind(),
- result, adaptor.getAcc());
+ result, adaptor.getAcc(),
+ op.getFastmathAttr());
} else if (isReductionLaneLocal(op)) {
// For lane-local reduction, lower to a sequence of vector.reduction ops
// over 1D slices extracted from the distributed source vector. This is
@@ -705,7 +708,7 @@ struct SgToLaneMultiDimReduction
result = xegpu::lowerToVectorReductions(
cast<TypedValue<VectorType>>(adaptor.getSource()),
cast<TypedValue<VectorType>>(adaptor.getAcc()), op.getKind(),
- reductionDim, op.getLoc(), rewriter);
+ reductionDim, op.getLoc(), rewriter, op.getFastmathAttr());
} else {
auto reductionDim = reductionDims[0];
VectorType sourceType = op.getSourceVectorType();
@@ -713,7 +716,8 @@ struct SgToLaneMultiDimReduction
result = xegpu::lowerCrossLaneReductionToShuffles(
cast<TypedValue<VectorType>>(adaptor.getSource()),
cast<TypedValue<VectorType>>(adaptor.getAcc()), op.getKind(),
- reductionDim, reductionDimSize, op.getLoc(), rewriter);
+ reductionDim, reductionDimSize, op.getLoc(), rewriter,
+ op.getFastmathAttr());
}
rewriter.replaceOp(op, result);
return success();
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index aaee71bb861822..00b5194ad59113 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -1008,7 +1008,8 @@ struct UnrollMultiReductionOp
Value reduced = tiles[0];
for (size_t i = 1; i < tiles.size(); ++i)
reduced =
- vector::makeArithReduction(rewriter, loc, kind, reduced, tiles[i]);
+ vector::makeArithReduction(rewriter, loc, kind, reduced, tiles[i],
+ reductionOp.getFastmathAttr());
// Step 3: Perform a single multi_reduction with the accumulator slice.
SmallVector<int64_t> accStrides(keptTileShape.size(), 1);
@@ -1016,7 +1017,8 @@ struct UnrollMultiReductionOp
rewriter, loc, acc, keptOffsets, keptTileShape, accStrides);
auto newReduction = vector::MultiDimReductionOp::create(
- rewriter, loc, reduced, accSlice, reductionMask, kind);
+ rewriter, loc, reduced, accSlice, reductionMask, kind,
+ reductionOp.getFastmath());
// Step 4: Insert the reduced result into the output vector.
SmallVector<int64_t> dstStrides(keptTileShape.size(), 1);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 9e2a0401900e08..69426d06533cd1 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -1129,11 +1129,11 @@ struct WgToSgMultiDimReductionOp
for (auto sgSrc : sgSrcs) {
// Create neutral accumulator for local reduction
Value neutralLocalAcc = xegpu::createReductionNeutralValue(
- rewriter, loc, sgDstType, op.getKind());
+ rewriter, loc, sgDstType, op.getKind(), op.getFastmathAttr());
// Local reduction with neutral accumulator
auto localReduce = vector::MultiDimReductionOp::create(
rewriter, loc, sgDstType, op.getKind(), sgSrc, neutralLocalAcc,
- reductionDims);
+ reductionDims, op.getFastmath());
localReductions.push_back(localReduce.getResult());
}
@@ -1154,7 +1154,8 @@ struct WgToSgMultiDimReductionOp
SmallVector<Value> results;
for (auto localResult : localReductions) {
auto finalResult = vector::makeArithReduction(
- rewriter, loc, op.getKind(), localResult, adaptor.getAcc()[0]);
+ rewriter, loc, op.getKind(), localResult, adaptor.getAcc()[0],
+ op.getFastmathAttr());
results.push_back(finalResult);
}
rewriter.replaceOpWithMultiple(op, {results});
@@ -1243,17 +1244,17 @@ struct WgToSgMultiDimReductionOp
// Step 5: Perform final reduction with neutral accumulator and add the
// original accumulator at the end
Value neutralFinalAcc = xegpu::createReductionNeutralValue(
- rewriter, loc, sgDstType, op.getKind());
+ rewriter, loc, sgDstType, op.getKind(), op.getFastmathAttr());
SmallVector<Value> finalResults;
for (size_t i = 0; i < slmLoadData.size(); ++i) {
auto loaded = slmLoadData[i];
auto finalReduce = vector::MultiDimReductionOp::create(
rewriter, loc, sgDstType, op.getKind(), loaded, neutralFinalAcc,
- reductionDims);
+ reductionDims, op.getFastmath());
finalResults.push_back(vector::makeArithReduction(
rewriter, loc, op.getKind(), finalReduce.getResult(),
- adaptor.getAcc()[i]));
+ adaptor.getAcc()[i], op.getFastmathAttr()));
}
rewriter.replaceOpWithMultiple(op, {finalResults});
return success();
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 56c8b440701e16..ccbe3efaf5ce6e 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -565,17 +565,25 @@ int xegpu::getLargestDivisor(T dim, ArrayRef<T> candidates,
return largest;
}
+/// Returns the flags held by `fastmath`, or none if it is null.
+static arith::FastMathFlags getFlags(arith::FastMathFlagsAttr fastmath) {
+ return fastmath ? fastmath.getValue() : arith::FastMathFlags::none;
+}
+
Value xegpu::subgroupReduction(Location loc, OpBuilder &builder, Value input,
- vector::CombiningKind kind, uint32_t size) {
+ vector::CombiningKind kind, uint32_t size,
+ arith::FastMathFlagsAttr fastmath) {
// First reduce on a single thread to get per lane reduction value.
- Value laneVal = vector::ReductionOp::create(builder, loc, kind, input);
+ Value laneVal = vector::ReductionOp::create(builder, loc, kind, input,
+ getFlags(fastmath));
// Parallel reduction using butterfly shuffles.
for (uint64_t i = 1; i < size; i <<= 1) {
Value shuffled =
gpu::ShuffleOp::create(builder, loc, laneVal, i, /** width = **/ size,
/** mode = **/ gpu::ShuffleMode::XOR)
.getShuffleResult();
- laneVal = makeArithReduction(builder, loc, kind, laneVal, shuffled);
+ laneVal =
+ makeArithReduction(builder, loc, kind, laneVal, shuffled, fastmath);
}
return laneVal;
}
@@ -584,7 +592,8 @@ Value xegpu::lowerToVectorReductions(TypedValue<VectorType> src,
TypedValue<VectorType> acc,
vector::CombiningKind kind,
int64_t reductionDim, Location loc,
- PatternRewriter &rewriter) {
+ PatternRewriter &rewriter,
+ arith::FastMathFlagsAttr fastmath) {
VectorType sourceType = src.getType();
int64_t sourceRank = sourceType.getRank();
// Expecting at least a 2D source vector. Leading dimensions (all except the
@@ -648,7 +657,7 @@ Value xegpu::lowerToVectorReductions(TypedValue<VectorType> src,
accIdx[accRank - 1] = i;
Value accExtract = vector::ExtractOp::create(rewriter, loc, acc, accIdx);
Value reduction = vector::ReductionOp::create(
- rewriter, loc, kind, slice.getResult(), accExtract);
+ rewriter, loc, kind, slice.getResult(), accExtract, getFlags(fastmath));
reductionResult = vector::InsertOp::create(rewriter, loc, reduction,
reductionResult, accIdx);
// Insert op should have the same layout as the accumulator.
@@ -660,7 +669,8 @@ Value xegpu::lowerToVectorReductions(TypedValue<VectorType> src,
Value xegpu::lowerCrossLaneReductionToShuffles(
TypedValue<VectorType> src, TypedValue<VectorType> acc,
vector::CombiningKind kind, int64_t reductionDim, int64_t reductionSize,
- Location loc, PatternRewriter &rewriter) {
+ Location loc, PatternRewriter &rewriter,
+ arith::FastMathFlagsAttr fastmath) {
VectorType sourceType = src.getType();
int64_t sourceRank = sourceType.getRank();
// Expecting at least a 2D source vector. Leading dimensions (all except the
@@ -715,10 +725,10 @@ Value xegpu::lowerCrossLaneReductionToShuffles(
SmallVector<int64_t> accIdx(accRank, 0);
accIdx[accRank - 1] = i;
Value accExtract = vector::ExtractOp::create(rewriter, loc, acc, accIdx);
- Value fullReduce =
- xegpu::subgroupReduction(loc, rewriter, slice, kind, reductionSize);
- fullReduce =
- vector::makeArithReduction(rewriter, loc, kind, fullReduce, accExtract);
+ Value fullReduce = xegpu::subgroupReduction(loc, rewriter, slice, kind,
+ reductionSize, fastmath);
+ fullReduce = vector::makeArithReduction(rewriter, loc, kind, fullReduce,
+ accExtract, fastmath);
reductionResult = vector::InsertOp::create(rewriter, loc, fullReduce,
reductionResult, accIdx);
}
@@ -726,37 +736,44 @@ Value xegpu::lowerCrossLaneReductionToShuffles(
}
/// Returns the neutral value of the float min or max reduction `kind` over
-/// elements of `semantics`. Builds only values the format can represent:
-/// - A kind that ignores NaN operands uses a quiet NaN, if the format has one.
-/// - Otherwise it uses the infinity that never wins.
-/// - A format without that infinity uses its lowest (for max) or highest (for
-/// min) value. That value never wins either.
+/// elements of `semantics`. The reduction combines this value under its
+/// fast-math `flags`, so build only a value that the format can represent and
+/// the flags allow:
+/// - A kind that ignores NaN operands uses a quiet NaN, if the format has one
+/// and `nnan` is not set.
+/// - Otherwise it uses the infinity that never wins, if the format has one and
+/// `ninf` is not set.
+/// - Otherwise it uses the lowest (for max) or highest (for min) finite value.
+/// That value never wins either: no number that the format and the flags
+/// allow is below the lowest or above the highest finite value.
static APFloat getMinMaxNeutral(const llvm::fltSemantics &semantics,
- vector::CombiningKind kind) {
+ vector::CombiningKind kind,
+ arith::FastMathFlags flags) {
bool isMax = kind == vector::CombiningKind::MAXIMUMF ||
kind == vector::CombiningKind::MAXNUMF ||
kind == vector::CombiningKind::MAXIMUMNUMF;
bool ignoresNaN = kind != vector::CombiningKind::MINIMUMF &&
kind != vector::CombiningKind::MAXIMUMF;
- if (ignoresNaN && APFloat::semanticsHasNaN(semantics))
+ if (ignoresNaN && APFloat::semanticsHasNaN(semantics) &&
+ !arith::bitEnumContainsAny(flags, arith::FastMathFlags::nnan))
return APFloat::getQNaN(semantics);
+ bool useInf = APFloat::semanticsHasInf(semantics) &&
+ !arith::bitEnumContainsAny(flags, arith::FastMathFlags::ninf);
if (!isMax)
- return APFloat::semanticsHasInf(semantics) ? APFloat::getInf(semantics)
- : APFloat::getLargest(semantics);
+ return useInf ? APFloat::getInf(semantics) : APFloat::getLargest(semantics);
// An unsigned format has no negative values. Its lowest value is zero, or
// its smallest positive value if it has no zero.
if (!APFloat::semanticsHasSignedRepr(semantics))
return APFloat::semanticsHasZero(semantics)
? APFloat::getZero(semantics)
: APFloat::getSmallest(semantics);
- return APFloat::semanticsHasInf(semantics)
- ? APFloat::getInf(semantics, /*Negative=*/true)
- : APFloat::getLargest(semantics, /*Negative=*/true);
+ return useInf ? APFloat::getInf(semantics, /*Negative=*/true)
+ : APFloat::getLargest(semantics, /*Negative=*/true);
}
Value xegpu::createReductionNeutralValue(OpBuilder &builder, Location loc,
- Type type,
- vector::CombiningKind kind) {
+ Type type, vector::CombiningKind kind,
+ arith::FastMathFlagsAttr fastmath) {
auto vecTy = dyn_cast<VectorType>(type);
Type elemTy = vecTy ? vecTy.getElementType() : type;
@@ -812,7 +829,8 @@ Value xegpu::createReductionNeutralValue(OpBuilder &builder, Location loc,
case vector::CombiningKind::MAXIMUMNUMF:
if (auto floatTy = dyn_cast<FloatType>(elemTy))
return makeConst(builder.getFloatAttr(
- elemTy, getMinMaxNeutral(floatTy.getFloatSemantics(), kind)));
+ elemTy, getMinMaxNeutral(floatTy.getFloatSemantics(), kind,
+ getFlags(fastmath))));
return nullptr;
}
return nullptr;
diff --git a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
index 0ca92d0fc972aa..b9d3e9da75a517 100644
--- a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
+++ b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
@@ -584,3 +584,60 @@ gpu.func @i4_transpose_memref_offset(%src: memref<512x2048xi4, strided<[2048, 1]
gpu.return %v : vector<32x64xi4>
}
}
+
+// -----
+// CHECK-LABEL: gpu.func @vector_reduce_2d_fastmath(
+// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <maximumf>, %{{.*}}, %{{.*}} fastmath<nnan,nsz> [0] : vector<4x16xf32> to vector<16xf32>
+// CHECK: vector.multi_reduction <maximumf>, %[[REDUCE_1]], %{{.*}} fastmath<nnan,nsz> [0] : vector<16xf32> to f32
+gpu.module @xevm_test {
+ gpu.func @vector_reduce_2d_fastmath(%src: memref<4x16xf32>, %dst: memref<256xf32>) {
+ %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 = [1, 1]>>
+ %load = xegpu.load_nd %tdesc[0, 0]
+ : !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+ -> vector<4x16xf32>
+ %reduce = vector.multi_reduction <maximumf>, %load, %cst fastmath<nnan,nsz>
+ {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 = [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 = [1, 1]>, dims = [0]>}> : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
+ gpu.return
+ }
+}
+
+// -----
+// The neutral accumulator of the intra-lane reduction must be a value that the
+// flags allow: no NaN under `nnan`, no infinity under `ninf`.
+// CHECK-LABEL: gpu.func @vector_reduce_2d_fastmath_neutral(
+// CHECK-DAG: %[[MIN_N:.*]] = arith.constant dense<3.40282347E+38> : vector<16xf32>
+// CHECK-DAG: %[[MAX_N:.*]] = arith.constant dense<0xFF800000> : vector<16xf32>
+// CHECK: vector.multi_reduction <maxnumf>, %{{.*}}, %[[MAX_N]] fastmath<nnan> [0] : vector<4x16xf32> to vector<16xf32>
+// CHECK: vector.multi_reduction <minimumf>, %{{.*}}, %[[MIN_N]] fastmath<fast> [0] : vector<4x16xf32> to vector<16xf32>
+gpu.module @xevm_test {
+ gpu.func @vector_reduce_2d_fastmath_neutral(%src: memref<4x16xf32>, %dst: memref<256xf32>) {
+ %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 = [1, 1]>>
+ %load = xegpu.load_nd %tdesc[0, 0]
+ : !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+ -> vector<4x16xf32>
+ %max = vector.multi_reduction <maxnumf>, %load, %cst fastmath<nnan>
+ {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0, 1]>}
+ [0, 1] : vector<4x16xf32> to f32
+ %min = vector.multi_reduction <minimumf>, %load, %max fastmath<fast>
+ {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 %min
+ {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 = [1, 1]>, dims = [0]>}> : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
+ gpu.return
+ }
+}
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
index 71e8f68f4dc542..6b322cf0d05b3e 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1847,4 +1847,49 @@ gpu.func @extract_strided_slice_packed_distributed_dim() {
}> : vector<1x64xbf16>
gpu.return
}
+
+// CHECK-LABEL: gpu.func @vector_reduction_fastmath
+// CHECK: vector.reduction <maximumf>, %{{.*}} fastmath<nnan,nsz> : vector<2xf32> into f32
+// CHECK-COUNT-5: arith.maximumf %{{.*}}, %{{.*}} fastmath<nnan,nsz> : f32
+gpu.func @vector_reduction_fastmath() {
+ %acc = arith.constant 1.0 : f32
+ %0 = "test.some_op"() : () -> vector<32xf32>
+ %2 = vector.reduction <maximumf>, %0, %acc fastmath<nnan,nsz> : vector<32xf32> into f32
+ %anchor = xegpu.convert_layout %2
+ <{
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [16], lane_data = [1]>, dims = [0]>
+ }> : f32
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @vector_multi_reduction_cross_lane_fastmath
+// CHECK: vector.reduction <maximumf>, %{{.*}} fastmath<nnan,nsz> : vector<1xf32> into f32
+// CHECK-COUNT-5: arith.maximumf %{{.*}}, %{{.*}} fastmath<nnan,nsz> : f32
+// CHECK: vector.reduction <maximumf>, %{{.*}} fastmath<nnan,nsz> : vector<1xf32> into f32
+// CHECK-COUNT-5: arith.maximumf %{{.*}}, %{{.*}} fastmath<nnan,nsz> : f32
+gpu.func @vector_multi_reduction_cross_lane_fastmath(%laneid: index) {
+ %src = "test.some_op"() : () -> vector<2x16xf32>
+ %acc = arith.constant dense<1.0> : vector<2xf32>
+ %1 = vector.multi_reduction <maximumf>, %src, %acc fastmath<nnan,nsz>
+ [1] : vector<2x16xf32> to vector<2xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>
+ }> : vector<2xf32>
+ gpu.return
+}
+
+// CHECK-LABEL: gpu.func @vector_multi_reduction_lane_local_fastmath
+// CHECK: vector.reduction <maximumf>, %{{.*}}, %{{.*}} fastmath<nnan,nsz> : vector<12xf32> into f32
+gpu.func @vector_multi_reduction_lane_local_fastmath(%laneid: index) {
+ %src = "test.some_op"() : () -> vector<16x12xf32>
+ %acc = arith.constant dense<1.0> : vector<16xf32>
+ %1 = vector.multi_reduction <maximumf>, %src, %acc fastmath<nnan,nsz>
+ [1] : vector<16x12xf32> to vector<16xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [1]>
+ }> : vector<16xf32>
+ gpu.return
+}
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir b/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
index 49a3f02dca58e2..3fa73d5dd74a74 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
@@ -341,5 +341,16 @@ gpu.module @test {
gpu.return
}
+
+//-----
+ // CHECK-LABEL: multi_reduction_2d_last_dim_fastmath
+ // CHECK-COUNT-4: arith.maximumf {{.*}} fastmath<nnan,nsz> : vector<16x16xf32>
+ // CHECK: vector.multi_reduction <maximumf>, {{.*}} fastmath<nnan,nsz> [1] : vector<16x16xf32> to vector<16xf32>
+ // CHECK-COUNT-4: arith.maximumf {{.*}} fastmath<nnan,nsz> : vector<16x16xf32>
+ // CHECK: vector.multi_reduction <maximumf>, {{.*}} fastmath<nnan,nsz> [1] : vector<16x16xf32> to vector<16xf32>
+ gpu.func @multi_reduction_2d_last_dim_fastmath(%src: vector<32x80xf32>, %acc: vector<32xf32>) -> vector<32xf32> {
+ %0 = vector.multi_reduction <maximumf>, %src, %acc fastmath<nnan,nsz> {layout_operand_0 = #xegpu.layout<inst_data = [16, 16]>} [1] : vector<32x80xf32> to vector<32xf32>
+ gpu.return %0 : vector<32xf32>
+ }
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index 0560c16591185c..2f5c037e07bb50 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -1574,3 +1574,183 @@ gpu.module @test_slice_layout {
gpu.return
}
}
+
+// -----
+
+// The fast-math flags are kept on the distributed reductions and on the
+// combine with the accumulator.
+gpu.module @test_distribution_fastmath {
+ // CHECK-LABEL: @vector_reduce_fastmath
+ // CHECK: %[[RED:.*]] = vector.multi_reduction <maximumf>, {{.*}}, {{.*}} fastmath<nnan,nsz> [1] : vector<16x64xf32> to vector<16xf32>
+ // CHECK: arith.maximumf %[[RED]], {{.*}} fastmath<nnan,nsz> : vector<16xf32>
+ gpu.func @vector_reduce_fastmath(%src: memref<256x64xf32>) {
+ %cst = arith.constant dense<1.0> : vector<256xf32>
+ %tdesc = xegpu.create_nd_tdesc %src : memref<256x64xf32>
+ -> !xegpu.tensor_desc<256x64xf32>
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>}>
+ : !xegpu.tensor_desc<256x64xf32>
+ -> vector<256x64xf32>
+ %reduce = vector.multi_reduction <maximumf>, %load, %cst fastmath<nnan,nsz> [1]
+ : vector<256x64xf32> to vector<256xf32>
+ %anchor = xegpu.convert_layout %reduce
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xf32>
+ gpu.return
+ }
+
+ // CHECK-LABEL: gpu.func @vector_reduce_cross_sg_fastmath
+ // CHECK: vector.multi_reduction <maximumf>, {{.*}}, {{.*}} fastmath<nnan,nsz> [0, 1] : vector<8x8xf32> to f32
+ // CHECK: gpu.barrier
+ // CHECK: %[[FINAL:.*]] = vector.multi_reduction <maximumf>, {{.*}}, {{.*}} fastmath<nnan,nsz> [0, 1] : vector<4x4xf32> to f32
+ // CHECK: arith.maximumf %[[FINAL]], {{.*}} fastmath<nnan,nsz> : f32
+ gpu.func @vector_reduce_cross_sg_fastmath(%src: memref<32x32xf32>) {
+ %cst = arith.constant 0.0 : f32
+ %tdesc = xegpu.create_nd_tdesc %src : memref<32x32xf32>
+ -> !xegpu.tensor_desc<32x32xf32>
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [8, 8]>}>
+ : !xegpu.tensor_desc<32x32xf32>
+ -> vector<32x32xf32>
+ %reduce = vector.multi_reduction <maximumf>, %load, %cst fastmath<nnan,nsz> [0, 1]
+ : vector<32x32xf32> to f32
+ %anchor = xegpu.convert_layout %reduce
+ <{
+ target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [4, 4], sg_data = [8, 8]>, dims = [0, 1]>
+ }> : f32
+ gpu.return
+ }
+
+ // The neutral accumulator of the local reductions must be a value that the
+ // flags allow: no NaN under `nnan`, no infinity under `ninf`.
+ // CHECK-LABEL: gpu.func @vector_reduce_fastmath_neutral
+ gpu.func @vector_reduce_fastmath_neutral(%src: memref<256x64xf32>) {
+ // CHECK: %[[LOAD:.*]] = xegpu.load_nd {{.*}} -> vector<16x64xf32>
+ // CHECK: %[[N0:.*]] = arith.constant dense<0xFF800000> : vector<16xf32>
+ // CHECK: vector.multi_reduction <maxnumf>, %[[LOAD]], %[[N0]] fastmath<nnan> [1]
+ // CHECK: %[[N1:.*]] = arith.constant dense<3.40282347E+38> : vector<16xf32>
+ // CHECK: vector.multi_reduction <minimumnumf>, %[[LOAD]], %[[N1]] fastmath<fast> [1]
+ // CHECK: %[[N2:.*]] = arith.constant dense<0x7FC00000> : vector<16xf32>
+ // CHECK: vector.multi_reduction <maxnumf>, %[[LOAD]], %[[N2]] fastmath<ninf> [1]
+ // CHECK: %[[N3:.*]] = arith.constant dense<-3.40282347E+38> : vector<16xf32>
+ // CHECK: vector.multi_reduction <maximumf>, %[[LOAD]], %[[N3]] fastmath<ninf> [1]
+ // CHECK: %[[N4:.*]] = arith.constant dense<0x7F800000> : vector<16xf32>
+ // CHECK: vector.multi_reduction <minimumf>, %[[LOAD]], %[[N4]] fastmath<nnan> [1]
+ %cst = arith.constant dense<1.0> : vector<256xf32>
+ %tdesc = xegpu.create_nd_tdesc %src : memref<256x64xf32>
+ -> !xegpu.tensor_desc<256x64xf32>
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>}>
+ : !xegpu.tensor_desc<256x64xf32>
+ -> vector<256x64xf32>
+ %r0 = vector.multi_reduction <maxnumf>, %load, %cst fastmath<nnan> [1]
+ : vector<256x64xf32> to vector<256xf32>
+ %a0 = xegpu.convert_layout %r0
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xf32>
+ %r1 = vector.multi_reduction <minimumnumf>, %load, %cst fastmath<fast> [1]
+ : vector<256x64xf32> to vector<256xf32>
+ %a1 = xegpu.convert_layout %r1
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xf32>
+ %r2 = vector.multi_reduction <maxnumf>, %load, %cst fastmath<ninf> [1]
+ : vector<256x64xf32> to vector<256xf32>
+ %a2 = xegpu.convert_layout %r2
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xf32>
+ %r3 = vector.multi_reduction <maximumf>, %load, %cst fastmath<ninf> [1]
+ : vector<256x64xf32> to vector<256xf32>
+ %a3 = xegpu.convert_layout %r3
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xf32>
+ %r4 = vector.multi_reduction <minimumf>, %load, %cst fastmath<nnan> [1]
+ : vector<256x64xf32> to vector<256xf32>
+ %a4 = xegpu.convert_layout %r4
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xf32>
+ gpu.return
+ }
+
+ // The neutral value must also be one the element type can represent. f8E4M3FN
+ // has no infinity, so under `nnan` the NaN-ignoring kinds use the lowest/highest
+ // finite value. Under `ninf` alone a quiet NaN is still allowed.
+ // CHECK-LABEL: gpu.func @vector_reduce_fastmath_neutral_no_inf
+ // CHECK: %[[N0:.*]] = arith.constant dense<-4.480000e+02> : vector<16xf8E4M3FN>
+ // CHECK: vector.multi_reduction <maxnumf>, %{{.*}}, %[[N0]] fastmath<nnan> [1] : vector<16x64xf8E4M3FN> to vector<16xf8E4M3FN>
+ // CHECK: %[[N1:.*]] = arith.constant dense<4.480000e+02> : vector<16xf8E4M3FN>
+ // CHECK: vector.multi_reduction <minnumf>, %{{.*}}, %[[N1]] fastmath<nnan> [1] : vector<16x64xf8E4M3FN> to vector<16xf8E4M3FN>
+ // CHECK: %[[N2:.*]] = arith.constant dense<0x7F> : vector<16xf8E4M3FN>
+ // CHECK: vector.multi_reduction <maxnumf>, %{{.*}}, %[[N2]] fastmath<ninf> [1] : vector<16x64xf8E4M3FN> to vector<16xf8E4M3FN>
+ gpu.func @vector_reduce_fastmath_neutral_no_inf(%src: memref<256x64xf8E4M3FN>) {
+ %cst = arith.constant dense<1.0> : vector<256xf8E4M3FN>
+ %tdesc = xegpu.create_nd_tdesc %src : memref<256x64xf8E4M3FN>
+ -> !xegpu.tensor_desc<256x64xf8E4M3FN>
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>}>
+ : !xegpu.tensor_desc<256x64xf8E4M3FN>
+ -> vector<256x64xf8E4M3FN>
+ %r0 = vector.multi_reduction <maxnumf>, %load, %cst fastmath<nnan> [1]
+ : vector<256x64xf8E4M3FN> to vector<256xf8E4M3FN>
+ %a0 = xegpu.convert_layout %r0
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xf8E4M3FN>
+ %r1 = vector.multi_reduction <minnumf>, %load, %cst fastmath<nnan> [1]
+ : vector<256x64xf8E4M3FN> to vector<256xf8E4M3FN>
+ %a1 = xegpu.convert_layout %r1
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xf8E4M3FN>
+ %r2 = vector.multi_reduction <maxnumf>, %load, %cst fastmath<ninf> [1]
+ : vector<256x64xf8E4M3FN> to vector<256xf8E4M3FN>
+ %a2 = xegpu.convert_layout %r2
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xf8E4M3FN>
+ gpu.return
+ }
+
+ // f4E2M1FN has neither NaN nor infinity: the lowest/highest finite value.
+ // CHECK-LABEL: gpu.func @vector_reduce_fastmath_neutral_finite_only
+ // CHECK: %[[N0:.*]] = arith.constant dense<-6.000000e+00> : vector<16xf4E2M1FN>
+ // CHECK: vector.multi_reduction <maxnumf>, %{{.*}}, %[[N0]] fastmath<fast> [1] : vector<16x64xf4E2M1FN> to vector<16xf4E2M1FN>
+ // CHECK: %[[N1:.*]] = arith.constant dense<6.000000e+00> : vector<16xf4E2M1FN>
+ // CHECK: vector.multi_reduction <minimumf>, %{{.*}}, %[[N1]] fastmath<fast> [1] : vector<16x64xf4E2M1FN> to vector<16xf4E2M1FN>
+ gpu.func @vector_reduce_fastmath_neutral_finite_only(%src: memref<256x64xf4E2M1FN>) {
+ %cst = arith.constant dense<1.0> : vector<256xf4E2M1FN>
+ %tdesc = xegpu.create_nd_tdesc %src : memref<256x64xf4E2M1FN>
+ -> !xegpu.tensor_desc<256x64xf4E2M1FN>
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>}>
+ : !xegpu.tensor_desc<256x64xf4E2M1FN>
+ -> vector<256x64xf4E2M1FN>
+ %r0 = vector.multi_reduction <maxnumf>, %load, %cst fastmath<fast> [1]
+ : vector<256x64xf4E2M1FN> to vector<256xf4E2M1FN>
+ %a0 = xegpu.convert_layout %r0
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xf4E2M1FN>
+ %r1 = vector.multi_reduction <minimumf>, %load, %cst fastmath<fast> [1]
+ : vector<256x64xf4E2M1FN> to vector<256xf4E2M1FN>
+ %a1 = xegpu.convert_layout %r1
+ <{target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
+ : vector<256xf4E2M1FN>
+ gpu.return
+ }
+
+ // Same for the neutral accumulator of the final cross-subgroup reduction.
+ // CHECK-LABEL: gpu.func @vector_reduce_cross_sg_fastmath_neutral
+ // CHECK: %[[ACC:.*]] = arith.constant 0.000000e+00 : f32
+ // CHECK: %[[LOCAL_N:.*]] = arith.constant 0xFF800000 : f32
+ // CHECK: vector.multi_reduction <maxnumf>, {{.*}}, %[[LOCAL_N]] fastmath<nnan> [0, 1] : vector<8x8xf32> to f32
+ // CHECK: gpu.barrier
+ // CHECK: %[[FINAL_N:.*]] = arith.constant 0xFF800000 : f32
+ // CHECK: %[[FINAL:.*]] = vector.multi_reduction <maxnumf>, {{.*}}, %[[FINAL_N]] fastmath<nnan> [0, 1] : vector<4x4xf32> to f32
+ // CHECK: arith.maxnumf %[[FINAL]], %[[ACC]] fastmath<nnan> : f32
+ gpu.func @vector_reduce_cross_sg_fastmath_neutral(%src: memref<32x32xf32>) {
+ %cst = arith.constant 0.0 : f32
+ %tdesc = xegpu.create_nd_tdesc %src : memref<32x32xf32>
+ -> !xegpu.tensor_desc<32x32xf32>
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [8, 8]>}>
+ : !xegpu.tensor_desc<32x32xf32>
+ -> vector<32x32xf32>
+ %reduce = vector.multi_reduction <maxnumf>, %load, %cst fastmath<nnan> [0, 1]
+ : vector<32x32xf32> to f32
+ %anchor = xegpu.convert_layout %reduce
+ <{
+ target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [4, 4], sg_data = [8, 8]>, dims = [0, 1]>
+ }> : f32
+ gpu.return
+ }
+}
More information about the Mlir-commits
mailing list