diff --git a/wrappers/rust/icicle-core/src/vec_ops/mod.rs b/wrappers/rust/icicle-core/src/vec_ops/mod.rs index ba22b776d..277846ee8 100644 --- a/wrappers/rust/icicle-core/src/vec_ops/mod.rs +++ b/wrappers/rust/icicle-core/src/vec_ops/mod.rs @@ -13,6 +13,8 @@ pub struct VecOpsConfig { pub is_b_on_device: bool, pub is_result_on_device: bool, pub is_async: bool, + pub batch_size: i32, + pub columns_batch: bool, pub ext: ConfigExtension, } @@ -24,6 +26,8 @@ impl VecOpsConfig { is_b_on_device: false, is_result_on_device: false, is_async: false, + batch_size: 1, + columns_batch: false, ext: ConfigExtension::new(), } } @@ -58,6 +62,46 @@ pub trait VecOps { cfg: &VecOpsConfig, ) -> Result<(), eIcicleError>; + fn div( + a: &(impl HostOrDeviceSlice + ?Sized), + b: &(impl HostOrDeviceSlice + ?Sized), + result: &mut (impl HostOrDeviceSlice + ?Sized), + cfg: &VecOpsConfig, + ) -> Result<(), eIcicleError>; + + fn sum( + a: &(impl HostOrDeviceSlice + ?Sized), + result: &mut (impl HostOrDeviceSlice + ?Sized), + cfg: &VecOpsConfig, + ) -> Result<(), eIcicleError>; + + fn product( + a: &(impl HostOrDeviceSlice + ?Sized), + result: &mut (impl HostOrDeviceSlice + ?Sized), + cfg: &VecOpsConfig, + ) -> Result<(), eIcicleError>; + + fn scalar_add( + a: &(impl HostOrDeviceSlice + ?Sized), + b: &(impl HostOrDeviceSlice + ?Sized), + result: &mut (impl HostOrDeviceSlice + ?Sized), + cfg: &VecOpsConfig, + ) -> Result<(), eIcicleError>; + + fn scalar_sub( + a: &(impl HostOrDeviceSlice + ?Sized), + b: &(impl HostOrDeviceSlice + ?Sized), + result: &mut (impl HostOrDeviceSlice + ?Sized), + cfg: &VecOpsConfig, + ) -> Result<(), eIcicleError>; + + fn scalar_mul( + a: &(impl HostOrDeviceSlice + ?Sized), + b: &(impl HostOrDeviceSlice + ?Sized), + result: &mut (impl HostOrDeviceSlice + ?Sized), + cfg: &VecOpsConfig, + ) -> Result<(), eIcicleError>; + fn transpose( input: &(impl HostOrDeviceSlice + ?Sized), nof_rows: u32, @@ -76,6 +120,16 @@ pub trait VecOps { input: &mut (impl HostOrDeviceSlice + ?Sized), cfg: &VecOpsConfig, ) -> Result<(), eIcicleError>; + + fn slice( + input: &(impl HostOrDeviceSlice + ?Sized), + offset: u64, + stride: u64, + size_in: u64, + size_out: u64, + cfg: &VecOpsConfig, + output: &mut (impl HostOrDeviceSlice + ?Sized), + ) -> Result<(), eIcicleError>; } fn check_vec_ops_args<'a, F>( @@ -166,6 +220,88 @@ where <::Config as VecOps>::mul(a, b, result, &cfg) } +pub fn div_scalars( + a: &(impl HostOrDeviceSlice + ?Sized), + b: &(impl HostOrDeviceSlice + ?Sized), + result: &mut (impl HostOrDeviceSlice + ?Sized), + cfg: &VecOpsConfig, +) -> Result<(), eIcicleError> +where + F: FieldImpl, + ::Config: VecOps, +{ + let cfg = check_vec_ops_args(a, b, result, cfg); + <::Config as VecOps>::div(a, b, result, &cfg) +} + +pub fn sum_scalars( + a: &(impl HostOrDeviceSlice + ?Sized), + result: &mut (impl HostOrDeviceSlice + ?Sized), + cfg: &VecOpsConfig, +) -> Result<(), eIcicleError> +where + F: FieldImpl, + ::Config: VecOps, +{ + let cfg = check_vec_ops_args(a, a, result, cfg); //TODO: emirsoyturk + <::Config as VecOps>::sum(a, result, &cfg) +} + +pub fn product_scalars( + a: &(impl HostOrDeviceSlice + ?Sized), + result: &mut (impl HostOrDeviceSlice + ?Sized), + cfg: &VecOpsConfig, +) -> Result<(), eIcicleError> +where + F: FieldImpl, + ::Config: VecOps, +{ + let cfg = check_vec_ops_args(a, a, result, cfg); //TODO: emirsoyturk + <::Config as VecOps>::product(a, result, &cfg) +} + +pub fn scalar_add( + a: &(impl HostOrDeviceSlice + ?Sized), + b: &(impl HostOrDeviceSlice + ?Sized), + result: &mut (impl HostOrDeviceSlice + ?Sized), + cfg: &VecOpsConfig, +) -> Result<(), eIcicleError> +where + F: FieldImpl, + ::Config: VecOps, +{ + let cfg = check_vec_ops_args(b, b, result, cfg); //TODO: emirsoyturk + <::Config as VecOps>::scalar_add(a, b, result, &cfg) +} + +pub fn scalar_sub( + a: &(impl HostOrDeviceSlice + ?Sized), + b: &(impl HostOrDeviceSlice + ?Sized), + result: &mut (impl HostOrDeviceSlice + ?Sized), + cfg: &VecOpsConfig, +) -> Result<(), eIcicleError> +where + F: FieldImpl, + ::Config: VecOps, +{ + let cfg = check_vec_ops_args(b, b, result, cfg); //TODO: emirsoyturk + <::Config as VecOps>::scalar_sub(a, b, result, &cfg) +} + +pub fn scalar_mul( + a: &(impl HostOrDeviceSlice + ?Sized), + b: &(impl HostOrDeviceSlice + ?Sized), + result: &mut (impl HostOrDeviceSlice + ?Sized), + cfg: &VecOpsConfig, +) -> Result<(), eIcicleError> +where + F: FieldImpl, + ::Config: VecOps, +{ + let cfg = check_vec_ops_args(b, b, result, cfg); //TODO: emirsoyturk + <::Config as VecOps>::scalar_mul(a, b, result, &cfg) +} + pub fn transpose_matrix( input: &(impl HostOrDeviceSlice + ?Sized), nof_rows: u32, @@ -205,6 +341,23 @@ where <::Config as VecOps>::bit_reverse_inplace(input, &cfg) } +pub fn slice( + input: &(impl HostOrDeviceSlice + ?Sized), + offset: u64, + stride: u64, + size_in: u64, + size_out: u64, + cfg: &VecOpsConfig, + output: &mut (impl HostOrDeviceSlice + ?Sized), +) -> Result<(), eIcicleError> +where + F: FieldImpl, + ::Config: VecOps, +{ + <::Config as VecOps>::slice(input, offset, stride, size_in, size_out, &cfg, output) +} + + #[macro_export] macro_rules! impl_vec_ops_field { ( @@ -255,6 +408,59 @@ macro_rules! impl_vec_ops_field { result: *mut $field, ) -> eIcicleError; + #[link_name = concat!($field_prefix, "_vector_div")] + pub(crate) fn vector_div_ffi( + a: *const $field, + b: *const $field, + size: u32, + cfg: *const VecOpsConfig, + result: *mut $field, + ) -> eIcicleError; + + #[link_name = concat!($field_prefix, "_vector_sum")] + pub(crate) fn vector_sum_ffi( + a: *const $field, + size: u32, + cfg: *const VecOpsConfig, + result: *mut $field, + ) -> eIcicleError; + + + #[link_name = concat!($field_prefix, "_vector_product")] + pub(crate) fn vector_product_ffi( + a: *const $field, + size: u32, + cfg: *const VecOpsConfig, + result: *mut $field, + ) -> eIcicleError; + + #[link_name = concat!($field_prefix, "_scalar_add_vec")] + pub(crate) fn scalar_add_ffi( + a: *const $field, + b: *const $field, + size: u32, + cfg: *const VecOpsConfig, + result: *mut $field, + ) -> eIcicleError; + + #[link_name = concat!($field_prefix, "_scalar_sub_vec")] + pub(crate) fn scalar_sub_ffi( + a: *const $field, + b: *const $field, + size: u32, + cfg: *const VecOpsConfig, + result: *mut $field, + ) -> eIcicleError; + + #[link_name = concat!($field_prefix, "_scalar_mul_vec")] + pub(crate) fn scalar_mul_ffi( + a: *const $field, + b: *const $field, + size: u32, + cfg: *const VecOpsConfig, + result: *mut $field, + ) -> eIcicleError; + #[link_name = concat!($field_prefix, "_matrix_transpose")] pub(crate) fn matrix_transpose_ffi( input: *const $field, @@ -271,6 +477,17 @@ macro_rules! impl_vec_ops_field { config: *const VecOpsConfig, output: *mut $field, ) -> eIcicleError; + + #[link_name = concat!($field_prefix, "_slice")] + pub(crate) fn slice_ffi( + input: *const $field, + offset: u64, + stride: u64, + size_in: u64, + size_out: u64, + cfg: *const VecOpsConfig, + output: *mut $field, + ) -> eIcicleError; } } @@ -344,6 +561,110 @@ macro_rules! impl_vec_ops_field { .wrap() } } + + fn div( + a: &(impl HostOrDeviceSlice<$field> + ?Sized), + b: &(impl HostOrDeviceSlice<$field> + ?Sized), + result: &mut (impl HostOrDeviceSlice<$field> + ?Sized), + cfg: &VecOpsConfig, + ) -> Result<(), eIcicleError> { + unsafe { + $field_prefix_ident::vector_div_ffi( + a.as_ptr(), + b.as_ptr(), + a.len() as u32, + cfg as *const VecOpsConfig, + result.as_mut_ptr(), + ) + .wrap() + } + } + + fn sum( + a: &(impl HostOrDeviceSlice<$field> + ?Sized), + result: &mut (impl HostOrDeviceSlice<$field> + ?Sized), + cfg: &VecOpsConfig, + ) -> Result<(), eIcicleError> { + unsafe { + $field_prefix_ident::vector_sum_ffi( + a.as_ptr(), + a.len() as u32, + cfg as *const VecOpsConfig, + result.as_mut_ptr(), + ) + .wrap() + } + } + + fn product( + a: &(impl HostOrDeviceSlice<$field> + ?Sized), + result: &mut (impl HostOrDeviceSlice<$field> + ?Sized), + cfg: &VecOpsConfig, + ) -> Result<(), eIcicleError> { + unsafe { + $field_prefix_ident::vector_sum_ffi( + a.as_ptr(), + a.len() as u32, + cfg as *const VecOpsConfig, + result.as_mut_ptr(), + ) + .wrap() + } + } + + fn scalar_add( + a: &(impl HostOrDeviceSlice<$field> + ?Sized), + b: &(impl HostOrDeviceSlice<$field> + ?Sized), + result: &mut (impl HostOrDeviceSlice<$field> + ?Sized), + cfg: &VecOpsConfig, + ) -> Result<(), eIcicleError> { + unsafe { + $field_prefix_ident::scalar_add_ffi( + a.as_ptr(), + b.as_ptr(), + b.len() as u32, + cfg as *const VecOpsConfig, + result.as_mut_ptr(), + ) + .wrap() + } + } + + fn scalar_sub( + a: &(impl HostOrDeviceSlice<$field> + ?Sized), + b: &(impl HostOrDeviceSlice<$field> + ?Sized), + result: &mut (impl HostOrDeviceSlice<$field> + ?Sized), + cfg: &VecOpsConfig, + ) -> Result<(), eIcicleError> { + unsafe { + $field_prefix_ident::scalar_sub_ffi( + a.as_ptr(), + b.as_ptr(), + b.len() as u32, + cfg as *const VecOpsConfig, + result.as_mut_ptr(), + ) + .wrap() + } + } + + fn scalar_mul( + a: &(impl HostOrDeviceSlice<$field> + ?Sized), + b: &(impl HostOrDeviceSlice<$field> + ?Sized), + result: &mut (impl HostOrDeviceSlice<$field> + ?Sized), + cfg: &VecOpsConfig, + ) -> Result<(), eIcicleError> { + unsafe { + $field_prefix_ident::scalar_mul_ffi( + a.as_ptr(), + b.as_ptr(), + b.len() as u32, + cfg as *const VecOpsConfig, + result.as_mut_ptr(), + ) + .wrap() + } + } fn transpose( input: &(impl HostOrDeviceSlice<$field> + ?Sized), @@ -394,6 +715,29 @@ macro_rules! impl_vec_ops_field { .wrap() } } + + fn slice( + input: &(impl HostOrDeviceSlice<$field> + ?Sized), + offset: u64, + stride: u64, + size_in: u64, + size_out: u64, + cfg: &VecOpsConfig, + output: &mut (impl HostOrDeviceSlice<$field> + ?Sized), + ) -> Result<(), eIcicleError> { + unsafe { + $field_prefix_ident::slice_ffi( + input.as_ptr(), + offset, + stride, + size_in, + size_out, + cfg as *const VecOpsConfig, + output.as_mut_ptr(), + ) + .wrap() + } + } } }; } @@ -436,6 +780,12 @@ macro_rules! impl_vec_ops_tests { initialize(); check_bit_reverse_inplace::<$field>() } + + #[test] + pub fn test_slice() { + initialize(); + check_slice::<$field>() + } } }; } diff --git a/wrappers/rust/icicle-core/src/vec_ops/tests.rs b/wrappers/rust/icicle-core/src/vec_ops/tests.rs index 6762f06c9..4a16fcb21 100644 --- a/wrappers/rust/icicle-core/src/vec_ops/tests.rs +++ b/wrappers/rust/icicle-core/src/vec_ops/tests.rs @@ -2,7 +2,7 @@ use crate::test_utilities; use crate::traits::GenerateRandom; use crate::vec_ops::{ - accumulate_scalars, add_scalars, bit_reverse, bit_reverse_inplace, mul_scalars, sub_scalars, transpose_matrix, + accumulate_scalars, add_scalars, bit_reverse, bit_reverse_inplace, mul_scalars, slice, div_scalars, sum_scalars, scalar_add, scalar_sub, scalar_mul, product_scalars, sub_scalars, transpose_matrix, FieldImpl, VecOps, VecOpsConfig, }; use icicle_runtime::device::Device; @@ -44,6 +44,12 @@ where check_vec_ops_scalars_add::(test_size); check_vec_ops_scalars_sub::(test_size); check_vec_ops_scalars_mul::(test_size); + check_vec_ops_scalars_div::(test_size); + check_vec_ops_scalars_sum::(test_size); + check_vec_ops_scalars_product::(test_size); + check_vec_ops_scalars_add_scalar::(test_size); + check_vec_ops_scalars_sub_scalar::(test_size); + check_vec_ops_scalars_mul_scalar::(test_size); check_vec_ops_scalars_accumulate::(test_size); } @@ -140,6 +146,191 @@ where .unwrap(); } +pub fn check_vec_ops_scalars_div(test_size: usize) +where + ::Config: VecOps + GenerateRandom, +{ + let a_main = F::Config::generate_random(test_size); + let b = F::Config::generate_random(test_size); + let mut result_main = vec![F::zero(); test_size]; + let mut result_ref = vec![F::zero(); test_size]; + + let a_main = HostSlice::from_slice(&a_main); + let b = HostSlice::from_slice(&b); + let result_main = HostSlice::from_mut_slice(&mut result_main); + let result_ref = HostSlice::from_mut_slice(&mut result_ref); + + let mut stream = IcicleStream::create().unwrap(); + let mut cfg = VecOpsConfig::default(); + cfg.stream_handle = *stream; + + test_utilities::test_set_main_device(); + div_scalars(a_main, b, result_main, &cfg).unwrap(); + + test_utilities::test_set_ref_device(); + div_scalars(a_main, b, result_ref, &cfg).unwrap(); + + assert_eq!(result_main.as_slice(), result_ref.as_slice()); + + stream + .destroy() + .unwrap(); +} + +pub fn check_vec_ops_scalars_sum(test_size: usize) +where + ::Config: VecOps + GenerateRandom, +{ + let a_main = F::Config::generate_random(test_size); + let mut result_main = vec![F::zero(); test_size]; + let mut result_ref = vec![F::zero(); test_size]; + + let a_main = HostSlice::from_slice(&a_main); + let result_main = HostSlice::from_mut_slice(&mut result_main); + let result_ref = HostSlice::from_mut_slice(&mut result_ref); + + let mut stream = IcicleStream::create().unwrap(); + let mut cfg = VecOpsConfig::default(); + cfg.stream_handle = *stream; + + test_utilities::test_set_main_device(); + sum_scalars(a_main, result_main, &cfg).unwrap(); + + test_utilities::test_set_ref_device(); + sum_scalars(a_main, result_ref, &cfg).unwrap(); + + assert_eq!(result_main.as_slice(), result_ref.as_slice()); + + stream + .destroy() + .unwrap(); +} + +pub fn check_vec_ops_scalars_product(test_size: usize) +where + ::Config: VecOps + GenerateRandom, +{ + let a_main = F::Config::generate_random(test_size); + let mut result_main = vec![F::zero(); test_size]; + let mut result_ref = vec![F::zero(); test_size]; + + let a_main = HostSlice::from_slice(&a_main); + let result_main = HostSlice::from_mut_slice(&mut result_main); + let result_ref = HostSlice::from_mut_slice(&mut result_ref); + + let mut stream = IcicleStream::create().unwrap(); + let mut cfg = VecOpsConfig::default(); + cfg.stream_handle = *stream; + + test_utilities::test_set_main_device(); + product_scalars(a_main, result_main, &cfg).unwrap(); + + test_utilities::test_set_ref_device(); + product_scalars(a_main, result_ref, &cfg).unwrap(); + + assert_eq!(result_main.as_slice(), result_ref.as_slice()); + + stream + .destroy() + .unwrap(); +} + +pub fn check_vec_ops_scalars_add_scalar(test_size: usize) +where + ::Config: VecOps + GenerateRandom, +{ + let a_main = F::Config::generate_random(1); + let b = F::Config::generate_random(test_size); + let mut result_main = vec![F::zero(); test_size]; + let mut result_ref = vec![F::zero(); test_size]; + + let a_main = HostSlice::from_slice(&a_main); + let b = HostSlice::from_slice(&b); + let result_main = HostSlice::from_mut_slice(&mut result_main); + let result_ref = HostSlice::from_mut_slice(&mut result_ref); + + let mut stream = IcicleStream::create().unwrap(); + let mut cfg = VecOpsConfig::default(); + cfg.stream_handle = *stream; + cfg.batch_size = 1; + + test_utilities::test_set_main_device(); + scalar_add(a_main, b, result_main, &cfg).unwrap(); + + test_utilities::test_set_ref_device(); + scalar_add(a_main, b, result_ref, &cfg).unwrap(); + + assert_eq!(result_main.as_slice(), result_ref.as_slice()); + + stream + .destroy() + .unwrap(); +} + +pub fn check_vec_ops_scalars_sub_scalar(test_size: usize) +where + ::Config: VecOps + GenerateRandom, +{ + let a_main = F::Config::generate_random(1); + let b = F::Config::generate_random(test_size); + let mut result_main = vec![F::zero(); test_size]; + let mut result_ref = vec![F::zero(); test_size]; + + let a_main = HostSlice::from_slice(&a_main); + let b = HostSlice::from_slice(&b); + let result_main = HostSlice::from_mut_slice(&mut result_main); + let result_ref = HostSlice::from_mut_slice(&mut result_ref); + + let mut stream = IcicleStream::create().unwrap(); + let mut cfg = VecOpsConfig::default(); + cfg.stream_handle = *stream; + cfg.batch_size = 1; + + test_utilities::test_set_main_device(); + scalar_sub(a_main, b, result_main, &cfg).unwrap(); + + test_utilities::test_set_ref_device(); + scalar_sub(a_main, b, result_ref, &cfg).unwrap(); + + assert_eq!(result_main.as_slice(), result_ref.as_slice()); + + stream + .destroy() + .unwrap(); +} + +pub fn check_vec_ops_scalars_mul_scalar(test_size: usize) +where + ::Config: VecOps + GenerateRandom, +{ + let a_main = F::Config::generate_random(1); + let b = F::Config::generate_random(test_size); + let mut result_main = vec![F::zero(); test_size]; + let mut result_ref = vec![F::zero(); test_size]; + + let a_main = HostSlice::from_slice(&a_main); + let b = HostSlice::from_slice(&b); + let result_main = HostSlice::from_mut_slice(&mut result_main); + let result_ref = HostSlice::from_mut_slice(&mut result_ref); + + let mut stream = IcicleStream::create().unwrap(); + let mut cfg = VecOpsConfig::default(); + cfg.stream_handle = *stream; + cfg.batch_size = 1; + + test_utilities::test_set_main_device(); + scalar_mul(a_main, b, result_main, &cfg).unwrap(); + + test_utilities::test_set_ref_device(); + scalar_mul(a_main, b, result_ref, &cfg).unwrap(); + + assert_eq!(result_main.as_slice(), result_ref.as_slice()); + + stream + .destroy() + .unwrap(); +} + pub fn check_vec_ops_scalars_accumulate(test_size: usize) where ::Config: VecOps + GenerateRandom, @@ -205,6 +396,47 @@ where assert_eq!(result_main, result_ref); } +pub fn check_slice() +where + ::Config: VecOps + GenerateRandom, +{ + let size_in: u64 = 1 << 10; + let offset: u64 = 10; + let stride: u64 = 3; + let size_out: u64 = ((size_in - offset) / stride) - 1; + + let input_matrix = F::Config::generate_random(size_in as usize); + let mut result_main = vec![F::zero(); size_out as usize]; + let mut result_ref = vec![F::zero(); size_out as usize]; + + let cfg = VecOpsConfig::default(); + test_utilities::test_set_main_device(); + slice( + HostSlice::from_slice(&input_matrix), + offset, + stride, + size_in, + size_out, + &cfg, + HostSlice::from_mut_slice(&mut result_main), + ) + .unwrap(); + + test_utilities::test_set_ref_device(); + slice( + HostSlice::from_slice(&input_matrix), + offset, + stride, + size_in, + size_out, + &cfg, + HostSlice::from_mut_slice(&mut result_ref), + ) + .unwrap(); + + assert_eq!(result_main, result_ref); +} + pub fn check_bit_reverse() where ::Config: VecOps + GenerateRandom,