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.

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

ML-DSA Parameter Set Monomorphization with Macros #732

Merged
merged 61 commits into from
Jan 9, 2025

Conversation

jschneider-bensch
Copy link
Collaborator

The motivation for this PR is to mostly get rid of allocations below the level of ml_dsa_generic. To this end, const generics are replaced by a macro-based manual monomorphization to the different parameter sets. In doing this, I've attempted to simplify parameter set specific constants, as well as platform instantiation and multiplexing.

Essentially ml_dsa_generic, constants, the instantiations in ml_dsa_generic::instantiation::{portable, neon, avx2} and ml_dsa_generic::multiplexing now have (macro-generated) submodules for each parameter set which should be guarded by their respective feature flag. They contain:

  • constants::ml_dsa_44/65/87: Parameter set specific constants go here. The constants module also defines const fns for deriving everything but the basic constants in the parameter set specific submodules.
  • ml_dsa_generic::ml_dsa_44/65/87: A version of the ml_dsa_generic::generic module where the parameter specific constants from constants are pulled in and derived constants are derived using the const fns from constants. The module also defines type aliases MLDSA{44/65/87}SigningKey, MLDSA{44/65/87}VerificationKey, etc. This is achieved by the ml_dsa_parameter_sets proc macro.
  • instantiations::{portable, neon, avx2}::ml_dsa_44/65/87: The platform specific instantiations of the versions from ml_dsa_generic. In the case of avx2that includes the application of thetarget_platform` cfg.
  • multiplexing::ml_dsa_44/65/87`: Runtime multiplexing, reverting to a portable version if platform-specific version was not compiled in.

I've verified that F* extraction still works and laxes. C extraction pending further updates.

@jschneider-bensch jschneider-bensch requested a review from a team as a code owner January 7, 2025 14:29
@jschneider-bensch jschneider-bensch self-assigned this Jan 7, 2025
Copy link
Contributor

@karthikbhargavan karthikbhargavan left a comment

Choose a reason for hiding this comment

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

This is a big one! Left a few comments, but happy for most of these being follow-ups.

However, let's document the general style guidelines implemented here, so that new code can also follow this pattern. I am adding a comment to the PR noting this.

libcrux-intrinsics/src/avx2.rs Show resolved Hide resolved
libcrux-ml-dsa/src/arithmetic.rs Show resolved Hide resolved
libcrux-ml-dsa/src/arithmetic.rs Outdated Show resolved Hide resolved
libcrux-ml-dsa/src/arithmetic.rs Outdated Show resolved Hide resolved
libcrux-ml-dsa/src/arithmetic.rs Show resolved Hide resolved
libcrux-ml-dsa/src/matrix.rs Show resolved Hide resolved
libcrux-ml-dsa/src/ml_dsa_generic.rs Outdated Show resolved Hide resolved
libcrux-ml-dsa/src/samplex4.rs Show resolved Hide resolved
libcrux-ml-dsa/src/simd/avx2.rs Outdated Show resolved Hide resolved
libcrux-ml-dsa/src/simd/portable/vector_type.rs Outdated Show resolved Hide resolved
@karthikbhargavan
Copy link
Contributor

This PR does a bit more cleanup than just the const generic monomorphization.
Let's try to document some of these style decisions so that we can be uniform in the codebase (and backport to ML-KEM and SHA-3 when we touch that code again.). Here are some things I noticed:

  • All struct and array arguments to functions are now borrows
  • We uniformly use &mut arguments instead of returning a new array or struct
  • We uniformly use slices (with implicit lengths) instead of fixed-size arrays
  • We eliminate const generic arguments and replace them by real arguments (when truly needed)
  • In several places, we replace calls to complex iterators like enumerate with simpler for loops over ranges.

@jschneider-bensch
Copy link
Collaborator Author

You're right, there's more here than the monomorphization. Thanks for noting down the style changes!

franziskuskiefer and others added 6 commits January 8, 2025 09:13
Eurydice can't handle associated types right now.
So this changes it such that the Operations trait is implemented
on top of a wrapper struct for the coefficients again.
This requires the implementations of the trait to point into the struct
again. But this is more self contained than propagating the generic type
for the trait.

I couldn't measure an impact on performance or stack size with this change
Co-authored-by: Jonas Schneider-Bensch <[email protected]>
Copy link
Member

@franziskuskiefer franziskuskiefer left a comment

Choose a reason for hiding this comment

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

I think we should simplify the plumbing for the variants a little to make it easier to add nice APIs. But that's for follow-ups. Let's get this in.

@franziskuskiefer franziskuskiefer added this pull request to the merge queue Jan 9, 2025
Merged via the queue into main with commit 31e77d0 Jan 9, 2025
57 checks passed
@franziskuskiefer franziskuskiefer deleted the franziskus/mldsa-cleanup branch January 9, 2025 10:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants