ml-dsa: reduce stack usage and add low-memory signing

.github/workflows/ml-dsa.yml

@@ -42,6 +42,7 @@ jobs:
       - run: cargo test --release
       - run: cargo test --all-features --release
       - run: cargo test --no-default-features --release
+      - run: cargo test --no-default-features --features low-memory,zeroize --release
       - run: cargo test --no-default-features --features alloc --lib --release
 
   no_std:
@@ -50,6 +51,7 @@ jobs:
       matrix:
         target:
           - thumbv7em-none-eabi
+          - thumbv8m.main-none-eabi
           - wasm32-unknown-unknown
         rust:
           - 1.85.0 # MSRV
@@ -63,3 +65,25 @@ jobs:
 
       - uses: RustCrypto/actions/cargo-hack-install@master
       - run: cargo hack build --target ${{ matrix.target }} --feature-powerset --exclude-features default,getrandom
+
+  cortex_m_stack:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        features:
+          - low-memory
+          - low-memory,ml-dsa-65
+          - low-memory,ml-dsa-87
+    steps:
+      - uses: actions/checkout@v7
+      - uses: dtolnay/rust-toolchain@stable
+        with:
+          targets: thumbv8m.main-none-eabi
+      - run: sudo apt-get update && sudo apt-get install -y qemu-system-arm
+      - run: >-
+          cargo build --locked --release --target thumbv8m.main-none-eabi
+          --manifest-path tools/cortex-m-stack/Cargo.toml --features ${{ matrix.features }}
+      - run: >-
+          timeout 60s qemu-system-arm -M mps2-an505 -display none -serial none -monitor none
+          -semihosting-config enable=on,target=native
+          -kernel tools/cortex-m-stack/target/thumbv8m.main-none-eabi/release/ml-dsa-cortex-m-stack

ml-dsa/Cargo.toml

