Skip to content

Commit a381495

Browse files
copybara-githubcopybara-github
committed
Merge pull request #1396 from j2kun:caratheodory-fejer
PiperOrigin-RevId: 726916132
2 parents fca1e70 + 5333eae commit a381495

File tree

7 files changed

+305
-15
lines changed

7 files changed

+305
-15
lines changed

‎lib/Utils/Approximation/BUILD

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,3 +28,27 @@ cc_test(
2828
"@llvm-project//mlir:Support",
2929
],
3030
)
31+
32+
cc_library(
33+
name = "CaratheodoryFejer",
34+
srcs = ["CaratheodoryFejer.cpp"],
35+
hdrs = ["CaratheodoryFejer.h"],
36+
deps = [
37+
":Chebyshev",
38+
"@eigen//:eigen3",
39+
"@heir//lib/Utils/Polynomial",
40+
"@llvm-project//llvm:Support",
41+
],
42+
)
43+
44+
cc_test(
45+
name = "CaratheodoryFejerTest",
46+
srcs = ["CaratheodoryFejerTest.cpp"],
47+
deps = [
48+
":CaratheodoryFejer",
49+
"@googletest//:gtest_main",
50+
"@heir//lib/Utils/Polynomial",
51+
"@llvm-project//llvm:Support",
52+
"@llvm-project//mlir:Support",
53+
],
54+
)

‎lib/Utils/Approximation/CaratheodoryFejer.cpp

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
#include "lib/Utils/Approximation/CaratheodoryFejer.h"
2+
3+
#include <cassert>
4+
#include <cstdint>
5+
#include <cstdlib>
6+
#include <functional>
7+
#include <optional>
8+
9+
#include "Eigen/Core" // from @eigen
10+
#include "Eigen/Dense" // from @eigen
11+
#include "Eigen/Eigenvalues" // from @eigen
12+
#include "lib/Utils/Approximation/Chebyshev.h"
13+
#include "lib/Utils/Polynomial/Polynomial.h"
14+
#include "llvm/include/llvm/ADT/APFloat.h" // from @llvm-project
15+
#include "llvm/include/llvm/ADT/SmallVector.h" // from @llvm-project
16+
17+
namespace mlir {
18+
namespace heir {
19+
namespace approximation {
20+
21+
using ::Eigen::MatrixXd;
22+
using ::Eigen::SelfAdjointEigenSolver;
23+
using ::Eigen::VectorXd;
24+
using ::llvm::APFloat;
25+
using ::llvm::SmallVector;
26+
using ::mlir::heir::polynomial::FloatPolynomial;
27+
28+
FloatPolynomial caratheodoryFejerApproximation(
29+
const std::function<APFloat(APFloat)> &func, int32_t degree,
30+
std::optional<int32_t> chebDegree) {
31+
// Construct the Chebyshev interpolant.
32+
SmallVector<APFloat> chebPts, chebEvalPts, chebCoeffs;
33+
int32_t actualChebDegree = chebDegree.value_or(2 * degree);
34+
int32_t numChebPts = 1 + actualChebDegree;
35+
assert(numChebPts >= 2 * degree + 1 &&
36+
"Chebyshev degree must be at least twice the CF degree plus 1");
37+
chebPts.reserve(numChebPts);
38+
getChebyshevPoints(numChebPts, chebPts);
39+
for (auto &pt : chebPts) {
40+
chebEvalPts.push_back(func(pt));
41+
}
42+
interpolateChebyshev(chebEvalPts, chebCoeffs);
43+
44+
// Use the tail coefficients to construct a Hankel matrix
45+
// where A[i, j] = c[i+j]
46+
// Cf. https://en.wikipedia.org/wiki/Hankel_matrix
47+
SmallVector<APFloat> tailChebCoeffs(chebCoeffs.begin() + (degree + 1),
48+
chebCoeffs.end());
49+
int32_t hankelSize = tailChebCoeffs.size();
50+
MatrixXd hankel(hankelSize, hankelSize);
51+
for (int i = 0; i < hankelSize; ++i) {
52+
for (int j = 0; j < hankelSize; ++j) {
53+
// upper left triangular region, including diagonal
54+
if (i + j < hankelSize)
55+
hankel(i, j) = tailChebCoeffs[i + j].convertToDouble();
56+
else
57+
hankel(i, j) = 0;
58+
}
59+
}
60+
61+
// Compute the eigenvalues and eigenvectors of the Hankel matrix
62+
SelfAdjointEigenSolver<MatrixXd> solver(hankel);
63+
64+
const VectorXd &eigenvalues = solver.eigenvalues();
65+
// Eigenvectors are columns of the matrix.
66+
const MatrixXd &eigenvectors = solver.eigenvectors();
67+
68+
// Extract the eigenvector for the (absolute value) largest eigenvalue.
69+
int32_t maxIndex = 0;
70+
double maxEigenvalue = eigenvalues(0);
71+
for (int32_t i = 1; i < eigenvalues.size(); ++i) {
72+
if (std::abs(eigenvalues(i)) > maxEigenvalue) {
73+
maxEigenvalue = std::abs(eigenvalues(i));
74+
maxIndex = i;
75+
}
76+
}
77+
VectorXd maxEigenvector = eigenvectors.col(maxIndex);
78+
79+
// A debug for comparing the eigenvalue solver with the reference
80+
// implementation.
81+
// std::cout << "Max eigenvector:" << std::endl;
82+
// for (int32_t i = 0; i < maxEigenvector.size(); ++i) {
83+
// std::cout << std::setprecision(18) << maxEigenvector(i) << ", ";
84+
// }
85+
// std::cout << std::endl;
86+
87+
double v1 = maxEigenvector(0);
88+
VectorXd vv = maxEigenvector.tail(maxEigenvector.size() - 1);
89+
90+
SmallVector<APFloat> b =
91+
SmallVector<APFloat>(tailChebCoeffs.begin(), tailChebCoeffs.end());
92+
93+
int32_t t = actualChebDegree - degree - 1;
94+
for (int32_t i = degree; i > -degree - 1; --i) {
95+
SmallVector<APFloat> sliceB(b.begin(), b.begin() + t);
96+
97+
APFloat sum = APFloat(0.0);
98+
for (int32_t j = 0; j < sliceB.size(); ++j) {
99+
double vvVal = vv(j);
100+
sum = sum + sliceB[j] * APFloat(vvVal);
101+
}
102+
103+
APFloat z = -sum / APFloat(v1);
104+
105+
// I suspect this insert is slow. Once it's working we can optimize this
106+
// loop to avoid the insert.
107+
b.insert(b.begin(), z);
108+
}
109+
110+
SmallVector<APFloat> bb(b.begin() + degree, b.begin() + (2 * degree + 1));
111+
for (int32_t i = 1; i < bb.size(); ++i) {
112+
bb[i] = bb[i] + b[degree - 1 - (i - 1)];
113+
}
114+
115+
SmallVector<APFloat> pk;
116+
pk.reserve(bb.size());
117+
for (int32_t i = 0; i < bb.size(); ++i) {
118+
pk.push_back(chebCoeffs[i] - bb[i]);
119+
}
120+
121+
return chebyshevToMonomial(pk);
122+
}
123+
124+
} // namespace approximation
125+
} // namespace heir
126+
} // namespace mlir

