[Mlir-commits] [mlir] [mlir] [arith] Fold trunc(extremum(ext(lhs), ext(rhs))) (PR #214658)

[Jakub Kuderski]

================
@@ -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(
----------------
kuhar wrote:

Could you also take care of the truncf(extf(x)) fold, in a separate PR?

https://github.com/llvm/llvm-project/pull/214658