Merge pull request #413 from kroma-network/feat/implement-poseidon2

feat: implement poseidon2

Cargo.toml

@@ -7,6 +7,8 @@ members = [
     "benchmark/fft/bellman",
     "benchmark/fft/halo2",
     "benchmark/poseidon/arkworks",
+    "benchmark/poseidon2/horizen",
+    "benchmark/poseidon2/plonky3",
     "tachyon/rs",
     "vendors/halo2",
 ]

README.md

@@ -53,9 +53,10 @@ Symbol Definitions:
 
 ### Hashes
 
-|          | CPU                | GPU |
-| -------- | ------------------ | --- |
-| Poseidon | :heavy_check_mark: | :x: |
+|           | CPU                | GPU |
+| --------- | ------------------ | --- |
+| Poseidon  | :heavy_check_mark: | :x: |
+| Poseidon2 | :heavy_check_mark: | :x: |
 
 ### Lookups
 

WORKSPACE

@@ -29,13 +29,10 @@ load("@rules_rust//rust:repositories.bzl", "rules_rust_dependencies", "rust_regi
 
 rules_rust_dependencies()
 
-# We need to change the default value of flag //tachyon/rs/base:rustc_version_ge_1.67.0
-# if we change the default rustc version.
-# See //tachyon/rs/base/BUILD.bazel.
 rust_register_toolchains(
     edition = "2021",
     versions = [
-        "1.66.1",
+        "1.77.1",
     ],
 )
 
@@ -58,6 +55,8 @@ crates_repository(
         "//benchmark/fft/bellman:Cargo.toml",
         "//benchmark/fft/halo2:Cargo.toml",
         "//benchmark/poseidon/arkworks:Cargo.toml",
+        "//benchmark/poseidon2/horizen:Cargo.toml",
+        "//benchmark/poseidon2/plonky3:Cargo.toml",
         "//tachyon/rs:Cargo.toml",
         "//vendors/halo2:Cargo.toml",
     ],

benchmark/fft/arkworks/Cargo.toml

@@ -9,7 +9,6 @@ Arkworks FFT Benchmark
 """
 license = "MIT OR Apache-2.0"
 repository = "https://github.com/kroma-network/tachyon"
-documentation = "https://docs.rs/halo2_benchmark"
 readme = "README.md"
 categories = ["cryptography"]
 keywords = ["tachyon", "benchmark", "arkworks"]

benchmark/fft/bellman/Cargo.toml

@@ -9,7 +9,6 @@ Bellman FFT Benchmark
 """
 license = "MIT OR Apache-2.0"
 repository = "https://github.com/kroma-network/tachyon"
-documentation = "https://docs.rs/bellman_benchmark"
 readme = "README.md"
 categories = ["cryptography"]
 keywords = ["tachyon", "benchmark", "bellman"]

benchmark/fft/halo2/Cargo.toml

@@ -9,7 +9,6 @@ Halo2 FFT Benchmark
 """
 license = "MIT OR Apache-2.0"
 repository = "https://github.com/kroma-network/tachyon"
-documentation = "https://docs.rs/halo2_benchmark"
 readme = "README.md"
 categories = ["cryptography"]
 keywords = ["tachyon", "benchmark", "halo2"]

benchmark/msm/arkworks/Cargo.toml

@@ -9,7 +9,6 @@ Arkworks MSM Benchmark
 """
 license = "MIT OR Apache-2.0"
 repository = "https://github.com/kroma-network/tachyon"
-documentation = "https://docs.rs/arkworks_benchmark"
 readme = "README.md"
 categories = ["cryptography"]
 keywords = ["tachyon", "benchmark", "arkworks"]

benchmark/msm/bellman/Cargo.toml

@@ -9,7 +9,6 @@ Bellman MSM Benchmark
 """
 license = "MIT OR Apache-2.0"
 repository = "https://github.com/kroma-network/tachyon"
-documentation = "https://docs.rs/bellman_benchmark"
 readme = "README.md"
 categories = ["cryptography"]
 keywords = ["tachyon", "benchmark", "bellman"]

benchmark/msm/halo2/Cargo.toml

@@ -9,7 +9,6 @@ Halo2 MSM Benchmark
 """
 license = "MIT OR Apache-2.0"
 repository = "https://github.com/kroma-network/tachyon"
-documentation = "https://docs.rs/halo2_benchmark"
 readme = "README.md"
 categories = ["cryptography"]
 keywords = ["tachyon", "benchmark", "halo2"]

benchmark/msm/msm_benchmark.cc

@@ -55,12 +55,16 @@ int RealMain(int argc, char** argv) {
   runner.Run(tachyon_bn254_g1_affine_msm, msm, point_nums, &results);
   for (const MSMConfig::Vendor vendor : config.vendors()) {
     std::vector<bn254::G1JacobianPoint> results_vendor;
-    if (vendor == MSMConfig::Vendor::kArkworks) {
-      runner.RunExternal(run_msm_arkworks, point_nums, &results_vendor);
-    } else if (vendor == MSMConfig::Vendor::kBellman) {
-      runner.RunExternal(run_msm_bellman, point_nums, &results_vendor);
-    } else if (vendor == MSMConfig::Vendor::kHalo2) {
-      runner.RunExternal(run_msm_halo2, point_nums, &results_vendor);
+    switch (vendor) {
+      case MSMConfig::Vendor::kArkworks:
+        runner.RunExternal(run_msm_arkworks, point_nums, &results_vendor);
+        break;
+      case MSMConfig::Vendor::kBellman:
+        runner.RunExternal(run_msm_bellman, point_nums, &results_vendor);
+        break;
+      case MSMConfig::Vendor::kHalo2:
+        runner.RunExternal(run_msm_halo2, point_nums, &results_vendor);
+        break;
     }
 
     if (config.check_results()) {

benchmark/poseidon/BUILD.bazel

@@ -4,6 +4,8 @@ load(
     "tachyon_cc_library",
 )
 
+package(default_visibility = ["//benchmark/poseidon2:__pkg__"])
+
 tachyon_cc_library(
     name = "poseidon_config",
     testonly = True,

benchmark/poseidon/README.md

@@ -10,21 +10,21 @@ CPU Caches:
 ```
 
 ```shell
-bazel run -c opt --//:has_openmp --//:has_rtti --//:has_matplotlib //benchmark/poseidon:poseidon_benchmark -- -s 10000 -a 10000
+bazel run -c opt --//:has_openmp --//:has_rtti --//:has_matplotlib //benchmark/poseidon:poseidon_benchmark
 ```
 
-| Repetition | Tachyon  | Arkworks     |
-| :--------: | -------- | ------------ |
-|     0      | 0.365703 | **0.26695**  |
-|     1      | 0.383136 | **0.267033** |
-|     2      | 0.383506 | **0.272285** |
-|     3      | 0.374358 | **0.257264** |
-|     4      | 0.370592 | **0.25442**  |
-|     5      | 0.371024 | **0.2545**   |
-|     6      | 0.373796 | **0.263127** |
-|     7      | 0.375338 | **0.27176**  |
-|     8      | 0.374944 | **0.267881** |
-|     9      | 0.374335 | **0.262994** |
-|    avg     | 0.374673 | **0.263821** |
+| Repetition | Tachyon  | Arkworks      |
+| :--------: | -------- | ------------- |
+|     0      | 0.000124 | **0.000114**  |
+|     1      | 0.000125 | **0.00011**   |
+|     2      | 0.000124 | **0.00011**   |
+|     3      | 0.000124 | **0.000106**  |
+|     4      | 0.000124 | **0.00011**   |
+|     5      | 0.000124 | **0.000109**  |
+|     6      | 0.000125 | **0.000108**  |
+|     7      | 0.000123 | **0.00011**   |
+|     8      | 0.000123 | **0.000109**  |
+|     9      | 0.000124 | **0.000106**  |
+|    avg     | 0.000124 | **0.0001092** |
 
 ![image](/benchmark/poseidon/Poseidon%20Benchmark.png)

benchmark/poseidon/arkworks/Cargo.toml

@@ -9,7 +9,6 @@ Arkworks Poseidon Hash Benchmark
 """
 license = "MIT OR Apache-2.0"
 repository = "https://github.com/kroma-network/tachyon"
-documentation = "https://docs.rs/arkworks_benchmark"
 readme = "README.md"
 categories = ["cryptography"]
 keywords = ["tachyon", "benchmark", "arkworks"]
@@ -18,9 +17,9 @@ publish = false
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-ark-ff = { version = "^0.4.0", default-features = false }
+ark-ff = { version = "^0.4.0", features = ["asm"] }
 ark-bn254 = "0.4.0"
-ark-crypto-primitives = { git = "https://github.com/arkworks-rs/crypto-primitives.git", features = [
+ark-crypto-primitives = { git = "https://github.com/kroma-network/crypto-primitives.git", features = [
     "sponge",
-], tag = "v0.4.0" }
+], rev = "99f5aff" }
 tachyon_rs = { path = "../../../tachyon/rs" }

benchmark/poseidon/arkworks/src/lib.rs

@@ -1,17 +1,12 @@
 use ark_bn254::Fr;
 use ark_crypto_primitives::sponge::poseidon::traits::find_poseidon_ark_and_mds;
 use ark_crypto_primitives::sponge::poseidon::{PoseidonConfig, PoseidonSponge};
-use ark_crypto_primitives::sponge::{CryptographicSponge, FieldBasedCryptographicSponge};
-use std::{mem, slice, time::Instant};
+use ark_crypto_primitives::sponge::CryptographicSponge;
+use std::time::Instant;
 use tachyon_rs::math::elliptic_curves::bn::bn254::Fr as CppFr;
 
 #[no_mangle]
-pub extern "C" fn run_poseidon_arkworks(
-    pre_images: *const CppFr,
-    absorbing_num: usize,
-    squeezing_num: usize,
-    duration: *mut u64,
-) -> *mut CppFr {
+pub extern "C" fn run_poseidon_arkworks(duration: *mut u64) -> *mut CppFr {
     let (ark, mds) = find_poseidon_ark_and_mds::<Fr>(254, 8, 8, 63, 0);
     let poseidon_config = PoseidonConfig {
         full_rounds: 8,
@@ -23,22 +18,13 @@ pub extern "C" fn run_poseidon_arkworks(
         capacity: 1,
     };
 
-    let pre_images = unsafe {
-        let pre_images: &[CppFr] = slice::from_raw_parts(pre_images, absorbing_num);
-        let pre_images: &[Fr] = mem::transmute(pre_images);
-
-        pre_images
-    };
     let mut sponge = PoseidonSponge::new(&poseidon_config);
 
     let start = Instant::now();
-    for pre_image in pre_images {
-        sponge.absorb(pre_image);
-    }
-    let squeezed_elements = sponge.squeeze_native_field_elements(squeezing_num);
+    sponge.permute();
     unsafe {
         duration.write(start.elapsed().as_micros() as u64);
     }
 
-    Box::into_raw(Box::new(squeezed_elements[0])) as *mut CppFr
+    Box::into_raw(Box::new(sponge.state[1])) as *mut CppFr
 }

benchmark/poseidon/poseidon_benchmark.cc

@@ -16,9 +16,7 @@ using namespace crypto;
 
 using Field = math::bn254::Fr;
 
-extern "C" tachyon_bn254_fr* run_poseidon_arkworks(
-    const tachyon_bn254_fr* pre_images, size_t aborbing_num,
-    size_t squeezing_num, uint64_t* duration);
+extern "C" tachyon_bn254_fr* run_poseidon_arkworks(uint64_t* duration);
 
 int RealMain(int argc, char** argv) {
   tachyon::PoseidonConfig config;
@@ -32,16 +30,11 @@ int RealMain(int argc, char** argv) {
   reporter.AddVendor("arkworks");
   PoseidonBenchmarkRunner<Field> runner(&reporter, &config);
 
-  std::vector<Field> results;
-  results.reserve(config.repeating_num());
-  runner.Run(&results);
-
-  std::vector<Field> results_vendor;
-  results_vendor.reserve(config.repeating_num());
-  runner.RunExternal(run_poseidon_arkworks, &results_vendor);
+  Field result = runner.Run();
+  Field result_arkworks = runner.RunExternal(run_poseidon_arkworks);
 
   if (config.check_results()) {
-    CHECK(results == results_vendor) << "Result not matched";
+    CHECK_EQ(result, result_arkworks) << "Result not matched";
   }
 
   reporter.AddAverageToLastRow();

benchmark/poseidon/poseidon_benchmark_runner.h

@@ -26,53 +26,43 @@ class PoseidonBenchmarkRunner {
  public:
   using CPrimeField = typename c::base::TypeTraits<Field>::CType;
 
-  typedef CPrimeField* (*PoseidonExternalFn)(const CPrimeField* pre_images,
-                                             size_t absorbing_num,
-                                             size_t squeezing_num,
-                                             uint64_t* duration);
+  typedef CPrimeField* (*PoseidonExternalFn)(uint64_t* duration);
 
   PoseidonBenchmarkRunner(SimplePoseidonBenchmarkReporter* reporter,
                           PoseidonConfig* config)
-      : reporter_(reporter), config_(config) {
-    pre_images_ = base::CreateVector(config->absorbing_num(),
-                                     []() { return Field::Random(); });
-  }
+      : reporter_(reporter), config_(config) {}
 
-  void Run(std::vector<Field>* results) {
+  Field Run() {
+    Field ret;
     for (size_t i = 0; i < config_->repeating_num(); ++i) {
       crypto::PoseidonConfig<Field> config =
           crypto::PoseidonConfig<Field>::CreateCustom(8, 5, 8, 63, 0);
       crypto::PoseidonSponge<Field> sponge(config);
       base::TimeTicks start = base::TimeTicks::Now();
-      for (size_t j = 0; j < config_->absorbing_num(); ++j) {
-        sponge.Absorb(pre_images_[j]);
-      }
-      std::vector<Field> squeezed_elements =
-          sponge.SqueezeFieldElements(config_->squeezing_num());
+      sponge.Permute();
       reporter_->AddTime(i, (base::TimeTicks::Now() - start).InSecondsF());
-      results->push_back(std::move(squeezed_elements[0]));
+      if (i == 0) {
+        ret = sponge.state.elements[1];
+      }
     }
+    return ret;
   }
 
-  void RunExternal(PoseidonExternalFn fn, std::vector<Field>* results) {
+  Field RunExternal(PoseidonExternalFn fn) {
+    std::unique_ptr<CPrimeField> ret;
     for (size_t i = 0; i < config_->repeating_num(); ++i) {
-      std::unique_ptr<CPrimeField> last_squeezed_element;
       uint64_t duration_in_us;
-      last_squeezed_element.reset(fn(
-          reinterpret_cast<const CPrimeField*>(pre_images_.data()),
-          config_->absorbing_num(), config_->squeezing_num(), &duration_in_us));
+      ret.reset(fn(&duration_in_us));
       reporter_->AddTime(i, base::Microseconds(duration_in_us).InSecondsF());
-      results->push_back(
-          *reinterpret_cast<Field*>(last_squeezed_element.get()));
     }
+    return *reinterpret_cast<Field*>(ret.get());
   }
 
  private:
   // not owned
   SimplePoseidonBenchmarkReporter* const reporter_;
   // not owned
   PoseidonConfig* const config_;
-  std::vector<Field> pre_images_;
 };
 
 }  // namespace tachyon

benchmark/poseidon/poseidon_config.cc

@@ -13,22 +13,14 @@ bool PoseidonConfig::Parse(int argc, char** argv) {
   parser.AddFlag<base::Flag<size_t>>(&repeating_num_)
       .set_short_name("-n")
       .set_help("Specify the number of repetition 'n'. By default, 10.");
-  parser.AddFlag<base::Flag<size_t>>(&absorbing_num_)
-      .set_short_name("-a")
-      .set_help("Specify the number of absorptions 'a'. By default, 100.");
-  parser.AddFlag<base::Flag<size_t>>(&squeezing_num_)
-      .set_short_name("-s")
-      .set_help("Specify the number of squeeze 's'. By default, 100.");
   parser.AddFlag<base::Flag<bool>>(&check_results_)
       .set_long_name("--check_results")
       .set_help("Whether checks results generated by each poseidon runner.");
 
-  {
-    std::string error;
-    if (!parser.Parse(argc, argv, &error)) {
-      tachyon_cerr << error << std::endl;
-      return false;
-    }
+  std::string error;
+  if (!parser.Parse(argc, argv, &error)) {
+    tachyon_cerr << error << std::endl;
+    return false;
   }
 
   return true;

benchmark/poseidon/poseidon_config.h

@@ -13,16 +13,12 @@ class PoseidonConfig {
 
   bool check_results() const { return check_results_; }
   size_t repeating_num() const { return repeating_num_; }
-  size_t absorbing_num() const { return absorbing_num_; }
-  size_t squeezing_num() const { return squeezing_num_; }
 
   bool Parse(int argc, char** argv);
 
  private:
   bool check_results_ = false;
   size_t repeating_num_ = 10;
-  size_t absorbing_num_ = 100;
-  size_t squeezing_num_ = 100;
 };
 
 }  // namespace tachyon