‎lib/Utils/Approximation/CaratheodoryFejer.h

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
#ifndef LIB_UTILS_APPROXIMATION_CARATHEODORYFEJER_H_
2+
#define LIB_UTILS_APPROXIMATION_CARATHEODORYFEJER_H_
3+
4+
#include <cstdint>
5+
#include <functional>
6+
#include <optional>
7+
8+
#include "lib/Utils/Polynomial/Polynomial.h"
9+
#include "llvm/include/llvm/ADT/APFloat.h" // from @llvm-project
10+
11+
namespace mlir {
12+
namespace heir {
13+
namespace approximation {
14+
15+
/// Construct the Caratheodory-Fejer approximation of a given function. The
16+
/// result polynomial is represented in the monomial basis and has degree
17+
/// `degree`.
18+
///
19+
/// A port of chebfun's cf.m, cf.
20+
/// https://github.com/chebfun/chebfun/blob/69c12cf75f93cb2f36fd4cfd5e287662cd2f1091/%40chebfun/cf.m
21+
/// Specifically, the path through `cf` that invokes `polynomialCF`
22+
/// (rational functions are not supported here).
23+
///
24+
/// Cf. https://doi.org/10.1007/s10543-011-0331-7 for a mathematical
25+
/// explanation.
26+
/// https://people.maths.ox.ac.uk/trefethen/publication/PDF/2011_138.pdf
27+
///
28+
/// The argument chebDegree provides control over the degree of the Chebyshev
29+
/// interpolation used to seed the CF approximation. If unset, it
30+
/// is heuristically chosen as twice the desired CF approximant degree.
31+
::mlir::heir::polynomial::FloatPolynomial caratheodoryFejerApproximation(
32+
const std::function<::llvm::APFloat(::llvm::APFloat)> &func, int32_t degree,
33+
std::optional<int32_t> chebDegree = std::nullopt);
34+
35+
} // namespace approximation
36+
} // namespace heir
37+
} // namespace mlir
38+
39+
#endif // LIB_UTILS_APPROXIMATION_CARATHEODORYFEJER_H_

