polynomial approx: implement square and multiply for math.fpowi

PiperOrigin-RevId: 961260919

Author: j2kun

lib/Transforms/PolynomialApproximation/PolynomialApproximation.cpp

@@ -21,6 +21,7 @@
 #include "mlir/include/mlir/Analysis/DataFlow/Utils.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"   // from @llvm-project
 #include "mlir/include/mlir/Dialect/Math/IR/Math.h"     // from @llvm-project
+#include "mlir/include/mlir/IR/Builders.h"              // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinAttributes.h"      // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
@@ -628,6 +629,47 @@ struct ReluViaCompositeSign : public OpRewritePattern<arith::MaximumFOp> {
   DataFlowSolver* solver;
 };
 
+// Use a square and multiply algorithm for x^n where n is a constant.
+struct SquareAndMultiplyForPowOp : public OpRewritePattern<math::FPowIOp> {
+  SquareAndMultiplyForPowOp(MLIRContext* context)
+      : OpRewritePattern<math::FPowIOp>(context, /*benefit=*/2) {}
+
+  LogicalResult matchAndRewrite(math::FPowIOp op,
+                                PatternRewriter& rewriter) const override {
+    ImplicitLocOpBuilder b(op.getLoc(), rewriter);
+    Value base = op.getLhs();
+    Value exp = op.getRhs();
+
+    APInt expVal;
+    if (!matchPattern(exp, m_ConstantInt(&expVal))) {
+      return rewriter.notifyMatchFailure(
+          op, "exponent is not a single-valued constant");
+    }
+    if (expVal.isNegative()) {
+      return op.emitOpError("negative exponent not supported");
+    }
+
+    int64_t expInt = static_cast<int64_t>(expVal.getSExtValue());
+    if (expInt == 0) {
+      rewriter.replaceOp(
+          op, arith::ConstantOp::create(b, b.getOneAttr(base.getType())));
+      return success();
+    }
+
+    auto res = base;
+    int highestBit = expVal.getActiveBits() - 1;
+    for (int i = highestBit - 1; i >= 0; --i) {
+      res = arith::MulFOp::create(b, res, res);
+      if ((expInt >> i) & 1) {
+        res = arith::MulFOp::create(b, res, base);
+      }
+    }
+
+    rewriter.replaceOp(op, res);
+    return success();
+  }
+};
+
 struct PolynomialApproximation
     : impl::PolynomialApproximationBase<PolynomialApproximation> {
   using PolynomialApproximationBase::PolynomialApproximationBase;
@@ -647,6 +689,7 @@ struct PolynomialApproximation
 
     // High priority patterns
     patterns.add<ExpOpTaylorApproximation>(context, &solver, /*k=*/7);
+    patterns.add<SquareAndMultiplyForPowOp>(context);
     if (useCompositeRelu) {
       patterns.add<ReluViaCompositeSign>(context, &solver);
     }

tests/Transforms/polynomial_approximation/fpowi.mlir

@@ -0,0 +1,45 @@
+// RUN: heir-opt --split-input-file --polynomial-approximation %s | FileCheck %s
+
+// CHECK: @test_fpowi_four
+func.func @test_fpowi_four(%x: tensor<4xf32>) -> tensor<4xf32> {
+  // CHECK-COUNT-2: arith.mulf
+  // CHECK: return
+  %c2 = arith.constant dense<4> : tensor<4xi64>
+  %0 = math.fpowi %x, %c2 : tensor<4xf32>, tensor<4xi64>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK: @test_fpowi_zero
+func.func @test_fpowi_zero(%x: tensor<4xf32>) -> tensor<4xf32> {
+  // CHECK: arith.constant dense<1
+  %c0 = arith.constant dense<0> : tensor<4xi64>
+  %0 = math.fpowi %x, %c0 : tensor<4xf32>, tensor<4xi64>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK: @test_fpowi_one
+// CHECK-SAME: %[[arg:.*]]: tensor<4xf32>
+func.func @test_fpowi_one(%x: tensor<4xf32>) -> tensor<4xf32> {
+  // CHECK: return %[[arg]]
+  %c1 = arith.constant dense<1> : tensor<4xi64>
+  %0 = math.fpowi %x, %c1 : tensor<4xf32>, tensor<4xi64>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK: @test_fpowi_five
+// CHECK-SAME: %[[arg:.*]]: tensor<4xf32>
+func.func @test_fpowi_five(%x: tensor<4xf32>) -> tensor<4xf32> {
+  // CHECK: %[[arg1:.*]] = arith.mulf %[[arg]], %[[arg]]
+  // CHECK: %[[arg2:.*]] = arith.mulf %[[arg1]], %[[arg1]]
+  // CHECK: %[[arg3:.*]] = arith.mulf %[[arg2]], %[[arg]]
+  // CHECK: return %[[arg3]]
+  %c1 = arith.constant dense<5> : tensor<4xi64>
+  %0 = math.fpowi %x, %c1 : tensor<4xf32>, tensor<4xi64>
+  return %0 : tensor<4xf32>
+}

tests/Transforms/polynomial_approximation/polynomial_approximation.mlir

@@ -115,8 +115,9 @@ func.func @test_sqrt_default_params(%x: f32 {secret.secret}) -> f32 {
 
 // CHECK: @test_fpowi_tensor
 func.func @test_fpowi_tensor(%x: tensor<1x5xf32> {secret.secret}) -> tensor<1x5xf32> {
-  // CHECK: polynomial.eval
+  // CHECK: arith.mulf
   // CHECK-NOT: math.fpowi
+  // CHECK-NOT: polynomial.eval
   %cst = arith.constant dense<2> : tensor<1x5xi64>
   %0 = math.fpowi %x, %cst : tensor<1x5xf32>, tensor<1x5xi64>
   return %0 : tensor<1x5xf32>
@@ -126,8 +127,9 @@ func.func @test_fpowi_tensor(%x: tensor<1x5xf32> {secret.secret}) -> tensor<1x5x
 
 // CHECK: @test_fpowi_scalar
 func.func @test_fpowi_scalar(%x: f32 {secret.secret}) -> f32 {
-  // CHECK: polynomial.eval
+  // CHECK: arith.mulf
   // CHECK-NOT: math.fpowi
+  // CHECK-NOT: polynomial.eval
   %cst = arith.constant 2 : i32
   %0 = math.fpowi %x, %cst : f32, i32
   return %0 : f32