Merge pull request #41 from okx/dev-dumi

update zeknox ref
okx · Nov 21, 2024 · 9e9f1a9 · 9e9f1a9
2 parents ff1b7ae + e380193
commit 9e9f1a9
Show file tree

Hide file tree

Showing 12 changed files with 36 additions and 52 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -3,7 +3,7 @@ members = ["field", "maybe_rayon", "plonky2", "starky", "util", "gen", "u32", "e
 resolver = "2"
 
 [workspace.dependencies]
-cryptography_cuda = { git = "ssh://git@github.com/okx/cryptography_cuda.git", rev = "547192b2ef42dc7519435059c86f88431b8de999" }
+zeknox = { git = "https://github.com/okx/zeknox.git", rev = "bfa744a8156850d8ee53d2aeec350981bb32ad4d" }
 ahash = { version = "0.8.7", default-features = false, features = [
     "compile-time-rng",
 ] } # NOTE: Be sure to keep this version the same as the dependency in `hashbrown`.

diff --git a/ecdsa/Cargo.toml b/ecdsa/Cargo.toml
@@ -11,8 +11,8 @@ categories.workspace = true
 
 [features]
 parallel = ["plonky2_maybe_rayon/parallel", "plonky2/parallel"]
-cuda = ["cryptography_cuda/cuda", "plonky2/cuda"]
-no_cuda = ["cryptography_cuda/no_cuda", "plonky2/no_cuda"]
+cuda = ["zeknox/cuda", "plonky2/cuda"]
+no_cuda = ["zeknox/no_cuda", "plonky2/no_cuda"]
 
 [dependencies]
 anyhow = { version = "1.0.40" }
@@ -22,7 +22,7 @@ num = { version = "0.4.0" }
 plonky2 = { path = "../plonky2" }
 plonky2_u32 = { path = "../u32" }
 serde = { version = "1.0", features = ["derive"] }
-cryptography_cuda = { workspace = true, optional = true }
+zeknox = { workspace = true, optional = true }
 
 [dev-dependencies]
 rand = { version = "0.8.4", features = ["getrandom"] }

diff --git a/ecgfp5/Cargo.toml b/ecgfp5/Cargo.toml
@@ -26,7 +26,7 @@ itertools = "0.10"
 serde = "1"
 rand = { version = "0.8.5", default-features = false, features = ["getrandom"] }
 hex = "0.4.3"
-cryptography_cuda = { workspace = true, optional = true }
+zeknox = { workspace = true, optional = true }
 
 [dev-dependencies]
 rand = { version = "0.8.5", features = ["min_const_gen"] }
@@ -48,5 +48,5 @@ name = "schnorr"
 harness = false
 
 [features]
-cuda = ["cryptography_cuda/cuda", "plonky2/cuda"]
-no_cuda = ["cryptography_cuda/no_cuda", "plonky2/no_cuda"]
+cuda = ["zeknox/cuda", "plonky2/cuda"]
+no_cuda = ["zeknox/no_cuda", "plonky2/no_cuda"]
diff --git a/field/Cargo.toml b/field/Cargo.toml
@@ -23,7 +23,7 @@ rand = { workspace = true, features = ["getrandom"] }
 serde = { workspace = true, features = ["alloc"] }
 static_assertions = { workspace = true }
 unroll = { workspace = true }
-cryptography_cuda = { workspace = true, optional = true }
+zeknox = { workspace = true, optional = true }
 
 [dev-dependencies]
 rand = { version = "0.8.5", default-features = false, features = ["getrandom"] }
@@ -37,9 +37,9 @@ quote = "1"
 
 [features]
 default = []
-cuda = ["cryptography_cuda/cuda"]
+cuda = ["zeknox/cuda"]
 precompile = []
-no_cuda = ["cryptography_cuda/no_cuda"]
+no_cuda = ["zeknox/no_cuda"]
 
 # Display math equations properly in documentation
 [package.metadata.docs.rs]

diff --git a/plonky2/Cargo.toml b/plonky2/Cargo.toml
@@ -17,8 +17,8 @@ gate_testing = []
 parallel = ["hashbrown/rayon", "plonky2_maybe_rayon/parallel"]
 std = ["anyhow/std", "rand/std", "itertools/use_std"]
 timing = ["std", "dep:web-time"]
-cuda = ["cryptography_cuda/cuda"]
-no_cuda = ["cryptography_cuda/no_cuda"]
+cuda = ["zeknox/cuda"]
+no_cuda = ["zeknox/no_cuda"]
 batch = []
 cuda_timing = []
 papi = []
@@ -44,7 +44,7 @@ papi-bindings = { version = "0.5.2" }
 plonky2_field = { version = "0.2.2", path = "../field", default-features = false }
 plonky2_maybe_rayon = { version = "0.2.0", path = "../maybe_rayon", default-features = false }
 plonky2_util = { version = "0.2.0", path = "../util", default-features = false }
-cryptography_cuda = { workspace = true, optional = true }
+zeknox = { workspace = true, optional = true }
 dyn-clone = "1.0.17"
 
 [target.'cfg(all(target_arch = "wasm32", target_os = "unknown"))'.dependencies]

diff --git a/plonky2/README.md b/plonky2/README.md
@@ -16,18 +16,9 @@ rustup override set nightly
 
 ## Poseidon Hash on GPU (CUDA)
 
-Build the shared library
-
-```
-cd cryptography_cuda/cuda/merkle
-make lib
-make libgpu
-```
-
 Run tests (in plonky2 folder)
 
 ```
-export LD_LIBRARY_PATH=<path-to cryptography_cuda/cuda/merkle>
 # CPU-only
 cargo test -- --nocapture merkle_trees
 # GPU
@@ -42,13 +33,6 @@ cargo bench merkle
 cargo bench --features=cuda merkle
 ```
 
-Run microbenchmarks
-
-```
-cd cryptography_cuda/cuda/merkle
-./run-benchmark.sh
-```
-
 ## License
 
 Licensed under either of

diff --git a/plonky2/benches/lde.rs b/plonky2/benches/lde.rs
@@ -1,8 +1,6 @@
 mod allocator;
 
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
-#[cfg(feature = "cuda")]
-use cryptography_cuda::init_cuda_degree_rs;
 use plonky2::field::extension::Extendable;
 use plonky2::field::goldilocks_field::GoldilocksField;
 use plonky2::field::polynomial::PolynomialCoeffs;
@@ -11,6 +9,8 @@ use plonky2::hash::hash_types::RichField;
 use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
 use plonky2::util::timing::TimingTree;
 use tynm::type_name;
+#[cfg(feature = "cuda")]
+use zeknox::init_cuda_degree_rs;
 
 pub(crate) fn bench_batch_lde<
     F: RichField + Extendable<D>,

diff --git a/plonky2/src/fri/oracle.rs b/plonky2/src/fri/oracle.rs
@@ -1,14 +1,14 @@
 #[cfg(not(feature = "std"))]
 use alloc::{format, vec::Vec};
 
+use itertools::Itertools;
+use plonky2_field::types::Field;
+use plonky2_maybe_rayon::*;
 #[cfg(feature = "cuda")]
-use cryptography_cuda::{
+use zeknox::{
     device::memory::HostOrDeviceSlice, lde_batch, lde_batch_multi_gpu, transpose_rev_batch,
     types::*,
 };
-use itertools::Itertools;
-use plonky2_field::types::Field;
-use plonky2_maybe_rayon::*;
 
 use crate::field::extension::Extendable;
 use crate::field::fft::FftRootTable;
@@ -33,7 +33,7 @@ pub static GPU_INIT: once_cell::sync::Lazy<std::sync::Arc<std::sync::Mutex<u64>>
 
 #[cfg(all(feature = "cuda", any(test, doctest)))]
 fn init_gpu() {
-    use cryptography_cuda::init_cuda_rs;
+    use zeknox::init_cuda_rs;
 
     let mut init = GPU_INIT.lock().unwrap();
     if *init == 0 {

diff --git a/plonky2/src/hash/merkle_tree.rs b/plonky2/src/hash/merkle_tree.rs
@@ -9,19 +9,19 @@ use std::collections::HashSet;
 use std::sync::Mutex;
 use std::time::Instant;
 
-#[cfg(feature = "cuda")]
-use cryptography_cuda::device::memory::HostOrDeviceSlice;
-#[cfg(feature = "cuda")]
-use cryptography_cuda::device::stream::CudaStream;
-#[cfg(feature = "cuda")]
-use cryptography_cuda::{
-    fill_digests_buf_linear_gpu_with_gpu_ptr, fill_digests_buf_linear_multigpu_with_gpu_ptr,
-};
 use num::range;
 #[cfg(feature = "cuda")]
 use once_cell::sync::Lazy;
 use plonky2_maybe_rayon::*;
 use serde::{Deserialize, Serialize};
+#[cfg(feature = "cuda")]
+use zeknox::device::memory::HostOrDeviceSlice;
+#[cfg(feature = "cuda")]
+use zeknox::device::stream::CudaStream;
+#[cfg(feature = "cuda")]
+use zeknox::{
+    fill_digests_buf_linear_gpu_with_gpu_ptr, fill_digests_buf_linear_multigpu_with_gpu_ptr,
+};
 
 #[cfg(all(target_feature = "avx2", target_feature = "avx512dq"))]
 use crate::hash::arch::x86_64::poseidon_goldilocks_avx512::{hash_leaf_avx512, hash_two_avx512};
@@ -552,7 +552,7 @@ fn fill_digests_buf_cpu<F: RichField, H: Hasher<F>>(
     leaf_size: usize,
     cap_height: usize,
 ) {
-    use cryptography_cuda::fill_digests_buf_linear_cpu;
+    use zeknox::fill_digests_buf_linear_cpu;
 
     let leaves_count = (leaves.len() / leaf_size) as u64;
     let digests_count: u64 = digests_buf.len().try_into().unwrap();

diff --git a/plonky2/src/util/test_utils.rs b/plonky2/src/util/test_utils.rs
@@ -1,8 +1,8 @@
 #[cfg(feature = "cuda")]
 pub fn init_cuda() {
-    use cryptography_cuda::{get_number_of_gpus_rs, init_coset_rs, init_twiddle_factors_rs};
     use plonky2_field::goldilocks_field::GoldilocksField;
     use plonky2_field::types::{Field, PrimeField64};
+    use zeknox::{get_number_of_gpus_rs, init_coset_rs, init_twiddle_factors_rs};
 
     let num_of_gpus = get_number_of_gpus_rs();
     println!("num of gpus: {:?}", num_of_gpus);

diff --git a/starky/Cargo.toml b/starky/Cargo.toml
@@ -16,8 +16,8 @@ default = ["parallel", "std", "timing"]
 parallel = ["plonky2/parallel", "plonky2_maybe_rayon/parallel"]
 std = ["anyhow/std", "plonky2/std"]
 timing = ["plonky2/timing"]
-cuda = ["cryptography_cuda/cuda", "plonky2/cuda"]
-no_cuda = ["cryptography_cuda/no_cuda", "plonky2/no_cuda"]
+cuda = ["zeknox/cuda", "plonky2/cuda"]
+no_cuda = ["zeknox/no_cuda", "plonky2/no_cuda"]
 
 [dependencies]
 ahash = { workspace = true }
@@ -26,7 +26,7 @@ hashbrown = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }
 num-bigint = { version = "0.4.3", default-features = false }
-cryptography_cuda = { workspace = true, optional = true }
+zeknox = { workspace = true, optional = true }
 
 
 # Local dependencies

diff --git a/u32/Cargo.toml b/u32/Cargo.toml
@@ -14,11 +14,11 @@ anyhow = { version = "1.0.40", default-features = false }
 itertools = { version = "0.10.0", default-features = false }
 num = { version = "0.4", default-features = false }
 plonky2 = { path = "../plonky2" }
-cryptography_cuda = { workspace = true, optional = true }
+zeknox = { workspace = true, optional = true }
 
 [dev-dependencies]
 rand = { version = "0.8.4", default-features = false, features = ["getrandom"] }
 
 [features]
-cuda = ["cryptography_cuda/cuda", "plonky2/cuda"]
-no_cuda = ["cryptography_cuda/no_cuda", "plonky2/no_cuda"]
+cuda = ["zeknox/cuda", "plonky2/cuda"]
+no_cuda = ["zeknox/no_cuda", "plonky2/no_cuda"]