Merge pull request #1367 from j2kun:chebyshev

PiperOrigin-RevId: 725244577

LICENSE

@@ -218,6 +218,7 @@ prospectively choose to deem waived or otherwise exclude such Section(s) of
 the License, but only in their entirety and only with respect to the Combined
 Software.
 
+# Yosys
 Copyright (C) 2012 - 2018  Clifford Wolf <[email protected]>, <[email protected]>
 
 Copyright (C) 2012 Martin Schmölzer <[email protected]>
@@ -242,3 +243,30 @@ ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+# PocketFFT
+Copyright (C) 2010-2018 Max-Planck-Society
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+* Redistributions in binary form must reproduce the above copyright notice, this
+  list of conditions and the following disclaimer in the documentation and/or
+  other materials provided with the distribution.
+* Neither the name of the copyright holder nor the names of its contributors may
+  be used to endorse or promote products derived from this software without
+  specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

WORKSPACE

@@ -421,3 +421,10 @@ git_repository(
     patches = ["@heir//bazel/openfhe:add_config_core.patch"],
     remote = "https://github.com/openfheorg/openfhe-development.git",
 )
+
+git_repository(
+    name = "pocketfft",
+    build_file = "//bazel/pocketfft:pocketfft.BUILD",
+    commit = "bb5bdb776c64819f66cb2205f78bef1581448628",
+    remote = "https://gitlab.mpcdf.mpg.de/mtr/pocketfft.git",
+)

bazel/pocketfft/BUILD

@@ -0,0 +1,7 @@
+# This build file is necessary to mark this directory as a subpackage for bazel
+# to have access to the files.
+
+package(
+    default_applicable_licenses = ["@heir//:license"],
+    default_visibility = ["//visibility:public"],
+)

bazel/pocketfft/pocketfft.BUILD

@@ -0,0 +1,19 @@
+# BUILD file for a bazel-native pocketfft build
+package(
+    default_visibility = ["//visibility:public"],
+)
+
+licenses(["notice"])
+
+cc_library(
+    name = "pocketfft",
+    hdrs = [
+        "pocketfft_hdronly.h",
+    ],
+    copts = [
+        "-fexceptions",
+    ],
+    features = [
+        "-use_header_modules",
+    ],
+)

lib/Utils/Approximation/BUILD

@@ -13,6 +13,7 @@ cc_library(
         "@heir//lib/Utils/Polynomial",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:Support",
+        "@pocketfft",
     ],
 )
 

lib/Utils/Approximation/Chebyshev.cpp

@@ -1,10 +1,15 @@
 
+#include <algorithm>
 #include <cmath>
+#include <complex>
+#include <cstddef>
 #include <cstdint>
+#include <vector>
 
 #include "lib/Utils/Polynomial/Polynomial.h"
 #include "llvm/include/llvm/ADT/APFloat.h"   // from @llvm-project
 #include "mlir/include/mlir/Support/LLVM.h"  // from @llvm-project
