Add precision tests for Taylor exponentiation approximation.

lib/Transforms/PolynomialApproximation/PolynomialApproximation.cpp

@@ -16,13 +16,15 @@
 #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/BuiltinAttributeInterfaces.h"  // from @llvm-project
-#include "mlir/include/mlir/IR/BuiltinAttributes.h"  // from @llvm-project
-#include "mlir/include/mlir/IR/BuiltinTypes.h"       // from @llvm-project
-#include "mlir/include/mlir/IR/MLIRContext.h"        // from @llvm-project
-#include "mlir/include/mlir/IR/Matchers.h"           // from @llvm-project
-#include "mlir/include/mlir/IR/PatternMatch.h"       // from @llvm-project
-#include "mlir/include/mlir/IR/Value.h"              // from @llvm-project
-#include "mlir/include/mlir/Support/LLVM.h"          // from @llvm-project
+#include "mlir/include/mlir/IR/BuiltinAttributes.h"      // from @llvm-project
+#include "mlir/include/mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/include/mlir/IR/BuiltinTypes.h"           // from @llvm-project
+#include "mlir/include/mlir/IR/MLIRContext.h"            // from @llvm-project
+#include "mlir/include/mlir/IR/Matchers.h"               // from @llvm-project
+#include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
+#include "mlir/include/mlir/IR/Types.h"                  // from @llvm-project
+#include "mlir/include/mlir/IR/Value.h"                  // from @llvm-project
+#include "mlir/include/mlir/Support/LLVM.h"              // from @llvm-project
 #include "mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
 
 // IWYU pragma: begin_keep
@@ -389,6 +391,75 @@ struct ConvertBinaryConstOp : public OpRewritePattern<OpTy> {
   double upper;
 };
 
+struct ExpOpTaylorApproximation : public OpRewritePattern<math::ExpOp> {
+  ExpOpTaylorApproximation(MLIRContext* context, int64_t k = 7)
+      : OpRewritePattern<math::ExpOp>(context, /*benefit=*/2), k(k) {}
+
+  LogicalResult matchAndRewrite(math::ExpOp op,
+                                PatternRewriter& rewriter) const override {
+    Location loc = op.getLoc();
+    Value operand = op.getOperand();
+    Type type = operand.getType();
+
+    double validLower = -static_cast<double>(1ULL << k);
+    double validUpper = 1.0;
+
+    if (op->hasAttr("domain_lower")) {
+      FloatAttr lowerAttr = dyn_cast<FloatAttr>(op->getAttr("domain_lower"));
+      if (!lowerAttr)
+        return op.emitOpError(
+            "domain_lower must be a floating-point attribute");
+      if (lowerAttr.getValueAsDouble() < validLower) {
+        return rewriter.notifyMatchFailure(
+            op, "domain_lower is less than valid interval bound -2^k");
+      }
+    }
+    if (op->hasAttr("domain_upper")) {
+      FloatAttr upperAttr = dyn_cast<FloatAttr>(op->getAttr("domain_upper"));
+      if (!upperAttr)
+        return op.emitOpError(
+            "domain_upper must be a floating-point attribute");
+      if (upperAttr.getValueAsDouble() > validUpper) {
+        return rewriter.notifyMatchFailure(
+            op, "domain_upper exceeds valid interval bound 1.0");
+      }
+    }
+
+    Type elemType =
+        isa<ShapedType>(type) ? cast<ShapedType>(type).getElementType() : type;
+
+    double inv2k = 1.0 / static_cast<double>(1ULL << k);
+    TypedAttr scaleAttr;
+    TypedAttr oneAttr;
+
+    if (ShapedType shapedType = dyn_cast<ShapedType>(type)) {
+      scaleAttr = DenseElementsAttr::get(
+          shapedType, rewriter.getFloatAttr(elemType, inv2k));
+      oneAttr = DenseElementsAttr::get(shapedType,
+                                       rewriter.getFloatAttr(elemType, 1.0));
+    } else {
+      scaleAttr = rewriter.getFloatAttr(elemType, inv2k);
+      oneAttr = rewriter.getFloatAttr(elemType, 1.0);
+    }
+
+    Value scaleConst = arith::ConstantOp::create(rewriter, loc, scaleAttr);
+    Value oneConst = arith::ConstantOp::create(rewriter, loc, oneAttr);
+
+    Value scaledX = arith::MulFOp::create(rewriter, loc, operand, scaleConst);
+    Value current = arith::AddFOp::create(rewriter, loc, scaledX, oneConst);
+
+    for (int64_t i = 0; i < k; ++i) {
+      current = arith::MulFOp::create(rewriter, loc, current, current);
+    }
+
+    rewriter.replaceOp(op, current);
+    return success();
+  }
+
+ private:
+  int64_t k;
+};
+
 struct PolynomialApproximation
     : impl::PolynomialApproximationBase<PolynomialApproximation> {
   using PolynomialApproximationBase::PolynomialApproximationBase;
@@ -397,6 +468,9 @@ struct PolynomialApproximation
     MLIRContext* context = &getContext();
     RewritePatternSet patterns(context);
 
+    // High priority patterns
+    patterns.add<ExpOpTaylorApproximation>(context, /*k=*/7);
+
     // Math unary ops
     patterns.add<ConvertUnaryOp<math::AbsFOp>>(context, absf);
     patterns.add<ConvertUnaryOp<math::AcosOp>>(context, acos);

lib/Transforms/PolynomialApproximation/PolynomialApproximation.td

@@ -59,7 +59,9 @@ def PolynomialApproximation : Pass<"polynomial-approximation"> {
     - `minnumf`
 
     These ops are replaced with `polynomial.eval` ops with a static polynomial
-    attribute.
+    attribute, with the exception of `exp` operations with domain in `[-2^k, 1.0]`
+    (default `k=7`), which are replaced using a Taylor approximation
+    `e^x = (1 + x/2^k)^(2^k)` evaluated via repeated squaring.
 
     (* example filepath=tests/Transforms/polynomial_approximation/doctest.mlir *)
   }];