‎lib/Utils/Approximation/CaratheodoryFejerTest.cpp

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
#include <cmath>
2+
3+
#include "gmock/gmock.h" // from @googletest
4+
#include "gtest/gtest.h" // from @googletest
5+
#include "lib/Utils/Approximation/CaratheodoryFejer.h"
6+
#include "lib/Utils/Polynomial/Polynomial.h"
7+
#include "llvm/include/llvm/ADT/APFloat.h" // from @llvm-project
8+
#include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project
9+
10+
#define EPSILON 1e-11
11+
12+
namespace mlir {
13+
namespace heir {
14+
namespace approximation {
15+
namespace {
16+
17+
using ::llvm::APFloat;
18+
using ::mlir::heir::polynomial::FloatPolynomial;
19+
using ::testing::DoubleNear;
20+
21+
TEST(CaratheodoryFejerTest, ApproximateExpDegree3) {
22+
auto func = [](const APFloat& x) {
23+
return APFloat(std::exp(x.convertToDouble()));
24+
};
25+
FloatPolynomial actual = caratheodoryFejerApproximation(func, 3);
26+
// Values taken from reference impl are exact.
27+
FloatPolynomial expected = FloatPolynomial::fromCoefficients(
28+
{0.9945811640427066, 0.9956553725361579, 0.5429702814725632,
29+
0.1795458211087378});
30+
EXPECT_EQ(actual, expected);
31+
}
32+
33+
TEST(CaratheodoryFejerTest, ApproximateReluDegree14) {
34+
auto relu = [](const APFloat& x) {
35+
APFloat zero = APFloat::getZero(x.getSemantics());
36+
return x > zero ? x : zero;
37+
};
38+
FloatPolynomial actual = caratheodoryFejerApproximation(relu, 14);
39+
40+
// The reference implementation prints coefficients that are ~1e-12 away from
41+
// our implementation, mainly because the eigenvalue solver details are
42+
// slightly different. For instance, the max eigenvalue in our implementation
43+
// is
44+
//
45+
//
46+
// -0.415033778742867843
47+
// 0.41503377874286651
48+
// 0.358041674766206519
49+
// -0.358041674766206408
50+
// -0.302283813201297547
51+
// 0.302283813201297602
52+
// 0.244610415109838275
53+
// -0.244610415109838636
54+
// -0.182890410854375879
55+
// 0.182890410854376101
56+
// 0.115325658064263425
57+
// -0.115325658064263148
58+
// -0.0399306015339729384
59+
// 0.0399306015339727718
60+
//
61+
// But in the reference implementation, the basis elements are permuted so
62+
// that the adjacent positive values and negative values are swapped. This is
63+
// still a valid eigenvector, but it results in slight difference in floating
64+
// point error accumulation as the rest of the algorithm proceeds.
65+
auto terms = actual.getTerms();
66+
EXPECT_THAT(terms[0].getCoefficient().convertToDouble(),
67+
DoubleNear(0.010384627976349288, EPSILON));
68+
EXPECT_THAT(terms[1].getCoefficient().convertToDouble(),
69+
DoubleNear(0.4999999999999994, EPSILON));
70+
EXPECT_THAT(terms[2].getCoefficient().convertToDouble(),
71+
DoubleNear(3.227328667600437, EPSILON));
72+
EXPECT_THAT(terms[3].getCoefficient().convertToDouble(),
73+
DoubleNear(2.1564570993799688e-14, EPSILON));
74+
EXPECT_THAT(terms[4].getCoefficient().convertToDouble(),
75+
DoubleNear(-27.86732536231614, EPSILON));
76+
EXPECT_THAT(terms[5].getCoefficient().convertToDouble(),
77+
DoubleNear(-1.965772591254676e-13, EPSILON));
78+
EXPECT_THAT(terms[6].getCoefficient().convertToDouble(),
79+
DoubleNear(139.12944753041404, EPSILON));
80+
EXPECT_THAT(terms[7].getCoefficient().convertToDouble(),
81+
DoubleNear(7.496488571843804e-13, EPSILON));
82+
EXPECT_THAT(terms[8].getCoefficient().convertToDouble(),
83+
DoubleNear(-363.6062351528312, EPSILON));
84+
EXPECT_THAT(terms[9].getCoefficient().convertToDouble(),
85+
DoubleNear(-1.3773783921527744e-12, EPSILON));
86+
EXPECT_THAT(terms[10].getCoefficient().convertToDouble(),
87+
DoubleNear(505.9489721657369, EPSILON));
88+
EXPECT_THAT(terms[11].getCoefficient().convertToDouble(),
89+
DoubleNear(1.2076732984649801e-12, EPSILON));
90+
EXPECT_THAT(terms[12].getCoefficient().convertToDouble(),
91+
DoubleNear(-355.4120699445272, EPSILON));
92+
EXPECT_THAT(terms[13].getCoefficient().convertToDouble(),
93+
DoubleNear(-4.050490139246503e-13, EPSILON));
94+
EXPECT_THAT(terms[14].getCoefficient().convertToDouble(),
95+
DoubleNear(99.07988219049058, EPSILON));
96+
}
97+
98+
} // namespace
99+
} // namespace approximation
100+
} // namespace heir
101+
} // namespace mlir

