[Mlir-commits] [mlir] [mlir] [arith] Fold trunc(extremum(ext(lhs), ext(rhs))) (PR #214658)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Aug 7 01:08:57 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Chuanqi Xu (ChuanqiXu9)
<details>
<summary>Changes</summary>
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.
---
Full diff: https://github.com/llvm/llvm-project/pull/214658.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Arith/IR/ArithOps.cpp (+65-4)
- (modified) mlir/test/Dialect/Arith/canonicalize.mlir (+111)
``````````diff
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
``````````
</details>
https://github.com/llvm/llvm-project/pull/214658
More information about the Mlir-commits
mailing list