Merge pull request #339 from kroma-network/feat/add-fft-benchmark

feat: add fft benchmark

Author: Ryan Kim

Cargo.Bazel.lock

@@ -3,6 +3,7 @@ members = [
     "benchmark/msm/arkworks",
     "benchmark/msm/bellman",
     "benchmark/msm/halo2",
+    "benchmark/fft/halo2",
     "tachyon/rs",
     "vendors/halo2",
 ]

Cargo.lock

@@ -54,6 +54,7 @@ crates_repository(
         "//benchmark/msm/arkworks:Cargo.toml",
         "//benchmark/msm/bellman:Cargo.toml",
         "//benchmark/msm/halo2:Cargo.toml",
+        "//benchmark/fft/halo2:Cargo.toml",
         "//tachyon/rs:Cargo.toml",
         "//vendors/halo2:Cargo.toml",
     ],

Cargo.toml

@@ -37,12 +37,14 @@ build --build_tag_filters -rust //enables cuda targets
 build --build_tag_filters="" //enables all targets
 ```
 
-To harness the plot chart feature, first ensure that matplotlib is installed (refer to the [installation guide](https://github.com/kroma-network/tachyon#matplotlib)). Then, append the `--//:has_rtti` and `--//:has_matplotlib` flags to your command:
+To harness the plot chart feature, first ensure that matplotlib is installed (refer to the [installation guide](/docs/how_to_use/how_to_build.md#matplotliboptional)). Then, append the `--//:has_rtti` and `--//:has_matplotlib` flags to your command:
 
 ```shell
 bazel run -c opt --//:has_rtti --//:has_matplotlib //benchmark/path/to/target:target_name -- -n <test_set_size>
 ```
 
+A build error with error message like `terminate called after throwing an instance of 'pybind11::error_already_set' what():  ImportError: ... : undefined symbol: PyTuple_Type` may occur when building the benchmark. This is because bazel does not know about updates to the python environment. To fix this, set environment varible [`PYTHON_LIB_PATH`](/third_party/py/python_configure.bzl) and rebuild the benchmark.
+
 For executing GPU benchmarks, make sure to configure [GPU config](https://github.com/kroma-network/tachyon#hardware-acceleration) for your environment. For instance, in CUDA:
 
 ```shell

WORKSPACE

@@ -17,6 +17,7 @@ tachyon_cc_library(
     visibility = ["//benchmark:__subpackages__"],
     deps = [
         "//benchmark:simple_benchmark_reporter",
+        "//tachyon/base/containers:container_util",
         "//tachyon/base/strings:string_number_conversions",
     ],
 )

benchmark/README.md

@@ -5,18 +5,18 @@
 
 #include "absl/strings/substitute.h"
 
+#include "tachyon/base/containers/container_util.h"
 #include "tachyon/base/strings/string_number_conversions.h"
 
 namespace tachyon {
 
 SimpleECBenchmarkReporter::SimpleECBenchmarkReporter(
-    std::string_view title, const std::vector<uint64_t>& nums) {
-  title_ = std::string(title);
+    std::string_view title, const std::vector<uint64_t>& nums)
+    : SimpleBenchmarkReporter(title) {
   column_headers_.push_back("CPU");
   column_headers_.push_back("GPU");
-  for (uint64_t num : nums) {
-    targets_.push_back(base::NumberToString(num));
-  }
+  targets_ =
+      base::Map(nums, [](uint64_t num) { return base::NumberToString(num); });
   results_.resize(nums.size());
 }
 

benchmark/ec/BUILD.bazel

@@ -0,0 +1,55 @@
+load(
+    "//bazel:tachyon_cc.bzl",
+    "tachyon_cc_binary",
+    "tachyon_cc_library",
+)
+
+tachyon_cc_library(
+    name = "fft_config",
+    testonly = True,
+    srcs = ["fft_config.cc"],
+    hdrs = ["fft_config.h"],
+    deps = [
+        "//tachyon/base/console",
+        "//tachyon/base/containers:container_util",
+        "//tachyon/base/flag:flag_parser",
+        "//tachyon/base/ranges:algorithm",
+    ],
+)
+
+tachyon_cc_library(
+    name = "fft_runner",
+    testonly = True,
+    hdrs = ["fft_runner.h"],
+    deps = [
+        ":simple_fft_benchmark_reporter",
+        "//tachyon/base/time",
+    ],
+)
+
+tachyon_cc_library(
+    name = "simple_fft_benchmark_reporter",
+    testonly = True,
+    srcs = ["simple_fft_benchmark_reporter.cc"],
+    hdrs = ["simple_fft_benchmark_reporter.h"],
+    deps = [
+        "//benchmark:simple_benchmark_reporter",
+        "//tachyon/base/containers:container_util",
+        "//tachyon/base/strings:string_number_conversions",
+    ],
+)
+
+tachyon_cc_binary(
+    name = "fft_benchmark",
+    testonly = True,
+    srcs = ["fft_benchmark.cc"],
+    deps = [
+        ":fft_config",
+        ":fft_runner",
+        "//benchmark/fft/halo2",
+        "//tachyon/c/math/polynomials/univariate:bn254_univariate_evaluation_domain",
+        "//tachyon/math/elliptic_curves/bn/bn254:fr",
+        "//tachyon/math/polynomials/univariate:univariate_evaluation_domain",
+        "//tachyon/math/polynomials/univariate:univariate_evaluation_domain_factory",
+    ],
+)

benchmark/ec/simple_ec_benchmark_reporter.cc

@@ -0,0 +1,48 @@
+# (I)FFT Benchmark
+
+```
+Run on 13th Gen Intel(R) Core(TM) i9-13900K (32 X 5500 MHz CPU s)
+CPU Caches:
+  L1 Data 48 KiB (x16)
+  L1 Instruction 32 KiB (x16)
+  L2 Unified 2048 KiB (x16)
+  L3 Unified 36864 KiB (x1)
+```
+
+## FFT
+
+```shell
+bazel run -c opt --config halo2 --//:has_openmp --//:polygon_zkevm_backend --//:has_rtti --//:has_matplotlib //benchmark/fft:fft_benchmark -- -k 16 -k 17 -k 18 -k 19 -k 20 -k 21 -k 22 -k 23
+```
+
+| Exponent | Tachyon      | Halo2        |
+| :------: | ------------ | ------------ |
+|    16    | **0.001964** | 0.005087     |
+|    17    | **0.006411** | 0.00705      |
+|    18    | **0.016552** | 0.023917     |
+|    19    | **0.027728** | 0.067646     |
+|    20    | 0.057617     | **0.056551** |
+|    21    | 0.123259     | **0.086488** |
+|    22    | 0.297385     | **0.18532**  |
+|    23    | 0.619081     | **0.399886** |
+
+![image](/benchmark/fft/FFT%20Benchmark.png)
+
+## IFFT
+
+```shell
+bazel run -c opt --config halo2 --//:has_openmp --//:polygon_zkevm_backend --//:has_rtti --//:has_matplotlib //benchmark/fft:fft_benchmark -- -k 16 -k 17 -k 18 -k 19 -k 20 -k 21 -k 22 -k 23 --run_ifft
+```
+
+| Exponent | Tachyon      | Halo2        |
+| :------: | ------------ | ------------ |
+|    16    | **0.002298** | 0.004557     |
+|    17    | **0.005008** | 0.005661     |
+|    18    | **0.009929** | 0.011304     |
+|    19    | **0.039518** | 0.044937     |
+|    20    | **0.045926** | 0.08374      |
+|    21    | 0.130692     | **0.108806** |
+|    22    | 0.316693     | **0.204709** |
+|    23    | 0.637058     | **0.451356** |
+
+![image](/benchmark/fft/IFFT%20Benchmark.png)

benchmark/fft/BUILD.bazel

@@ -0,0 +1,106 @@
+#include <stddef.h>
+#include <stdint.h>
+
+#include <iostream>
+
+// clang-format off
+#include "benchmark/fft/fft_config.h"
+#include "benchmark/fft/fft_runner.h"
+#include "benchmark/fft/simple_fft_benchmark_reporter.h"
+// clang-format on
+#include "tachyon/c/math/polynomials/univariate/bn254_univariate_evaluation_domain.h"
+#include "tachyon/math/elliptic_curves/bn/bn254/fr.h"
+#include "tachyon/math/polynomials/univariate/univariate_evaluation_domain.h"
+#include "tachyon/math/polynomials/univariate/univariate_evaluation_domain_factory.h"
+
+namespace tachyon {
+
+using namespace math;
+
+extern "C" tachyon_bn254_fr* run_fft_halo2(const tachyon_bn254_fr* coeffs,
+                                           size_t n,
+                                           const tachyon_bn254_fr* omega,
+                                           uint32_t k,
+                                           uint64_t* duration_in_us);
+
+extern "C" tachyon_bn254_fr* run_ifft_halo2(const tachyon_bn254_fr* coeffs,
+                                            size_t n,
+                                            const tachyon_bn254_fr* omega_inv,
+                                            uint32_t k,
+                                            uint64_t* duration_in_us);
+
+template <typename PolyOrEvals>
+void CheckResults(bool check_results, const std::vector<PolyOrEvals>& results,
+                  const std::vector<PolyOrEvals>& results_halo2) {
+  if (check_results) {
+    CHECK(results == results_halo2) << "Results not matched";
+  }
+}
+
+template <typename Domain, typename PolyOrEvals,
+          typename RetPoly = std::conditional_t<
+              std::is_same_v<PolyOrEvals, typename Domain::Evals>,
+              typename Domain::DensePoly, typename Domain::Evals>>
+void Run(const FFTConfig& config) {
+  // NOTE(TomTaehoonKim): To remove code duplication, we named it as "(I)FFT
+  // Benchmark" in the code, but for the plots in benchmark.md, I used "FFT
+  // Benchmark" and "IFFT Benchmark" for better readability.
+  SimpleFFTBenchmarkReporter reporter("(I)FFT Benchmark", config.exponents());
+  reporter.AddVendor("halo2");
+
+  std::vector<uint64_t> degrees = config.GetDegrees();
+
+  std::cout << "Generating evaluation domain and random polys..." << std::endl;
+  std::vector<std::unique_ptr<Domain>> domains = base::Map(
+      degrees, [](uint64_t degree) { return Domain::Create(degree); });
+  std::vector<PolyOrEvals> polys = base::Map(
+      degrees, [](uint64_t degree) { return PolyOrEvals::Random(degree); });
+  std::cout << "Generation completed" << std::endl;
+
+  FFTRunner<Domain, PolyOrEvals> runner(&reporter);
+  runner.SetInputs(&polys, std::move(domains));
+
+  std::vector<RetPoly> results;
+  std::vector<RetPoly> results_halo2;
+  if constexpr (std::is_same_v<PolyOrEvals, typename Domain::Evals>) {
+    runner.Run(tachyon_bn254_univariate_evaluation_domain_ifft, degrees,
+               &results);
+    runner.RunExternal(run_ifft_halo2, config.exponents(), &results_halo2);
+    // NOLINTNEXTLINE(readability/braces)
+  } else if constexpr (std::is_same_v<PolyOrEvals,
+                                      typename Domain::DensePoly>) {
+    runner.Run(tachyon_bn254_univariate_evaluation_domain_fft, degrees,
+               &results);
+    runner.RunExternal(run_fft_halo2, config.exponents(), &results_halo2);
+  }
+  CheckResults(config.check_results(), results, results_halo2);
+
+  reporter.Show();
+}
+
+int RealMain(int argc, char** argv) {
+  using Field = bn254::Fr;
+  constexpr size_t kMaxDegree = SIZE_MAX - 1;
+  using Domain = UnivariateEvaluationDomain<Field, kMaxDegree>;
+  using DensePoly = Domain::DensePoly;
+  using Evals = Domain::Evals;
+
+  Field::Init();
+
+  FFTConfig config;
+  if (!config.Parse(argc, argv)) {
+    return 1;
+  }
+
+  if (config.run_ifft()) {
+    Run<Domain, Evals>(config);
+  } else {
+    Run<Domain, DensePoly>(config);
+  }
+
+  return 0;
+}
+
+}  // namespace tachyon
+
+int main(int argc, char** argv) { return tachyon::RealMain(argc, argv); }

benchmark/fft/FFT Benchmark.png

@@ -0,0 +1,43 @@
+#include "benchmark/fft/fft_config.h"
+
+#include <string>
+
+#include "tachyon/base/console/iostream.h"
+#include "tachyon/base/containers/container_util.h"
+#include "tachyon/base/flag/flag_parser.h"
+#include "tachyon/base/ranges/algorithm.h"
+
+namespace tachyon {
+
+bool FFTConfig::Parse(int argc, char** argv) {
+  base::FlagParser parser;
+  // clang-format off
+  parser.AddFlag<base::Flag<std::vector<uint64_t>>>(&exponents_)
+      .set_short_name("-k")
+      .set_required()
+      .set_help("Specify the exponent 'k's where the degree of poly to test is 2ᵏ.");
+  // clang-format on
+  parser.AddFlag<base::BoolFlag>(&run_ifft_)
+      .set_long_name("--run_ifft")
+      .set_help("Run IFFT benchmark. Default is FFT benchmark.");
+  parser.AddFlag<base::BoolFlag>(&check_results_)
+      .set_long_name("--check_results")
+      .set_help("Whether checks results generated by each fft runner.");
+  {
+    std::string error;
+    if (!parser.Parse(argc, argv, &error)) {
+      tachyon_cerr << error << std::endl;
+      return false;
+    }
+  }
+
+  base::ranges::sort(exponents_);  // NOLINT
+  return true;
+}
+
+std::vector<uint64_t> FFTConfig::GetDegrees() const {
+  return base::Map(exponents_,
+                   [](uint64_t exponent) { return uint64_t{1} << exponent; });
+}
+
+}  // namespace tachyon

benchmark/fft/IFFT Benchmark.png

@@ -0,0 +1,32 @@
+#ifndef BENCHMARK_FFT_FFT_CONFIG_H_
+#define BENCHMARK_FFT_FFT_CONFIG_H_
+
+#include <stdint.h>
+
+#include <vector>
+
+namespace tachyon {
+
+class FFTConfig {
+ public:
+  FFTConfig() = default;
+  FFTConfig(const FFTConfig& other) = delete;
+  FFTConfig& operator=(const FFTConfig& other) = delete;
+
+  const std::vector<uint64_t>& exponents() const { return exponents_; }
+  bool run_ifft() const { return run_ifft_; }
+  bool check_results() const { return check_results_; }
+
+  bool Parse(int argc, char** argv);
+
+  std::vector<uint64_t> GetDegrees() const;
+
+ private:
+  std::vector<uint64_t> exponents_;
+  bool run_ifft_ = false;
+  bool check_results_ = false;
+};
+
+}  // namespace tachyon
+
+#endif  // BENCHMARK_FFT_FFT_CONFIG_H_