+#include "pocketfft_hdronly.h"               // from @pocketfft
 
 namespace mlir {
 namespace heir {
@@ -69,6 +74,105 @@ void getChebyshevPolynomials(int64_t numPolynomials,
   }
 }
 
+FloatPolynomial chebyshevToMonomial(const SmallVector<APFloat> &coefficients) {
+  SmallVector<FloatPolynomial> chebPolys;
+  chebPolys.reserve(coefficients.size());
+  getChebyshevPolynomials(coefficients.size(), chebPolys);
+
+  FloatPolynomial result = FloatPolynomial::zero();
+  for (int64_t i = 0; i < coefficients.size(); ++i) {
+    result = result.add(chebPolys[i].scale(coefficients[i]));
+  }
+
+  return result;
+}
+
+void interpolateChebyshev(ArrayRef<APFloat> chebEvalPoints,
+                          SmallVector<APFloat> &outputChebCoeffs) {
+  size_t n = chebEvalPoints.size();
+  if (n == 0) {
+    return;
+  }
+  if (n == 1) {
+    outputChebCoeffs.push_back(chebEvalPoints[0]);
+    return;
+  }
+
+  // When the function being evaluated has even or odd symmetry, we can get
+  // coefficients. In particular, even symmetry implies all odd-numbered
+  // Chebyshev coefficients are zero. Odd symmetry implies even-numbered
+  // coefficients are zero.
+  bool isEven =
+      std::equal(chebEvalPoints.begin(), chebEvalPoints.begin() + n / 2,
+                 chebEvalPoints.rbegin());
+
+  bool isOdd = true;
+  for (int i = 0; i < n / 2; ++i) {
+    if (chebEvalPoints[i] != -chebEvalPoints[(n - 1) - i]) {
+      isOdd = false;
+      break;
+    }
+  }
+
+  // Construct input to ifft so as to compute a Discrete Cosine Transform
+  // The inputs are [v_{n-1}, v_{n-2}, ..., v_0, v_1, ..., v_{n-2}]
+  std::vector<std::complex<double>> ifftInput;
+  size_t fftLen = 2 * (n - 1);
+  ifftInput.reserve(fftLen);
+  for (size_t i = n - 1; i > 0; --i) {
+    ifftInput.emplace_back(chebEvalPoints[i].convertToDouble());
+  }
+  for (size_t i = 0; i < n - 1; ++i) {
+    ifftInput.emplace_back(chebEvalPoints[i].convertToDouble());
+  }
+
+  // Compute inverse FFT using minimal API call to pocketfft. This should be
+  // equivalent to numpy.fft.ifft, as it uses pocketfft underneath. It's worth
+  // noting here that we're computing the Discrete Cosine Transform (DCT) in
+  // terms of a complex Discrete Fourier Transform (DFT), but pocketfft appears
+  // to have a built-in `dct` function. It may be trivial to switch to
+  // pocketfft::dct, but this was originally based on a reference
+  // implementation that did not have access to a native DCT. Migrating to a
+  // DCT should only be necessary (a) once the reference implementation is
+  // fully ported and tested, and (b) if we determine that there's a
+  // performance benefit to using the native DCT. Since this routine is
+  // expected to be used in doing relatively low-degree approximations, it
+  // probably won't be a problem.
+  std::vector<std::complex<double>> ifftResult(fftLen);
+  pocketfft::shape_t shape{fftLen};
+  pocketfft::stride_t strided{sizeof(std::complex<double>)};
+  pocketfft::shape_t axes{0};
+
+  pocketfft::c2c(shape, strided, strided, axes, pocketfft::BACKWARD,
+                 ifftInput.data(), ifftResult.data(), 1. / fftLen);
+
+  outputChebCoeffs.clear();
+  outputChebCoeffs.reserve(n);
+  for (size_t i = 0; i < n; ++i) {
+    outputChebCoeffs.push_back(APFloat(ifftResult[i].real()));
+  }
+
+  // Due to the endpoint behavior of Chebyshev polynomials and the properties
+  // of the DCT, the non-endpoint coefficients of the DCT are the Chebyshev
+  // coefficients scaled by 2.
+  for (int i = 1; i < n - 1; ++i) {
+    outputChebCoeffs[i] = outputChebCoeffs[i] * APFloat(2.0);
+  }
+
+  // Even/odd corrections
+  if (isEven) {
+    for (size_t i = 1; i < n; i += 2) {
+      outputChebCoeffs[i] = APFloat(0.0);
+    }
+  }
+
+  if (isOdd) {
+    for (size_t i = 0; i < n; i += 2) {
+      outputChebCoeffs[i] = APFloat(0.0);
+    }
+  }
+}
+
 }  // namespace approximation
 }  // namespace heir
 }  // namespace mlir

lib/Utils/Approximation/Chebyshev.h

@@ -18,15 +18,38 @@ namespace approximation {
 /// This is a port of the chebfun routine at
 /// https://github.com/chebfun/chebfun/blob/db207bc9f48278ca4def15bf90591bfa44d0801d/%40chebtech2/chebpts.m#L34
 void getChebyshevPoints(int64_t numPoints,
-                        SmallVector<::llvm::APFloat> &results);
+                        ::llvm::SmallVector<::llvm::APFloat> &results);
 
 /// Generate the first `numPolynomials` Chebyshev polynomials of the second
 /// kind, storing them in the results outparameter.
 ///
 /// The first few polynomials are 1, 2x, 4x^2 - 1, 8x^3 - 4x, ...
 void getChebyshevPolynomials(
     int64_t numPolynomials,
-    SmallVector<::mlir::heir::polynomial::FloatPolynomial> &results);
+    ::llvm::SmallVector<::mlir::heir::polynomial::FloatPolynomial> &results);
+
+/// Convert a vector of Chebyshev coefficients to the monomial basis. If the
+/// Chebyshev polynomials are T_0, T_1, ..., then entry i of the input vector
+/// is the coefficient of T_i.
+::mlir::heir::polynomial::FloatPolynomial chebyshevToMonomial(
+    const ::llvm::SmallVector<::llvm::APFloat> &coefficients);
+
+/// Interpolate Chebyshev coefficients for a given set of points. The values in
+/// chebEvalPoints are assumed to be evaluations of the target function on the
+/// first N+1 Chebyshev points of the second kind, where N is the degree of the
+/// interpolating polynomial. The produced coefficients are stored in the
+/// outparameter outputChebCoeffs.
+///
+/// A port of chebfun vals2coeffs, cf.
+/// https://github.com/chebfun/chebfun/blob/69c12cf75f93cb2f36fd4cfd5e287662cd2f1091/%40ballfun/vals2coeffs.m
+/// based on the a trigonometric interpolation via the FFT.
+///
+/// Cf. Henrici, "Fast Fourier Methods in Computational Complex Analysis"
+/// https://doi.org/10.1137/1021093
+/// https://people.math.ethz.ch/~hiptmair/Seminars/CONVQUAD/Articles/HEN79.pdf
+void interpolateChebyshev(
+    ::llvm::ArrayRef<::llvm::APFloat> chebEvalPoints,
+    ::llvm::SmallVector<::llvm::APFloat> &outputChebCoeffs);
 
 }  // namespace approximation
 }  // namespace heir