lib/Utils/Approximation/BUILD

@@ -55,3 +55,20 @@ cc_test(
         "@llvm-project//mlir:Support",
     ],
 )
+
+cc_library(
+    name = "Taylor",
+    hdrs = ["Taylor.h"],
+)
+
+cc_test(
+    name = "TaylorTest",
+    srcs = ["TaylorTest.cpp"],
+    deps = [
+        ":CaratheodoryFejer",
+        ":Taylor",
+        "@googletest//:gtest_main",
+        "@heir//lib/Utils/Polynomial",
+        "@llvm-project//llvm:Support",
+    ],
+)

lib/Utils/Approximation/Taylor.h

@@ -0,0 +1,25 @@
+#ifndef LIB_UTILS_APPROXIMATION_TAYLOR_H_
+#define LIB_UTILS_APPROXIMATION_TAYLOR_H_
+
+#include <cstdint>
+
+namespace mlir {
+namespace heir {
+namespace approximation {
+
+/// Evaluates the Taylor exponential approximation (1 + x / 2^k)^(2^k)
+/// via repeated squaring for a given input x and parameter k.
+inline double expTaylorApproximation(double x, int64_t k = 7) {
+  double scale = 1.0 / static_cast<double>(1ULL << k);
+  double val = 1.0 + x * scale;
+  for (int64_t i = 0; i < k; ++i) {
+    val = val * val;
+  }
+  return val;
+}
+
+}  // namespace approximation
+}  // namespace heir
+}  // namespace mlir
+
+#endif  // LIB_UTILS_APPROXIMATION_TAYLOR_H_

lib/Utils/Approximation/TaylorTest.cpp

@@ -0,0 +1,82 @@
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+
+#include "gtest/gtest.h"  // from @googletest
+#include "lib/Utils/Approximation/CaratheodoryFejer.h"
+#include "lib/Utils/Approximation/Taylor.h"
+#include "lib/Utils/Polynomial/Polynomial.h"
+#include "llvm/include/llvm/ADT/APFloat.h"  // from @llvm-project
+
+namespace mlir {
+namespace heir {
+namespace approximation {
+namespace {
+
+struct ExpTaylorTestParams {
+  int64_t k;
+  double domainLower;
+  double domainUpper;
+  double step;
+  double maxInfinityNormError;
+};
+
+class TaylorApproximationTest
+    : public ::testing::TestWithParam<ExpTaylorTestParams> {};
+
+TEST_P(TaylorApproximationTest, ExpAccuracyAcrossWholeDomain) {
+  const ExpTaylorTestParams& params = GetParam();
+  auto expFunc = [](const ::llvm::APFloat& x) {
+    return ::llvm::APFloat(std::exp(x.convertToDouble()));
+  };
+  // Compare Taylor with k squarings against a Chebyshev polynomial of degree k.
+  polynomial::FloatPolynomial chebyshevPoly =
+      caratheodoryFejerApproximation(expFunc, params.k, params.domainLower,
+                                     params.domainUpper)
+          .toStandardBasis();
+
+  double taylorInfinityNormError = 0.0;
+  double taylorSumSqDiff = 0.0;
+  double chebyshevSumSqDiff = 0.0;
+
+  for (double x = params.domainLower; x <= params.domainUpper;
+       x += params.step) {
+    double expected = std::exp(x);
+    double taylorActual = expTaylorApproximation(x, params.k);
+
+    double chebyshevActual = 0.0;
+    for (const auto& term : chebyshevPoly.getTerms()) {
+      chebyshevActual += term.getCoefficient().convertToDouble() *
+                         std::pow(x, term.getExponent().getZExtValue());
+    }
+
+    double taylorAbsError = std::abs(taylorActual - expected);
+    double chebyshevAbsError = std::abs(chebyshevActual - expected);
+
+    taylorInfinityNormError = std::max(taylorInfinityNormError, taylorAbsError);
+
+    taylorSumSqDiff += taylorAbsError * taylorAbsError;
+    chebyshevSumSqDiff += chebyshevAbsError * chebyshevAbsError;
+  }
+
+  // Verify infinity norm (max absolute error) for Taylor.
+  EXPECT_LT(taylorInfinityNormError, params.maxInfinityNormError);
+
+  // Compare L2 norm against Chebyshev of degree k.
+  EXPECT_LT(taylorSumSqDiff, chebyshevSumSqDiff);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    TaylorTests, TaylorApproximationTest,
+    ::testing::Values(ExpTaylorTestParams{/*k=*/7, /*domainLower=*/-128.0,
+                                          /*domainUpper=*/1.0, /*step=*/0.1,
+                                          /*maxInfinityNormError=*/0.015},
+                      ExpTaylorTestParams{/*k=*/14, /*domainLower=*/-16384.0,
+                                          /*domainUpper=*/1.0, /*step=*/1.0,
+                                          /*maxInfinityNormError=*/1e-4}));
+
+}  // namespace
+}  // namespace approximation
+}  // namespace heir
+}  // namespace mlir

