NEON backend for aarch64

[rubdos]

As promised in #449! This is joint work with @Tarinn. In fact, @Tarinn did most of the work :-)

Fixes #147 for Aarch64, and you'll see an ARMv7 PR coming in after this.

I'm working on benchmarks across several devices now (including an unpublished very hacky armv7 version), will update here when they become available. We're seeing speedups of 20-30% in relevant benchmarks. This is the same code as in zkcrypto#19.

I'll be cleaning up the commits in the next few minutes (trailing white space fixes, mostly), and I'll run a comparison benchmark for 343be3a, because that was only introduced for future ARMv7 support, and has not been tested under load.

TODO:

• Cleanup whitespace and commits
• Test shuffle! call vs vqtbx1q_u8 performance
• Add Aarch64 cross build to CI
• Figure out how-the-💥 we get rid of the ARMv8 TBL/TBX issue
• Add NEON documentation to README
  Explain that ARMv7 is not yet supported

[rubdos]

The shuffle! macro use in 26c494f seems to win 1.5% to 3% over the use of the vqtbx1q_u8 intrinsic on a Raspberry Pi 4. Seems like the LLVM compiler people know what they're doing.

I'm testing another point where we can use the shuffle! macro instead of some manual lane swapping, will report back.

[rubdos]

2d343e9 does not have an influence on performance, but the code is quite a bit cleaner, so let's leave it in.

[tarcieri]

How does this work without cross?

[rubdos]

Rust has a built-in cross-compiler. I've never really used cross to cross-compile myself; always had to roll my own due to circumstances.

Probably, the underlaying Ubuntu machine needs to install a gcc toolchain to actually work, but I hoped that the CI would give me some feedback already.

[tarcieri]

A cross-compiler works for builds, but how does it run the tests on a non-native architecture?

cross uses QEMU

[rubdos]

Oh okay, I did not spot that. The new version is only with cargo build, not with test, based on the new build-simd CI snippet.

[tarcieri]

I'd suggest using cross for this as it should be fine for testing NEON in CI.