‎lib/Utils/Approximation/Chebyshev.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,8 @@ void getChebyshevPolynomials(int64_t numPolynomials,
6161
results.push_back(FloatPolynomial::fromCoefficients({1.}));
6262
}
6363
if (numPolynomials >= 2) {
64-
// 2x
65-
results.push_back(FloatPolynomial::fromCoefficients({0., 2.}));
64+
// x
65+
results.push_back(FloatPolynomial::fromCoefficients({0., 1.}));
6666
}
6767

6868
if (numPolynomials <= 2) return;

‎lib/Utils/Approximation/Chebyshev.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,10 @@ namespace approximation {
2020
void getChebyshevPoints(int64_t numPoints,
2121
::llvm::SmallVector<::llvm::APFloat> &results);
2222

23-
/// Generate the first `numPolynomials` Chebyshev polynomials of the second
23+
/// Generate the first `numPolynomials` Chebyshev polynomials of the first
2424
/// kind, storing them in the results outparameter.
2525
///
26-
/// The first few polynomials are 1, 2x, 4x^2 - 1, 8x^3 - 4x, ...
26+
/// The first few polynomials are 1, x, 2x^2 - 1, 4x^3 - 3x, ...
2727
void getChebyshevPolynomials(
2828
int64_t numPolynomials,
2929
::llvm::SmallVector<::mlir::heir::polynomial::FloatPolynomial> &results);

‎lib/Utils/Approximation/ChebyshevTest.cpp

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -55,25 +55,25 @@ TEST(ChebyshevTest, TestGetChebyshevPolynomials) {
5555
chebPolys,
5656
ElementsAre(
5757
FloatPolynomial::fromCoefficients({1.}),
58-
FloatPolynomial::fromCoefficients({0., 2.}),
59-
FloatPolynomial::fromCoefficients({-1., 0., 4.}),
60-
FloatPolynomial::fromCoefficients({0., -4., 0., 8.}),
61-
FloatPolynomial::fromCoefficients({1., 0., -12., 0., 16.}),
62-
FloatPolynomial::fromCoefficients({0., 6., 0., -32., 0., 32.}),
63-
FloatPolynomial::fromCoefficients({-1., 0., 24., 0., -80., 0., 64.}),
58+
FloatPolynomial::fromCoefficients({0., 1.}),
59+
FloatPolynomial::fromCoefficients({-1., 0., 2.}),
60+
FloatPolynomial::fromCoefficients({0., -3., 0., 4.}),
61+
FloatPolynomial::fromCoefficients({1., 0., -8., 0., 8.}),
62+
FloatPolynomial::fromCoefficients({0., 5., 0., -20., 0., 16.}),
63+
FloatPolynomial::fromCoefficients({-1., 0., 18., 0., -48., 0., 32.}),
6464
FloatPolynomial::fromCoefficients(
65-
{0., -8., 0., 80., 0., -192., 0., 128.}),
65+
{0., -7., 0., 56., 0., -112., 0., 64.}),
6666
FloatPolynomial::fromCoefficients(
67-
{1., 0., -40., 0., 240., 0., -448., 0., 256.})));
67+
{1., 0., -32., 0., 160., 0., -256., 0., 128.})));
6868
}
6969

7070
TEST(ChebyshevTest, TestChebyshevToMonomial) {
71-
// 1 (1) - 1 (-1 + 4x^2) + 2 (-4x + 8x^3)
71+
// 1 (1) - 1 (-1 + 2x^2) + 2 (-3x + 4x^3)
7272
SmallVector<APFloat> chebCoeffs = {APFloat(1.0), APFloat(0.0), APFloat(-1.0),
7373
APFloat(2.0)};
74-
// 2 - 8 x - 4 x^2 + 16 x^3
74+
// 2 - 6 x - 2 x^2 + 8 x^3
7575
FloatPolynomial expected =
76-
FloatPolynomial::fromCoefficients({2.0, -8.0, -4.0, 16.0});
76+
FloatPolynomial::fromCoefficients({2.0, -6.0, -2.0, 8.0});
7777
FloatPolynomial actual = chebyshevToMonomial(chebCoeffs);
7878
EXPECT_EQ(actual, expected);
7979
}

0 commit comments

Comments
 (0)