diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 88c6a717..e2f4f354 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -655,6 +655,10 @@ jobs: strategy: fail-fast: false matrix: + # `--all-features` includes `x86_64-asm`, so this lane (and the + # platform-smoke jobs below) executes the assembly field backend, + # which requires ADX/BMI2 on the runner's CPU — true of GitHub's + # current fleet. A SIGILL here means the runner lacks ADX. features: [--all-features, --no-default-features] steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 diff --git a/pasta_curves/CHANGELOG.md b/pasta_curves/CHANGELOG.md index 495efebd..b10d3d85 100644 --- a/pasta_curves/CHANGELOG.md +++ b/pasta_curves/CHANGELOG.md @@ -41,6 +41,24 @@ and this project adheres to Rust's notion of - Moved the portable (pure Rust) `Fp` and `Fq` wide-squaring routine into a shared `fields::portable` module. Both fields delegate to the same limb-array implementation; the algorithm and generated operations are unchanged. +- Added an `x86_64-asm` feature: MULX/ADCX/ADOX Montgomery multiplication + and a dedicated squaring for the Pasta fields on x86-64, a + transcription of the `aarch64-asm` backend's five-limb CIOS rounds with + the same canonicity contract. Requires BMI2 and ADX (Intel Broadwell / + AMD Zen or newer; enabling it on an older CPU faults at runtime), and is + a no-op on other architectures. Measured on Skylake-X: field + multiplication ~1.25x faster than the portable path and a dedicated + assembly squaring ~1.05-1.10x (2-5% ahead of squaring through the + multiplication, mirroring the AArch64 backend's own square-over-mul + margin). Two slower schedulings are pinned in the module docs so they + are not retried: squaring routed through the multiplication, and + interleaved-ADCX/ADOX Montgomery reduction sweeps (~10% slower than the + two short sequential sweeps on Skylake-X). +- Fixed the `fp`/`fq` benches' `square` cell to call `Field::square` + explicitly: the inherent (always-portable) `square` shadowed the + runtime-dispatched trait method, so the cell measured the portable path + even under the assembly features and once mis-sized an assembly squaring + decision. - All of this release's new MSM machinery — the Eisenstein-orbit backend (`glv::orbit`), the magnitude-profiled backend planner, the prepared zero-checks (`glv::zero`), and the `arithmetic::PreparedZeroCheck` / diff --git a/pasta_curves/Cargo.toml b/pasta_curves/Cargo.toml index 4247565e..03f84a8a 100644 --- a/pasta_curves/Cargo.toml +++ b/pasta_curves/Cargo.toml @@ -46,6 +46,11 @@ rustdoc-args = [ [features] aarch64-asm = ["dep:cc"] +# MULX/ADCX/ADOX Montgomery multiplication for the Pasta fields on x86-64. +# Requires a CPU with the BMI2 and ADX extensions (Intel Broadwell / AMD Zen +# or later); enabling it on an older CPU faults at runtime. A no-op on other +# architectures. +x86_64-asm = ["dep:cc"] alloc = [ "group/alloc", "blake2b_simd", diff --git a/pasta_curves/benches/fp.rs b/pasta_curves/benches/fp.rs index d3d4384b..be11821d 100644 --- a/pasta_curves/benches/fp.rs +++ b/pasta_curves/benches/fp.rs @@ -117,7 +117,12 @@ fn bench_fp_square(b: &mut Bencher) { let mut count = 0; b.iter(|| { let mut tmp = v[count]; - tmp = tmp.square(); + // The inherent (always-portable) `Fp::square` shadows + // `Field::square`, the runtime-dispatched production path; call the + // trait method explicitly so the assembly backends are what gets + // measured. A shadowed cell here once mis-sized an assembly + // squaring decision. + tmp = Field::square(&tmp); count = (count + 1) % SAMPLES; tmp }); diff --git a/pasta_curves/benches/fq.rs b/pasta_curves/benches/fq.rs index 90c37363..a2ac1bd4 100644 --- a/pasta_curves/benches/fq.rs +++ b/pasta_curves/benches/fq.rs @@ -117,7 +117,12 @@ fn bench_fq_square(b: &mut Bencher) { let mut count = 0; b.iter(|| { let mut tmp = v[count]; - tmp = tmp.square(); + // The inherent (always-portable) `Fq::square` shadows + // `Field::square`, the runtime-dispatched production path; call the + // trait method explicitly so the assembly backends are what gets + // measured. A shadowed cell here once mis-sized an assembly + // squaring decision. + tmp = Field::square(&tmp); count = (count + 1) % SAMPLES; tmp }); diff --git a/pasta_curves/build.rs b/pasta_curves/build.rs index ea809125..c25a0f36 100644 --- a/pasta_curves/build.rs +++ b/pasta_curves/build.rs @@ -1,11 +1,27 @@ -#[cfg(feature = "aarch64-asm")] +#[cfg(any(feature = "aarch64-asm", feature = "x86_64-asm"))] use std::env; fn main() { println!("cargo:rerun-if-changed=src/asm/pasta_mul-armv8.S"); + println!("cargo:rerun-if-changed=src/asm/pasta_lazy_square-x86_64.S"); #[cfg(feature = "aarch64-asm")] build_aarch64_asm(); + + #[cfg(feature = "x86_64-asm")] + build_x86_64_asm(); +} + +#[cfg(feature = "x86_64-asm")] +fn build_x86_64_asm() { + let target_arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap(); + let target_family = env::var("CARGO_CFG_TARGET_FAMILY").unwrap(); + + if target_arch == "x86_64" && target_family == "unix" { + cc::Build::new() + .file("src/asm/pasta_lazy_square-x86_64.S") + .compile("pasta_curves_x86_64_lazy_square"); + } } #[cfg(feature = "aarch64-asm")] diff --git a/pasta_curves/src/asm/pasta_lazy_square-x86_64.S b/pasta_curves/src/asm/pasta_lazy_square-x86_64.S new file mode 100644 index 00000000..c8fa44b3 --- /dev/null +++ b/pasta_curves/src/asm/pasta_lazy_square-x86_64.S @@ -0,0 +1,416 @@ +// Fused x86-64 lazy-squaring chain for the Pasta fields. +// +// The loop is the instruction schedule generated for the interleaved +// `portable::square_reduce_lazy_pasta` implementation. Its four-limb +// accumulator remains in registers across iterations. When `by` is non-null, +// the closing Montgomery multiplication follows in this same leaf function +// and consumes those four accumulator limbs directly from registers. The only +// conditional subtraction is the one at the end of that multiplication. + +#if defined(__APPLE__) +#define C_SYMBOL(name) _##name +#else +#define C_SYMBOL(name) name +#endif + + .text + .p2align 4 + .globl C_SYMBOL(pasta_x86_64_sqr_n) +#if defined(__APPLE__) + .private_extern C_SYMBOL(pasta_x86_64_sqr_n) +#else + .hidden C_SYMBOL(pasta_x86_64_sqr_n) + .type C_SYMBOL(pasta_x86_64_sqr_n), @function +#endif +C_SYMBOL(pasta_x86_64_sqr_n): + .cfi_startproc + pushq %rbp + .cfi_def_cfa_offset 16 + .cfi_offset %rbp, -16 + movq %rsp, %rbp + .cfi_def_cfa_register %rbp + pushq %r15 + pushq %r14 + pushq %r13 + pushq %r12 + pushq %rbx + .cfi_offset %rbx, -56 + .cfi_offset %r12, -48 + .cfi_offset %r13, -40 + .cfi_offset %r14, -32 + .cfi_offset %r15, -24 + + // Fixed state in the red zone. This is a leaf function on every target + // for which it is built (Unix x86-64). + movq %rdi, -88(%rbp) // output pointer + movq %rcx, -128(%rbp) // optional closing multiplier + movq (%r8), %rdi + movq %rdi, -64(%rbp) // modulus[0] + movq 8(%r8), %rdi + movq %rdi, -56(%rbp) // modulus[1] + movq %r9, -48(%rbp) // Montgomery inverse + + // The lazy accumulator stays in r11:rbx:r12:rax (little-endian). + movq (%rsi), %r11 + movq 8(%rsi), %rbx + movq 16(%rsi), %r12 + movq 24(%rsi), %rax + testl %edx, %edx + je .Lsquare_done + + .p2align 4 +.Lsquare_loop: + movl %edx, -68(%rbp) + + // Low off-diagonal products and the low diagonal terms. + movq %r11, %rdx + mulxq %rbx, %r9, %rcx + mulxq %r12, %r14, %rsi + addq %rcx, %r14 + mulxq %rax, %r10, %r8 + movq %rbx, %rdx + mulxq %r12, %rdi, %r15 + adcq %rsi, %r10 + mulxq %rax, %rcx, %rdx + adcq %r8, %rcx + adcq $0, %rdx + movq %rdx, %rsi + movq %r11, %rdx + mulxq %r11, %r13, %r8 + movq %rbx, %rdx + mulxq %rbx, %rdx, %rbx + addq %r10, %rdi + adcq %r15, %rcx + movq %rcx, -80(%rbp) + adcq $0, %rsi + movq %rsi, -120(%rbp) + movq %rdi, %r11 + shldq $1, %r14, %r11 + shldq $1, %r9, %r14 + addq %r9, %r9 + addq %rdx, %r14 + adcq $0, %rbx + addq %r8, %r9 + adcq $0, %r14 + adcq $0, %rbx + + // Montgomery round 0. + movq -48(%rbp), %r8 + imulq %r13, %r8 + movq -64(%rbp), %rdx + mulxq %r8, %rcx, %rsi + movq -56(%rbp), %rdx + mulxq %r8, %r10, %r15 + addq %r9, %r10 + adcq $0, %r15 + addq %r13, %rcx + adcq %rsi, %r10 + adcq $0, %r15 + movq %r8, %rcx + shlq $62, %rcx + shrq $2, %r8 + + // Independent upper diagonal work overlaps the reduction. + movq -80(%rbp), %rdx + shrdq $63, %rdx, %rdi + movq %r12, %rdx + mulxq %r12, %rdx, %rsi + addq %rdi, %rdx + adcq $0, %rsi + addq %r11, %rbx + adcq $0, %rdx + movq %rdx, -112(%rbp) + adcq $0, %rsi + movq %rsi, -104(%rbp) + addq %r14, %r15 + adcq %rcx, %rbx + adcq $0, %r8 + + // Montgomery round 1. + movq -48(%rbp), %r9 + movq %r9, %r11 + imulq %r10, %r11 + movq %r11, %rdx + movq -64(%rbp), %r13 + mulxq %r13, %rsi, %rcx + movq %rcx, -96(%rbp) + movq -56(%rbp), %rcx + mulxq %rcx, %rdi, %r14 + addq %r15, %rdi + adcq $0, %r14 + addq %r10, %rsi + adcq -96(%rbp), %rdi + adcq $0, %r14 + movq %r11, %r10 + shlq $62, %r10 + shrq $2, %r11 + addq %rbx, %r14 + adcq %r8, %r10 + adcq $0, %r11 + + // Montgomery round 2. + movq %r9, %rbx + imulq %rdi, %rbx + movq %rbx, %rdx + mulxq %r13, %rsi, %r9 + mulxq %rcx, %r8, %r13 + addq %r14, %r8 + adcq $0, %r13 + addq %rdi, %rsi + adcq %r9, %r8 + adcq $0, %r13 + movq %rbx, %r15 + shlq $62, %r15 + shrq $2, %rbx + addq %r10, %r13 + adcq %r11, %r15 + adcq $0, %rbx + + // Montgomery round 3. + movq -48(%rbp), %r14 + imulq %r8, %r14 + movq %r14, %rdx + mulxq -64(%rbp), %rsi, %rcx + mulxq -56(%rbp), %r11, %rdi + addq %r13, %r11 + adcq $0, %rdi + addq %r8, %rsi + adcq %rcx, %r11 + adcq $0, %rdi + movq %r14, %r8 + shlq $62, %r8 + shrq $2, %r14 + addq %r15, %rdi + adcq %rbx, %r8 + + // Finish the upper off-diagonal products and fold the high half. + movq %r12, %rdx + mulxq %rax, %rbx, %rcx + adcq $0, %r14 + addq -120(%rbp), %rbx + adcq $0, %rcx + movq %rcx, %rsi + shldq $1, %rbx, %rcx + movq %rax, %rdx + mulxq %rax, %r12, %rax + movl -68(%rbp), %edx + movq -80(%rbp), %r10 + shldq $1, %r10, %rbx + shrq $63, %rsi + addq %rcx, %r12 + adcq $0, %rax + addq -104(%rbp), %rbx + adcq $0, %r12 + adcq %rsi, %rax + addq -112(%rbp), %r11 + adcq %rdi, %rbx + adcq %r8, %r12 + adcq %r14, %rax + decl %edx + jne .Lsquare_loop + +.Lsquare_done: + cmpq $0, -128(%rbp) + je .Lstore_lazy + + // Reassign the live lazy accumulator as fixed multiplier limbs: + // rax:rbx:rcx:rdi. This avoids a boundary spill and all sixteen lhs + // reloads in the four multiplication rows. + movq %rax, %rdi + movq %r12, %rcx + movq %r11, %rax + movq -128(%rbp), %rsi + + // Round 0: initialize a five-limb accumulator with lazy * by[0]. + movq (%rsi), %rdx + mulxq %rax, %r8, %r9 + mulxq %rbx, %r13, %r10 + addq %r13, %r9 + mulxq %rcx, %r13, %r11 + adcq %r13, %r10 + mulxq %rdi, %r13, %r12 + adcq %r13, %r11 + adcq $0, %r12 + + // Montgomery step 0. + movq %r8, %rdx + imulq -48(%rbp), %rdx + mulxq -56(%rbp), %r13, %r14 + movq %rdx, %r15 + shlq $62, %r15 + negq %r8 + adcq %r13, %r9 + adcq $0, %r10 + adcq %r15, %r11 + adcq $0, %r12 + mulxq -64(%rbp), %r15, %r13 + movq %rdx, %r15 + shrq $2, %r15 + movq $0, %r8 + addq %r13, %r9 + adcq %r14, %r10 + adcq $0, %r11 + adcq %r15, %r12 + adcq $0, %r8 + + // Round 1, using CF and OF as independent product carry chains. + movq 8(%rsi), %rdx + xorq %r13, %r13 + mulxq %rax, %r13, %r14 + adcxq %r13, %r9 + adoxq %r14, %r10 + mulxq %rbx, %r13, %r14 + adcxq %r13, %r10 + adoxq %r14, %r11 + mulxq %rcx, %r13, %r14 + adcxq %r13, %r11 + adoxq %r14, %r12 + mulxq %rdi, %r13, %r14 + adcxq %r13, %r12 + adoxq %r14, %r8 + movq $0, %r13 + adcxq %r13, %r8 + adoxq %r13, %r8 + + // Montgomery step 1. + movq %r9, %rdx + imulq -48(%rbp), %rdx + mulxq -56(%rbp), %r13, %r14 + movq %rdx, %r15 + shlq $62, %r15 + negq %r9 + adcq %r13, %r10 + adcq $0, %r11 + adcq %r15, %r12 + adcq $0, %r8 + mulxq -64(%rbp), %r15, %r13 + movq %rdx, %r15 + shrq $2, %r15 + movq $0, %r9 + addq %r13, %r10 + adcq %r14, %r11 + adcq $0, %r12 + adcq %r15, %r8 + adcq $0, %r9 + + // Round 2. + movq 16(%rsi), %rdx + xorq %r13, %r13 + mulxq %rax, %r13, %r14 + adcxq %r13, %r10 + adoxq %r14, %r11 + mulxq %rbx, %r13, %r14 + adcxq %r13, %r11 + adoxq %r14, %r12 + mulxq %rcx, %r13, %r14 + adcxq %r13, %r12 + adoxq %r14, %r8 + mulxq %rdi, %r13, %r14 + adcxq %r13, %r8 + adoxq %r14, %r9 + movq $0, %r13 + adcxq %r13, %r9 + adoxq %r13, %r9 + + // Montgomery step 2. + movq %r10, %rdx + imulq -48(%rbp), %rdx + mulxq -56(%rbp), %r13, %r14 + movq %rdx, %r15 + shlq $62, %r15 + negq %r10 + adcq %r13, %r11 + adcq $0, %r12 + adcq %r15, %r8 + adcq $0, %r9 + mulxq -64(%rbp), %r15, %r13 + movq %rdx, %r15 + shrq $2, %r15 + movq $0, %r10 + addq %r13, %r11 + adcq %r14, %r12 + adcq $0, %r8 + adcq %r15, %r9 + adcq $0, %r10 + + // Round 3. + movq 24(%rsi), %rdx + xorq %r13, %r13 + mulxq %rax, %r13, %r14 + adcxq %r13, %r11 + adoxq %r14, %r12 + mulxq %rbx, %r13, %r14 + adcxq %r13, %r12 + adoxq %r14, %r8 + mulxq %rcx, %r13, %r14 + adcxq %r13, %r8 + adoxq %r14, %r9 + mulxq %rdi, %r13, %r14 + adcxq %r13, %r9 + adoxq %r14, %r10 + movq $0, %r13 + adcxq %r13, %r10 + adoxq %r13, %r10 + + // Montgomery step 3. + movq %r11, %rdx + imulq -48(%rbp), %rdx + mulxq -56(%rbp), %r13, %r14 + movq %rdx, %r15 + shlq $62, %r15 + negq %r11 + adcq %r13, %r12 + adcq $0, %r8 + adcq %r15, %r9 + adcq $0, %r10 + mulxq -64(%rbp), %r15, %r13 + movq %rdx, %r15 + shrq $2, %r15 + addq %r13, %r12 + adcq %r14, %r8 + adcq $0, %r9 + adcq %r15, %r10 + + // The chain's sole conditional subtraction. + movabsq $0x4000000000000000, %rdx + movq %r12, %r13 + movq %r8, %r14 + movq %r9, %r15 + movq %r10, %r11 + subq -64(%rbp), %r13 + sbbq -56(%rbp), %r14 + sbbq $0, %r15 + sbbq %rdx, %r11 + cmovncq %r13, %r12 + cmovncq %r14, %r8 + cmovncq %r15, %r9 + cmovncq %r11, %r10 + + movq -88(%rbp), %rdi + movq %r12, (%rdi) + movq %r8, 8(%rdi) + movq %r9, 16(%rdi) + movq %r10, 24(%rdi) + jmp .Lreturn + +.Lstore_lazy: + movq -88(%rbp), %rdi + movq %r11, (%rdi) + movq %rbx, 8(%rdi) + movq %r12, 16(%rdi) + movq %rax, 24(%rdi) + +.Lreturn: + popq %rbx + popq %r12 + popq %r13 + popq %r14 + popq %r15 + popq %rbp + retq + .cfi_endproc + +#if !defined(__APPLE__) + .size C_SYMBOL(pasta_x86_64_sqr_n), .-C_SYMBOL(pasta_x86_64_sqr_n) + .section .note.GNU-stack,"",@progbits +#endif diff --git a/pasta_curves/src/fields.rs b/pasta_curves/src/fields.rs index 77048e0a..b97d8c0d 100644 --- a/pasta_curves/src/fields.rs +++ b/pasta_curves/src/fields.rs @@ -16,6 +16,11 @@ mod portable; ))] mod aarch64_asm; +// Same containment for the x86-64 inline-assembly backend. +#[allow(unsafe_code)] +#[cfg(all(feature = "x86_64-asm", target_arch = "x86_64"))] +mod x86_64_asm; + pub use fp::*; pub use fq::*; diff --git a/pasta_curves/src/fields/fp.rs b/pasta_curves/src/fields/fp.rs index 7620d1d0..be6a4d47 100644 --- a/pasta_curves/src/fields/fp.rs +++ b/pasta_curves/src/fields/fp.rs @@ -395,10 +395,18 @@ impl Fp { Fp(super::aarch64_asm::mul(&self.0, &rhs.0, &MODULUS.0, INV)) } - #[cfg(not(all( - feature = "aarch64-asm", - target_arch = "aarch64", - target_vendor = "apple" + #[cfg(all(feature = "x86_64-asm", target_arch = "x86_64"))] + { + Fp(super::x86_64_asm::mul(&self.0, &rhs.0, &MODULUS.0, INV)) + } + + #[cfg(not(any( + all( + feature = "aarch64-asm", + target_arch = "aarch64", + target_vendor = "apple" + ), + all(feature = "x86_64-asm", target_arch = "x86_64") )))] { self.mul(rhs) @@ -416,10 +424,18 @@ impl Fp { Fp(super::aarch64_asm::square(&self.0, &MODULUS.0, INV)) } - #[cfg(not(all( - feature = "aarch64-asm", - target_arch = "aarch64", - target_vendor = "apple" + #[cfg(all(feature = "x86_64-asm", target_arch = "x86_64"))] + { + Fp(super::x86_64_asm::square(&self.0, &MODULUS.0, INV)) + } + + #[cfg(not(any( + all( + feature = "aarch64-asm", + target_arch = "aarch64", + target_vendor = "apple" + ), + all(feature = "x86_64-asm", target_arch = "x86_64") )))] { self.square() @@ -427,8 +443,8 @@ impl Fp { } /// Squares `self` `n` times (`n` must be at least 1), then multiplies the - /// result by `by`. The assembly backend keeps the accumulator in - /// registers for the whole chain. + /// result by `by`. Assembly backends leave the accumulator lazily reduced + /// until the closing multiplication. #[inline] fn sqr_n_mul_runtime(&self, n: u32, by: &Self) -> Self { assert!(n >= 1); @@ -444,15 +460,30 @@ impl Fp { )) } - #[cfg(not(all( - feature = "aarch64-asm", - target_arch = "aarch64", - target_vendor = "apple" + #[cfg(all(feature = "x86_64-asm", target_arch = "x86_64"))] + { + if n >= super::x86_64_asm::LAZY_SQUARE_THRESHOLD { + Fp(super::x86_64_asm::sqr_n_mul( + &self.0, n, &by.0, &MODULUS.0, INV, + )) + } else { + Fp(portable::sqr_n_lazy(&self.0, n, &MODULUS.0, INV)).mul_runtime(by) + } + } + + #[cfg(not(any( + all( + feature = "aarch64-asm", + target_arch = "aarch64", + target_vendor = "apple" + ), + all(feature = "x86_64-asm", target_arch = "x86_64") )))] { // Leave the accumulator unreduced between squarings. The closing - // multiplication canonicalizes its result. - Fp(portable::sqr_n_lazy(&self.0, n, &MODULUS.0, INV)).mul(by) + // multiplication canonicalizes its result through the selected + // runtime backend. + Fp(portable::sqr_n_lazy(&self.0, n, &MODULUS.0, INV)).mul_runtime(by) } } @@ -472,10 +503,23 @@ impl Fp { (0..n).fold(*self, |acc, _| acc.square_runtime()) } - #[cfg(not(all( - feature = "aarch64-asm", - target_arch = "aarch64", - target_vendor = "apple" + #[cfg(all(feature = "x86_64-asm", target_arch = "x86_64"))] + { + let lazy = if n >= super::x86_64_asm::LAZY_SQUARE_THRESHOLD { + super::x86_64_asm::sqr_n_lazy(&self.0, n, &MODULUS.0, INV) + } else { + portable::sqr_n_lazy(&self.0, n, &MODULUS.0, INV) + }; + Fp(portable::canonicalize(&lazy, &MODULUS.0)) + } + + #[cfg(not(any( + all( + feature = "aarch64-asm", + target_arch = "aarch64", + target_vendor = "apple" + ), + all(feature = "x86_64-asm", target_arch = "x86_64") )))] { Fp(portable::canonicalize( @@ -1676,3 +1720,137 @@ fn aarch64_asm_mul_rejects_non_canonical_rhs_in_debug() { // The modulus itself is the smallest non-canonical value. let _ = Fp::one().mul_runtime(&MODULUS); } + +#[cfg(all(test, feature = "x86_64-asm", target_arch = "x86_64"))] +#[test] +fn x86_64_asm_mul_matches_portable() { + use rand::{Rng, SeedableRng}; + + // Random canonical pairs through the inline block against the portable + // multiplication, and the squaring route against the portable squaring. + let mut rng = rand_xorshift::XorShiftRng::from_seed([0x2a; 16]); + for case in 0..100_000 { + let a = Fp::from_raw([ + rng.next_u64(), + rng.next_u64(), + rng.next_u64(), + rng.next_u64(), + ]); + let b = Fp::from_raw([ + rng.next_u64(), + rng.next_u64(), + rng.next_u64(), + rng.next_u64(), + ]); + let asm = a.mul_runtime(&b); + assert_eq!(asm, Fp::mul(&a, &b), "lhs {:x?} rhs {:x?}", a.0, b.0); + assert!(is_canonical(&asm)); + assert_eq!(a.square_runtime(), Fp::square(&a), "value {:x?}", a.0); + + if case < 256 { + for n in [1, 2, 7, 129] { + let asm_lazy = super::x86_64_asm::sqr_n_lazy(&a.0, n, &MODULUS.0, INV); + let portable_lazy = portable::sqr_n_lazy(&a.0, n, &MODULUS.0, INV); + assert_eq!(asm_lazy, portable_lazy, "lazy square, n = {n}"); + assert_eq!( + a.sqr_n_mul_runtime(n, &b), + Fp(portable::canonicalize(&portable_lazy, &MODULUS.0)).mul(&b), + "lazy square-and-multiply, n = {n}" + ); + } + } + } + + // Edge operands: zero, one, the largest canonical value, and dense + // all-ones-shaped canonical limbs. + let mut max_canonical = MODULUS; + max_canonical.0[0] -= 1; + let mut dense_canonical = Fp([u64::MAX - 3; 4]); + dense_canonical.0[3] = MODULUS.0[3] - 1; + let edges = [ + Fp::zero(), + Fp::one(), + max_canonical, + dense_canonical, + R2, + R3, + ]; + for a in edges { + for n in [1, 2, 7, 129, 512] { + assert_eq!( + super::x86_64_asm::sqr_n_lazy(&a.0, n, &MODULUS.0, INV), + portable::sqr_n_lazy(&a.0, n, &MODULUS.0, INV), + "edge lazy square, value {:x?}, n = {n}", + a.0 + ); + } + for b in edges { + assert_eq!( + a.mul_runtime(&b), + Fp::mul(&a, &b), + "lhs {:x?} rhs {:x?}", + a.0, + b.0 + ); + for n in [1, 129] { + let eager = (0..n).fold(a, |acc, _| Fp::square(&acc)); + assert_eq!( + a.sqr_n_mul_runtime(n, &b), + Fp::mul(&eager, &b), + "edge lazy square-and-multiply, n = {n}" + ); + } + } + assert_eq!(a.square_runtime(), Fp::square(&a), "value {:x?}", a.0); + } +} + +#[cfg(all(test, feature = "x86_64-asm", target_arch = "x86_64"))] +#[test] +fn x86_64_asm_mul_unreduced_lhs_near_modulus_rhs_matches_portable() { + use rand::{Rng, SeedableRng}; + + // Same contract and five-limb structure as the AArch64 block: the final + // shift omits the fifth candidate limb because + // `(lhs * rhs + m * modulus) / R < 2 * modulus < R` once the rhs is + // canonical. Stress the bound where it is tightest: lhs with its top bit + // set, rhs within a few limbs of the modulus (kept canonical, and within + // the per-limb no-wrap condition that an unreduced lhs separately + // requires). + let mut rng = rand_xorshift::XorShiftRng::from_seed([0x35; 16]); + let mut n = 0u32; + while n < 100_000 { + let lhs = Fp([ + rng.next_u64(), + rng.next_u64(), + rng.next_u64(), + rng.next_u64() | (1 << 63), + ]); + let mut rhs = MODULUS; + rhs.0[0] = rhs.0[0].wrapping_sub(rng.next_u64() >> (rng.next_u32() % 64)); + if rng.next_u32() & 1 == 1 { + rhs.0[1] = rhs.0[1].wrapping_sub(rng.next_u64() >> 60); + } + if !is_canonical(&rhs) || rhs.0.iter().any(|&l| l > u64::MAX - 3) { + continue; + } + n += 1; + let asm = lhs.mul_runtime(&rhs); + assert_eq!( + asm, + Fp::mul(&lhs, &rhs), + "lhs {:x?} rhs {:x?}", + lhs.0, + rhs.0 + ); + assert!(is_canonical(&asm)); + } +} + +#[cfg(all(test, debug_assertions, feature = "x86_64-asm", target_arch = "x86_64"))] +#[test] +#[should_panic(expected = "requires a canonical rhs")] +fn x86_64_asm_mul_rejects_non_canonical_rhs_in_debug() { + // The modulus itself is the smallest non-canonical value. + let _ = Fp::one().mul_runtime(&MODULUS); +} diff --git a/pasta_curves/src/fields/fq.rs b/pasta_curves/src/fields/fq.rs index a86faf05..bb330146 100644 --- a/pasta_curves/src/fields/fq.rs +++ b/pasta_curves/src/fields/fq.rs @@ -395,10 +395,18 @@ impl Fq { Fq(super::aarch64_asm::mul(&self.0, &rhs.0, &MODULUS.0, INV)) } - #[cfg(not(all( - feature = "aarch64-asm", - target_arch = "aarch64", - target_vendor = "apple" + #[cfg(all(feature = "x86_64-asm", target_arch = "x86_64"))] + { + Fq(super::x86_64_asm::mul(&self.0, &rhs.0, &MODULUS.0, INV)) + } + + #[cfg(not(any( + all( + feature = "aarch64-asm", + target_arch = "aarch64", + target_vendor = "apple" + ), + all(feature = "x86_64-asm", target_arch = "x86_64") )))] { self.mul(rhs) @@ -416,10 +424,18 @@ impl Fq { Fq(super::aarch64_asm::square(&self.0, &MODULUS.0, INV)) } - #[cfg(not(all( - feature = "aarch64-asm", - target_arch = "aarch64", - target_vendor = "apple" + #[cfg(all(feature = "x86_64-asm", target_arch = "x86_64"))] + { + Fq(super::x86_64_asm::square(&self.0, &MODULUS.0, INV)) + } + + #[cfg(not(any( + all( + feature = "aarch64-asm", + target_arch = "aarch64", + target_vendor = "apple" + ), + all(feature = "x86_64-asm", target_arch = "x86_64") )))] { self.square() @@ -427,8 +443,8 @@ impl Fq { } /// Squares `self` `n` times (`n` must be at least 1), then multiplies the - /// result by `by`. The assembly backend keeps the accumulator in - /// registers for the whole chain. + /// result by `by`. Assembly backends leave the accumulator lazily reduced + /// until the closing multiplication. #[inline] fn sqr_n_mul_runtime(&self, n: u32, by: &Self) -> Self { assert!(n >= 1); @@ -444,15 +460,30 @@ impl Fq { )) } - #[cfg(not(all( - feature = "aarch64-asm", - target_arch = "aarch64", - target_vendor = "apple" + #[cfg(all(feature = "x86_64-asm", target_arch = "x86_64"))] + { + if n >= super::x86_64_asm::LAZY_SQUARE_THRESHOLD { + Fq(super::x86_64_asm::sqr_n_mul( + &self.0, n, &by.0, &MODULUS.0, INV, + )) + } else { + Fq(portable::sqr_n_lazy(&self.0, n, &MODULUS.0, INV)).mul_runtime(by) + } + } + + #[cfg(not(any( + all( + feature = "aarch64-asm", + target_arch = "aarch64", + target_vendor = "apple" + ), + all(feature = "x86_64-asm", target_arch = "x86_64") )))] { // Leave the accumulator unreduced between squarings. The closing - // multiplication canonicalizes its result. - Fq(portable::sqr_n_lazy(&self.0, n, &MODULUS.0, INV)).mul(by) + // multiplication canonicalizes its result through the selected + // runtime backend. + Fq(portable::sqr_n_lazy(&self.0, n, &MODULUS.0, INV)).mul_runtime(by) } } @@ -472,10 +503,23 @@ impl Fq { (0..n).fold(*self, |acc, _| acc.square_runtime()) } - #[cfg(not(all( - feature = "aarch64-asm", - target_arch = "aarch64", - target_vendor = "apple" + #[cfg(all(feature = "x86_64-asm", target_arch = "x86_64"))] + { + let lazy = if n >= super::x86_64_asm::LAZY_SQUARE_THRESHOLD { + super::x86_64_asm::sqr_n_lazy(&self.0, n, &MODULUS.0, INV) + } else { + portable::sqr_n_lazy(&self.0, n, &MODULUS.0, INV) + }; + Fq(portable::canonicalize(&lazy, &MODULUS.0)) + } + + #[cfg(not(any( + all( + feature = "aarch64-asm", + target_arch = "aarch64", + target_vendor = "apple" + ), + all(feature = "x86_64-asm", target_arch = "x86_64") )))] { Fq(portable::canonicalize( @@ -1675,3 +1719,137 @@ fn aarch64_asm_mul_rejects_non_canonical_rhs_in_debug() { // The modulus itself is the smallest non-canonical value. let _ = Fq::one().mul_runtime(&MODULUS); } + +#[cfg(all(test, feature = "x86_64-asm", target_arch = "x86_64"))] +#[test] +fn x86_64_asm_mul_matches_portable() { + use rand::{Rng, SeedableRng}; + + // Random canonical pairs through the inline block against the portable + // multiplication, and the squaring route against the portable squaring. + let mut rng = rand_xorshift::XorShiftRng::from_seed([0x2a; 16]); + for case in 0..100_000 { + let a = Fq::from_raw([ + rng.next_u64(), + rng.next_u64(), + rng.next_u64(), + rng.next_u64(), + ]); + let b = Fq::from_raw([ + rng.next_u64(), + rng.next_u64(), + rng.next_u64(), + rng.next_u64(), + ]); + let asm = a.mul_runtime(&b); + assert_eq!(asm, Fq::mul(&a, &b), "lhs {:x?} rhs {:x?}", a.0, b.0); + assert!(is_canonical(&asm)); + assert_eq!(a.square_runtime(), Fq::square(&a), "value {:x?}", a.0); + + if case < 256 { + for n in [1, 2, 7, 129] { + let asm_lazy = super::x86_64_asm::sqr_n_lazy(&a.0, n, &MODULUS.0, INV); + let portable_lazy = portable::sqr_n_lazy(&a.0, n, &MODULUS.0, INV); + assert_eq!(asm_lazy, portable_lazy, "lazy square, n = {n}"); + assert_eq!( + a.sqr_n_mul_runtime(n, &b), + Fq(portable::canonicalize(&portable_lazy, &MODULUS.0)).mul(&b), + "lazy square-and-multiply, n = {n}" + ); + } + } + } + + // Edge operands: zero, one, the largest canonical value, and dense + // all-ones-shaped canonical limbs. + let mut max_canonical = MODULUS; + max_canonical.0[0] -= 1; + let mut dense_canonical = Fq([u64::MAX - 3; 4]); + dense_canonical.0[3] = MODULUS.0[3] - 1; + let edges = [ + Fq::zero(), + Fq::one(), + max_canonical, + dense_canonical, + R2, + R3, + ]; + for a in edges { + for n in [1, 2, 7, 129, 512] { + assert_eq!( + super::x86_64_asm::sqr_n_lazy(&a.0, n, &MODULUS.0, INV), + portable::sqr_n_lazy(&a.0, n, &MODULUS.0, INV), + "edge lazy square, value {:x?}, n = {n}", + a.0 + ); + } + for b in edges { + assert_eq!( + a.mul_runtime(&b), + Fq::mul(&a, &b), + "lhs {:x?} rhs {:x?}", + a.0, + b.0 + ); + for n in [1, 129] { + let eager = (0..n).fold(a, |acc, _| Fq::square(&acc)); + assert_eq!( + a.sqr_n_mul_runtime(n, &b), + Fq::mul(&eager, &b), + "edge lazy square-and-multiply, n = {n}" + ); + } + } + assert_eq!(a.square_runtime(), Fq::square(&a), "value {:x?}", a.0); + } +} + +#[cfg(all(test, feature = "x86_64-asm", target_arch = "x86_64"))] +#[test] +fn x86_64_asm_mul_unreduced_lhs_near_modulus_rhs_matches_portable() { + use rand::{Rng, SeedableRng}; + + // Same contract and five-limb structure as the AArch64 block: the final + // shift omits the fifth candidate limb because + // `(lhs * rhs + m * modulus) / R < 2 * modulus < R` once the rhs is + // canonical. Stress the bound where it is tightest: lhs with its top bit + // set, rhs within a few limbs of the modulus (kept canonical, and within + // the per-limb no-wrap condition that an unreduced lhs separately + // requires). + let mut rng = rand_xorshift::XorShiftRng::from_seed([0x35; 16]); + let mut n = 0u32; + while n < 100_000 { + let lhs = Fq([ + rng.next_u64(), + rng.next_u64(), + rng.next_u64(), + rng.next_u64() | (1 << 63), + ]); + let mut rhs = MODULUS; + rhs.0[0] = rhs.0[0].wrapping_sub(rng.next_u64() >> (rng.next_u32() % 64)); + if rng.next_u32() & 1 == 1 { + rhs.0[1] = rhs.0[1].wrapping_sub(rng.next_u64() >> 60); + } + if !is_canonical(&rhs) || rhs.0.iter().any(|&l| l > u64::MAX - 3) { + continue; + } + n += 1; + let asm = lhs.mul_runtime(&rhs); + assert_eq!( + asm, + Fq::mul(&lhs, &rhs), + "lhs {:x?} rhs {:x?}", + lhs.0, + rhs.0 + ); + assert!(is_canonical(&asm)); + } +} + +#[cfg(all(test, debug_assertions, feature = "x86_64-asm", target_arch = "x86_64"))] +#[test] +#[should_panic(expected = "requires a canonical rhs")] +fn x86_64_asm_mul_rejects_non_canonical_rhs_in_debug() { + // The modulus itself is the smallest non-canonical value. + let _ = Fq::one().mul_runtime(&MODULUS); +} diff --git a/pasta_curves/src/fields/x86_64_asm.rs b/pasta_curves/src/fields/x86_64_asm.rs new file mode 100644 index 00000000..c8232b55 --- /dev/null +++ b/pasta_curves/src/fields/x86_64_asm.rs @@ -0,0 +1,583 @@ +//! Private x86-64 backend for the Pasta fields. +//! +//! Montgomery multiplication and single squaring are implemented as inline +//! `asm!` blocks using MULX (BMI2) with ADCX/ADOX dual carry chains (ADX) in +//! the multiplication rows. Repeated lazy squaring uses the interleaved #218 +//! schedule in `pasta_lazy_square-x86_64.S`; it retains the accumulator across +//! iterations and fuses the optional closing multiplication. Two negative +//! scheduling results are pinned here so +//! they are not retried on this microarchitecture family: routing squaring +//! through the multiplication measured 2–5% *slower* (run-dependent) than +//! the dedicated squaring below (21.0 vs 20.0–20.7 ns on Skylake-X — +//! mirroring the AArch64 +//! backend, whose inline square also beats its multiplication; an earlier +//! contrary reading came from a benchmark cell in which the inherent +//! portable `square` shadowed `Field::square`), and merging each +//! Montgomery step's two carry sweeps into interleaved ADCX/ADOX chains +//! (staging all five `q*p` operands flag-free, `TEST` to clear both +//! chains) measured ~10% slower than the two short sequential sweeps +//! (22.3 vs 20.2 ns) despite the shorter nominal dependency length. +//! +//! The round structure is a transcription of the AArch64 backend +//! (`aarch64_asm.rs`), which is itself the upstream Semolina +//! `mul_mont_pasta`: a five-limb CIOS accumulator, one Montgomery +//! cancellation per round, and the shared Pasta modulus shape — +//! `modulus[2] = 0` and `modulus[3] = 2^62` — materialized as shifts, so +//! only `modulus[0]`, `modulus[1]`, and `inv` distinguish Fp from Fq. +//! Because the mathematical structure is identical, the AArch64 module's +//! bounds analysis carries over verbatim; see its module docs for the +//! five-limb no-wrap argument. +//! +//! Unlike the AArch64 block, operand limbs are addressed through pointers +//! (`readonly` memory operands) rather than individual registers: the +//! interleaved rounds plus staging temporaries do not fit x86-64's fourteen +//! allocatable registers with twelve limbs pinned. The loads are L1 hits off +//! the multiplier's critical path. +//! +//! Canonicity contract (same as the AArch64 backend): `rhs` in `mul` and +//! the input of `square` must be canonical (below the modulus) — the +//! five-limb accumulator drops the candidate's would-be fifth limb, and +//! for `rhs >= R - p` the result would be an incorrect residue that still +//! looks canonical. Both routines debug-assert that precondition, and with +//! both operands canonical they are always safe. `lhs` in `mul` may be an +//! unreduced 256-bit value only if every `rhs` limb is at most `2^64 - 4` +//! (the accumulator no-wrap bound) — a condition that is *not* asserted. +//! The lazy squaring chain is the one narrower exception: its `lhs` stays +//! below `2p < 2^255`, so its top bit is clear and the same five-limb +//! accumulator cannot wrap even for an arbitrary canonical `rhs`. +//! `from_u512`, the other place that produces an unreduced value, continues +//! to use the portable path. The +//! `x86_64_asm_mul_unreduced_lhs_near_modulus_rhs_matches_portable` tests +//! in `fp.rs`/`fq.rs` pin the allowance. Outputs are canonical. +//! +//! The direct-operation blocks are straight-line. The lazy chain branches only +//! on its public iteration count and whether a public closing multiplier was +//! supplied. There are no data-dependent branches or memory addresses, and +//! final conditional subtraction uses CMOV, so the code is constant-time. +//! +//! ISA requirement: MULX needs BMI2 and ADCX/ADOX need ADX (Intel Broadwell +//! / AMD Zen or newer). The feature is opt-in precisely because this is not +//! checked at runtime; enabling it on an older CPU faults with an illegal +//! instruction. + +use core::arch::asm; + +type Limbs = [u64; 4]; + +/// The fused assembly path wins even for the shortest measured chains, so all +/// nonempty chains use it. +pub(super) const LAZY_SQUARE_THRESHOLD: u32 = 1; + +/// Whether `value < modulus` as little-endian 256-bit integers. +#[inline(always)] +fn is_canonical(value: &Limbs, modulus: &Limbs) -> bool { + for i in (0..4).rev() { + if value[i] != modulus[i] { + return value[i] < modulus[i]; + } + } + false +} +/// Multiplies two Montgomery residues for a Pasta modulus. `rhs` must be +/// canonical (debug-asserted; a violation yields an incorrect residue, see +/// the module docs). `lhs` may be unreduced only if every `rhs` limb is at +/// most `2^64 - 4`; see the AArch64 module docs for the carry-chain bound +/// behind this. +#[inline(always)] +pub(super) fn mul(lhs: &Limbs, rhs: &Limbs, modulus: &Limbs, inv: u64) -> Limbs { + debug_assert!( + is_canonical(rhs, modulus), + "x86_64_asm::mul requires a canonical rhs" + ); + let (o0, o1, o2, o3): (u64, u64, u64, u64); + // SAFETY: straight-line arithmetic reading only the twelve limbs behind + // the three passed references (`readonly`); no stack use, and outputs + // depend only on the declared inputs. + // + // Register roles: the five-limb accumulator lives in {ae}/{be}/{ce}/ + // {de}/{ee} and its window rotates down one register per round (the + // register cancelled by the round's Montgomery step becomes the next + // round's fifth limb), so after four rounds the candidate sits in + // {ee},{ae},{be},{ce}. {s1}/{s2}/{s3} stage multiplier halves and + // shifted `q * modulus[3]` terms so no flag-writing instruction lands + // inside a carry chain. RDX is the implicit MULX source: each round's + // `rhs` limb, then the round's Montgomery factor `q`. + unsafe { + asm!( + // Round 0: initialize the accumulator with lhs * rhs[0]. + "mov rdx, qword ptr [{b}]", // rdx = b[0]. + "mulx {be}, {ae}, qword ptr [{a}]", // ae = low(a[0]*b[0]), be = high. + "mulx {ce}, {s1}, qword ptr [{a} + 8]", + "add {be}, {s1}", // Fold low(a[1]*b[0]) into limb 1. + "mulx {de}, {s1}, qword ptr [{a} + 16]", + "adc {ce}, {s1}", // Fold low(a[2]*b[0]) and carry. + "mulx {ee}, {s1}, qword ptr [{a} + 24]", + "adc {de}, {s1}", // Fold low(a[3]*b[0]) and carry. + "adc {ee}, 0", // Fifth limb of lhs * b[0]. + + // Montgomery step 0: q = limb0 * inv; add q*p; shift one limb. + "mov rdx, {ae}", + "imul rdx, {inv}", // rdx = q (low 64 bits only). + "mulx {s2}, {s1}, qword ptr [{p} + 8]", // s1 = low(q*p[1]), s2 = high (kept). + "mov {s3}, rdx", + "shl {s3}, 62", // s3 = low(q*p[3]); p[2] contributes nothing. + // low(q*p[0]) cancels limb 0; its carry is one exactly when the + // limb is nonzero, which NEG leaves in CF. + "neg {ae}", // CF = (limb0 != 0); ae is dead. + "adc {be}, {s1}", // Add low(q*p[1]) and the cancellation carry. + "adc {ce}, 0", // Propagate across zero p[2]. + "adc {de}, {s3}", // Add low(q*p[3]) and carry. + "adc {ee}, 0", // Propagate into the fifth limb. + "mulx {s1}, {s3}, qword ptr [{p}]", // s1 = high(q*p[0]); the low half is spent. + "mov {s3}, rdx", + "shr {s3}, 2", // s3 = high(q*p[3]). + "mov {ae}, 0", // Next round's fifth limb (MOV keeps flags). + "add {be}, {s1}", // New limb 0 includes high(q*p[0]). + "adc {ce}, {s2}", // New limb 1 includes high(q*p[1]). + "adc {de}, 0", // New limb 2; p[2] contributes zero. + "adc {ee}, {s3}", // New limb 3 includes high(q*p[3]). + "adc {ae}, 0", // Capture the reduction carry as limb 4. + + // Round 1: accumulator window is [be,ce,de,ee,ae]; add lhs*b[1] + // on dual carry chains (CF: low halves, OF: high halves). + "mov rdx, qword ptr [{b} + 8]", // rdx = b[1]. + "xor {s1}, {s1}", // Clear CF and OF. + "mulx {s2}, {s1}, qword ptr [{a}]", + "adcx {be}, {s1}", + "adox {ce}, {s2}", + "mulx {s2}, {s1}, qword ptr [{a} + 8]", + "adcx {ce}, {s1}", + "adox {de}, {s2}", + "mulx {s2}, {s1}, qword ptr [{a} + 16]", + "adcx {de}, {s1}", + "adox {ee}, {s2}", + "mulx {s2}, {s1}, qword ptr [{a} + 24]", + "adcx {ee}, {s1}", + "adox {ae}, {s2}", + "mov {s1}, 0", + "adcx {ae}, {s1}", // Close the CF chain into limb 4. + "adox {ae}, {s1}", // Close the OF chain into limb 4. + + // Montgomery step 1. + "mov rdx, {be}", + "imul rdx, {inv}", + "mulx {s2}, {s1}, qword ptr [{p} + 8]", + "mov {s3}, rdx", + "shl {s3}, 62", + "neg {be}", + "adc {ce}, {s1}", + "adc {de}, 0", + "adc {ee}, {s3}", + "adc {ae}, 0", + "mulx {s1}, {s3}, qword ptr [{p}]", + "mov {s3}, rdx", + "shr {s3}, 2", + "mov {be}, 0", + "add {ce}, {s1}", + "adc {de}, {s2}", + "adc {ee}, 0", + "adc {ae}, {s3}", + "adc {be}, 0", + + // Round 2: window [ce,de,ee,ae,be]; add lhs*b[2]. + "mov rdx, qword ptr [{b} + 16]", + "xor {s1}, {s1}", + "mulx {s2}, {s1}, qword ptr [{a}]", + "adcx {ce}, {s1}", + "adox {de}, {s2}", + "mulx {s2}, {s1}, qword ptr [{a} + 8]", + "adcx {de}, {s1}", + "adox {ee}, {s2}", + "mulx {s2}, {s1}, qword ptr [{a} + 16]", + "adcx {ee}, {s1}", + "adox {ae}, {s2}", + "mulx {s2}, {s1}, qword ptr [{a} + 24]", + "adcx {ae}, {s1}", + "adox {be}, {s2}", + "mov {s1}, 0", + "adcx {be}, {s1}", + "adox {be}, {s1}", + + // Montgomery step 2. + "mov rdx, {ce}", + "imul rdx, {inv}", + "mulx {s2}, {s1}, qword ptr [{p} + 8]", + "mov {s3}, rdx", + "shl {s3}, 62", + "neg {ce}", + "adc {de}, {s1}", + "adc {ee}, 0", + "adc {ae}, {s3}", + "adc {be}, 0", + "mulx {s1}, {s3}, qword ptr [{p}]", + "mov {s3}, rdx", + "shr {s3}, 2", + "mov {ce}, 0", + "add {de}, {s1}", + "adc {ee}, {s2}", + "adc {ae}, 0", + "adc {be}, {s3}", + "adc {ce}, 0", + + // Round 3: window [de,ee,ae,be,ce]; add lhs*b[3]. + "mov rdx, qword ptr [{b} + 24]", + "xor {s1}, {s1}", + "mulx {s2}, {s1}, qword ptr [{a}]", + "adcx {de}, {s1}", + "adox {ee}, {s2}", + "mulx {s2}, {s1}, qword ptr [{a} + 8]", + "adcx {ee}, {s1}", + "adox {ae}, {s2}", + "mulx {s2}, {s1}, qword ptr [{a} + 16]", + "adcx {ae}, {s1}", + "adox {be}, {s2}", + "mulx {s2}, {s1}, qword ptr [{a} + 24]", + "adcx {be}, {s1}", + "adox {ce}, {s2}", + "mov {s1}, 0", + "adcx {ce}, {s1}", + "adox {ce}, {s1}", + + // Montgomery step 3. Canonical rhs bounds the candidate below + // 2p < R, so the final shift produces no fifth limb (see the + // AArch64 module docs); the shift's carry adc is omitted. + "mov rdx, {de}", + "imul rdx, {inv}", + "mulx {s2}, {s1}, qword ptr [{p} + 8]", + "mov {s3}, rdx", + "shl {s3}, 62", + "neg {de}", + "adc {ee}, {s1}", + "adc {ae}, 0", + "adc {be}, {s3}", + "adc {ce}, 0", + "mulx {s1}, {s3}, qword ptr [{p}]", + "mov {s3}, rdx", + "shr {s3}, 2", + "add {ee}, {s1}", // Final candidate limb 0. + "adc {ae}, {s2}", // Final candidate limb 1. + "adc {be}, 0", // Final candidate limb 2. + "adc {ce}, {s3}", // Final candidate limb 3. + + // Conditional subtraction of p = [p0, p1, 0, 2^62]. + "movabs rdx, 0x4000000000000000", // Materialize p[3] = 2^62. + "mov {s1}, {ee}", + "mov {s2}, {ae}", + "mov {s3}, {be}", + "mov {de}, {ce}", + "sub {s1}, qword ptr [{p}]", // Tentative limb 0 = candidate - p[0]. + "sbb {s2}, qword ptr [{p} + 8]", // Tentative limb 1 minus p[1]. + "sbb {s3}, 0", // Tentative limb 2; p[2] is zero. + "sbb {de}, rdx", // Tentative limb 3 minus p[3]. + // No borrow (CF clear) means the candidate is at least p, so the + // subtracted value is the canonical output. + "cmovnc {ee}, {s1}", + "cmovnc {ae}, {s2}", + "cmovnc {be}, {s3}", + "cmovnc {ce}, {de}", + a = in(reg) lhs.as_ptr(), + b = in(reg) rhs.as_ptr(), + p = in(reg) modulus.as_ptr(), + inv = in(reg) inv, + ae = out(reg) o1, + be = out(reg) o2, + ce = out(reg) o3, + de = out(reg) _, + ee = out(reg) o0, + s1 = out(reg) _, + s2 = out(reg) _, + s3 = out(reg) _, + out("rdx") _, + options(pure, readonly, nostack), + ); + } + [o0, o1, o2, o3] +} + +/// Squares a canonical Montgomery residue for a Pasta modulus (the input's +/// canonicity is debug-asserted). +/// +/// A transcription of the AArch64 backend's dedicated squaring: the 512-bit +/// square as cross products, one doubling pass, and the diagonals (ten MULX +/// against the multiplication's sixteen), then four Montgomery +/// cancellations on a rotating four-limb window with a carried fifth limb, +/// the high product half folded in (the sum stays below `2p`, so no carry +/// escapes — see the AArch64 module's bounds), and a CMOV conditional +/// subtraction. Measured 2–5% ahead of squaring through [`mul`] on +/// Skylake-X (20.0–20.7 vs 21.0 ns across runs), mirroring the AArch64 +/// backend's own square-over-mul margin. +#[inline(always)] +pub(super) fn square(value: &Limbs, modulus: &Limbs, inv: u64) -> Limbs { + debug_assert!( + is_canonical(value, modulus), + "x86_64_asm::square requires a canonical input" + ); + let (o0, o1, o2, o3): (u64, u64, u64, u64); + // SAFETY: straight-line arithmetic reading only the limbs behind the two + // passed references (`readonly`); no stack use, and outputs depend only + // on the declared inputs. The input pointer's register is reclaimed as + // the reduction's carry limb once phase 1 has consumed the last load. + unsafe { + asm!( + // Phase 1: the 512-bit square in z0..z7. + // Cross products a[i]*a[j] (i < j), accumulated as they stream. + "xor {z5:e}, {z5:e}", + "xor {z6:e}, {z6:e}", + "xor {z7:e}, {z7:e}", + "mov rdx, qword ptr [{a}]", + "mulx {t1}, {z1}, qword ptr [{a} + 8]", // a0*a1. + "mulx {t2}, {z2}, qword ptr [{a} + 16]", // a0*a2. + "mulx {z4}, {z3}, qword ptr [{a} + 24]", // a0*a3. + "add {z2}, {t1}", // Fold high(a0*a1). + "adc {z3}, {t2}", // Fold high(a0*a2) and carry. + "adc {z4}, 0", + "mov rdx, qword ptr [{a} + 8]", + "mulx {t2}, {t1}, qword ptr [{a} + 16]", // a1*a2. + "add {z3}, {t1}", + "adc {z4}, {t2}", + "adc {z5}, 0", + "mulx {t2}, {t1}, qword ptr [{a} + 24]", // a1*a3. + "add {z4}, {t1}", + "adc {z5}, {t2}", + "adc {z6}, 0", + "mov rdx, qword ptr [{a} + 16]", + "mulx {t2}, {t1}, qword ptr [{a} + 24]", // a2*a3. + "add {z5}, {t1}", + "adc {z6}, {t2}", + "adc {z7}, 0", + // Double the cross products. The doubled sum is below 2^512, so + // no carry leaves z7. + "add {z1}, {z1}", + "adc {z2}, {z2}", + "adc {z3}, {z3}", + "adc {z4}, {z4}", + "adc {z5}, {z5}", + "adc {z6}, {z6}", + "adc {z7}, {z7}", + // Add the diagonal squares in one carry chain (MOV and MULX + // preserve flags). + "mov rdx, qword ptr [{a}]", + "mulx {t2}, {z0}, rdx", // z0 = low(a0^2). + "add {z1}, {t2}", // High(a0^2). + "mov rdx, qword ptr [{a} + 8]", + "mulx {t2}, {t1}, rdx", + "adc {z2}, {t1}", + "adc {z3}, {t2}", + "mov rdx, qword ptr [{a} + 16]", + "mulx {t2}, {t1}, rdx", + "adc {z4}, {t1}", + "adc {z5}, {t2}", + "mov rdx, qword ptr [{a} + 24]", + "mulx {t2}, {t1}, rdx", + "adc {z6}, {t1}", + "adc {z7}, {t2}", // a^2 < 2^510: no carry out. + + // Phase 2: four Montgomery cancellations on the low half, the + // same two-sweep step as [`mul`]'s. The window rotates down one + // register per step; {a} (its loads are done) serves as the + // first carried fifth limb. + // Step 0: window [z0, z1, z2, z3], carry into {a}. + "mov rdx, {z0}", + "imul rdx, {inv}", // rdx = q. + "mulx {t2}, {t1}, qword ptr [{p} + 8]", // t1/t2 = low/high(q*p1). + "mov {a}, rdx", + "shl {a}, 62", // low(q*p3); p2 is zero. + "neg {z0}", // CF = (limb0 != 0). + "adc {z1}, {t1}", + "adc {z2}, 0", + "adc {z3}, {a}", + "mov {a}, 0", + "adc {a}, 0", // Carry above limb 3. + "mulx {t1}, {z0}, qword ptr [{p}]", // t1 = high(q*p0); low is spent. + "mov {z0}, rdx", + "shr {z0}, 2", // high(q*p3). + "add {z1}, {t1}", // New limb 0. + "adc {z2}, {t2}", // New limb 1 += high(q*p1). + "adc {z3}, 0", // New limb 2. + "adc {a}, {z0}", // New limb 3 += high(q*p3). + // Step 1: window [z1, z2, z3, a], carry into z0. + "mov rdx, {z1}", + "imul rdx, {inv}", + "mulx {t2}, {t1}, qword ptr [{p} + 8]", + "mov {z0}, rdx", + "shl {z0}, 62", + "neg {z1}", + "adc {z2}, {t1}", + "adc {z3}, 0", + "adc {a}, {z0}", + "mov {z0}, 0", + "adc {z0}, 0", + "mulx {t1}, {z1}, qword ptr [{p}]", + "mov {z1}, rdx", + "shr {z1}, 2", + "add {z2}, {t1}", + "adc {z3}, {t2}", + "adc {a}, 0", + "adc {z0}, {z1}", + // Step 2: window [z2, z3, a, z0], carry into z1. + "mov rdx, {z2}", + "imul rdx, {inv}", + "mulx {t2}, {t1}, qword ptr [{p} + 8]", + "mov {z1}, rdx", + "shl {z1}, 62", + "neg {z2}", + "adc {z3}, {t1}", + "adc {a}, 0", + "adc {z0}, {z1}", + "mov {z1}, 0", + "adc {z1}, 0", + "mulx {t1}, {z2}, qword ptr [{p}]", + "mov {z2}, rdx", + "shr {z2}, 2", + "add {z3}, {t1}", + "adc {a}, {t2}", + "adc {z0}, 0", + "adc {z1}, {z2}", + // Step 3: window [z3, a, z0, z1], carry into z2. + "mov rdx, {z3}", + "imul rdx, {inv}", + "mulx {t2}, {t1}, qword ptr [{p} + 8]", + "mov {z2}, rdx", + "shl {z2}, 62", + "neg {z3}", + "adc {a}, {t1}", + "adc {z0}, 0", + "adc {z1}, {z2}", + "mov {z2}, 0", + "adc {z2}, 0", + "mulx {t1}, {z3}, qword ptr [{p}]", + "mov {z3}, rdx", + "shr {z3}, 2", + "add {a}, {t1}", + "adc {z0}, {t2}", + "adc {z1}, 0", + "adc {z2}, {z3}", + + // Fold in the high product half; the sum stays below 2p, so no + // carry escapes and a four-limb conditional subtraction suffices. + "add {a}, {z4}", + "adc {z0}, {z5}", + "adc {z1}, {z6}", + "adc {z2}, {z7}", + "movabs rdx, 0x4000000000000000", // p3 = 2^62. + "mov {t1}, {a}", + "mov {t2}, {z0}", + "mov {z3}, {z1}", + "mov {z4}, {z2}", + "sub {t1}, qword ptr [{p}]", + "sbb {t2}, qword ptr [{p} + 8]", + "sbb {z3}, 0", + "sbb {z4}, rdx", + "cmovnc {a}, {t1}", + "cmovnc {z0}, {t2}", + "cmovnc {z1}, {z3}", + "cmovnc {z2}, {z4}", + a = inout(reg) value.as_ptr() => o0, + p = in(reg) modulus.as_ptr(), + inv = in(reg) inv, + z0 = out(reg) o1, + z1 = out(reg) o2, + z2 = out(reg) o3, + z3 = out(reg) _, + z4 = out(reg) _, + z5 = out(reg) _, + z6 = out(reg) _, + z7 = out(reg) _, + t1 = out(reg) _, + t2 = out(reg) _, + out("rdx") _, + options(pure, readonly, nostack), + ); + } + [o0, o1, o2, o3] +} + +#[cfg(target_family = "unix")] +extern "C" { + fn pasta_x86_64_sqr_n( + out: *mut u64, + value: *const u64, + n: u32, + by: *const u64, + modulus: *const u64, + inv: u64, + ); +} + +/// Runs the interleaved lazy-squaring schedule in one assembly leaf function. +/// A null `by` returns the lazy accumulator; otherwise the same assembly block +/// performs the closing multiplication and its single conditional correction. +#[cfg(target_family = "unix")] +#[inline(always)] +fn run_lazy_square_chain( + value: &Limbs, + n: u32, + by: *const u64, + modulus: &Limbs, + inv: u64, +) -> Limbs { + let mut out = [0; 4]; + // SAFETY: the assembly reads four limbs from `value`, `modulus`, and from + // `by` when it is non-null, then initializes all four limbs of `out`. + // Its stack-relative accesses stay within the SysV red zone; the build + // script only links it on Unix x86-64 targets where that ABI applies. + unsafe { + pasta_x86_64_sqr_n( + out.as_mut_ptr(), + value.as_ptr(), + n, + by, + modulus.as_ptr(), + inv, + ); + } + out +} + +/// Squares a canonical starting value `n` times with #218's interleaved lazy +/// schedule, keeping its four-limb accumulator live across loop iterations. +#[inline(always)] +pub(super) fn sqr_n_lazy(value: &Limbs, n: u32, modulus: &Limbs, inv: u64) -> Limbs { + debug_assert!(n >= 1); + debug_assert!( + is_canonical(value, modulus), + "x86_64_asm::sqr_n_lazy requires a canonical starting value" + ); + + #[cfg(target_family = "unix")] + { + run_lazy_square_chain(value, n, core::ptr::null(), modulus, inv) + } + + #[cfg(not(target_family = "unix"))] + { + super::portable::sqr_n_lazy(value, n, modulus, inv) + } +} + +/// Runs the fused x86 assembly lazy-squaring chain, then performs the closing +/// multiplication and the chain's only conditional correction in that same +/// assembly leaf function. +#[inline(always)] +pub(super) fn sqr_n_mul(value: &Limbs, n: u32, by: &Limbs, modulus: &Limbs, inv: u64) -> Limbs { + debug_assert!(n >= 1); + debug_assert!( + is_canonical(value, modulus), + "x86_64_asm::sqr_n_mul requires a canonical starting value" + ); + debug_assert!( + is_canonical(by, modulus), + "x86_64_asm::sqr_n_mul requires a canonical multiplier" + ); + + #[cfg(target_family = "unix")] + { + run_lazy_square_chain(value, n, by.as_ptr(), modulus, inv) + } + + #[cfg(not(target_family = "unix"))] + { + let acc = super::portable::sqr_n_lazy(value, n, modulus, inv); + mul(&acc, by, modulus, inv) + } +}