Skip to content

Commit

Permalink
Merge pull request kroma-network#435 from kroma-network/feat/enable-p…
Browse files Browse the repository at this point in the history 
…oseidon-to-hash-packed-prime-fields

feat(math): enable poseidon to hash packed prime fields
  • Loading branch information
@chokobole
chokobole authored Jun 18, 2024
2 parents 24db7b5 + 7cc8138 commit c9e6f93
Show file tree
Hide file tree
Showing 56 changed files with 919 additions and 393 deletions.
Binary file added benchmark/poseidon/Poseidon Benchmark MacM3.png
View file
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file modified benchmark/poseidon/Poseidon Benchmark.png
View file
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
52 changes: 39 additions & 13 deletions benchmark/poseidon/README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,24 +7,50 @@ CPU Caches:
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)
```

```shell
bazel run -c opt --//:has_openmp --//:has_rtti --//:has_matplotlib //benchmark/poseidon:poseidon_benchmark
```

| 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** |
## On Intel i9-13900K

| Repetition | Tachyon | Arkworks |
| :--------: | ------------ | --------- |
| 0 | **9.9e-05** | 0.000122 |
| 1 | **9.9e-05** | 0.000119 |
| 2 | **9.9e-05** | 0.000117 |
| 3 | **9.9e-05** | 0.000119 |
| 4 | **9.9e-05** | 0.000118 |
| 5 | **0.000101** | 0.000116 |
| 6 | **9.8e-05** | 0.000118 |
| 7 | **9.9e-05** | 0.000118 |
| 8 | **9.9e-05** | 0.000118 |
| 9 | **9.9e-05** | 0.000118 |
| avg | **9.91e-05** | 0.0001183 |

![image](/benchmark/poseidon/Poseidon%20Benchmark.png)

## On Mac M3 Pro

| Repetition | Tachyon | Arkworks |
| :--------: | ------------- | --------- |
| 0 | **0.000112** | 0.000131 |
| 1 | **0.000118** | 0.000125 |
| 2 | **0.000111** | 0.000121 |
| 3 | **0.000103** | 0.000121 |
| 4 | **0.000104** | 0.00012 |
| 5 | **0.000103** | 0.000117 |
| 6 | **0.000103** | 0.000118 |
| 7 | **0.000104** | 0.000117 |
| 8 | **0.0001** | 0.000118 |
| 9 | **0.000113** | 0.000118 |
| avg | **0.0001071** | 0.0001206 |

![image](/benchmark/poseidon/Poseidon%20Benchmark%20MacM3.png)
19 changes: 19 additions & 0 deletions tachyon/crypto/hashes/sponge/BUILD.bazel
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,25 @@ tachyon_cc_library(
],
)

tachyon_cc_library(
name = "sponge_config",
hdrs = ["sponge_config.h"],
deps = [
"//tachyon:export",
"//tachyon/base/buffer:copyable",
],
)

tachyon_cc_library(
name = "sponge_state",
hdrs = ["sponge_state.h"],
deps = [
":duplex_sponge_mode",
":sponge_config",
"//tachyon/math/matrix:matrix_types",
],
)

tachyon_cc_library(
name = "truncated_permutation",
hdrs = ["truncated_permutation.h"],
Expand Down
20 changes: 9 additions & 11 deletions tachyon/crypto/hashes/sponge/poseidon/BUILD.bazel
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@ tachyon_cc_library(
deps = [
":poseidon_config",
":poseidon_sponge_base",
":poseidon_state",
"//tachyon/base/buffer:copyable",
"//tachyon/crypto/hashes/sponge:sponge_state",
],
)

Expand All @@ -31,6 +31,8 @@ tachyon_cc_library(
hdrs = ["poseidon_config_base.h"],
deps = [
"//tachyon/base/buffer:copyable",
"//tachyon/crypto/hashes/sponge:sponge_config",
"//tachyon/math/finite_fields:finite_field_traits",
"//tachyon/math/matrix:matrix_types",
],
)
Expand All @@ -44,14 +46,18 @@ tachyon_cc_library(
tachyon_cc_library(
name = "poseidon_config_entry_base",
hdrs = ["poseidon_config_entry_base.h"],
deps = [":poseidon_grain_lfsr"],
deps = [
":poseidon_grain_lfsr",
"//tachyon/math/finite_fields:finite_field_traits",
],
)

tachyon_cc_library(
name = "poseidon_grain_lfsr",
hdrs = ["poseidon_grain_lfsr.h"],
deps = [
"//tachyon/math/base:big_int",
"//tachyon/math/finite_fields:finite_field_traits",
"//tachyon/math/matrix:matrix_types",
"//tachyon/math/matrix:prime_field_num_traits",
],
Expand All @@ -68,15 +74,6 @@ tachyon_cc_library(
],
)

tachyon_cc_library(
name = "poseidon_state",
hdrs = ["poseidon_state.h"],
deps = [
"//tachyon/crypto/hashes/sponge:duplex_sponge_mode",
"//tachyon/math/matrix:matrix_types",
],
)

tachyon_cc_unittest(
name = "poseidon_unittests",
srcs = [
Expand All @@ -90,6 +87,7 @@ tachyon_cc_unittest(
"//tachyon/base/buffer:vector_buffer",
"//tachyon/math/elliptic_curves/bls12/bls12_381:fr",
"//tachyon/math/elliptic_curves/bn/bn254:fr",
"//tachyon/math/finite_fields/baby_bear:packed_baby_bear",
"//tachyon/math/finite_fields/test:finite_field_test",
],
)
74 changes: 52 additions & 22 deletions tachyon/crypto/hashes/sponge/poseidon/poseidon.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
#include "tachyon/base/buffer/copyable.h"
#include "tachyon/crypto/hashes/sponge/poseidon/poseidon_config.h"
#include "tachyon/crypto/hashes/sponge/poseidon/poseidon_sponge_base.h"
#include "tachyon/crypto/hashes/sponge/poseidon/poseidon_state.h"
#include "tachyon/crypto/hashes/sponge/sponge_state.h"

namespace tachyon {
namespace crypto {
Expand All @@ -26,31 +26,62 @@ namespace crypto {
// Squeeze: Squeeze elements out of the sponge.
// This implementation of Poseidon is entirely Fractal's implementation in
// [COS20][cos] with small syntax changes. See https://eprint.iacr.org/2019/1076
template <typename PrimeField>
struct PoseidonSponge final
: public PoseidonSpongeBase<PoseidonSponge<PrimeField>> {
using F = PrimeField;

template <typename F>
struct PoseidonSponge final : public PoseidonSpongeBase<PoseidonSponge<F>> {
// Sponge Config
PoseidonConfig<F> config;

// Sponge State
PoseidonState<F> state;
SpongeState<F> state;

PoseidonSponge() = default;
explicit PoseidonSponge(const PoseidonConfig<F>& config)
: config(config), state(config.rate + config.capacity) {}
PoseidonSponge(const PoseidonConfig<F>& config, const PoseidonState<F>& state)
PoseidonSponge(const PoseidonConfig<F>& config, const SpongeState<F>& state)
: config(config), state(state) {}
PoseidonSponge(const PoseidonConfig<F>& config, PoseidonState<F>&& state)
PoseidonSponge(const PoseidonConfig<F>& config, SpongeState<F>&& state)
: config(config), state(std::move(state)) {}

// PoseidonSpongeBase methods
void ApplyARK(Eigen::Index round_number, bool) {
state.elements += config.ark.row(round_number);
}

void ApplyMix(bool) { state.elements = config.mds * state.elements; }
void ApplyMix(bool) {
// NOTE (chokobole): Eigen matrix multiplication has a computational
// overhead unlike naive matrix multiplication.
//
// Example: Multiplying a 2x2 matrix by a 2x1 vector:
//
// +----+----+ +----+
// | m₀ | m₁ | * | v₀ |
// +----+----+ +----+
// | m₂ | m₃ | | v₁ |
// +----+----+ +----+
//
// The operations involved in this multiplication are as follows:
//
// 1 * 1
// 1 * 1
// m₀ * v₀
// m₀v₀ + 0
// m₂ * v₀
// m₂v₀ + 0
// m₁ * v₁
// m₁v₁ + m₀v₀
// m₃ * v₁
// m₃v₁ + m₂v₀
// m₁v₁ + m₀v₀ * 1
// m₁v₁ + m₀v₀ + 0
// m₃v₁ + m₂v₀ * 1
// m₃v₁ + m₂v₀ + 0
math::Vector<F> elements(state.elements.size());
for (Eigen::Index i = 0; i < config.mds.rows(); ++i) {
for (Eigen::Index j = 0; j < config.mds.cols(); ++j) {
elements[i] += config.mds(i, j) * state.elements[j];
}
}
state.elements = std::move(elements);
}

bool operator==(const PoseidonSponge& other) const {
return config == other.config && state == other.state;
Expand All @@ -60,28 +91,28 @@ struct PoseidonSponge final
}
};

template <typename PrimeField>
struct CryptographicSpongeTraits<PoseidonSponge<PrimeField>> {
using F = PrimeField;
template <typename Field>
struct CryptographicSpongeTraits<PoseidonSponge<Field>> {
using F = Field;
constexpr static bool kApplyMixAtFront = false;
};

} // namespace crypto

namespace base {

template <typename PrimeField>
class Copyable<crypto::PoseidonSponge<PrimeField>> {
template <typename F>
class Copyable<crypto::PoseidonSponge<F>> {
public:
static bool WriteTo(const crypto::PoseidonSponge<PrimeField>& poseidon,
static bool WriteTo(const crypto::PoseidonSponge<F>& poseidon,
Buffer* buffer) {
return buffer->WriteMany(poseidon.config, poseidon.state);
}

static bool ReadFrom(const ReadOnlyBuffer& buffer,
crypto::PoseidonSponge<PrimeField>* poseidon) {
crypto::PoseidonConfig<PrimeField> config;
crypto::PoseidonState<PrimeField> state;
crypto::PoseidonSponge<F>* poseidon) {
crypto::PoseidonConfig<F> config;
crypto::SpongeState<F> state;
if (!buffer.ReadMany(&config, &state)) {
return false;
}
Expand All @@ -90,8 +121,7 @@ class Copyable<crypto::PoseidonSponge<PrimeField>> {
return true;
}

static size_t EstimateSize(
const crypto::PoseidonSponge<PrimeField>& poseidon) {
static size_t EstimateSize(const crypto::PoseidonSponge<F>& poseidon) {
return base::EstimateSize(poseidon.config, poseidon.state);
}
};
Expand Down
Loading

0 comments on commit c9e6f93

Please sign in to comment.