diff --git a/pasta_curves/src/asm/pasta_mul-armv8.S b/pasta_curves/src/asm/pasta_mul-armv8.S index 64b074e5..6f87d643 100644 --- a/pasta_curves/src/asm/pasta_mul-armv8.S +++ b/pasta_curves/src/asm/pasta_mul-armv8.S @@ -5,8 +5,8 @@ // Adapted from the pre-generated Mach-O assembly in Semolina v0.1.4: // https://github.com/supranational/semolina/blob/v0.1.4/src/mach-o/pasta_mul-armv8.S // -// Only the fused repeated-squaring multiplication, the conversion out of -// Montgomery form, and their shared reduction helper are retained here. +// Only the fused repeated-squaring chains, the conversion out of Montgomery +// form, and their shared reduction helper are retained here. // Single Montgomery multiplication and squaring live as inline `asm!` // transcriptions of the same upstream routines in // src/fields/aarch64_asm.rs. Symbols are renamed for this crate. The @@ -20,12 +20,14 @@ // // Apple AArch64 argument registers: // +// - sqr_n: x0 = out, x1 = value, x2 = count, x3 = modulus, x4 = inv // - sqr_n_mul: x0 = out, x1 = value, x2 = count, x3 = rhs, x4 = modulus, // x5 = inv // - from_mont: x0 = out, x1 = value, x2 = modulus, x3 = inv // -// The code has no secret-dependent branches or memory accesses. Conditional -// reductions use `csel` after a full-width subtraction. +// The code has no secret-dependent branches or memory accesses. The chains +// branch only on their public counts. Conditional reductions use `csel` +// after a full-width subtraction. .text @@ -504,6 +506,7 @@ L$pasta_curves_sqr_n_mul_loop: adcs x12,x13,xzr // Final candidate limb 2. adc x13,x14,x22 // Final candidate limb 3. +L$pasta_curves_sqr_n_canonicalize: // Subtract p = [p0,p1,0,2^62]; the upper limbs are materialized inline. subs x19,x10,x23 // Tentative result limb 0 = candidate - p[0]. mov x22,#1<<62 // Materialize p[3] = 2^62. @@ -525,3 +528,186 @@ L$pasta_curves_sqr_n_mul_loop: ldp x23,x24,[x29,#48] // Restore callee-saved x23 and x24. ldr x29,[sp],#80 // Restore frame pointer and release the frame. ret // Return; x30 still holds the caller's address. + +.globl _pasta_curves_sqr_n_mont_pasta +.private_extern _pasta_curves_sqr_n_mont_pasta + +// Squares the input `count` times (count must be at least 1), retaining the +// lazy accumulator in registers across the entire chain and canonicalizing +// only once. This is a separate emitted loop so the existing fused +// square-and-multiply path remains instruction-for-instruction unchanged. + +.align 5 +_pasta_curves_sqr_n_mont_pasta: + stp x29,x30,[sp,#-80]! // Allocate frame; save frame pointer and LR. + add x29,sp,#0 // Set the frame pointer to the new stack top. + stp x19,x20,[sp,#16] // Preserve callee-saved scratch x19 and x20. + stp x21,x22,[sp,#32] // Preserve callee-saved scratch x21 and x22. + stp x23,x24,[sp,#48] // Preserve callee-saved modulus limbs 0 and 1. + + mov x5,x4 // Move inv into the loop's expected register. + mov x4,x3 // Move modulus into the expected register. + + ldp x6,x7,[x1] // Load value limbs a[0] and a[1]. + ldp x8,x9,[x1,#16] // Load value limbs a[2] and a[3]. + + ldp x23,x24,[x4] // Keep p[0] and p[1] resident across the loop. + // p[2] = 0 and p[3] = 2^62 are implicit throughout. + +L$pasta_curves_sqr_n_loop: + sub x2,x2,#1 // Consume one squaring from the count. + + // Square a (in x6..x9) exactly as the inline `square` in + // src/fields/aarch64_asm.rs does; the 512-bit + // product limbs A0..A7 map to x10,x11,x12,x13,x14,x15,x16,x17. + + mul x11,x7,x6 // x11 = low(a[1] * a[0]). + umulh x20,x7,x6 // x20 = high(a[1] * a[0]). + mul x12,x8,x6 // x12 = low(a[2] * a[0]). + umulh x21,x8,x6 // x21 = high(a[2] * a[0]). + mul x13,x9,x6 // x13 = low(a[3] * a[0]). + umulh x14,x9,x6 // x14 = high(a[3] * a[0]). + + adds x12,x12,x20 // Fold high(a[1] * a[0]) into product limb 2. + mul x19,x8,x7 // x19 = low(a[2] * a[1]). + umulh x20,x8,x7 // x20 = high(a[2] * a[1]). + adcs x13,x13,x21 // Fold high(a[2] * a[0]) into product limb 3. + mul x21,x9,x7 // x21 = low(a[3] * a[1]). + umulh x22,x9,x7 // x22 = high(a[3] * a[1]). + adc x14,x14,xzr // Propagate carry into product limb 4. + + mul x15,x9,x8 // x15 = low(a[3] * a[2]). + umulh x16,x9,x8 // x16 = high(a[3] * a[2]). + + adds x20,x20,x21 // Combine terms contributing to product limb 4. + mul x10,x6,x6 // x10 = low(a[0]^2), product limb 0. + adc x21,x22,xzr // Combine terms contributing to product limb 5. + + adds x13,x13,x19 // Add low(a[2] * a[1]) into product limb 3. + umulh x6,x6,x6 // x6 = high(a[0]^2). + adcs x14,x14,x20 // Accumulate cross terms into product limb 4. + mul x20,x7,x7 // x20 = low(a[1]^2). + adcs x15,x15,x21 // Accumulate cross terms into product limb 5. + umulh x7,x7,x7 // x7 = high(a[1]^2). + adc x16,x16,xzr // Propagate carry into product limb 6. + + adds x11,x11,x11 // Double cross-term product limb 1. + mul x21,x8,x8 // x21 = low(a[2]^2). + adcs x12,x12,x12 // Double cross-term product limb 2. + umulh x8,x8,x8 // x8 = high(a[2]^2). + adcs x13,x13,x13 // Double cross-term product limb 3. + mul x22,x9,x9 // x22 = low(a[3]^2). + adcs x14,x14,x14 // Double cross-term product limb 4. + umulh x9,x9,x9 // x9 = high(a[3]^2). + adcs x15,x15,x15 // Double cross-term product limb 5. + adcs x16,x16,x16 // Double cross-term product limb 6. + adc x17,xzr,xzr // Capture the doubled cross-term carry in limb 7. + + mul x4,x5,x10 // q = product limb 0 * inv mod 2^64. + + // Add diagonal squares to obtain a^2 in x10..x17. + adds x11,x11,x6 // Add high(a[0]^2) to product limb 1. + adcs x12,x12,x20 // Add low(a[1]^2) to product limb 2. + adcs x13,x13,x7 // Add high(a[1]^2) to product limb 3. + adcs x14,x14,x21 // Add low(a[2]^2) to product limb 4. + adcs x15,x15,x8 // Add high(a[2]^2) to product limb 5. + adcs x16,x16,x22 // Add low(a[3]^2) to product limb 6. + adc x17,x17,x9 // Add high(a[3]^2) to product limb 7. + + // Montgomery cancellation 0 on the low half, as in the shared helper. + // low(q * p[0]) cancels x10 and is discarded by the limb shift. + mul x20,x24,x4 // x20 = low(q * p[1]). + // q * p[2] is zero because p[2] = 0. + lsl x22,x4,#62 // x22 = low(q * p[3]). + // Carry from x10 + low(q*p[0]) is one exactly when x10 is nonzero. + subs xzr,x10,#1 // Set that cancellation carry. + umulh x19,x23,x4 // x19 = high(q * p[0]). + adcs x11,x11,x20 // Add low(q * p[1]) and cancellation carry. + umulh x20,x24,x4 // x20 = high(q * p[1]). + adcs x12,x12,xzr // Propagate carry across zero p[2]. + // high(q * p[2]) is zero. + adcs x13,x13,x22 // Add low(q * p[3]) and carry. + lsr x22,x4,#2 // x22 = high(q * p[3]). + adc x1,xzr,xzr // Save the carry above limb 3. + + // Shift out cancelled limb 0 and start cancellation 1. + adds x10,x11,x19 // New limb 0 includes high(q * p[0]). + adcs x11,x12,x20 // New limb 1 includes high(q * p[1]). + adcs x12,x13,xzr // New limb 2; p[2] contributes zero. + mul x4,x5,x10 // Next q = new limb 0 * inv mod 2^64. + adc x13,x1,x22 // New limb 3 includes high(q * p[3]). + // low(q * p[0]) cancels x10 and is discarded. + mul x20,x24,x4 // x20 = low(next q * p[1]). + // next q * p[2] is zero. + lsl x22,x4,#62 // x22 = low(next q * p[3]). + subs xzr,x10,#1 // Set the low-limb cancellation carry. + umulh x19,x23,x4 // x19 = high(next q * p[0]). + adcs x11,x11,x20 // Add low(next q * p[1]) and carry. + umulh x20,x24,x4 // x20 = high(next q * p[1]). + adcs x12,x12,xzr // Propagate carry across zero p[2]. + // high(next q * p[2]) is zero. + adcs x13,x13,x22 // Add low(next q * p[3]) and carry. + lsr x22,x4,#2 // x22 = high(next q * p[3]). + adc x1,xzr,xzr // Save the carry above limb 3. + + // Shift out cancelled limb 1 and start cancellation 2. + adds x10,x11,x19 // New limb 0 includes high(q * p[0]). + adcs x11,x12,x20 // New limb 1 includes high(q * p[1]). + adcs x12,x13,xzr // New limb 2; p[2] contributes zero. + mul x4,x5,x10 // Next q = new limb 0 * inv mod 2^64. + adc x13,x1,x22 // New limb 3 includes high(q * p[3]). + // low(q * p[0]) cancels x10 and is discarded. + mul x20,x24,x4 // x20 = low(next q * p[1]). + // next q * p[2] is zero. + lsl x22,x4,#62 // x22 = low(next q * p[3]). + subs xzr,x10,#1 // Set the low-limb cancellation carry. + umulh x19,x23,x4 // x19 = high(next q * p[0]). + adcs x11,x11,x20 // Add low(next q * p[1]) and carry. + umulh x20,x24,x4 // x20 = high(next q * p[1]). + adcs x12,x12,xzr // Propagate carry across zero p[2]. + // high(next q * p[2]) is zero. + adcs x13,x13,x22 // Add low(next q * p[3]) and carry. + lsr x22,x4,#2 // x22 = high(next q * p[3]). + adc x1,xzr,xzr // Save the carry above limb 3. + + // Shift out cancelled limb 2 and start cancellation 3. + adds x10,x11,x19 // New limb 0 includes high(q * p[0]). + adcs x11,x12,x20 // New limb 1 includes high(q * p[1]). + adcs x12,x13,xzr // New limb 2; p[2] contributes zero. + mul x4,x5,x10 // Final q = new limb 0 * inv mod 2^64. + adc x13,x1,x22 // New limb 3 includes high(q * p[3]). + // low(q * p[0]) cancels x10 and is discarded. + mul x20,x24,x4 // x20 = low(final q * p[1]). + // final q * p[2] is zero. + lsl x22,x4,#62 // x22 = low(final q * p[3]). + subs xzr,x10,#1 // Set the low-limb cancellation carry. + umulh x19,x23,x4 // x19 = high(final q * p[0]). + adcs x11,x11,x20 // Add low(final q * p[1]) and carry. + umulh x20,x24,x4 // x20 = high(final q * p[1]). + adcs x12,x12,xzr // Propagate carry across zero p[2]. + // high(final q * p[2]) is zero. + adcs x13,x13,x22 // Add low(final q * p[3]) and carry. + lsr x22,x4,#2 // x22 = high(final q * p[3]). + adc x1,xzr,xzr // Save the carry above limb 3. + + // Shift out cancelled limb 3 to finish dividing the low half by R. + adds x10,x11,x19 // Reduced limb 0 includes high(q * p[0]). + adcs x11,x12,x20 // Reduced limb 1 includes high(q * p[1]). + adcs x12,x13,xzr // Reduced limb 2; p[2] contributes zero. + adc x13,x1,x22 // Reduced limb 3 includes high(q * p[3]). + // Add the upper product half; the sum stays below 2p (see above), so no + // carry escapes and no conditional subtraction is needed mid-loop. + adds x6,x10,x14 // Next-iteration a[0]. + adcs x7,x11,x15 // Next-iteration a[1]. + adcs x8,x12,x16 // Next-iteration a[2]. + adc x9,x13,x17 // Next-iteration a[3]. + + cbnz x2,L$pasta_curves_sqr_n_loop // Square again if count remains. + + // The loop leaves the lazy accumulator in x6..x9. Move it into the + // shared canonicalization/output registers without a dummy multiply. + mov x10,x6 + mov x11,x7 + mov x12,x8 + mov x13,x9 + b L$pasta_curves_sqr_n_canonicalize diff --git a/pasta_curves/src/fields/aarch64_asm.rs b/pasta_curves/src/fields/aarch64_asm.rs index 27a9631d..c1c45af7 100644 --- a/pasta_curves/src/fields/aarch64_asm.rs +++ b/pasta_curves/src/fields/aarch64_asm.rs @@ -1,7 +1,7 @@ //! Private Apple AArch64 backend for the Pasta fields. //! //! Montgomery multiplication and squaring are implemented as inline `asm!` -//! blocks below; the fused repeated-squaring chain and the canonical-form +//! blocks below; the fused repeated-squaring chains and the canonical-form //! conversion remain in `src/asm/pasta_mul-armv8.S` and are reached through //! `extern "C"`. //! @@ -58,14 +58,22 @@ //! construction; `mul` and `square` debug-assert the precondition so a future //! caller that breaks it fails loudly under test instead of silently. //! -//! There are no branches and no memory accesses inside the blocks, so the -//! code is constant-time. +//! The inline blocks have no branches or memory accesses. The out-of-line +//! repeated-squaring chains branch only on their public counts, so the code +//! is constant-time. use core::arch::asm; type Limbs = [u64; 4]; extern "C" { + fn pasta_curves_sqr_n_mont_pasta( + out: *mut Limbs, + value: *const Limbs, + count: usize, + modulus: *const Limbs, + inv: u64, + ); fn pasta_curves_sqr_n_mul_mont_pasta( out: *mut Limbs, value: *const Limbs, @@ -512,6 +520,26 @@ pub(super) fn square(value: &Limbs, modulus: &Limbs, inv: u64) -> Limbs { [a0, a1, a2, a3] } +/// Squares a canonical Montgomery residue `count` times, keeping the lazy +/// accumulator in registers and canonicalizing it once at the end. +#[inline] +pub(super) fn sqr_n(value: &Limbs, count: usize, modulus: &Limbs, inv: u64) -> Limbs { + // The assembly decrements the count before testing it, so a zero count + // would wrap around and effectively never terminate. + assert!(count >= 1); + debug_assert!( + is_canonical(value, modulus), + "aarch64_asm::sqr_n requires a canonical starting value" + ); + let mut out = Limbs::default(); + // SAFETY: All pointers refer to four initialized `u64` limbs for the + // duration of the call. The backend writes exactly four limbs to `out`. + unsafe { + pasta_curves_sqr_n_mont_pasta(&mut out, value, count, modulus, inv); + } + out +} + /// Squares a canonical Montgomery residue `count` times, then multiplies the /// result by the canonical Montgomery residue `rhs`, keeping the accumulator /// in registers throughout. @@ -526,6 +554,14 @@ pub(super) fn sqr_n_mul( // The assembly decrements the count before testing it, so a zero count // would wrap around and effectively never terminate. assert!(count >= 1); + debug_assert!( + is_canonical(value, modulus), + "aarch64_asm::sqr_n_mul requires a canonical starting value" + ); + debug_assert!( + is_canonical(rhs, modulus), + "aarch64_asm::sqr_n_mul requires a canonical multiplier" + ); let mut out = Limbs::default(); // SAFETY: All pointers refer to four initialized `u64` limbs for the // duration of the call. The backend writes exactly four limbs to `out`. diff --git a/pasta_curves/src/fields/fp.rs b/pasta_curves/src/fields/fp.rs index 85baaa04..94aebd8f 100644 --- a/pasta_curves/src/fields/fp.rs +++ b/pasta_curves/src/fields/fp.rs @@ -486,7 +486,15 @@ impl Fp { target_vendor = "apple" ))] { - (0..n).fold(*self, |acc, _| acc.square_runtime()) + // Calling the dedicated single-square routine is faster than + // entering the repeated-squaring loop for a one-element chain. + if n == 1 { + self.square_runtime() + } else { + Fp(super::aarch64_asm::sqr_n( + &self.0, n as usize, &MODULUS.0, INV, + )) + } } #[cfg(not(all( @@ -1047,9 +1055,16 @@ fn aarch64_asm_matches_portable_arithmetic() { (0..n).fold(value, |acc, _| Fp::square(&acc)).mul(&by) } + fn portable_sqr_n(value: Fp, n: u32) -> Fp { + (0..n).fold(value, |acc, _| Fp::square(&acc)) + } + for lhs in boundaries { aarch64_asm_check_repr(lhs); assert_eq!(::square(&lhs), Fp::square(&lhs)); + for n in [1, 2, 7, 129] { + assert_eq!(lhs.sqr_n_runtime(n), portable_sqr_n(lhs, n)); + } for rhs in boundaries { assert_eq!(lhs.cmp(&rhs), aarch64_asm_portable_cmp(lhs, rhs)); assert_eq!(&lhs * &rhs, Fp::mul(&lhs, &rhs)); @@ -1082,6 +1097,7 @@ fn aarch64_asm_matches_portable_arithmetic() { assert_eq!(&lhs * &rhs, Fp::mul(&lhs, &rhs)); assert_eq!(::square(&lhs), Fp::square(&lhs)); for n in [1, 129] { + assert_eq!(lhs.sqr_n_runtime(n), portable_sqr_n(lhs, n)); assert_eq!( lhs.sqr_n_mul_runtime(n, &rhs), portable_sqr_n_mul(lhs, n, rhs) diff --git a/pasta_curves/src/fields/fq.rs b/pasta_curves/src/fields/fq.rs index 8065d772..027620bb 100644 --- a/pasta_curves/src/fields/fq.rs +++ b/pasta_curves/src/fields/fq.rs @@ -486,7 +486,15 @@ impl Fq { target_vendor = "apple" ))] { - (0..n).fold(*self, |acc, _| acc.square_runtime()) + // Calling the dedicated single-square routine is faster than + // entering the repeated-squaring loop for a one-element chain. + if n == 1 { + self.square_runtime() + } else { + Fq(super::aarch64_asm::sqr_n( + &self.0, n as usize, &MODULUS.0, INV, + )) + } } #[cfg(not(all( @@ -1046,9 +1054,16 @@ fn aarch64_asm_matches_portable_arithmetic() { (0..n).fold(value, |acc, _| Fq::square(&acc)).mul(&by) } + fn portable_sqr_n(value: Fq, n: u32) -> Fq { + (0..n).fold(value, |acc, _| Fq::square(&acc)) + } + for lhs in boundaries { aarch64_asm_check_repr(lhs); assert_eq!(::square(&lhs), Fq::square(&lhs)); + for n in [1, 2, 7, 129] { + assert_eq!(lhs.sqr_n_runtime(n), portable_sqr_n(lhs, n)); + } for rhs in boundaries { assert_eq!(lhs.cmp(&rhs), aarch64_asm_portable_cmp(lhs, rhs)); assert_eq!(&lhs * &rhs, Fq::mul(&lhs, &rhs)); @@ -1081,6 +1096,7 @@ fn aarch64_asm_matches_portable_arithmetic() { assert_eq!(&lhs * &rhs, Fq::mul(&lhs, &rhs)); assert_eq!(::square(&lhs), Fq::square(&lhs)); for n in [1, 129] { + assert_eq!(lhs.sqr_n_runtime(n), portable_sqr_n(lhs, n)); assert_eq!( lhs.sqr_n_mul_runtime(n, &rhs), portable_sqr_n_mul(lhs, n, rhs)