[Mlir-commits] [mlir] [mlir] [arith] Fold trunc(extremum(ext(lhs), ext(rhs))) (PR #214658)
Chuanqi Xu
llvmlistbot at llvm.org
Fri Aug 7 01:08:11 PDT 2026
https://github.com/ChuanqiXu9 created https://github.com/llvm/llvm-project/pull/214658
For example, now we can't fold the following case:
```
%lhs = arith.extf %arg0 : f16 to f32
%rhs = arith.extf %arg1 : f16 to f32
%maximum = arith.maximumf %lhs, %rhs fastmath<nnan> : f32
%maximumTrunc = arith.truncf %maximum : f32 to f16
```
to
```
%maximum = arith.maximumf %arg0, %arg1 fastmath<nnan> : f16
```
And this patch made this optimization. We can extend the op list to arith.select, arith.cmpf and so on in the future.
I was told this may be intentional as some old hardware doesn't have instructions to deal with maximum with fp16. But I think in the higher level we should do such optimizations. And specific hardware can make their own transformations according to their specific details.
>From 8a0d589e84c92c98823d84a47946ebb8b09d0117 Mon Sep 17 00:00:00 2001
From: Chuanqi Xu <chuanqi.xcq at alibaba-inc.com>
Date: Fri, 7 Aug 2026 15:54:51 +0800
Subject: [PATCH] [mlir] [arith] Fold trunc(extremum(ext(lhs), ext(rhs)))
For example, now we can't fold the following case:
```
%lhs = arith.extf %arg0 : f16 to f32
%rhs = arith.extf %arg1 : f16 to f32
%maximum = arith.maximumf %lhs, %rhs fastmath<nnan> : f32
%maximumTrunc = arith.truncf %maximum : f32 to f16
```
to
```
%maximum = arith.maximumf %arg0, %arg1 fastmath<nnan> : f16
```
And this patch made this optimization. We can extend the op list to
arith.select, arith.cmpf and so on in the future.
I was told this may be intentional as some old hardware doesn't have
instructions to deal with maximum with fp16. But I think in the higher
level we should do such optimizations. And specific hardware can
make their own transformations according to their specific details.
---
mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 69 +++++++++++++-
mlir/test/Dialect/Arith/canonicalize.mlir | 111 ++++++++++++++++++++++
2 files changed, 176 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index ff6a5d4a0c29a..eadce07b93a3a 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -1364,6 +1364,61 @@ void arith::SubFOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
patterns.add<SubFOfNegZero>(context);
}
+namespace {
+
+/// Narrow an extremum whose operands were extended from the result type:
+///
+/// trunc(extremum(ext(lhs), ext(rhs))) -> extremum(lhs, rhs)
+///
+/// The concrete extension is part of the pattern so each extremum is only
+/// registered with extensions that preserve its ordering.
+/// For floating-point types, also require the extension to represent every
+/// source value exactly.
+template <typename TruncOp, typename ExtOp, typename ExtremumOp>
+struct NarrowExtremum final : OpRewritePattern<TruncOp> {
+ using OpRewritePattern<TruncOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(TruncOp truncOp,
+ PatternRewriter &rewriter) const override {
+ auto extremumOp = truncOp.getIn().template getDefiningOp<ExtremumOp>();
+ if (!extremumOp || !extremumOp->hasOneUse())
+ return failure();
+
+ auto lhsExt = extremumOp.getLhs().template getDefiningOp<ExtOp>();
+ auto rhsExt = extremumOp.getRhs().template getDefiningOp<ExtOp>();
+ if (!lhsExt || !rhsExt)
+ return failure();
+
+ Value lhs = lhsExt.getIn();
+ Value rhs = rhsExt.getIn();
+ Type narrowType = truncOp.getType();
+ if (lhs.getType() != narrowType || rhs.getType() != narrowType)
+ return failure();
+
+ // A floating-point extension is not necessarily lossless between arbitrary
+ // floating-point semantics, even when the destination has a larger bit
+ // width. Require every narrow value to be exactly representable in the
+ // wide type so moving the extremum before the extension cannot change
+ // which value is selected.
+ if (auto narrowFloatType =
+ dyn_cast<FloatType>(getElementTypeOrSelf(narrowType))) {
+ auto wideFloatType =
+ dyn_cast<FloatType>(getElementTypeOrSelf(extremumOp.getType()));
+ if (!wideFloatType || !llvm::APFloatBase::isRepresentableBy(
+ narrowFloatType.getFloatSemantics(),
+ wideFloatType.getFloatSemantics()))
+ return failure();
+ }
+
+ rewriter.replaceOpWithNewOp<ExtremumOp>(truncOp, TypeRange{narrowType},
+ ValueRange{lhs, rhs},
+ extremumOp->getAttrs());
+ return success();
+ }
+};
+
+} // namespace
+
//===----------------------------------------------------------------------===//
// MaximumFOp
//===----------------------------------------------------------------------===//
@@ -1876,9 +1931,11 @@ bool arith::TruncIOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
void arith::TruncIOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
MLIRContext *context) {
- patterns
- .add<TruncIExtSIToExtSI, TruncIExtUIToExtUI, TruncIShrSIToTrunciShrUI>(
- context);
+ patterns.add<NarrowExtremum<TruncIOp, ExtSIOp, MaxSIOp>,
+ NarrowExtremum<TruncIOp, ExtSIOp, MinSIOp>,
+ NarrowExtremum<TruncIOp, ExtUIOp, MaxUIOp>,
+ NarrowExtremum<TruncIOp, ExtUIOp, MinUIOp>, TruncIExtSIToExtSI,
+ TruncIExtUIToExtUI, TruncIShrSIToTrunciShrUI>(context);
}
LogicalResult arith::TruncIOp::verify() {
@@ -1932,7 +1989,11 @@ OpFoldResult arith::TruncFOp::fold(FoldAdaptor adaptor) {
void arith::TruncFOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
MLIRContext *context) {
- patterns.add<TruncFSIToFPToSIToFP, TruncFUIToFPToUIToFP>(context);
+ patterns.add<NarrowExtremum<TruncFOp, ExtFOp, MaximumFOp>,
+ NarrowExtremum<TruncFOp, ExtFOp, MaxNumFOp>,
+ NarrowExtremum<TruncFOp, ExtFOp, MinimumFOp>,
+ NarrowExtremum<TruncFOp, ExtFOp, MinNumFOp>,
+ TruncFSIToFPToSIToFP, TruncFUIToFPToUIToFP>(context);
}
bool arith::TruncFOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 0c06aa6e861a4..8acd3167a9bd3 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -1159,6 +1159,117 @@ func.func @truncExtf3(%arg0: f32) -> f16 {
return %truncf : f16
}
+// CHECK-LABEL: @narrowExtremaOfExtf
+// CHECK-NOT: arith.extf
+// CHECK-NOT: arith.truncf
+// CHECK: %[[MAXIMUM:.*]] = arith.maximumf %arg0, %arg1 fastmath<nnan> : f16
+// CHECK: %[[MAXNUM:.*]] = arith.maxnumf %arg0, %arg1 : f16
+// CHECK: %[[MINIMUM:.*]] = arith.minimumf %arg0, %arg1 : f16
+// CHECK: %[[MINNUM:.*]] = arith.minnumf %arg0, %arg1 : f16
+// CHECK: return %[[MAXIMUM]], %[[MAXNUM]], %[[MINIMUM]], %[[MINNUM]]
+func.func @narrowExtremaOfExtf(%arg0: f16, %arg1: f16)
+ -> (f16, f16, f16, f16) {
+ %lhs = arith.extf %arg0 : f16 to f32
+ %rhs = arith.extf %arg1 : f16 to f32
+ %maximum = arith.maximumf %lhs, %rhs fastmath<nnan> : f32
+ %maxnum = arith.maxnumf %lhs, %rhs : f32
+ %minimum = arith.minimumf %lhs, %rhs : f32
+ %minnum = arith.minnumf %lhs, %rhs : f32
+ %maximumTrunc = arith.truncf %maximum : f32 to f16
+ %maxnumTrunc = arith.truncf %maxnum : f32 to f16
+ %minimumTrunc = arith.truncf %minimum : f32 to f16
+ %minnumTrunc = arith.truncf %minnum : f32 to f16
+ return %maximumTrunc, %maxnumTrunc, %minimumTrunc, %minnumTrunc
+ : f16, f16, f16, f16
+}
+
+// CHECK-LABEL: @narrowIntegerExtrema
+// CHECK-NOT: arith.extsi
+// CHECK-NOT: arith.extui
+// CHECK-NOT: arith.trunci
+// CHECK: %[[SMAX:.*]] = arith.maxsi %arg0, %arg1 : i8
+// CHECK: %[[SMIN:.*]] = arith.minsi %arg0, %arg1 : i8
+// CHECK: %[[UMAX:.*]] = arith.maxui %arg0, %arg1 : i8
+// CHECK: %[[UMIN:.*]] = arith.minui %arg0, %arg1 : i8
+// CHECK: return %[[SMAX]], %[[SMIN]], %[[UMAX]], %[[UMIN]]
+func.func @narrowIntegerExtrema(%arg0: i8, %arg1: i8)
+ -> (i8, i8, i8, i8) {
+ %slhs = arith.extsi %arg0 : i8 to i32
+ %srhs = arith.extsi %arg1 : i8 to i32
+ %ulhs = arith.extui %arg0 : i8 to i32
+ %urhs = arith.extui %arg1 : i8 to i32
+ %smax = arith.maxsi %slhs, %srhs : i32
+ %smin = arith.minsi %slhs, %srhs : i32
+ %umax = arith.maxui %ulhs, %urhs : i32
+ %umin = arith.minui %ulhs, %urhs : i32
+ %smaxTrunc = arith.trunci %smax : i32 to i8
+ %sminTrunc = arith.trunci %smin : i32 to i8
+ %umaxTrunc = arith.trunci %umax : i32 to i8
+ %uminTrunc = arith.trunci %umin : i32 to i8
+ return %smaxTrunc, %sminTrunc, %umaxTrunc, %uminTrunc : i8, i8, i8, i8
+}
+
+// CHECK-LABEL: @doNotNarrowIntegerExtremumWithWideUse
+// CHECK: %[[LHS:.*]] = arith.extsi %arg0 : i8 to i32
+// CHECK: %[[RHS:.*]] = arith.extsi %arg1 : i8 to i32
+// CHECK: %[[MAX:.*]] = arith.maxsi %[[LHS]], %[[RHS]] : i32
+// CHECK: %[[TRUNC:.*]] = arith.trunci %[[MAX]] : i32 to i8
+// CHECK: return %[[TRUNC]], %[[MAX]]
+func.func @doNotNarrowIntegerExtremumWithWideUse(%arg0: i8, %arg1: i8)
+ -> (i8, i32) {
+ %lhs = arith.extsi %arg0 : i8 to i32
+ %rhs = arith.extsi %arg1 : i8 to i32
+ %max = arith.maxsi %lhs, %rhs : i32
+ %trunc = arith.trunci %max : i32 to i8
+ return %trunc, %max : i8, i32
+}
+
+// Zero extension does not preserve signed ordering across the sign boundary.
+// CHECK-LABEL: @doNotNarrowSignedExtremumOfZeroExtension
+// CHECK: %[[LHS:.*]] = arith.extui %arg0 : i8 to i32
+// CHECK: %[[RHS:.*]] = arith.extui %arg1 : i8 to i32
+// CHECK: %[[MAX:.*]] = arith.maxsi %[[LHS]], %[[RHS]] : i32
+// CHECK: %[[TRUNC:.*]] = arith.trunci %[[MAX]] : i32 to i8
+// CHECK: return %[[TRUNC]]
+func.func @doNotNarrowSignedExtremumOfZeroExtension(
+ %arg0: i8, %arg1: i8) -> i8 {
+ %lhs = arith.extui %arg0 : i8 to i32
+ %rhs = arith.extui %arg1 : i8 to i32
+ %max = arith.maxsi %lhs, %rhs : i32
+ %trunc = arith.trunci %max : i32 to i8
+ return %trunc : i8
+}
+
+// CHECK-LABEL: @doNotNarrowExtremumWithWideUse
+// CHECK: %[[LHS:.*]] = arith.extf %arg0 : f16 to f32
+// CHECK: %[[RHS:.*]] = arith.extf %arg1 : f16 to f32
+// CHECK: %[[MAXIMUM:.*]] = arith.maximumf %[[LHS]], %[[RHS]] : f32
+// CHECK: %[[TRUNC:.*]] = arith.truncf %[[MAXIMUM]] : f32 to f16
+// CHECK: return %[[TRUNC]], %[[MAXIMUM]]
+func.func @doNotNarrowExtremumWithWideUse(%arg0: f16, %arg1: f16)
+ -> (f16, f32) {
+ %lhs = arith.extf %arg0 : f16 to f32
+ %rhs = arith.extf %arg1 : f16 to f32
+ %maximum = arith.maximumf %lhs, %rhs : f32
+ %trunc = arith.truncf %maximum : f32 to f16
+ return %trunc, %maximum : f16, f32
+}
+
+// CHECK-LABEL: @doNotNarrowExtremumThroughInexactExtension
+// CHECK: %[[LHS:.*]] = arith.extf %arg0 : f8E8M0FNU to f16
+// CHECK: %[[RHS:.*]] = arith.extf %arg1 : f8E8M0FNU to f16
+// CHECK: %[[MAXIMUM:.*]] = arith.maximumf %[[LHS]], %[[RHS]] : f16
+// CHECK: %[[TRUNC:.*]] = arith.truncf %[[MAXIMUM]] : f16 to f8E8M0FNU
+// CHECK: return %[[TRUNC]]
+func.func @doNotNarrowExtremumThroughInexactExtension(
+ %arg0: f8E8M0FNU, %arg1: f8E8M0FNU) -> f8E8M0FNU {
+ %lhs = arith.extf %arg0 : f8E8M0FNU to f16
+ %rhs = arith.extf %arg1 : f8E8M0FNU to f16
+ %maximum = arith.maximumf %lhs, %rhs : f16
+ %trunc = arith.truncf %maximum : f16 to f8E8M0FNU
+ return %trunc : f8E8M0FNU
+}
+
// CHECK-LABEL: @truncSitofp
// CHECK: %[[SITOFP:.*]] = arith.sitofp %[[ARG0:.*]] : i32 to f32
// CHECK-NOT: truncf
More information about the Mlir-commits
mailing list