tests/Transforms/polynomial_approximation/doctest.mlir

@@ -2,8 +2,7 @@
 
 // CHECK: @test_exp
 func.func @test_exp(%x: f32) -> f32 {
-  // CHECK: polynomial.eval
-  // CHECK-SAME: [{{.*}}, {{.*}}, {{.*}}, {{.*}}]
+  // CHECK: arith.mulf
   %0 = math.exp %x {
       degree = 3 : i32,
       domain_lower = -1.0 : f64,

tests/Transforms/polynomial_approximation/polynomial_approximation.mlir

@@ -2,16 +2,37 @@
 
 // CHECK: @test_exp
 func.func @test_exp(%x: f32) -> f32 {
-  // Don't assert the quality of the approximation, just that it was applied
-  // and has the right degree. Leave quality-of-approximation for unit testing.
-  // CHECK: polynomial.eval
-  // CHECK-SAME: [{{.*}}, {{.*}}, {{.*}}, {{.*}}]
+  // CHECK: %[[SCALE:.*]] = arith.constant 7.812500e-03 : f32
+  // CHECK: %[[ONE:.*]] = arith.constant 1.000000e+00 : f32
+  // CHECK: %[[SCALED:.*]] = arith.mulf %{{.*}}, %[[SCALE]] : f32
+  // CHECK: %[[V0:.*]] = arith.addf %[[SCALED]], %[[ONE]] : f32
+  // CHECK: %[[V1:.*]] = arith.mulf %[[V0]], %[[V0]] : f32
+  // CHECK: %[[V2:.*]] = arith.mulf %[[V1]], %[[V1]] : f32
+  // CHECK: %[[V3:.*]] = arith.mulf %[[V2]], %[[V2]] : f32
+  // CHECK: %[[V4:.*]] = arith.mulf %[[V3]], %[[V3]] : f32
+  // CHECK: %[[V5:.*]] = arith.mulf %[[V4]], %[[V4]] : f32
+  // CHECK: %[[V6:.*]] = arith.mulf %[[V5]], %[[V5]] : f32
+  // CHECK: %[[V7:.*]] = arith.mulf %[[V6]], %[[V6]] : f32
+  // CHECK: return %[[V7]] : f32
   %0 = math.exp %x {degree = 3 : i32, domain_lower = -1.0 : f64, domain_upper = 1.0 : f64} : f32
   return %0 : f32
 }
 
 // -----
 
+// CHECK: @test_exp_tensor
+func.func @test_exp_tensor(%x: tensor<4xf32>) -> tensor<4xf32> {
+  // CHECK: %[[SCALE:.*]] = arith.constant dense<7.812500e-03> : tensor<4xf32>
+  // CHECK: %[[ONE:.*]] = arith.constant dense<1.000000e+00> : tensor<4xf32>
+  // CHECK: %[[SCALED:.*]] = arith.mulf %{{.*}}, %[[SCALE]] : tensor<4xf32>
+  // CHECK: %[[V0:.*]] = arith.addf %[[SCALED]], %[[ONE]] : tensor<4xf32>
+  // CHECK: %[[V1:.*]] = arith.mulf %[[V0]], %[[V0]] : tensor<4xf32>
+  %0 = math.exp %x : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
 // CHECK: @test_domain
 func.func @test_domain(%x: f32) -> f32 {
   // CHECK: polynomial.eval