AtHeartEngineer/icicle
1
0
Fork 0
mirror of https://github.com/pseXperiments/icicle.git synced 2026-01-08 23:17:54 -05:00
Code Issues Packages Projects Releases Wiki Activity

add keccak tree builder (#555)

Browse Source
This commit is contained in:
ChickenLover
2024-07-15 15:31:12 +07:00
committed by GitHub
parent 7fd9ed1b49
commit ea71faf1fa
74 changed files with 777 additions and 642 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
wrappers/rust/icicle-core/src/hash.rs
View File

@@ -11,44 +11,28 @@ use crate::ntt::IcicleResult;
/// Struct that encodes Sponge hash parameters.
#[repr(C)]
#[derive(Debug, Clone)]
pub struct SpongeConfig<'a> {
pub struct HashConfig<'a> {
/// Details related to the device such as its id and stream id. See [DeviceContext](@ref device_context::DeviceContext).
pub ctx: DeviceContext<'a>,
pub(crate) are_inputs_on_device: bool,
pub(crate) are_outputs_on_device: bool,
pub input_rate: u32,
pub output_rate: u32,
pub offset: u32,
/// If true - input should be already aligned for poseidon permutation.
/// Aligned format: [0, A, B, 0, C, D, ...] (as you might get by using loop_state)
/// not aligned format: [A, B, 0, C, D, 0, ...] (as you might get from cudaMemcpy2D)
pub recursive_squeeze: bool,
/// If true, hash results will also be copied in the input pointer in aligned format
pub aligned: bool,
pub are_inputs_on_device: bool,
pub are_outputs_on_device: bool,
/// Whether to run the sponge operations asynchronously. If set to `true`, the functions will be non-blocking and you'd need to synchronize
/// it explicitly by running `stream.synchronize()`. If set to false, the functions will block the current CPU thread.
pub is_async: bool,
}
impl<'a> Default for SpongeConfig<'a> {
impl<'a> Default for HashConfig<'a> {
fn default() -> Self {
Self::default_for_device(DEFAULT_DEVICE_ID)
}
}
impl<'a> SpongeConfig<'a> {
impl<'a> HashConfig<'a> {
pub(crate) fn default_for_device(device_id: usize) -> Self {
SpongeConfig {
HashConfig {
ctx: DeviceContext::default_for_device(device_id),
are_inputs_on_device: false,
are_outputs_on_device: false,
input_rate: 0,
output_rate: 0,
offset: 0,
recursive_squeeze: false,
aligned: false,
is_async: false,
}
}
@@ -62,10 +46,10 @@ pub trait SpongeHash<PreImage, Image> {
number_of_states: usize,
input_block_len: usize,
output_len: usize,
cfg: &SpongeConfig,
cfg: &HashConfig,
) -> IcicleResult<()>;
fn default_config<'a>(&self) -> SpongeConfig<'a>;
fn default_config<'a>(&self) -> HashConfig<'a>;
fn get_handle(&self) -> *const c_void;
}

19
wrappers/rust/icicle-core/src/poseidon/mod.rs
View File

@@ -7,7 +7,7 @@ use icicle_cuda_runtime::{device_context::DeviceContext, memory::HostOrDeviceSli
use crate::{
error::IcicleResult,
hash::{sponge_check_input, sponge_check_outputs, SpongeConfig, SpongeHash},
hash::{sponge_check_input, sponge_check_outputs, HashConfig, SpongeHash},
traits::FieldImpl,
};
@@ -87,7 +87,7 @@ where
number_of_states: usize,
input_block_len: usize,
output_len: usize,
cfg: &SpongeConfig,
cfg: &HashConfig,
) -> IcicleResult<()> {
sponge_check_input(inputs, number_of_states, input_block_len, self.width - 1, &cfg.ctx);
sponge_check_outputs(output, number_of_states, output_len, self.width, false, &cfg.ctx);
@@ -107,11 +107,8 @@ where
)
}
fn default_config<'a>(&self) -> SpongeConfig<'a> {
let mut cfg = SpongeConfig::default();
cfg.input_rate = self.width as u32 - 1;
cfg.output_rate = self.width as u32;
cfg
fn default_config<'a>(&self) -> HashConfig<'a> {
HashConfig::default()
}
}
@@ -148,7 +145,7 @@ pub trait PoseidonImpl<F: FieldImpl> {
input_block_len: u32,
output_len: u32,
poseidon: PoseidonHandle,
cfg: &SpongeConfig,
cfg: &HashConfig,
) -> IcicleResult<()>;
fn delete(poseidon: PoseidonHandle) -> IcicleResult<()>;
@@ -163,7 +160,7 @@ macro_rules! impl_poseidon {
$field_config:ident
) => {
mod $field_prefix_ident {
use crate::poseidon::{$field, $field_config, CudaError, DeviceContext, PoseidonHandle, SpongeConfig};
use crate::poseidon::{$field, $field_config, CudaError, DeviceContext, HashConfig, PoseidonHandle};
extern "C" {
#[link_name = concat!($field_prefix, "_poseidon_create_cuda")]
pub(crate) fn create(
@@ -194,7 +191,7 @@ macro_rules! impl_poseidon {
number_of_states: u32,
input_block_len: u32,
output_len: u32,
cfg: &SpongeConfig,
cfg: &HashConfig,
) -> CudaError;
}
}
@@ -248,7 +245,7 @@ macro_rules! impl_poseidon {
input_block_len: u32,
output_len: u32,
poseidon: PoseidonHandle,
cfg: &SpongeConfig,
cfg: &HashConfig,
) -> IcicleResult<()> {
unsafe {
$field_prefix_ident::hash_many(

32
wrappers/rust/icicle-core/src/poseidon2/mod.rs
View File

@@ -7,7 +7,7 @@ use icicle_cuda_runtime::{device_context::DeviceContext, memory::HostOrDeviceSli
use crate::{
error::IcicleResult,
hash::{sponge_check_input, sponge_check_outputs, SpongeConfig, SpongeHash},
hash::{sponge_check_input, sponge_check_outputs, HashConfig, SpongeHash},
traits::FieldImpl,
};
@@ -32,6 +32,7 @@ where
<F as FieldImpl>::Config: Poseidon2Impl<F>,
{
width: usize,
rate: usize,
handle: Poseidon2Handle,
phantom: PhantomData<F>,
}
@@ -52,6 +53,7 @@ where
.and_then(|handle| {
Ok(Self {
width,
rate,
handle,
phantom: PhantomData,
})
@@ -85,6 +87,7 @@ where
.and_then(|handle| {
Ok(Self {
width,
rate,
handle,
phantom: PhantomData,
})
@@ -108,15 +111,9 @@ where
number_of_states: usize,
input_block_len: usize,
output_len: usize,
cfg: &SpongeConfig,
cfg: &HashConfig,
) -> IcicleResult<()> {
sponge_check_input(
inputs,
number_of_states,
input_block_len,
cfg.input_rate as usize,
&cfg.ctx,
);
sponge_check_input(inputs, number_of_states, input_block_len, self.rate, &cfg.ctx);
sponge_check_outputs(output, number_of_states, output_len, self.width, false, &cfg.ctx);
let mut local_cfg = cfg.clone();
@@ -134,11 +131,8 @@ where
)
}
fn default_config<'a>(&self) -> SpongeConfig<'a> {
let mut cfg = SpongeConfig::default();
cfg.input_rate = self.width as u32;
cfg.output_rate = self.width as u32;
cfg
fn default_config<'a>(&self) -> HashConfig<'a> {
HashConfig::default()
}
}
@@ -181,7 +175,7 @@ pub trait Poseidon2Impl<F: FieldImpl> {
input_block_len: u32,
output_len: u32,
poseidon: Poseidon2Handle,
cfg: &SpongeConfig,
cfg: &HashConfig,
) -> IcicleResult<()>;
fn delete(poseidon: Poseidon2Handle) -> IcicleResult<()>;
@@ -197,8 +191,8 @@ macro_rules! impl_poseidon2 {
) => {
mod $field_prefix_ident {
use crate::poseidon2::{
$field, $field_config, CudaError, DeviceContext, DiffusionStrategy, MdsType, Poseidon2Handle,
SpongeConfig,
$field, $field_config, CudaError, DeviceContext, DiffusionStrategy, HashConfig, MdsType,
Poseidon2Handle,
};
use icicle_core::error::IcicleError;
extern "C" {
@@ -238,7 +232,7 @@ macro_rules! impl_poseidon2 {
number_of_states: u32,
input_block_len: u32,
output_len: u32,
cfg: &SpongeConfig,
cfg: &HashConfig,
) -> CudaError;
}
}
@@ -298,7 +292,7 @@ macro_rules! impl_poseidon2 {
input_block_len: u32,
output_len: u32,
poseidon: Poseidon2Handle,
cfg: &SpongeConfig,
cfg: &HashConfig,
) -> IcicleResult<()> {
unsafe {
$field_prefix_ident::hash_many(

2
wrappers/rust/icicle-curves/icicle-bls12-377/src/poseidon/mod.rs
View File

@@ -3,7 +3,7 @@ use crate::curve::{BaseCfg, BaseField};
use crate::curve::{ScalarCfg, ScalarField};
use icicle_core::error::IcicleResult;
use icicle_core::hash::SpongeConfig;
use icicle_core::hash::HashConfig;
use icicle_core::impl_poseidon;
use icicle_core::poseidon::{PoseidonHandle, PoseidonImpl};
use icicle_core::traits::IcicleResultWrap;

2
wrappers/rust/icicle-curves/icicle-bls12-381/src/poseidon/mod.rs
View File

@@ -1,7 +1,7 @@
use crate::curve::{ScalarCfg, ScalarField};
use icicle_core::error::IcicleResult;
use icicle_core::hash::SpongeConfig;
use icicle_core::hash::HashConfig;
use icicle_core::impl_poseidon;
use icicle_core::poseidon::{PoseidonHandle, PoseidonImpl};
use icicle_core::traits::IcicleResultWrap;

2
wrappers/rust/icicle-curves/icicle-bn254/src/poseidon/mod.rs
View File

@@ -1,7 +1,7 @@
use crate::curve::{ScalarCfg, ScalarField};
use icicle_core::error::IcicleResult;
use icicle_core::hash::SpongeConfig;
use icicle_core::hash::HashConfig;
use icicle_core::impl_poseidon;
use icicle_core::poseidon::{PoseidonHandle, PoseidonImpl};
use icicle_core::traits::IcicleResultWrap;

2
wrappers/rust/icicle-curves/icicle-bn254/src/poseidon2/mod.rs
View File

@@ -1,7 +1,7 @@
use crate::curve::{ScalarCfg, ScalarField};
use icicle_core::error::IcicleResult;
use icicle_core::hash::SpongeConfig;
use icicle_core::hash::HashConfig;
use icicle_core::impl_poseidon2;
use icicle_core::poseidon2::{DiffusionStrategy, MdsType, Poseidon2Handle, Poseidon2Impl};
use icicle_core::traits::IcicleResultWrap;

2
wrappers/rust/icicle-curves/icicle-grumpkin/src/poseidon/mod.rs
View File

@@ -1,7 +1,7 @@
use crate::curve::{ScalarCfg, ScalarField};
use icicle_core::error::IcicleResult;
use icicle_core::hash::SpongeConfig;
use icicle_core::hash::HashConfig;
use icicle_core::impl_poseidon;
use icicle_core::poseidon::{PoseidonHandle, PoseidonImpl};
use icicle_core::traits::IcicleResultWrap;

2
wrappers/rust/icicle-fields/icicle-babybear/src/poseidon2/mod.rs
View File

@@ -1,7 +1,7 @@
use crate::field::{ScalarCfg, ScalarField};
use icicle_core::error::IcicleResult;
use icicle_core::hash::SpongeConfig;
use icicle_core::hash::HashConfig;
use icicle_core::impl_poseidon2;
use icicle_core::poseidon2::{DiffusionStrategy, MdsType, Poseidon2Handle, Poseidon2Impl};
use icicle_core::traits::IcicleResultWrap;

133
wrappers/rust/icicle-hash/src/keccak/mod.rs
View File

@@ -1,84 +1,64 @@
use icicle_core::hash::HashConfig;
use icicle_core::tree::TreeBuilderConfig;
use icicle_cuda_runtime::error::CudaError;
use icicle_cuda_runtime::{
device_context::{DeviceContext, DEFAULT_DEVICE_ID},
memory::HostOrDeviceSlice,
};
use icicle_cuda_runtime::memory::HostOrDeviceSlice;
use icicle_core::error::IcicleResult;
use icicle_core::traits::IcicleResultWrap;
pub mod tests;
#[repr(C)]
#[derive(Debug, Clone)]
pub struct KeccakConfig<'a> {
/// Details related to the device such as its id and stream id. See [DeviceContext](@ref device_context::DeviceContext).
pub ctx: DeviceContext<'a>,
/// True if inputs are on device and false if they're on host. Default value: false.
pub are_inputs_on_device: bool,
/// If true, output is preserved on device, otherwise on host. Default value: false.
pub are_outputs_on_device: bool,
/// Whether to run the Keccak asynchronously. If set to `true`, the keccak_hash function will be
/// non-blocking and you'd need to synchronize it explicitly by running
/// `cudaStreamSynchronize` or `cudaDeviceSynchronize`. If set to false, keccak_hash
/// function will block the current CPU thread.
pub is_async: bool,
}
impl<'a> Default for KeccakConfig<'a> {
fn default() -> Self {
Self::default_for_device(DEFAULT_DEVICE_ID)
}
}
impl<'a> KeccakConfig<'a> {
pub fn default_for_device(device_id: usize) -> Self {
KeccakConfig {
ctx: DeviceContext::default_for_device(device_id),
are_inputs_on_device: false,
are_outputs_on_device: false,
is_async: false,
}
}
}
extern "C" {
pub(crate) fn keccak256_cuda(
input: *const u8,
input_block_size: i32,
number_of_blocks: i32,
input_block_size: u32,
number_of_blocks: u32,
output: *mut u8,
config: &KeccakConfig,
config: &HashConfig,
) -> CudaError;
pub(crate) fn keccak512_cuda(
input: *const u8,
input_block_size: i32,
number_of_blocks: i32,
input_block_size: u32,
number_of_blocks: u32,
output: *mut u8,
config: &KeccakConfig,
config: &HashConfig,
) -> CudaError;
pub(crate) fn build_keccak256_merkle_tree_cuda(
leaves: *const u8,
digests: *mut u64,
height: u32,
input_block_len: u32,
config: &TreeBuilderConfig,
) -> CudaError;
pub(crate) fn build_keccak512_merkle_tree_cuda(
leaves: *const u8,
digests: *mut u64,
height: u32,
input_block_len: u32,
config: &TreeBuilderConfig,
) -> CudaError;
}
pub fn keccak256(
input: &(impl HostOrDeviceSlice<u8> + ?Sized),
input_block_size: i32,
number_of_blocks: i32,
input_block_size: u32,
number_of_blocks: u32,
output: &mut (impl HostOrDeviceSlice<u8> + ?Sized),
config: &mut KeccakConfig,
config: &HashConfig,
) -> IcicleResult<()> {
config.are_inputs_on_device = input.is_on_device();
config.are_outputs_on_device = output.is_on_device();
let mut local_cfg = config.clone();
local_cfg.are_inputs_on_device = input.is_on_device();
local_cfg.are_outputs_on_device = output.is_on_device();
unsafe {
keccak256_cuda(
input.as_ptr(),
input_block_size,
number_of_blocks,
output.as_mut_ptr(),
config,
&local_cfg,
)
.wrap()
}
@@ -86,19 +66,58 @@ pub fn keccak256(
pub fn keccak512(
input: &(impl HostOrDeviceSlice<u8> + ?Sized),
input_block_size: i32,
number_of_blocks: i32,
input_block_size: u32,
number_of_blocks: u32,
output: &mut (impl HostOrDeviceSlice<u8> + ?Sized),
config: &mut KeccakConfig,
config: &HashConfig,
) -> IcicleResult<()> {
config.are_inputs_on_device = input.is_on_device();
config.are_outputs_on_device = output.is_on_device();
let mut local_cfg = config.clone();
local_cfg.are_inputs_on_device = input.is_on_device();
local_cfg.are_outputs_on_device = output.is_on_device();
unsafe {
keccak512_cuda(
input.as_ptr(),
input_block_size,
number_of_blocks,
output.as_mut_ptr(),
&local_cfg,
)
.wrap()
}
}
pub fn build_keccak256_merkle_tree(
leaves: &(impl HostOrDeviceSlice<u8> + ?Sized),
digests: &mut (impl HostOrDeviceSlice<u64> + ?Sized),
height: usize,
input_block_len: usize,
config: &TreeBuilderConfig,
) -> IcicleResult<()> {
unsafe {
build_keccak256_merkle_tree_cuda(
leaves.as_ptr(),
digests.as_mut_ptr(),
height as u32,
input_block_len as u32,
config,
)
.wrap()
}
}
pub fn build_keccak512_merkle_tree(
leaves: &(impl HostOrDeviceSlice<u8> + ?Sized),
digests: &mut (impl HostOrDeviceSlice<u64> + ?Sized),
height: usize,
input_block_len: usize,
config: &TreeBuilderConfig,
) -> IcicleResult<()> {
unsafe {
build_keccak512_merkle_tree_cuda(
leaves.as_ptr(),
digests.as_mut_ptr(),
height as u32,
input_block_len as u32,
config,
)
.wrap()

47
wrappers/rust/icicle-hash/src/keccak/tests.rs
View File

@@ -1 +1,48 @@
#[cfg(test)]
pub(crate) mod tests {
use icicle_core::{
hash::HashConfig,
tree::{merkle_tree_digests_len, TreeBuilderConfig},
};
use icicle_cuda_runtime::memory::HostSlice;
use crate::keccak::{build_keccak256_merkle_tree, keccak256};
#[test]
fn keccak_hash_test() {
let config = HashConfig::default();
let input_block_len = 136;
let number_of_hashes = 1024;
let preimages = vec![1u8; number_of_hashes * input_block_len];
let mut digests = vec![0u8; number_of_hashes * 64];
let preimages_slice = HostSlice::from_slice(&preimages);
let digests_slice = HostSlice::from_mut_slice(&mut digests);
keccak256(
preimages_slice,
input_block_len as u32,
number_of_hashes as u32,
digests_slice,
&config,
)
.unwrap();
}
#[test]
fn keccak_merkle_tree_test() {
let mut config = TreeBuilderConfig::default();
config.arity = 2;
let height = 22;
let input_block_len = 136;
let leaves = vec![1u8; (1 << height) * input_block_len];
let mut digests = vec![0u64; merkle_tree_digests_len((height + 1) as u32, 2, 1)];
let leaves_slice = HostSlice::from_slice(&leaves);
let digests_slice = HostSlice::from_mut_slice(&mut digests);
build_keccak256_merkle_tree(leaves_slice, digests_slice, height, input_block_len, &config).unwrap();
println!("Root: {:?}", digests_slice[0]);
}
}