use crate::core_crypto::gpu::CudaStreams;
use crate::core_crypto::prelude::LweBskGroupingFactor;
use crate::integer::gpu::ciphertext::{CudaIntegerRadixCiphertext, CudaUnsignedRadixCiphertext};
use crate::integer::gpu::server_key::{CudaBootstrappingKey, CudaDynamicKeyswitchingKey};
use crate::integer::gpu::{
cuda_backend_get_full_propagate_assign_size_on_gpu, cuda_backend_get_rotate_left_size_on_gpu,
cuda_backend_get_rotate_right_size_on_gpu, cuda_backend_unchecked_rotate_left_assign,
cuda_backend_unchecked_rotate_right_assign, CudaServerKey, PBSType,
};
impl CudaServerKey {
pub fn unchecked_rotate_right_assign<T>(
&self,
ct: &mut T,
rotate: &CudaUnsignedRadixCiphertext,
streams: &CudaStreams,
) where
T: CudaIntegerRadixCiphertext,
{
let lwe_ciphertext_count = ct.as_ref().d_blocks.lwe_ciphertext_count();
let is_signed = T::IS_SIGNED;
let CudaDynamicKeyswitchingKey::Standard(computing_ks_key) = &self.key_switching_key else {
panic!("Only the standard atomic pattern is supported on GPU")
};
unsafe {
match &self.bootstrapping_key {
CudaBootstrappingKey::Classic(d_bsk) => {
cuda_backend_unchecked_rotate_right_assign(
streams,
ct.as_mut(),
rotate.as_ref(),
&d_bsk.d_vec,
&computing_ks_key.d_vec,
self.message_modulus,
self.carry_modulus,
d_bsk.glwe_dimension,
d_bsk.polynomial_size,
computing_ks_key.input_key_lwe_size().to_lwe_dimension(),
computing_ks_key.output_key_lwe_size().to_lwe_dimension(),
computing_ks_key.decomposition_level_count(),
computing_ks_key.decomposition_base_log(),
d_bsk.decomp_level_count,
d_bsk.decomp_base_log,
lwe_ciphertext_count.0 as u32,
is_signed,
PBSType::Classical,
LweBskGroupingFactor(0),
d_bsk.ms_noise_reduction_configuration.as_ref(),
);
}
CudaBootstrappingKey::MultiBit(d_multibit_bsk) => {
cuda_backend_unchecked_rotate_right_assign(
streams,
ct.as_mut(),
rotate.as_ref(),
&d_multibit_bsk.d_vec,
&computing_ks_key.d_vec,
self.message_modulus,
self.carry_modulus,
d_multibit_bsk.glwe_dimension,
d_multibit_bsk.polynomial_size,
computing_ks_key.input_key_lwe_size().to_lwe_dimension(),
computing_ks_key.output_key_lwe_size().to_lwe_dimension(),
computing_ks_key.decomposition_level_count(),
computing_ks_key.decomposition_base_log(),
d_multibit_bsk.decomp_level_count,
d_multibit_bsk.decomp_base_log,
lwe_ciphertext_count.0 as u32,
is_signed,
PBSType::MultiBit,
d_multibit_bsk.grouping_factor,
None,
);
}
}
}
}
pub fn unchecked_rotate_right<T>(
&self,
ct: &T,
rotate: &CudaUnsignedRadixCiphertext,
streams: &CudaStreams,
) -> T
where
T: CudaIntegerRadixCiphertext,
{
let mut result = ct.duplicate(streams);
self.unchecked_rotate_right_assign(&mut result, rotate, streams);
result
}
pub fn unchecked_rotate_left_assign<T>(
&self,
ct: &mut T,
rotate: &CudaUnsignedRadixCiphertext,
streams: &CudaStreams,
) where
T: CudaIntegerRadixCiphertext,
{
let lwe_ciphertext_count = ct.as_ref().d_blocks.lwe_ciphertext_count();
let is_signed = T::IS_SIGNED;
let CudaDynamicKeyswitchingKey::Standard(computing_ks_key) = &self.key_switching_key else {
panic!("Only the standard atomic pattern is supported on GPU")
};
unsafe {
match &self.bootstrapping_key {
CudaBootstrappingKey::Classic(d_bsk) => {
cuda_backend_unchecked_rotate_left_assign(
streams,
ct.as_mut(),
rotate.as_ref(),
&d_bsk.d_vec,
&computing_ks_key.d_vec,
self.message_modulus,
self.carry_modulus,
d_bsk.glwe_dimension,
d_bsk.polynomial_size,
computing_ks_key.input_key_lwe_size().to_lwe_dimension(),
computing_ks_key.output_key_lwe_size().to_lwe_dimension(),
computing_ks_key.decomposition_level_count(),
computing_ks_key.decomposition_base_log(),
d_bsk.decomp_level_count,
d_bsk.decomp_base_log,
lwe_ciphertext_count.0 as u32,
is_signed,
PBSType::Classical,
LweBskGroupingFactor(0),
d_bsk.ms_noise_reduction_configuration.as_ref(),
);
}
CudaBootstrappingKey::MultiBit(d_multibit_bsk) => {
cuda_backend_unchecked_rotate_left_assign(
streams,
ct.as_mut(),
rotate.as_ref(),
&d_multibit_bsk.d_vec,
&computing_ks_key.d_vec,
self.message_modulus,
self.carry_modulus,
d_multibit_bsk.glwe_dimension,
d_multibit_bsk.polynomial_size,
computing_ks_key.input_key_lwe_size().to_lwe_dimension(),
computing_ks_key.output_key_lwe_size().to_lwe_dimension(),
computing_ks_key.decomposition_level_count(),
computing_ks_key.decomposition_base_log(),
d_multibit_bsk.decomp_level_count,
d_multibit_bsk.decomp_base_log,
lwe_ciphertext_count.0 as u32,
is_signed,
PBSType::MultiBit,
d_multibit_bsk.grouping_factor,
None,
);
}
}
}
}
pub fn unchecked_rotate_left<T>(
&self,
ct: &T,
rotate: &CudaUnsignedRadixCiphertext,
streams: &CudaStreams,
) -> T
where
T: CudaIntegerRadixCiphertext,
{
let mut result = ct.duplicate(streams);
self.unchecked_rotate_left_assign(&mut result, rotate, streams);
result
}
pub fn rotate_right<T>(
&self,
ct: &T,
rotate: &CudaUnsignedRadixCiphertext,
streams: &CudaStreams,
) -> T
where
T: CudaIntegerRadixCiphertext,
{
let mut tmp_lhs: T;
let mut tmp_rhs: CudaUnsignedRadixCiphertext;
let (lhs, rhs) = match (
ct.block_carries_are_empty(),
rotate.block_carries_are_empty(),
) {
(true, true) => (ct, rotate),
(true, false) => {
tmp_rhs = rotate.duplicate(streams);
self.full_propagate_assign(&mut tmp_rhs, streams);
(ct, &tmp_rhs)
}
(false, true) => {
tmp_lhs = ct.duplicate(streams);
self.full_propagate_assign(&mut tmp_lhs, streams);
(&tmp_lhs, rotate)
}
(false, false) => {
tmp_lhs = ct.duplicate(streams);
tmp_rhs = rotate.duplicate(streams);
self.full_propagate_assign(&mut tmp_lhs, streams);
self.full_propagate_assign(&mut tmp_rhs, streams);
(&tmp_lhs, &tmp_rhs)
}
};
let mut result = lhs.duplicate(streams);
self.unchecked_rotate_right_assign(&mut result, rhs, streams);
result
}
pub fn rotate_right_assign<T>(
&self,
ct: &mut T,
rotate: &CudaUnsignedRadixCiphertext,
streams: &CudaStreams,
) where
T: CudaIntegerRadixCiphertext,
{
let mut tmp_lhs: T;
let mut tmp_rhs: CudaUnsignedRadixCiphertext;
let (lhs, rhs) = match (
ct.block_carries_are_empty(),
rotate.block_carries_are_empty(),
) {
(true, true) => (ct, rotate),
(true, false) => {
tmp_rhs = rotate.duplicate(streams);
self.full_propagate_assign(&mut tmp_rhs, streams);
(ct, &tmp_rhs)
}
(false, true) => {
tmp_lhs = ct.duplicate(streams);
self.full_propagate_assign(&mut tmp_lhs, streams);
(&mut tmp_lhs, rotate)
}
(false, false) => {
tmp_lhs = ct.duplicate(streams);
tmp_rhs = rotate.duplicate(streams);
self.full_propagate_assign(&mut tmp_lhs, streams);
self.full_propagate_assign(&mut tmp_rhs, streams);
(&mut tmp_lhs, &tmp_rhs)
}
};
self.unchecked_rotate_right_assign(lhs, rhs, streams);
}
pub fn rotate_left<T>(
&self,
ct: &T,
rotate: &CudaUnsignedRadixCiphertext,
streams: &CudaStreams,
) -> T
where
T: CudaIntegerRadixCiphertext,
{
let mut tmp_lhs: T;
let mut tmp_rhs: CudaUnsignedRadixCiphertext;
let (lhs, rhs) = match (
ct.block_carries_are_empty(),
rotate.block_carries_are_empty(),
) {
(true, true) => (ct, rotate),
(true, false) => {
tmp_rhs = rotate.duplicate(streams);
self.full_propagate_assign(&mut tmp_rhs, streams);
(ct, &tmp_rhs)
}
(false, true) => {
tmp_lhs = ct.duplicate(streams);
self.full_propagate_assign(&mut tmp_lhs, streams);
(&tmp_lhs, rotate)
}
(false, false) => {
tmp_lhs = ct.duplicate(streams);
tmp_rhs = rotate.duplicate(streams);
self.full_propagate_assign(&mut tmp_lhs, streams);
self.full_propagate_assign(&mut tmp_rhs, streams);
(&tmp_lhs, &tmp_rhs)
}
};
let mut result = lhs.duplicate(streams);
self.unchecked_rotate_left_assign(&mut result, rhs, streams);
result
}
pub fn rotate_left_assign<T>(
&self,
ct: &mut T,
rotate: &CudaUnsignedRadixCiphertext,
streams: &CudaStreams,
) where
T: CudaIntegerRadixCiphertext,
{
let mut tmp_lhs: T;
let mut tmp_rhs: CudaUnsignedRadixCiphertext;
let (lhs, rhs) = match (
ct.block_carries_are_empty(),
rotate.block_carries_are_empty(),
) {
(true, true) => (ct, rotate),
(true, false) => {
tmp_rhs = rotate.duplicate(streams);
self.full_propagate_assign(&mut tmp_rhs, streams);
(ct, &tmp_rhs)
}
(false, true) => {
tmp_lhs = ct.duplicate(streams);
self.full_propagate_assign(&mut tmp_lhs, streams);
(&mut tmp_lhs, rotate)
}
(false, false) => {
tmp_lhs = ct.duplicate(streams);
tmp_rhs = rotate.duplicate(streams);
self.full_propagate_assign(&mut tmp_lhs, streams);
self.full_propagate_assign(&mut tmp_rhs, streams);
(&mut tmp_lhs, &tmp_rhs)
}
};
self.unchecked_rotate_left_assign(lhs, rhs, streams);
}
pub fn get_rotate_left_size_on_gpu<T: CudaIntegerRadixCiphertext>(
&self,
ct_left: &T,
ct_right: &CudaUnsignedRadixCiphertext,
streams: &CudaStreams,
) -> u64 {
assert_eq!(
ct_left.as_ref().d_blocks.lwe_dimension(),
ct_right.as_ref().d_blocks.lwe_dimension()
);
assert_eq!(
ct_left.as_ref().d_blocks.lwe_ciphertext_count(),
ct_right.as_ref().d_blocks.lwe_ciphertext_count()
);
let CudaDynamicKeyswitchingKey::Standard(computing_ks_key) = &self.key_switching_key else {
panic!("Only the standard atomic pattern is supported on GPU")
};
let full_prop_mem = match &self.bootstrapping_key {
CudaBootstrappingKey::Classic(d_bsk) => {
cuda_backend_get_full_propagate_assign_size_on_gpu(
streams,
d_bsk.input_lwe_dimension(),
d_bsk.glwe_dimension(),
d_bsk.polynomial_size(),
computing_ks_key.decomposition_level_count(),
computing_ks_key.decomposition_base_log(),
d_bsk.decomp_level_count(),
d_bsk.decomp_base_log(),
self.message_modulus,
self.carry_modulus,
PBSType::Classical,
LweBskGroupingFactor(0),
d_bsk.ms_noise_reduction_configuration.as_ref(),
)
}
CudaBootstrappingKey::MultiBit(d_multibit_bsk) => {
cuda_backend_get_full_propagate_assign_size_on_gpu(
streams,
d_multibit_bsk.input_lwe_dimension(),
d_multibit_bsk.glwe_dimension(),
d_multibit_bsk.polynomial_size(),
computing_ks_key.decomposition_level_count(),
computing_ks_key.decomposition_base_log(),
d_multibit_bsk.decomp_level_count(),
d_multibit_bsk.decomp_base_log(),
self.message_modulus,
self.carry_modulus,
PBSType::MultiBit,
d_multibit_bsk.grouping_factor,
None,
)
}
};
let actual_full_prop_mem = match (
ct_left.block_carries_are_empty(),
ct_right.block_carries_are_empty(),
) {
(true, true) => 0,
(true, false) => self.get_ciphertext_size_on_gpu(ct_right) + full_prop_mem,
(false, true) => full_prop_mem,
(false, false) => self.get_ciphertext_size_on_gpu(ct_right) + full_prop_mem,
};
let lwe_ciphertext_count = ct_left.as_ref().d_blocks.lwe_ciphertext_count();
let rotate_mem = match &self.bootstrapping_key {
CudaBootstrappingKey::Classic(d_bsk) => cuda_backend_get_rotate_left_size_on_gpu(
streams,
self.message_modulus,
self.carry_modulus,
d_bsk.glwe_dimension,
d_bsk.polynomial_size,
computing_ks_key.input_key_lwe_size().to_lwe_dimension(),
computing_ks_key.output_key_lwe_size().to_lwe_dimension(),
computing_ks_key.decomposition_level_count(),
computing_ks_key.decomposition_base_log(),
d_bsk.decomp_level_count,
d_bsk.decomp_base_log,
lwe_ciphertext_count.0 as u32,
T::IS_SIGNED,
PBSType::Classical,
LweBskGroupingFactor(0),
d_bsk.ms_noise_reduction_configuration.as_ref(),
),
CudaBootstrappingKey::MultiBit(d_multibit_bsk) => {
cuda_backend_get_rotate_left_size_on_gpu(
streams,
self.message_modulus,
self.carry_modulus,
d_multibit_bsk.glwe_dimension,
d_multibit_bsk.polynomial_size,
computing_ks_key.input_key_lwe_size().to_lwe_dimension(),
computing_ks_key.output_key_lwe_size().to_lwe_dimension(),
computing_ks_key.decomposition_level_count(),
computing_ks_key.decomposition_base_log(),
d_multibit_bsk.decomp_level_count,
d_multibit_bsk.decomp_base_log,
lwe_ciphertext_count.0 as u32,
T::IS_SIGNED,
PBSType::MultiBit,
d_multibit_bsk.grouping_factor,
None,
)
}
};
actual_full_prop_mem.max(rotate_mem)
}
pub fn get_rotate_right_size_on_gpu<T: CudaIntegerRadixCiphertext>(
&self,
ct_left: &T,
ct_right: &CudaUnsignedRadixCiphertext,
streams: &CudaStreams,
) -> u64 {
assert_eq!(
ct_left.as_ref().d_blocks.lwe_dimension(),
ct_right.as_ref().d_blocks.lwe_dimension()
);
assert_eq!(
ct_left.as_ref().d_blocks.lwe_ciphertext_count(),
ct_right.as_ref().d_blocks.lwe_ciphertext_count()
);
let CudaDynamicKeyswitchingKey::Standard(computing_ks_key) = &self.key_switching_key else {
panic!("Only the standard atomic pattern is supported on GPU")
};
let full_prop_mem = match &self.bootstrapping_key {
CudaBootstrappingKey::Classic(d_bsk) => {
cuda_backend_get_full_propagate_assign_size_on_gpu(
streams,
d_bsk.input_lwe_dimension(),
d_bsk.glwe_dimension(),
d_bsk.polynomial_size(),
computing_ks_key.decomposition_level_count(),
computing_ks_key.decomposition_base_log(),
d_bsk.decomp_level_count(),
d_bsk.decomp_base_log(),
self.message_modulus,
self.carry_modulus,
PBSType::Classical,
LweBskGroupingFactor(0),
d_bsk.ms_noise_reduction_configuration.as_ref(),
)
}
CudaBootstrappingKey::MultiBit(d_multibit_bsk) => {
cuda_backend_get_full_propagate_assign_size_on_gpu(
streams,
d_multibit_bsk.input_lwe_dimension(),
d_multibit_bsk.glwe_dimension(),
d_multibit_bsk.polynomial_size(),
computing_ks_key.decomposition_level_count(),
computing_ks_key.decomposition_base_log(),
d_multibit_bsk.decomp_level_count(),
d_multibit_bsk.decomp_base_log(),
self.message_modulus,
self.carry_modulus,
PBSType::MultiBit,
d_multibit_bsk.grouping_factor,
None,
)
}
};
let actual_full_prop_mem = match (
ct_left.block_carries_are_empty(),
ct_right.block_carries_are_empty(),
) {
(true, true) => 0,
(true, false) => self.get_ciphertext_size_on_gpu(ct_right) + full_prop_mem,
(false, true) => full_prop_mem,
(false, false) => self.get_ciphertext_size_on_gpu(ct_right) + full_prop_mem,
};
let lwe_ciphertext_count = ct_left.as_ref().d_blocks.lwe_ciphertext_count();
let rotate_mem = match &self.bootstrapping_key {
CudaBootstrappingKey::Classic(d_bsk) => cuda_backend_get_rotate_right_size_on_gpu(
streams,
self.message_modulus,
self.carry_modulus,
d_bsk.glwe_dimension,
d_bsk.polynomial_size,
computing_ks_key.input_key_lwe_size().to_lwe_dimension(),
computing_ks_key.output_key_lwe_size().to_lwe_dimension(),
computing_ks_key.decomposition_level_count(),
computing_ks_key.decomposition_base_log(),
d_bsk.decomp_level_count,
d_bsk.decomp_base_log,
lwe_ciphertext_count.0 as u32,
T::IS_SIGNED,
PBSType::Classical,
LweBskGroupingFactor(0),
d_bsk.ms_noise_reduction_configuration.as_ref(),
),
CudaBootstrappingKey::MultiBit(d_multibit_bsk) => {
cuda_backend_get_rotate_right_size_on_gpu(
streams,
self.message_modulus,
self.carry_modulus,
d_multibit_bsk.glwe_dimension,
d_multibit_bsk.polynomial_size,
computing_ks_key.input_key_lwe_size().to_lwe_dimension(),
computing_ks_key.output_key_lwe_size().to_lwe_dimension(),
computing_ks_key.decomposition_level_count(),
computing_ks_key.decomposition_base_log(),
d_multibit_bsk.decomp_level_count,
d_multibit_bsk.decomp_base_log,
lwe_ciphertext_count.0 as u32,
T::IS_SIGNED,
PBSType::MultiBit,
d_multibit_bsk.grouping_factor,
None,
)
}
};
actual_full_prop_mem.max(rotate_mem)
}
}