[Mlir-commits] [mlir] [mlir][math] Fold FPowIOp with square-and-multiply to match powi expansion (PR #210982)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-math

Author: Alexey Bataev (alexey-bataev)

<details>
<summary>Changes</summary>

Fold math.fpowi using the base's own floating-point semantics via iterative
square-and-multiply, matching the multiply sequence ExpandPowI builds in
SelectionDAGBuilder. The prior pow/powf-based fold computed in float/double
and rounded differently from the runtime expansion, disagreeing with x**n by ~1 ULP.


---
Full diff: https://github.com/llvm/llvm-project/pull/210982.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Math/IR/MathOps.cpp (+18-12) 
- (modified) mlir/test/Dialect/Math/canonicalize.mlir (+16-4) 


``````````diff
diff --git a/mlir/lib/Dialect/Math/IR/MathOps.cpp b/mlir/lib/Dialect/Math/IR/MathOps.cpp
index b900cb1911759..2fd081282c2aa 100644
--- a/mlir/lib/Dialect/Math/IR/MathOps.cpp
+++ b/mlir/lib/Dialect/Math/IR/MathOps.cpp
@@ -831,23 +831,29 @@ OpFoldResult math::FPowIOp::fold(FoldAdaptor adaptor) {
   return constFoldBinaryOpConditional<FloatAttr, IntegerAttr>(
       adaptor.getOperands(),
       [](const APFloat &base, const APInt &exp) -> std::optional<APFloat> {
-        const llvm::fltSemantics &sem = base.getSemantics();
-        // Fold when the exponent is exactly representable in the
-        // floating-point type of the base.
-        APFloat fExp(sem);
-        if (fExp.convertFromAPInt(exp, /*isSigned=*/true,
-                                  APFloat::rmNearestTiesToEven) !=
-            APFloat::opOK)
-          return {};
-
-        switch (APFloat::SemanticsToEnum(sem)) {
+        switch (APFloat::SemanticsToEnum(base.getSemantics())) {
         case APFloat::Semantics::S_IEEEdouble:
-          return APFloat(pow(base.convertToDouble(), fExp.convertToDouble()));
         case APFloat::Semantics::S_IEEEsingle:
-          return APFloat(powf(base.convertToFloat(), fExp.convertToFloat()));
+          break;
         default:
           return {};
         }
+
+        // Square-and-multiply using the base's own semantics, matching the
+        // multiply sequence ExpandPowI builds in SelectionDAGBuilder.cpp.
+        const llvm::fltSemantics &sem = base.getSemantics();
+        APInt magnitude = exp.abs();
+        APFloat res = APFloat::getOne(sem);
+        APFloat curSquare = base;
+        while (!magnitude.isZero()) {
+          if (magnitude[0])
+            res = res * curSquare;
+          curSquare = curSquare * curSquare;
+          magnitude.lshrInPlace(1);
+        }
+        if (exp.isNegative())
+          res = APFloat::getOne(sem) / res;
+        return res;
       });
 }
 
diff --git a/mlir/test/Dialect/Math/canonicalize.mlir b/mlir/test/Dialect/Math/canonicalize.mlir
index 3459164c5c0a7..a73a15f34e577 100644
--- a/mlir/test/Dialect/Math/canonicalize.mlir
+++ b/mlir/test/Dialect/Math/canonicalize.mlir
@@ -638,10 +638,22 @@ func.func @fpowi_fold_vec() -> vector<4xf32> {
   return %0 : vector<4xf32>
 }
 
-// 16777217 is not exactly representable in f32.
-// CHECK-LABEL: @fpowi_fold_failed
-// CHECK:       math.fpowi
-func.func @fpowi_fold_failed() -> f32 {
+// CHECK-LABEL: @fpowi_fold_const
+// CHECK: %[[cst:.+]] = arith.constant 29.1229134 : f32
+// CHECK: return %[[cst]]
+func.func @fpowi_fold_const() -> f32 {
+  %cst = arith.constant 1.234567 : f32
+  %c16_i32 = arith.constant 16 : i32
+  %0 = math.fpowi %cst, %c16_i32 : f32, i32
+  return %0 : f32
+}
+
+// The fold no longer requires the exponent to be exactly representable as
+// f32: it is used as an integer, so this folds and overflows to +inf.
+// CHECK-LABEL: @fpowi_fold_overflow
+// CHECK: %[[cst:.+]] = arith.constant 0x7F800000 : f32
+// CHECK: return %[[cst]]
+func.func @fpowi_fold_overflow() -> f32 {
   %cst = arith.constant 2.000000e+00 : f32
   %c16777217_i32 = arith.constant 16777217 : i32
   %0 = math.fpowi %cst, %c16777217_i32 : f32, i32

``````````

</details>


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