Merge pull request #3552 from autonomys/pot-avx512f+vaes

Faster PoT verification for CPUs that support AVX512F+VAES

Author: nazar-pc

crates/subspace-proof-of-time/src/aes.rs

@@ -15,6 +15,7 @@ pub(crate) fn create(seed: PotSeed, key: PotKey, checkpoint_iterations: u32) ->
     {
         cpufeatures::new!(has_aes, "aes");
         if has_aes::get() {
+            // SAFETY: Checked `aes` feature
             return unsafe { x86_64::create(seed.as_ref(), key.as_ref(), checkpoint_iterations) };
         }
     }
@@ -51,6 +52,47 @@ pub(crate) fn verify_sequential(
 ) -> bool {
     assert_eq!(checkpoint_iterations % 2, 0);
 
+    #[cfg(target_arch = "x86_64")]
+    {
+        cpufeatures::new!(has_avx512f_vaes, "avx512f", "vaes");
+        if has_avx512f_vaes::get() {
+            // SAFETY: Checked `avx512f` and `vaes` features
+            return unsafe {
+                x86_64::verify_sequential_avx512f_vaes(
+                    &seed,
+                    &key,
+                    checkpoints,
+                    checkpoint_iterations,
+                )
+            };
+        }
+
+        cpufeatures::new!(has_avx2_vaes, "avx2", "vaes");
+        if has_avx2_vaes::get() {
+            // SAFETY: Checked `avx2` and `vaes` features
+            return unsafe {
+                x86_64::verify_sequential_avx2_vaes(&seed, &key, checkpoints, checkpoint_iterations)
+            };
+        }
+
+        cpufeatures::new!(has_aes_sse41, "aes", "sse4.1");
+        if has_aes_sse41::get() {
+            // SAFETY: Checked `aes` and `sse4.1` features
+            return unsafe {
+                x86_64::verify_sequential_aes_sse41(&seed, &key, checkpoints, checkpoint_iterations)
+            };
+        }
+    }
+
+    verify_sequential_generic(seed, key, checkpoints, checkpoint_iterations)
+}
+
+fn verify_sequential_generic(
+    seed: PotSeed,
+    key: PotKey,
+    checkpoints: &PotCheckpoints,
+    checkpoint_iterations: u32,
+) -> bool {
     let key = Array::from(*key);
     let cipher = Aes128::new(&key);
 
@@ -94,6 +136,65 @@ mod tests {
     ];
     const BAD_CIPHER: [u8; 16] = [22; 16];
 
+    fn verify_test(
+        seed: PotSeed,
+        key: PotKey,
+        checkpoints: &PotCheckpoints,
+        checkpoint_iterations: u32,
+    ) -> bool {
+        let sequential = verify_sequential(seed, key, checkpoints, checkpoint_iterations);
+        let sequential_generic =
+            verify_sequential_generic(seed, key, checkpoints, checkpoint_iterations);
+        assert_eq!(sequential, sequential_generic);
+
+        #[cfg(target_arch = "x86_64")]
+        {
+            cpufeatures::new!(has_avx512f_vaes, "avx512f", "vaes");
+            if has_avx512f_vaes::get() {
+                // SAFETY: Checked `avx512f` and `vaes` features
+                let avx512f_vaes = unsafe {
+                    x86_64::verify_sequential_avx512f_vaes(
+                        &seed,
+                        &key,
+                        checkpoints,
+                        checkpoint_iterations,
+                    )
+                };
+                assert_eq!(sequential, avx512f_vaes);
+            }
+
+            cpufeatures::new!(has_avx2_vaes, "avx2", "vaes");
+            if has_avx2_vaes::get() {
+                // SAFETY: Checked `avx2` and `vaes` features
+                let avx2_vaes = unsafe {
+                    x86_64::verify_sequential_avx2_vaes(
+                        &seed,
+                        &key,
+                        checkpoints,
+                        checkpoint_iterations,
+                    )
+                };
+                assert_eq!(sequential, avx2_vaes);
+            }
+
+            cpufeatures::new!(has_aes_sse41, "aes", "sse4.1");
+            if has_aes_sse41::get() {
+                // SAFETY: Checked `aes` and `sse4.1` features
+                let aes = unsafe {
+                    x86_64::verify_sequential_aes_sse41(
+                        &seed,
+                        &key,
+                        checkpoints,
+                        checkpoint_iterations,
+                    )
+                };
+                assert_eq!(sequential, aes);
+            }
+        }
+
+        sequential
+    }
+
     #[test]
     fn test_create_verify() {
         let seed = PotSeed::from(SEED);
@@ -107,47 +208,42 @@ mod tests {
             assert_eq!(checkpoints, generic_checkpoints);
         }
 
-        assert!(verify_sequential(
-            seed,
-            key,
-            &checkpoints,
-            checkpoint_iterations,
-        ));
+        assert!(verify_test(seed, key, &checkpoints, checkpoint_iterations,));
 
         // Decryption of invalid cipher text fails.
         let mut checkpoints_1 = checkpoints;
         checkpoints_1[0] = PotOutput::from(BAD_CIPHER);
-        assert!(!verify_sequential(
+        assert!(!verify_test(
             seed,
             key,
             &checkpoints_1,
             checkpoint_iterations,
         ));
 
         // Decryption with wrong number of iterations fails.
-        assert!(!verify_sequential(
+        assert!(!verify_test(
             seed,
             key,
             &checkpoints,
             checkpoint_iterations + 2,
         ));
-        assert!(!verify_sequential(
+        assert!(!verify_test(
             seed,
             key,
             &checkpoints,
             checkpoint_iterations - 2,
         ));
 
         // Decryption with wrong seed fails.
-        assert!(!verify_sequential(
+        assert!(!verify_test(
             PotSeed::from(SEED_1),
             key,
             &checkpoints,
             checkpoint_iterations,
         ));
 
         // Decryption with wrong key fails.
-        assert!(!verify_sequential(
+        assert!(!verify_test(
             seed,
             PotKey::from(KEY_1),
             &checkpoints,