test(benchmark): introduce FFTBatch and CosetLDEBatch benchmark

Cargo.Bazel.lock

@@ -1,5 +1,5 @@
 {
- "checksum": "e580ed4d35bc3975a2a92b01f0b25307c5311abb39e26d03b4f436f05ad77d8a",
+ "checksum": "fe9ae6f4561ad04bed9e5ebe9cba69153f8f8ff4c819e68d1aa1ffc9128967e2",
  "crates": {
  "addchain 0.2.0": {
  "name": "addchain",
@@ -12906,6 +12906,52 @@
  },
  "license": "Apache-2.0 OR MIT"
  },
+ "plonky3_batch_fft_benchmark 0.0.1": {
+ "name": "plonky3_batch_fft_benchmark",
+ "version": "0.0.1",
+ "repository": null,
+ "targets": [
+ {
+ "Library": {
+ "crate_name": "plonky3_batch_fft_benchmark",
+ "crate_root": "src/lib.rs",
+ "srcs": [
+ "**/*.rs"
+ ]
+ }
+ }
+ ],
+ "library_target_name": "plonky3_batch_fft_benchmark",
+ "common_attrs": {
+ "compile_data_glob": [
+ "**"
+ ],
+ "deps": {
+ "common": [
+ {
+ "id": "p3-baby-bear 0.1.3-succinct",
+ "target": "p3_baby_bear"
+ },
+ {
+ "id": "p3-dft 0.1.3-succinct",
+ "target": "p3_dft"
+ },
+ {
+ "id": "p3-field 0.1.3-succinct",
+ "target": "p3_field"
+ },
+ {
+ "id": "p3-matrix 0.1.3-succinct",
+ "target": "p3_matrix"
+ }
+ ],
+ "selects": {}
+ },
+ "edition": "2021",
+ "version": "0.0.1"
+ },
+ "license": "MIT OR Apache-2.0"
+ },
  "plonky3_poseidon2_benchmark 0.0.1": {
  "name": "plonky3_poseidon2_benchmark",
  "version": "0.0.1",
@@ -21494,6 +21540,7 @@
  "halo2_fft_benchmark 0.0.1": "benchmark/fft/halo2",
  "halo2_msm_benchmark 0.0.1": "benchmark/msm/halo2",
  "horizen_poseidon2_benchmark 0.0.1": "benchmark/poseidon2/horizen",
+ "plonky3_batch_fft_benchmark 0.0.1": "benchmark/fft_batch/plonky3",
  "plonky3_poseidon2_benchmark 0.0.1": "benchmark/poseidon2/plonky3",
  "tachyon_plonky3 0.0.1": "vendors/plonky3",
  "tachyon_rs 0.0.1": "tachyon/rs",

Cargo.toml

@@ -6,6 +6,7 @@ members = [
  "benchmark/fft/arkworks",
  "benchmark/fft/bellman",
  "benchmark/fft/halo2",
+ "benchmark/fft_batch/plonky3",
  "benchmark/poseidon/arkworks",
  "benchmark/poseidon2/horizen",
  "benchmark/poseidon2/plonky3",

WORKSPACE

@@ -57,6 +57,7 @@ crates_repository(
  "//benchmark/fft/arkworks:Cargo.toml",
  "//benchmark/fft/bellman:Cargo.toml",
  "//benchmark/fft/halo2:Cargo.toml",
+ "//benchmark/fft_batch/plonky3:Cargo.toml",
  "//benchmark/poseidon/arkworks:Cargo.toml",
  "//benchmark/poseidon2/horizen:Cargo.toml",
  "//benchmark/poseidon2/plonky3:Cargo.toml",

benchmark/fft_batch/BUILD.bazel

@@ -0,0 +1,50 @@
+load(
+ "//bazel:tachyon_cc.bzl",
+ "tachyon_cc_binary",
+ "tachyon_cc_library",
+)
+
+tachyon_cc_library(
+ name = "fft_batch_config",
+ testonly = True,
+ srcs = ["fft_batch_config.cc"],
+ hdrs = ["fft_batch_config.h"],
+ deps = [
+ "//benchmark:config",
+ "//benchmark:field_type",
+ "//tachyon/base/console",
+ "//tachyon/base/containers:container_util",
+ "//tachyon/base/flag:flag_parser",
+ ],
+)
+
+tachyon_cc_library(
+ name = "fft_batch_runner",
+ testonly = True,
+ hdrs = ["fft_batch_runner.h"],
+ deps = [
+ ":fft_batch_config",
+ "//benchmark:simple_reporter",
+ "//tachyon/base/containers:container_util",
+ "//tachyon/base/time",
+ "//tachyon/c/math/matrix",
+ "//tachyon/math/matrix:matrix_types",
+ "//tachyon/math/polynomials/univariate:radix2_evaluation_domain",
+ ],
+)
+
+tachyon_cc_binary(
+ name = "fft_batch_benchmark",
+ testonly = True,
+ srcs = ["fft_batch_benchmark.cc"],
+ deps = [
+ ":fft_batch_config",
+ ":fft_batch_runner",
+ "//benchmark:simple_reporter",
+ "//benchmark/fft_batch/plonky3",
+ "//tachyon/c/math/finite_fields/baby_bear",
+ "//tachyon/math/matrix:matrix_types",
+ "//tachyon/math/polynomials/univariate:radix2_evaluation_domain",
+ "@com_google_absl//absl/strings",
+ ],
+)

benchmark/fft_batch/README.md

@@ -0,0 +1,90 @@
+# FFTBatch/CosetLDEBatch Benchmark
+
+## CPU
+
+```
+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)
+
+Run on Apple M3 Pro (12 X 4050 MHz)
+CPU Caches:
+ L1 Data 64 KiB (x12)
+ L1 Instruction 128 KiB (x12)
+ L2 Unified 4096 KiB (x12)
+```
+
+### FFTBatch
+
+```shell
+bazel run -c opt --//:has_openmp --//:has_rtti --//:has_matplotlib //benchmark/fft_batch:fft_batch_benchmark -- -k 16 -k 17 -k 18 -k 19 -k 20 -k 21 -k 22 -k 23 --vendor plonky3 -p baby_bear --check_results
+```
+
+#### On Intel i9-13900K
+
+| Exponent | TachyonCPU | Plonky3 |
+| :------- | ------------ | ------------ |
+| 16 | **0.003543** | 0.005809 |
+| 17 | **0.007168** | 0.009239 |
+| 18 | 0.027791 | **0.023848** |
+| 19 | 0.063468 | **0.049085** |
+| 20 | 0.133178 | **0.102343** |
+| 21 | 0.238817 | **0.208557** |
+| 22 | 0.507061 | **0.427260** |
+| 23 | 1.11136 | **0.922439** |
+
+![image](/benchmark/fft_batch/fft_batch_benchmark_ubuntu_i9.png)
+
+#### On Mac M3 Pro
+
+| Exponent | TachyonCPU | Plonky3 |
+| :------- | ------------ | -------- |
+| 16 | **0.007926** | 0.01639 |
+| 17 | **0.016391** | 0.024321 |
+| 18 | **0.035098** | 0.085961 |
+| 19 | **0.076266** | 0.096928 |
+| 20 | **0.145975** | 0.151024 |
+| 21 | **0.309752** | 0.339549 |
+| 22 | **0.674991** | 2.66605 |
+| 23 | **1.727520** | 7.79002 |
+
+![image](/benchmark/fft_batch/fft_batch_benchmark_mac_m3.png)
+
+### CosetLDEBatch
+
+```shell
+bazel run -c opt --//:has_openmp --//:has_rtti --//:has_matplotlib //benchmark/fft_batch:fft_batch_benchmark -- -k 16 -k 17 -k 18 -k 19 -k 20 -k 21 -k 22 -k 23 --vendor plonky3 -p baby_bear --run_coset_lde --check_results
+```
+
+#### On Intel i9-13900K
+
+| Exponent | TachyonCPU | Plonky3 |
+| :------- | ------------ | -------- |
+| 16 | **0.008384** | 0.018529 |
+| 17 | **0.017164** | 0.043266 |
+| 18 | **0.052999** | 0.093348 |
+| 19 | **0.128624** | 0.19531 |
+| 20 | **0.246412** | 0.418079 |
+| 21 | **0.508587** | 0.816136 |
+| 22 | **1.071360** | 1.63289 |
+| 23 | **2.225130** | 3.53348 |
+
+![image](/benchmark/fft_batch/coset_lde_batch_benchmark_ubuntu_i9.png)
+
+#### On Mac M3 Pro
+
+| Exponent | TachyonCPU | Plonky3 |
+| :------- | ------------ | ------------ |
+| 16 | **0.013503** | 0.020251 |
+| 17 | **0.028850** | 0.044731 |
+| 18 | 0.061675 | **0.055980** |
+| 19 | 0.132670 | **0.113062** |
+| 20 | **0.263582** | 0.300933 |
+| 21 | **0.548267** | 0.573641 |
+| 22 | **1.164400** | 1.19051 |
+| 23 | **2.818220** | 5.11368 |
+
+![image](/benchmark/fft_batch/coset_lde_batch_benchmark_mac_m3.png)

benchmark/fft_batch/fft_batch_benchmark.cc

@@ -0,0 +1,112 @@
+#include <iostream>
+
+#include "absl/strings/substitute.h"
+
+// clang-format off
+#include "benchmark/fft_batch/fft_batch_config.h"
+#include "benchmark/fft_batch/fft_batch_runner.h"
+#include "benchmark/simple_reporter.h"
+// clang-format on
+#include "tachyon/base/console/iostream.h"
+#include "tachyon/c/math/finite_fields/baby_bear/baby_bear_type_traits.h"
+
+namespace tachyon::benchmark {
+
+extern "C" tachyon_baby_bear* run_fft_batch_plonky3_baby_bear(
+ const tachyon_baby_bear* data, uint32_t n_log, size_t batch_size,
+ uint64_t* duration);
+
+extern "C" tachyon_baby_bear* run_coset_lde_batch_plonky3_baby_bear(
+ const tachyon_baby_bear* data, uint32_t n_log, size_t batch_size,
+ uint64_t* duration);
+
+template <typename F>
+void CheckResults(bool check_results,
+ const std::vector<math::RowMajorMatrix<F>>& results,
+ const std::vector<math::RowMajorMatrix<F>>& results_vendor) {
+ if (check_results) {
+ CHECK(results == results_vendor) << "Results not matched";
+ }
+}
+
+template <typename F>
+int Run(const FFTBatchConfig& config) {
+ using Domain = math::Radix2EvaluationDomain<F, SIZE_MAX - 1>;
+
+ F::Init();
+
+ std::string name;
+ if (config.run_coset_lde()) {
+ name = absl::Substitute("CosetLDEBatch Benchmark (Batch Size: $0)",
+ config.batch_size());
+ } else {
+ name = absl::Substitute("FFTBatch Benchmark (Batch Size: $0)",
+ config.batch_size());
+ }
+
+ SimpleReporter reporter(name);
+
+ reporter.SetXLabel("Degree (2ˣ)");
+ reporter.SetColumnLabels(base::Map(config.exponents(), [](uint32_t exponent) {
+ return base::NumberToString(exponent);
+ }));
+
+ std::vector<size_t> degrees = config.GetDegrees();
+
+ std::cout << "Generating evaluation domain and random matrices..."
+ << std::endl;
+ std::vector<std::unique_ptr<Domain>> domains =
+ base::Map(degrees, [](size_t degree) { return Domain::Create(degree); });
+ std::vector<math::RowMajorMatrix<F>> inputs =
+ base::Map(degrees, [&config](size_t degree) {
+ math::RowMajorMatrix<F> matrix =
+ math::RowMajorMatrix<F>::Random(degree, config.batch_size());
+ return matrix;
+ });
+
+ FFTBatchRunner<Domain> runner(reporter, config);
+ runner.set_domains(absl::MakeSpan(domains));
+ runner.set_inputs(absl::MakeSpan(inputs));
+
+ std::vector<math::RowMajorMatrix<F>> results;
+ runner.Run(Vendor::TachyonCPU(), results, config.run_coset_lde());
+ for (Vendor vendor : config.vendors()) {
+ std::vector<math::RowMajorMatrix<F>> results_vendor;
+ if (vendor.value() == Vendor::kPlonky3) {
+ if (config.run_coset_lde()) {
+ runner.RunExternal(vendor, run_coset_lde_batch_plonky3_baby_bear,
+ results_vendor);
+ } else {
+ runner.RunExternal(vendor, run_fft_batch_plonky3_baby_bear,
+ results_vendor);
+ }
+ CheckResults(config.check_results(), results, results_vendor);
+ } else {
+ tachyon_cerr << "Unsupported vendor\n";
+ return 1;
+ }
+ }
+
+ reporter.Show();
+ return 0;
+}
+
+int RealMain(int argc, char** argv) {
+ FFTBatchConfig config;
+ if (!config.Parse(argc, argv)) {
+ return 1;
+ }
+
+ if (config.prime_field().value() == FieldType::kBabyBear) {
+ return Run<math::BabyBear>(config);
+ } else {
+ tachyon_cerr << "Unsupported prime field\n";
+ return 1;
+ }
+}
+
+} // namespace tachyon::benchmark
+
+int main(int argc, char** argv) {
+ return tachyon::benchmark::RealMain(argc, argv);
+}