@@ -27,6 +27,8 @@ exclude = [
 [features]
 default = ["alloc", "getrandom", "pkcs8"]
 alloc = ["hybrid-array/alloc", "module-lattice/alloc", "pkcs8?/alloc", "signature/alloc"]
+# Regenerate public matrix entries and secret NTTs to reduce retained RAM and stack usage.
+low-memory = []
 
 getrandom = ["common/getrandom", "rand_core"]
 pkcs8 = ["dep:const-oid", "dep:pkcs8"]

ml-dsa/README.md

@@ -14,6 +14,14 @@ Pure Rust implementation of the Module-Lattice-Based Digital Signature Standard
 
 ML-DSA was formerly known as [CRYSTALS-Dilithium].
 
+With the `alloc` feature, expanded matrix rows and cached key components are stored on the heap to
+reduce construction and cloning stack usage. The opt-in `low-memory` feature regenerates public
+matrix entries and secret NTTs instead of retaining them, and provides a reusable `SigningWorkspace`.
+With `default-features = false`, this path needs no allocator. Enable `zeroize` to erase scratch after
+signing. This trades additional SHAKE and NTT work for lower RAM usage. A caller-owned workspace can
+reside in a static SRAM cell; existing signing methods construct it locally. See the
+[Cortex-M33 stack probe](tools/cortex-m-stack) for reproduction instructions.
+
 ## ⚠️ Security Warning
 
 The implementation contained in this crate has never been independently audited!

ml-dsa/benches/ml_dsa.rs

@@ -11,9 +11,10 @@ use ml_dsa::{
 /// ML-DSA benchmarks.
 #[allow(deprecated)] // TODO(tarcieri): stop using expanded signing keys
 fn criterion_benchmark(c: &mut Criterion) {
-    let xi = B32::generate();
-    let m = B32::generate();
-    let ctx = B32::generate();
+    // Keep the signing rejection sequence identical between benchmark runs.
+    let xi = B32::default();
+    let m = B32::from([1; 32]);
+    let ctx = B32::from([2; 32]);
 
     let kp = SigningKey::<MlDsa65>::from_seed(&xi);
     let sk = kp.expanded_key();
@@ -24,6 +25,19 @@ fn criterion_benchmark(c: &mut Criterion) {
     let vk_bytes = vk.encode();
     let sig_bytes = sig.encode();
 
+    // Reuse the expanded key so this measures signing independently of key import.
+    c.bench_function("sign_reused_key", |b| {
+        #[cfg(feature = "low-memory")]
+        let mut workspace = ml_dsa::SigningWorkspace::<MlDsa65>::new();
+        b.iter(|| {
+            #[cfg(feature = "low-memory")]
+            let sig = sk.sign_deterministic_with_workspace(&m, &ctx, &mut workspace);
+            #[cfg(not(feature = "low-memory"))]
+            let sig = sk.sign_deterministic(&m, &ctx);
+            core::hint::black_box(sig)
+        });
+    });
+
     // Key generation
     c.bench_function("keygen", |b| {
         b.iter(|| {

ml-dsa/examples/stack_usage.rs

@@ -0,0 +1,93 @@
+//! Probe stack requirements in a separate process: a stack overflow aborts the process.
+//!
+//! Run with `cargo run -p ml-dsa --release --example stack_usage -- 87 keygen 524288`.
+//! Stack sizes and compiler optimizations vary between platforms; compare with the same toolchain.
+
+use core::hint::black_box;
+use ml_dsa::{
+    ExpandedSigningKey, Keypair, MlDsa44, MlDsa65, MlDsa87, MlDsaParams, Seed, SigningKey,
+    VerifyingKey,
+};
+use std::{env, thread};
+
+#[inline(never)]
+fn keygen<P: MlDsaParams>() {
+    black_box(SigningKey::<P>::from_seed(black_box(&Seed::default())));
+}
+
+fn probe<P: MlDsaParams + Send + Sync + 'static>(operation: &str, stack_size: usize) {
+    let worker = thread::Builder::new().stack_size(stack_size);
+    let handle = match operation {
+        "keygen" => worker.spawn(keygen::<P>),
+        "clone" => {
+            let sk = SigningKey::<P>::from_seed(&Seed::default());
+            worker.spawn(move || {
+                black_box(black_box(&sk).clone());
+            })
+        }
+        "decode" => {
+            let sk = SigningKey::<P>::from_seed(&Seed::default());
+            let enc = sk.verifying_key().encode();
+            worker.spawn(move || {
+                black_box(VerifyingKey::<P>::decode(black_box(&enc)));
+            })
+        }
+        "derive" => {
+            let sk = SigningKey::<P>::from_seed(&Seed::default());
+            worker.spawn(move || {
+                black_box(sk.expanded_key().verifying_key());
+            })
+        }
+        "import" => {
+            let sk = SigningKey::<P>::from_seed(&Seed::default());
+            #[allow(deprecated)]
+            let enc = sk.expanded_key().to_expanded();
+            worker.spawn(move || {
+                #[allow(deprecated)]
+                black_box(ExpandedSigningKey::<P>::from_expanded(black_box(&enc)));
+            })
+        }
+        "sign" => {
+            let sk = SigningKey::<P>::from_seed(&Seed::default());
+            worker.spawn(move || {
+                black_box(
+                    sk.expanded_key()
+                        .sign_deterministic(black_box(b"stack probe"), &[]),
+                )
+                .expect("signing failed");
+            })
+        }
+        "verify" => {
+            let sk = SigningKey::<P>::from_seed(&Seed::default());
+            let vk = sk.verifying_key().clone();
+            let sig = sk
+                .expanded_key()
+                .sign_deterministic(b"stack probe", &[])
+                .expect("signing failed");
+            worker.spawn(move || {
+                assert!(vk.verify_with_context(black_box(b"stack probe"), &[], black_box(&sig)));
+            })
+        }
+        _ => panic!("operation must be keygen, clone, decode, derive, import, sign, or verify"),
+    };
+    handle
+        .expect("spawning thread failed")
+        .join()
+        .expect("worker failed");
+}
+
+fn main() {
+    let args: Vec<_> = env::args().collect();
+    assert_eq!(
+        args.len(),
+        4,
+        "usage: stack_usage <44|65|87> <operation> <stack bytes>"
+    );
+    let stack_size = args[3].parse().expect("invalid stack size");
+    match args[1].as_str() {
+        "44" => probe::<MlDsa44>(&args[2], stack_size),
+        "65" => probe::<MlDsa65>(&args[2], stack_size),
+        "87" => probe::<MlDsa87>(&args[2], stack_size),
+        _ => panic!("parameter set must be 44, 65, or 87"),
+    }
+}

ml-dsa/src/algebra.rs

@@ -1,9 +1,17 @@
+use core::ops::Mul;
 use ctutils::{CtEq, CtGt, CtLt, CtSelect};
 use hybrid_array::{
-    ArraySize,
+    Array, ArraySize,
     typenum::{Shleft, U1, U13, Unsigned},
 };
+#[cfg(not(feature = "low-memory"))]
+use module_lattice::MaybeBox;
 use module_lattice::{Field, Truncate};
+#[cfg(feature = "low-memory")]
+use {
+    crate::{B32, sampling::rej_ntt_poly},
+    core::marker::PhantomData,
+};
 
 module_lattice::define_field!(BaseField, u32, u64, u128, 8_380_417);
 
@@ -14,7 +22,64 @@ pub(crate) type Polynomial = module_lattice::Polynomial<BaseField>;
 pub(crate) type Vector<K> = module_lattice::Vector<BaseField, K>;
 pub(crate) type NttPolynomial = module_lattice::NttPolynomial<BaseField>;
 pub(crate) type NttVector<K> = module_lattice::NttVector<BaseField, K>;
-pub(crate) type NttMatrix<K, L> = module_lattice::NttMatrix<BaseField, K, L>;
+
+/// An expanded public matrix, with each row independently heap-allocated when `alloc` is enabled.
+///
+/// With `alloc`, building and cloning one row at a time avoids a full matrix temporary on the stack.
+/// Row sizes, allocation counts, and multiplication's iteration order depend only on public parameters.
+#[cfg(not(feature = "low-memory"))]
+#[derive(Clone, Debug, PartialEq)]
+pub(crate) struct NttMatrix<K: ArraySize, L: ArraySize>(Array<MaybeBox<NttVector<L>>, K>);
+
+#[cfg(not(feature = "low-memory"))]
+impl<K: ArraySize, L: ArraySize> NttMatrix<K, L> {
+    pub(crate) fn from_fn(mut f: impl FnMut(usize) -> NttVector<L>) -> Self {
+        Self(Array::from_fn(|r| MaybeBox::new(f(r))))
+    }
+}
+
+/// In low-memory builds only the public seed is retained. Entries are regenerated in row order.
+#[cfg(feature = "low-memory")]
+#[derive(Clone, Debug, PartialEq)]
+pub(crate) struct NttMatrix<K: ArraySize, L: ArraySize>(B32, PhantomData<fn() -> (K, L)>);
+
+#[cfg(feature = "low-memory")]
+impl<K: ArraySize, L: ArraySize> NttMatrix<K, L> {
+    pub(crate) fn from_seed(rho: &B32) -> Self {
+        Self(rho.clone(), PhantomData)
+    }
+
+    /// One polynomial of scratch, regardless of the public matrix dimensions.
+    pub(crate) fn multiply_row_into(&self, r: usize, rhs: &NttVector<L>, out: &mut Polynomial) {
+        out.0.fill(Elem::new(0));
+        for s in 0..L::USIZE {
+            // Rejection sampling depends solely on rho and these public indices.
+            let a = rej_ntt_poly(&self.0, Truncate::truncate(r), Truncate::truncate(s));
+            for j in 0..256 {
+                out.0[j] = out.0[j] + a.0[j] * rhs.0[s].0[j];
+            }
+        }
+    }
+}
+
+impl<K: ArraySize, L: ArraySize> Mul<&NttVector<L>> for &NttMatrix<K, L> {
+    type Output = NttVector<K>;
+
+    fn mul(self, rhs: &NttVector<L>) -> Self::Output {
+        #[cfg(not(feature = "low-memory"))]
+        {
+            NttVector::new(Array::from_fn(|r| &*self.0[r] * rhs))
+        }
+        #[cfg(feature = "low-memory")]
+        {
+            NttVector::new(Array::from_fn(|r| {
+                let mut out = Polynomial::default();
+                self.multiply_row_into(r, rhs, &mut out);
+                NttPolynomial::new(out.0)
+            }))
+        }
+    }
+}
 
 // We require modular reduction for three moduli: q, 2^d, and 2 * gamma2.  All three of these are
 // greater than sqrt(q), which means that a number reduced mod q will always be less than M^2,
@@ -349,3 +414,31 @@ mod test {
         }
     }
 }
+
+#[cfg(all(test, feature = "low-memory"))]
+mod streamed_matrix_tests {
+    use super::*;
+    use hybrid_array::typenum::{U5, U6};
+
+    #[test]
+    fn streamed_matrix_matches_explicit_expansion() {
+        let rho = B32::from([0x55; 32]);
+        let a = NttMatrix::<U6, U5>::from_seed(&rho);
+        let input = NttVector::new(Array::from_fn(|s| {
+            NttPolynomial::new(Array::from_fn(|j| {
+                Elem::new(u32::try_from(s * 256 + j).expect("small index"))
+            }))
+        }));
+        let actual = &a * &input;
+        for r in 0..6 {
+            let row = NttVector::new(Array::from_fn(|s| {
+                rej_ntt_poly(
+                    &rho,
+                    <u8 as Truncate<_>>::truncate(r),
+                    <u8 as Truncate<_>>::truncate(s),
+                )
+            }));
+            assert_eq!(actual.0[r], &row * &input);
+        }
+    }
+}

ml-dsa/src/hint.rs

@@ -11,7 +11,7 @@ use module_lattice::{Field, MaybeBox, Truncate};
 
 /// Algorithm 39 `MakeHint`: computes hint bit indicating whether adding `z` to `r` alters the high
 /// bits of `r`.
-fn make_hint<TwoGamma2: Unsigned>(z: Elem, r: Elem) -> bool {
+pub(crate) fn make_hint<TwoGamma2: Unsigned>(z: Elem, r: Elem) -> bool {
     let r1 = r.high_bits::<TwoGamma2>();
     let v1 = (r + z).high_bits::<TwoGamma2>();
     r1 != v1
@@ -61,6 +61,7 @@ impl<P> Hint<P>
 where
     P: SignatureParams,
 {
+    #[cfg(any(not(feature = "low-memory"), test))]
     pub(crate) fn new(z: &Vector<P::K>, r: &Vector<P::K>) -> Self {
         let zi = z.0.iter();
         let ri = r.0.iter();
@@ -79,6 +80,7 @@ where
         ))
     }
 
+    #[cfg(any(not(feature = "low-memory"), test))]
     pub(crate) fn hamming_weight(&self) -> usize {
         self.0
             .iter()

ml-dsa/src/lib.rs

@@ -55,6 +55,11 @@ mod param;
 mod sampling;
 mod signing;
 mod verifying;
+#[cfg(feature = "low-memory")]
+mod workspace;
+
+#[cfg(feature = "low-memory")]
+pub use workspace::SigningWorkspace;
 
 pub use crate::{
     param::{EncodedSignature, EncodedVerifyingKey, ExpandedSigningKeyBytes, MlDsaParams},