[Mlir-commits] [mlir] [mlir][math] Propagate fast math attrs in AlgebraicSimplification (PR #166802)

[Aleksei Nurmukhametov]

https://github.com/nurmukhametov created https://github.com/llvm/llvm-project/pull/166802

Fix missing propagation of fast-math flags in algebraic simplification patterns of the MLIR math dialect.

>From 5378706533f3202a566d373c7c63410140ee9872 Mon Sep 17 00:00:00 2001
From: Aleksei Nurmukhametov <anurmukh at amd.com>
Date: Thu, 6 Nov 2025 16:44:12 +0000
Subject: [PATCH 1/2] [NFC][mlir][math] Minor code cleanup in
 AlgebraicSimplification

---
 .../Dialect/Math/Transforms/AlgebraicSimplification.cpp  | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp b/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
index 77b10cec48d8e..7fb26f487e3af 100644
--- a/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
@@ -66,7 +66,7 @@ PowFStrengthReduction::matchAndRewrite(math::PowFOp op,
   // Maybe broadcasts scalar value into vector type compatible with `op`.
   auto bcast = [&](Value value) -> Value {
     if (auto vec = dyn_cast<VectorType>(op.getType()))
-      return vector::BroadcastOp::create(rewriter, op.getLoc(), vec, value);
+      return vector::BroadcastOp::create(rewriter, loc, vec, value);
     return value;
   };
 
@@ -84,8 +84,7 @@ PowFStrengthReduction::matchAndRewrite(math::PowFOp op,
 
   // Replace `pow(x, 3.0)` with `x * x * x`.
   if (isExponentValue(3.0)) {
-    Value square =
-        arith::MulFOp::create(rewriter, op.getLoc(), ValueRange({x, x}));
+    Value square = arith::MulFOp::create(rewriter, loc, ValueRange({x, x}));
     rewriter.replaceOpWithNewOp<arith::MulFOp>(op, ValueRange({x, square}));
     return success();
   }
@@ -113,8 +112,8 @@ PowFStrengthReduction::matchAndRewrite(math::PowFOp op,
 
   // Replace `pow(x, 0.75)` with `sqrt(sqrt(x)) * sqrt(x)`.
   if (isExponentValue(0.75)) {
-    Value powHalf = math::SqrtOp::create(rewriter, op.getLoc(), x);
-    Value powQuarter = math::SqrtOp::create(rewriter, op.getLoc(), powHalf);
+    Value powHalf = math::SqrtOp::create(rewriter, loc, x);
+    Value powQuarter = math::SqrtOp::create(rewriter, loc, powHalf);
     rewriter.replaceOpWithNewOp<arith::MulFOp>(op,
                                                ValueRange{powHalf, powQuarter});
     return success();

>From 4b463cec6ccaecae403b1aa7b7ac7cfb345d8169 Mon Sep 17 00:00:00 2001
From: Aleksei Nurmukhametov <anurmukh at amd.com>
Date: Thu, 6 Nov 2025 16:51:04 +0000
Subject: [PATCH 2/2] [mlir][math] Propagate fast math attrs in
 AlgebraicSimplification

---
 .../Transforms/AlgebraicSimplification.cpp    | 20 +++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp b/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
index 7fb26f487e3af..677d7505662a0 100644
--- a/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
@@ -43,6 +43,7 @@ PowFStrengthReduction::matchAndRewrite(math::PowFOp op,
                                        PatternRewriter &rewriter) const {
   Location loc = op.getLoc();
   Value x = op.getLhs();
+  auto fmf = op.getFastmathAttr().getValue();
 
   FloatAttr scalarExponent;
   DenseFPElementsAttr vectorExponent;
@@ -78,14 +79,14 @@ PowFStrengthReduction::matchAndRewrite(math::PowFOp op,
 
   // Replace `pow(x, 2.0)` with `x * x`.
   if (isExponentValue(2.0)) {
-    rewriter.replaceOpWithNewOp<arith::MulFOp>(op, ValueRange({x, x}));
+    rewriter.replaceOpWithNewOp<arith::MulFOp>(op, x, x, fmf);
     return success();
   }
 
   // Replace `pow(x, 3.0)` with `x * x * x`.
   if (isExponentValue(3.0)) {
-    Value square = arith::MulFOp::create(rewriter, loc, ValueRange({x, x}));
-    rewriter.replaceOpWithNewOp<arith::MulFOp>(op, ValueRange({x, square}));
+    Value square = arith::MulFOp::create(rewriter, loc, x, x, fmf);
+    rewriter.replaceOpWithNewOp<arith::MulFOp>(op, x, square, fmf);
     return success();
   }
 
@@ -94,28 +95,27 @@ PowFStrengthReduction::matchAndRewrite(math::PowFOp op,
     Value one = arith::ConstantOp::create(
         rewriter, loc,
         rewriter.getFloatAttr(getElementTypeOrSelf(op.getType()), 1.0));
-    rewriter.replaceOpWithNewOp<arith::DivFOp>(op, ValueRange({bcast(one), x}));
+    rewriter.replaceOpWithNewOp<arith::DivFOp>(op, bcast(one), x, fmf);
     return success();
   }
 
   // Replace `pow(x, 0.5)` with `sqrt(x)`.
   if (isExponentValue(0.5)) {
-    rewriter.replaceOpWithNewOp<math::SqrtOp>(op, x);
+    rewriter.replaceOpWithNewOp<math::SqrtOp>(op, x, fmf);
     return success();
   }
 
   // Replace `pow(x, -0.5)` with `rsqrt(x)`.
   if (isExponentValue(-0.5)) {
-    rewriter.replaceOpWithNewOp<math::RsqrtOp>(op, x);
+    rewriter.replaceOpWithNewOp<math::RsqrtOp>(op, x, fmf);
     return success();
   }
 
   // Replace `pow(x, 0.75)` with `sqrt(sqrt(x)) * sqrt(x)`.
   if (isExponentValue(0.75)) {
-    Value powHalf = math::SqrtOp::create(rewriter, loc, x);
-    Value powQuarter = math::SqrtOp::create(rewriter, loc, powHalf);
-    rewriter.replaceOpWithNewOp<arith::MulFOp>(op,
-                                               ValueRange{powHalf, powQuarter});
+    Value powHalf = math::SqrtOp::create(rewriter, loc, x, fmf);
+    Value powQuarter = math::SqrtOp::create(rewriter, loc, powHalf, fmf);
+    rewriter.replaceOpWithNewOp<arith::MulFOp>(op, powHalf, powQuarter, fmf);
     return success();
   }