Here are some examples (and it's probably fine to reuse the @RustCrypto cross-install workflow too)

https://github.com/RustCrypto/block-ciphers/blob/master/.github/workflows/aes.yml#L201-L249

[rubdos]

Seems like that's done in release/4.0 now, so I'll just drop my own commit.

I'd still like to point out that when I submit the armv7hl NEON work, it will work on nightly only for a while, so you might want to add armv7-nightly to the test matrix.

[tarcieri]

@rubdos we're working off the releases/4.0 branch for development, so you'll definitely want to rebase this

[jrose-signal]

Don't forget to make sure it's still constant-time! I don't actually know what that entails; it is not something I've looked into before. But a good amount of effort was put into the initial version of this crate using constant-time operations, and new backends should adhere to that too.

[rubdos]

@jrose-signal Indeed, that's pretty important. As far as @Tarinn and I know, this code should be constant time. It is a direct port of the AVX2 work, and there should be no time-dependent calls in there. The most sketchy call that I could see is the shuffle!, but that does not take secret-dependent indexes.

I'm not sure how this is generally done. Should we ask some independent reviewer to have a look at the code?

[pinkforest]

fyi - we're continuing all development via in main -

Could this be re-based to main which is basically from release/4.0 - Thanks!

Also would we have any benchmarks between u64, neon and fiat ?
e.g. a table like this would be great: https://github.com/pinkforest/ed25519-dalek/tree/docs-docsrs

For constant time, e.g. subtle CsOption etc.

[rubdos]

fyi - we're continuing all development via in main -

Could this be re-based to main which is basically from release/4.0 - Thanks!

Of course - Just to make sure: is that going to remain? I did the rebase four days ago on release/4.0.

Do we have any benchmarks between u64, neon and fiat ?
e.g. a table like this would be great: https://github.com/pinkforest/ed25519-dalek/tree/docs-docsrs

I'm working on this across a bunch of devices (including armv7). I'd like these results to appear in a small publication, but I'll certainly dump a summary in the README when I have them. I should have something by the end of the week.

[pinkforest]

Of course - Just to make sure: is that going to remain? I did the rebase four days ago on release/4.0.

Yeah this is going to stay in main

• Development / Maintenance / Release workflow #466

I'm working on this across a bunch of devices

Awesome! 🥳 would love to see a link to write-up in This-Week-in-Rust TWiR if you can share there too:
https://github.com/rust-lang/this-week-in-rust/

[rubdos]

Awesome! partying_face would love to see a link to write-up in This-Week-in-Rust TWiR if you can share there too: https://github.com/rust-lang/this-week-in-rust/

After publication, I'll make sure to make a huge amount of noise about it. That should be pretty short notice, I have some very strict deadlines about that now. Just to clarify: the deadline is there because this work is complete; there was no pressure on completion of this work that would have affected the quality.

[pinkforest]

Suggested change

	let mut out = [FieldElement51::zero(); 4];
	let mut out = [FieldElement51::ZERO; 4];

This was changed to ZERO

• Use inherent constants for ZERO, ONE, and MINUS_ONE #470

[pinkforest]

Suggested change

	println!("{} = {:?}", stringify!($x), $x.to_bytes());
	println!("{} = {:?}", stringify!($x), $x.as_bytes());

These were renamed to as_bytes()

[pinkforest]

Suggested change

	use backend::serial::u64::field::FieldElement51;
	use crate::backend::serial::u64::field::FieldElement51;

[pinkforest]

Instead of flooding some changes here is some .patches:

src/backend/vector/neon/field.rs

diff --git a/src/backend/vector/neon/field.rs b/src/backend/vector/neon/field.rs
index 5cf57ee..85630cd 100644
--- a/src/backend/vector/neon/field.rs
+++ b/src/backend/vector/neon/field.rs
@@ -21,7 +21,7 @@
 //! arm instructions.
 
 use core::ops::{Add, Mul, Neg};
-use packed_simd::{u32x4, u32x2, i32x4, u8x16, u64x4, u64x2, IntoBits};
+use packed_simd::{u32x4, u32x2, i32x4, u64x4, u64x2, IntoBits};
 
 use crate::backend::vector::neon::constants::{P_TIMES_16_HI, P_TIMES_16_LO, P_TIMES_2_HI, P_TIMES_2_LO};
 use crate::backend::serial::u64::field::FieldElement51;
@@ -141,7 +141,7 @@ impl ConditionallySelectable for FieldElement2625x4 {
 impl FieldElement2625x4 {
 
     pub fn split(&self) -> [FieldElement51; 4] {
-        let mut out = [FieldElement51::zero(); 4];
+        let mut out = [FieldElement51::ZERO; 4];
         for i in 0..5 {
             let a_2i   = self.0[i].0.extract(0) as u64;
             let b_2i   = self.0[i].0.extract(1) as u64;
@@ -367,7 +367,9 @@ impl FieldElement2625x4 {
     #[inline]
     fn reduce64(mut z: [(u64x2, u64x2); 10]) -> FieldElement2625x4 {
 
+        #[allow(non_snake_case)]
         let LOW_25_BITS: u64x2 = u64x2::splat((1 << 25) - 1);
+        #[allow(non_snake_case)]
         let LOW_26_BITS: u64x2 = u64x2::splat((1 << 26) - 1);
 
         let carry = |z: &mut [(u64x2, u64x2); 10], i: usize| {
@@ -419,6 +421,7 @@ impl FieldElement2625x4 {
         ])
     }
 
+    #[allow(non_snake_case)]
     pub fn square_and_negate_D(&self) -> FieldElement2625x4 {
         #[inline(always)]
         fn m(x: (u32x2, u32x2), y: (u32x2, u32x2)) -> u64x4 {
@@ -463,16 +466,16 @@ impl FieldElement2625x4 {
         let x8_19  = m_lo(v19, x8);
         let x9_19  = m_lo(v19, x9);
 
-        let mut z0 = m(x0,  x0) + m(x2_2,x8_19) + m(x4_2,x6_19) + ((m(x1_2,x9_19) +  m(x3_2,x7_19) +    m(x5,x5_19)) << 1);
-        let mut z1 = m(x0_2,x1) + m(x3_2,x8_19) + m(x5_2,x6_19) +                  ((m(x2,x9_19)   +    m(x4,x7_19)) << 1);
-        let mut z2 = m(x0_2,x2) + m(x1_2,x1)    + m(x4_2,x8_19) + m(x6,x6_19)    + ((m(x3_2,x9_19) +  m(x5_2,x7_19)) << 1);
-        let mut z3 = m(x0_2,x3) + m(x1_2,x2)    + m(x5_2,x8_19) +                  ((m(x4,x9_19)   +    m(x6,x7_19)) << 1);
-        let mut z4 = m(x0_2,x4) + m(x1_2,x3_2)  + m(x2,  x2)    + m(x6_2,x8_19)  + ((m(x5_2,x9_19) +    m(x7,x7_19)) << 1);
-        let mut z5 = m(x0_2,x5) + m(x1_2,x4)    + m(x2_2,x3)    + m(x7_2,x8_19)                    +  ((m(x6,x9_19)) << 1);
-        let mut z6 = m(x0_2,x6) + m(x1_2,x5_2)  + m(x2_2,x4)    + m(x3_2,x3) + m(x8,x8_19)        + ((m(x7_2,x9_19)) << 1);
-        let mut z7 = m(x0_2,x7) + m(x1_2,x6)    + m(x2_2,x5)    + m(x3_2,x4)                      +   ((m(x8,x9_19)) << 1);
-        let mut z8 = m(x0_2,x8) + m(x1_2,x7_2)  + m(x2_2,x6)    + m(x3_2,x5_2) + m(x4,x4)         +   ((m(x9,x9_19)) << 1);
-        let mut z9 = m(x0_2,x9) + m(x1_2,x8)    + m(x2_2,x7)    + m(x3_2,x6) + m(x4_2,x5);
+        let z0 = m(x0,  x0) + m(x2_2,x8_19) + m(x4_2,x6_19) + ((m(x1_2,x9_19) +  m(x3_2,x7_19) +    m(x5,x5_19)) << 1);
+        let z1 = m(x0_2,x1) + m(x3_2,x8_19) + m(x5_2,x6_19) +                  ((m(x2,x9_19)   +    m(x4,x7_19)) << 1);
+        let z2 = m(x0_2,x2) + m(x1_2,x1)    + m(x4_2,x8_19) + m(x6,x6_19)    + ((m(x3_2,x9_19) +  m(x5_2,x7_19)) << 1);
+        let z3 = m(x0_2,x3) + m(x1_2,x2)    + m(x5_2,x8_19) +                  ((m(x4,x9_19)   +    m(x6,x7_19)) << 1);
+        let z4 = m(x0_2,x4) + m(x1_2,x3_2)  + m(x2,  x2)    + m(x6_2,x8_19)  + ((m(x5_2,x9_19) +    m(x7,x7_19)) << 1);
+        let z5 = m(x0_2,x5) + m(x1_2,x4)    + m(x2_2,x3)    + m(x7_2,x8_19)                    +  ((m(x6,x9_19)) << 1);
+        let z6 = m(x0_2,x6) + m(x1_2,x5_2)  + m(x2_2,x4)    + m(x3_2,x3) + m(x8,x8_19)        + ((m(x7_2,x9_19)) << 1);
+        let z7 = m(x0_2,x7) + m(x1_2,x6)    + m(x2_2,x5)    + m(x3_2,x4)                      +   ((m(x8,x9_19)) << 1);
+        let z8 = m(x0_2,x8) + m(x1_2,x7_2)  + m(x2_2,x6)    + m(x3_2,x5_2) + m(x4,x4)         +   ((m(x9,x9_19)) << 1);
+        let z9 = m(x0_2,x9) + m(x1_2,x8)    + m(x2_2,x7)    + m(x3_2,x6) + m(x4_2,x5);
 
 
         let low__p37 = u64x4::splat(0x3ffffed << 37);
@@ -675,7 +678,7 @@ mod test {
 
     #[test]
     fn scale_by_curve_constants() {
-        let mut x = FieldElement2625x4::splat(&FieldElement51::one());
+        let mut x = FieldElement2625x4::splat(&FieldElement51::ONE);
 
         x = x * (121666, 121666, 2*121666, 2*121665);
 

src/backend/vector/neon/edwards.rs

diff --git a/src/backend/vector/neon/edwards.rs b/src/backend/vector/neon/edwards.rs
index 8073c3a..357d4e3 100644
--- a/src/backend/vector/neon/edwards.rs
+++ b/src/backend/vector/neon/edwards.rs
@@ -321,14 +321,14 @@ mod test {
     use super::*;
 
     fn serial_add(P: edwards::EdwardsPoint, Q: edwards::EdwardsPoint) -> edwards::EdwardsPoint {
-        use backend::serial::u64::field::FieldElement51;
+        use crate::backend::serial::u64::field::FieldElement51;
 
         let (X1, Y1, Z1, T1) = (P.X, P.Y, P.Z, P.T);
         let (X2, Y2, Z2, T2) = (Q.X, Q.Y, Q.Z, Q.T);
 
         macro_rules! print_var {
             ($x:ident) => {
-                println!("{} = {:?}", stringify!($x), $x.to_bytes());
+                println!("{} = {:?}", stringify!($x), $x.as_bytes());
             };
         }
 
@@ -411,8 +411,8 @@ mod test {
 
     #[test]
     fn vector_addition_vs_serial_addition_vs_edwards_extendedpoint() {
-        use constants;
-        use scalar::Scalar;
+        use crate::constants;
+        use crate::scalar::Scalar;
 
         println!("Testing id +- id");
         let P = edwards::EdwardsPoint::identity();
@@ -440,7 +440,7 @@ mod test {
 
         macro_rules! print_var {
             ($x:ident) => {
-                println!("{} = {:?}", stringify!($x), $x.to_bytes());
+                println!("{} = {:?}", stringify!($x), $x.as_bytes());
             };
         }
 
@@ -498,8 +498,8 @@ mod test {
 
     #[test]
     fn vector_doubling_vs_serial_doubling_vs_edwards_extendedpoint() {
-        use constants;
-        use scalar::Scalar;
+        use crate::constants;
+        use crate::scalar::Scalar;
 
         println!("Testing [2]id");
         let P = edwards::EdwardsPoint::identity();
@@ -516,8 +516,8 @@ mod test {
 
     #[test]
     fn basepoint_odd_lookup_table_verify() {
-        use constants;
-        use backend::vector::neon::constants::{BASEPOINT_ODD_LOOKUP_TABLE};
+        use crate::constants;
+        use crate::backend::vector::neon::constants::{BASEPOINT_ODD_LOOKUP_TABLE};
 
         let basepoint_odd_table = NafLookupTable8::<CachedPoint>::from(&constants::ED25519_BASEPOINT_POINT);
         println!("Testing basepoint table");

clippy is still whining but I'll send a .patch later - we are using the 2021 edition

$ env RUSTFLAGS='-C target_feature=+neon --cfg curve25519_dalek_backend="simd"' cargo +nightly build --features rand_core

$ env RUSTFLAGS='-C target_feature=+neon --cfg curve25519_dalek_backend="simd"' cargo +nightly clippy --features rand_core

$ env RUSTFLAGS='-C target_feature=+neon --cfg curve25519_dalek_backend="simd"' cargo +nightly bench --features rand_core

[rubdos]

I've done the rebase together with your patches, and squashed the NEON work into the first commit. I had not tested the rebase on 4.0, sorry about that. I'll do that now.

[rubdos]

Working on the rustfmt attributes now.

[pinkforest]

started a benchmark repo here -

I ran some benchmarks earlier (will re-run now) on Mac M1 Max Aarch64 comparing between backends FWIW

I'm going to re-organise it to cover wider range benchmarks between various impl's not just dalek's

[rubdos]

I'll see whether I'll have all the data in that format. I currently only retain the raw outputs from Criterion for further automatic processing for our paper.

[rubdos]

This took a bit longer than anticipated. We got a publication out, so academia-wise, we're all happy now. In between writing that paper and getting it out somewhere, there was quite some time, but I'm trying to sum up the current remaining issues with this branch in this post.

Results are especially spectacular on ARMv7/32-bit ARM CPUs (~50% speedups). Results are still great on ARMv8/aarch64 (~20%) until ARMv8.2-ish: when run in A64 mode, we get slow-downs relative to base. This is because ARM decided to nerf the TBL/TBX instructions relative to previous versions. The throughput remained the same, but the total latency of the instruction is now a function of the input size.

(shamelessly stolen from https://www.stonybrook.edu/commcms/ookami/support/_docs/A64FX_Microarchitecture_Manual_en_1.3.pdf)

We'll require a bit of in-depth looking into the assembly to figure out a way around this.
The TBL/TBX instruction comes from various vqtbx1q_u8 calls in FieldElement2625x4::blend. @Tarinn made a note that we can probably use a combination of vset/vget to get blend faster, but we don't currently have the hardware to reproduce this.

We sent some executables to @direc85 over in Finland to get the benchmark run on an Xperia 10 III, which is how we noted the issue. That's not really a tight debuggable loop. I should soon™ have an Xperia 10 IV to close the loop a bit.

[rubdos]

(that's a disfunctional rebase, we still have to reintegrate in the new backend-selection)

[pinkforest]

Awesome -

For CI there is ARM already for cross - .github/workflows/cross.yml#L19

We also have simd/x86_64-AVX related tests in .github/worksflows/curve25519-dalek.yml#L118

But I think we could split backend tests into a separate file - qemu should support neon at this level - backends.yml for the sake of clarity and discoverability.

I'm re-working the build script so I could include a placeholder separately there so it doesn't cause conflict.

• curve25519: Remove build-dep:platforms in favor of CARGO_CFG_TARGET_POINTER_WIDTH #636

Also nightly CI is currently failing - I've fixed it here:

• Fix new nightly redundant import lint warns #638

[jrose-signal]

Today I idly tested this on my M1 Mac and the current state of the branch is a regression from the u64 backend. Changing the implementation of blend to use core::simd::Mask::select (with constant masks) got it back to about par with the u64 backend. I wasn't being super scientific with this (I was doing other things while the benchmark ran, for example, though nothing CPU-intensive), but it was a big enough gap to be indicative. There's probably something better than Mask, but just confirming that the TBL slowdown is serious. 😭

[rubdos]

There's probably something better than Mask, but just confirming that the TBL slowdown is serious. 😭

Yep, it's quite devastating indeed... If we've got news on that front, we'll post our findings as soon as we have them. Thanks for testing it out! ❤️