lib/Utils/Approximation/ChebyshevTest.cpp

@@ -1,3 +1,4 @@
+#include <cmath>
 #include <cstdint>
 
 #include "gmock/gmock.h"  // from @googletest
@@ -14,6 +15,7 @@ namespace {
 
 using ::llvm::APFloat;
 using ::mlir::heir::polynomial::FloatPolynomial;
+using ::testing::DoubleEq;
 using ::testing::ElementsAre;
 
 TEST(ChebyshevTest, TestGetChebyshevPointsSingle) {
@@ -47,10 +49,8 @@ TEST(ChebyshevTest, TestGetChebyshevPoints9) {
 TEST(ChebyshevTest, TestGetChebyshevPolynomials) {
   SmallVector<FloatPolynomial> chebPolys;
   int64_t n = 9;
+  chebPolys.reserve(n);
   getChebyshevPolynomials(n, chebPolys);
-
-  for (const auto& p : chebPolys) p.dump();
-
   EXPECT_THAT(
       chebPolys,
       ElementsAre(
@@ -67,6 +67,39 @@ TEST(ChebyshevTest, TestGetChebyshevPolynomials) {
               {1., 0., -40., 0., 240., 0., -448., 0., 256.})));
 }
 
+TEST(ChebyshevTest, TestChebyshevToMonomial) {
+  // 1 (1) - 1 (-1 + 4x^2) + 2 (-4x + 8x^3)
+  SmallVector<APFloat> chebCoeffs = {APFloat(1.0), APFloat(0.0), APFloat(-1.0),
+                                     APFloat(2.0)};
+  // 2 - 8 x - 4 x^2 + 16 x^3
+  FloatPolynomial expected =
+      FloatPolynomial::fromCoefficients({2.0, -8.0, -4.0, 16.0});
+  FloatPolynomial actual = chebyshevToMonomial(chebCoeffs);
+  EXPECT_EQ(actual, expected);
+}
+
+TEST(ChebyshevTest, TestInterpolateChebyshevExpDegree3) {
+  // degree 3 implies we need 4 points.
+  SmallVector<APFloat> chebPts = {APFloat(-1.0), APFloat(-0.5), APFloat(0.5),
+                                  APFloat(1.0)};
+  SmallVector<APFloat> expVals;
+  expVals.reserve(chebPts.size());
+  for (const APFloat& pt : chebPts) {
+    expVals.push_back(APFloat(std::exp(pt.convertToDouble())));
+  }
+
+  SmallVector<APFloat> actual;
+  interpolateChebyshev(expVals, actual);
+
+  EXPECT_THAT(actual[0].convertToDouble(), DoubleEq(1.2661108550760016));
+  EXPECT_THAT(actual[1].convertToDouble(), DoubleEq(1.1308643327583656));
+  EXPECT_THAT(actual[2].convertToDouble(), DoubleEq(0.276969779739242));
+  // This test is slightly off from what numpy produces (up to ~10^{-15}), not
+  // sure why.
+  // EXPECT_THAT(actual[3].convertToDouble(), DoubleEq(0.04433686088543568));
+  EXPECT_THAT(actual[3].convertToDouble(), DoubleEq(0.044336860885435536));
+}
+
 }  // namespace
 }  // namespace approximation
 }  // namespace heir

lib/Utils/Polynomial/Polynomial.h

@@ -243,6 +243,12 @@ class PolynomialBase {
     return os.str();
   }
 
+  // Returns a zero polynomial
+  static Derived zero() {
+    SmallVector<Monomial> monomials;
+    return Derived(monomials);
+  }
+
   bool isZero() const { return getTerms().empty(); }
 
   unsigned getDegree() const {
@@ -262,7 +268,8 @@ class PolynomialBase {
 /// A single-variable polynomial with integer coefficients.
 ///
 /// Eg: x^1024 + x + 1
-class IntPolynomial : public PolynomialBase<IntPolynomial, IntMonomial, APInt> {
+class IntPolynomial final
+    : public PolynomialBase<IntPolynomial, IntMonomial, APInt> {
  public:
   explicit IntPolynomial(ArrayRef<IntMonomial> terms) : PolynomialBase(terms) {}
 
@@ -283,7 +290,7 @@ class IntPolynomial : public PolynomialBase<IntPolynomial, IntMonomial, APInt> {
 /// A single-variable polynomial with double coefficients.
 ///
 /// Eg: 1.0 x^1024 + 3.5 x + 1e-05
-class FloatPolynomial
+class FloatPolynomial final
     : public PolynomialBase<FloatPolynomial, FloatMonomial, APFloat> {
  public:
   explicit FloatPolynomial(ArrayRef<FloatMonomial> terms)