From 419b7c28e62142c0530e935c37db11782e5e43a4 Mon Sep 17 00:00:00 2001 From: HEIR Team Date: Mon, 20 Jul 2026 00:13:44 -0700 Subject: [PATCH] Add precision tests for Taylor exponentiation approximation. We tests for k=7 and k=14 that the approximation performs better than the default Chebyshev pass on the chosen domain for the L2 norm. PiperOrigin-RevId: 950641161 --- .../PolynomialApproximation.cpp | 88 +++++++++++++++++-- .../PolynomialApproximation.td | 4 +- lib/Utils/Approximation/BUILD | 17 ++++ lib/Utils/Approximation/Taylor.h | 25 ++++++ lib/Utils/Approximation/TaylorTest.cpp | 82 +++++++++++++++++ .../polynomial_approximation/doctest.mlir | 3 +- .../polynomial_approximation.mlir | 29 +++++- 7 files changed, 234 insertions(+), 14 deletions(-) create mode 100644 lib/Utils/Approximation/Taylor.h create mode 100644 lib/Utils/Approximation/TaylorTest.cpp diff --git a/lib/Transforms/PolynomialApproximation/PolynomialApproximation.cpp b/lib/Transforms/PolynomialApproximation/PolynomialApproximation.cpp index 60ca2059be..59c9c8d100 100644 --- a/lib/Transforms/PolynomialApproximation/PolynomialApproximation.cpp +++ b/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 { double upper; }; +struct ExpOpTaylorApproximation : public OpRewritePattern { + ExpOpTaylorApproximation(MLIRContext* context, int64_t k = 7) + : OpRewritePattern(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(1ULL << k); + double validUpper = 1.0; + + if (op->hasAttr("domain_lower")) { + FloatAttr lowerAttr = dyn_cast(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(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(type) ? cast(type).getElementType() : type; + + double inv2k = 1.0 / static_cast(1ULL << k); + TypedAttr scaleAttr; + TypedAttr oneAttr; + + if (ShapedType shapedType = dyn_cast(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 { using PolynomialApproximationBase::PolynomialApproximationBase; @@ -397,6 +468,9 @@ struct PolynomialApproximation MLIRContext* context = &getContext(); RewritePatternSet patterns(context); + // High priority patterns + patterns.add(context, /*k=*/7); + // Math unary ops patterns.add>(context, absf); patterns.add>(context, acos); diff --git a/lib/Transforms/PolynomialApproximation/PolynomialApproximation.td b/lib/Transforms/PolynomialApproximation/PolynomialApproximation.td index 4eb36a0598..1cf28b4e4d 100644 --- a/lib/Transforms/PolynomialApproximation/PolynomialApproximation.td +++ b/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 *) }]; diff --git a/lib/Utils/Approximation/BUILD b/lib/Utils/Approximation/BUILD index c031ee9c01..b289124aa6 100644 --- a/lib/Utils/Approximation/BUILD +++ b/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", + ], +) diff --git a/lib/Utils/Approximation/Taylor.h b/lib/Utils/Approximation/Taylor.h new file mode 100644 index 0000000000..6bba38531d --- /dev/null +++ b/lib/Utils/Approximation/Taylor.h @@ -0,0 +1,25 @@ +#ifndef LIB_UTILS_APPROXIMATION_TAYLOR_H_ +#define LIB_UTILS_APPROXIMATION_TAYLOR_H_ + +#include + +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(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_ diff --git a/lib/Utils/Approximation/TaylorTest.cpp b/lib/Utils/Approximation/TaylorTest.cpp new file mode 100644 index 0000000000..b21e9c16f8 --- /dev/null +++ b/lib/Utils/Approximation/TaylorTest.cpp @@ -0,0 +1,82 @@ +#include +#include +#include +#include + +#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 {}; + +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 diff --git a/tests/Transforms/polynomial_approximation/doctest.mlir b/tests/Transforms/polynomial_approximation/doctest.mlir index 34758babde..66b36fc8da 100644 --- a/tests/Transforms/polynomial_approximation/doctest.mlir +++ b/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, diff --git a/tests/Transforms/polynomial_approximation/polynomial_approximation.mlir b/tests/Transforms/polynomial_approximation/polynomial_approximation.mlir index 6cd771d48c..b62aa92508 100644 --- a/tests/Transforms/polynomial_approximation/polynomial_approximation.mlir +++ b/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