Skip to content

chore(pasta): fuse AArch64 repeated squaring #240

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

Closed
TalDerei wants to merge 1 commit into from
Closed

chore(pasta): fuse AArch64 repeated squaring #240

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
194 changes: 190 additions & 4 deletions pasta_curves/src/asm/pasta_mul-armv8.S
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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.
Expand All @@ -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
42 changes: 39 additions & 3 deletions pasta_curves/src/fields/aarch64_asm.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"`.
//!
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand All @@ -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`.
Expand Down
18 changes: 17 additions & 1 deletion pasta_curves/src/fields/fp.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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!(<Fp as Field>::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));
Expand Down Expand Up @@ -1082,6 +1097,7 @@ fn aarch64_asm_matches_portable_arithmetic() {
assert_eq!(&lhs * &rhs, Fp::mul(&lhs, &rhs));
assert_eq!(<Fp as Field>::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)
Expand Down
18 changes: 17 additions & 1 deletion pasta_curves/src/fields/fq.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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!(<Fq as Field>::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));
Expand Down Expand Up @@ -1081,6 +1096,7 @@ fn aarch64_asm_matches_portable_arithmetic() {
assert_eq!(&lhs * &rhs, Fq::mul(&lhs, &rhs));
assert_eq!(<Fq as Field>::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)
Expand Down