Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

refactor(core): start refactoring pbs code #627

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

refactor(core): start refactoring pbs code #627

wants to merge 1 commit into from

Conversation

sarah-ek
Copy link
Contributor

closes: please link all relevant issues

PR content/description

Check-list:

  • Tests for the changes have been added (for bug fixes / features)
  • Docs have been added / updated (for bug fixes / features)
  • Relevant issues are marked as resolved/closed, related issues are linked in the description
  • Check for breaking changes (including serialization changes) and add them to commit message following the conventional commit specification

@cla-bot cla-bot bot added the cla-signed label Oct 11, 2023
@github-actions
Copy link

@slab-ci cpu_fast_test

IceTDrinker
IceTDrinker reviewed Oct 12, 2023
View reviewed changes
Copy link
Member

@IceTDrinker IceTDrinker left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

First batch of comments, thanks for the work, I'll check the remainder after lunch :)

tfhe/src/core_crypto/algorithms/ggsw_conversion.rs Outdated
@@ -3,14 +3,14 @@
//! like the Fourier domain.

use crate::core_crypto::commons::computation_buffers::ComputationBuffers;
use crate::core_crypto::commons::math::fft::{Fft, FftView};
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should we maybe plan for the fft64/fft128 split as we had in the fft_impl module ? I believe we want to also move the f128 at some point, we can check with JB

tfhe/src/core_crypto/algorithms/lwe_bootstrap_key_conversion.rs Outdated
use dyn_stack::{PodStack, SizeOverflow, StackReq};
use dyn_stack::{PodStack, ReborrowMut, SizeOverflow, StackReq};

impl<'a> FourierLweBootstrapKeyMutView<'a> {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should this be in the entities module for now ?

My reasoning : the algorithm for now is tied to the entity, what we would do is the same you did in the ggsw_conversion module i.e. rewrite the algorithm in this file as a function not in an impl block ?

tfhe/src/core_crypto/commons/math/fft/mod.rs Outdated
let mut stack = PodStack::new(&mut stack);

for (fourier_poly, standard_poly) in izip!(
fourier_poly_chunk.chunks_exact_mut(f_polynomial_size),
standard_poly_chunk.chunks_exact(polynomial_size.0)
) {
fft.forward_as_torus(
FourierPolynomialMutView { data: fourier_poly },
FourierPolynomialMutView::from_container(fourier_poly),
PolynomialView::from_container(standard_poly),
stack.rb_mut(),
);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could we put the mod tests that is at the the end of the file towards the top ?

tfhe/src/core_crypto/fft_impl/fft64/math/mod.rs
@@ -1,3 +0,0 @@
pub mod decomposition;
pub mod fft;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks like fft has not been deleted ?

tfhe/src/core_crypto/entities/fourier_ggsw_ciphertext.rs Outdated
) -> Self {
assert_eq!(
data.container_len(),
polynomial_size.to_fourier_polynomial_size().0 * glwe_size.0 * row_count
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

row_count should be equal to glwe_size you can check ggsw_ciphertext.rs I think there are some size functions :)

IceTDrinker
IceTDrinker reviewed Oct 12, 2023
View reviewed changes
tfhe/src/core_crypto/entities/fourier_lwe_bootstrap_key.rs Outdated
}

pub fn output_lwe_dimension(&self) -> LweDimension {
LweDimension((self.glwe_size.0 - 1) * self.polynomial_size().0)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you can now use self.glwe_size.to_glwe_dimension().to_equivalent_lwe_dimension(self.polynomial_size())

IceTDrinker
IceTDrinker reviewed Oct 12, 2023
View reviewed changes
tfhe/src/core_crypto/entities/fourier_ggsw_ciphertext.rs Outdated
data: C,
glwe_size: GlweSize,
polynomial_size: PolynomialSize,
row_count: usize,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should not be needed row_count == glwe_size

IceTDrinker
IceTDrinker reviewed Oct 12, 2023
View reviewed changes
tfhe/src/core_crypto/entities/fourier_ggsw_ciphertext.rs Outdated
Comment on lines 171 to 174
pub fn row_count(&self) -> usize {
self.row_count
}

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should not be required

IceTDrinker
IceTDrinker reviewed Oct 12, 2023
View reviewed changes
tfhe/src/core_crypto/algorithms/lwe_programmable_bootstrapping.rs
Comment on lines +1831 to +1554
impl<Scalar> FourierBootstrapKey<Scalar> for FourierLweBootstrapKeyOwned
where
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this impl should be in the entity file for now

IceTDrinker
IceTDrinker reviewed Oct 12, 2023
View reviewed changes
tfhe/src/core_crypto/algorithms/lwe_programmable_bootstrapping.rs Outdated
}

#[cfg_attr(__profiling, inline(never))]
pub(crate) fn update_with_fmadd(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe the update_with_fmadd could be seen as FFT primitives ? As this allows to make polynomial mulitplications in an efficient way ? what do you think ?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

makes sense to me. i can put them with the fft

IceTDrinker
IceTDrinker reviewed Oct 12, 2023
View reviewed changes
tfhe/src/core_crypto/algorithms/lwe_programmable_bootstrapping.rs Outdated
})
}

pub(crate) fn update_with_fmadd_factor(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same comment :)

IceTDrinker
IceTDrinker reviewed Oct 12, 2023
View reviewed changes
tfhe/src/core_crypto/algorithms/lwe_programmable_bootstrapping.rs Outdated
@@ -335,6 +480,155 @@ pub fn add_external_product_assign<Scalar, OutputGlweCont, InputGlweCont, GgswCo
add_external_product_assign_mem_optimized(out, ggsw, glwe, fft, buffers.stack());
}

#[cfg_attr(__profiling, inline(never))]
pub(crate) fn collect_next_term<Scalar: UnsignedTorus>(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could it be a method on the TensorSignedDecomposition struct ?

IceTDrinker
IceTDrinker reviewed Oct 12, 2023
View reviewed changes
tfhe/src/core_crypto/algorithms/lwe_programmable_bootstrapping.rs Outdated
Comment on lines 32 to 33
impl<'a> FourierLweBootstrapKeyView<'a> {
// CastInto required for PBS modulus switch which returns a usize
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

for now I believe this should be in the entity module

@github-actions
Copy link

@slab-ci cpu_fast_test

@sarah-ek
Copy link
Contributor Author

i pushed a new commit with your suggestions. lemme know if there's anything else to fix
i'm also gonna run the benchmarks locally, to check that there are no regressions

@sarah-ek
Copy link
Contributor Author

no regressions found on my machine

@IceTDrinker
Copy link
Member

created a follow up issue https://github.com/zama-ai/tfhe-rs-internal/issues/287

This looks very good, finishing reading up !

IceTDrinker
IceTDrinker reviewed Oct 13, 2023
View reviewed changes
tfhe/src/core_crypto/entities/fourier_lwe_bootstrap_key.rs
Comment on lines -366 to -367

impl<Scalar> FourierBootstrapKey<Scalar> for FourierLweBootstrapKeyOwned
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any reason to remove this implementation ?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

that might have been a mistake on my part

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@IceTDrinker IceTDrinker IceTDrinker left review comments

Assignees
No one assigned
Labels
cla-signed
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
2 participants
@sarah-ek @IceTDrinker