Devnet4 fix-avx512 (#194)

* poseidon avx2 avx512 (#192)

Co-authored-by: Tom Wambsgans <TomWambsgans@users.noreply.github.com>

* fix avx512 (panicked on small instances) (#193)

* fix avx512 (panicked on small instances)

* add `test_aggregation`

---------

Co-authored-by: Tom Wambsgans <TomWambsgans@users.noreply.github.com>

* fmt

---------

Co-authored-by: Tom Wambsgans <TomWambsgans@users.noreply.github.com>

Author: TomWambsgans

crates/backend/koala-bear/src/monty_31/aarch64_neon/poseidon_helpers.rs

@@ -8,6 +8,9 @@ use core::mem::transmute;
 use super::exp_small;
 use crate::{FieldParameters, MontyParameters, PackedMontyField31Neon, PackedMontyParameters, RelativelyPrimePower};
 
+// Convenience alias to match the naming used for the AVX2/AVX512 helpers.
+pub(crate) use convert_to_vec_neg_form_neon as convert_to_vec_neg_form;
+
 /// A specialized representation of the Poseidon state for a width of 16.
 ///
 /// Splits the state into `s0` (undergoes S-box) and `s_hi` (undergoes only linear transforms),
@@ -79,3 +82,16 @@ where
         *val = PackedMontyField31Neon::<PMP>::from_vector(output);
     }
 }
+
+/// Applies the S-Box `x -> x^D` to a packed vector. Output is in canonical form.
+#[inline(always)]
+pub(crate) fn sbox<PMP, const D: u64>(val: PackedMontyField31Neon<PMP>) -> PackedMontyField31Neon<PMP>
+where
+    PMP: PackedMontyParameters + FieldParameters + RelativelyPrimePower<D>,
+{
+    unsafe {
+        let signed = val.to_signed_vector();
+        let out = exp_small::<PMP, D>(signed);
+        PackedMontyField31Neon::<PMP>::from_vector(out)
+    }
+}

crates/backend/koala-bear/src/monty_31/x86_64_avx2/poseidon_helpers.rs

@@ -5,10 +5,8 @@
 use core::arch::x86_64::{self, __m256i};
 use core::mem::transmute;
 
-use crate::{
-    MontyParameters, PackedMontyField31AVX2, PackedMontyParameters, apply_func_to_even_odd, packed_exp_3, packed_exp_5,
-    packed_exp_7,
-};
+use super::{apply_func_to_even_odd, packed_exp_3, packed_exp_5, packed_exp_7};
+use crate::{MontyParameters, PackedMontyField31AVX2, PackedMontyParameters};
 
 /// A specialized representation of the Poseidon state for a width of 16.
 ///
@@ -34,29 +32,6 @@ impl<PMP: PackedMontyParameters> InternalLayer16<PMP> {
     }
 }
 
-/// A specialized representation of the Poseidon state for a width of 24.
-///
-/// Same split as `InternalLayer16` but for width 24.
-#[derive(Clone, Copy)]
-#[repr(C)]
-pub struct InternalLayer24<PMP: PackedMontyParameters> {
-    pub(crate) s0: PackedMontyField31AVX2<PMP>,
-    pub(crate) s_hi: [__m256i; 23],
-}
-
-impl<PMP: PackedMontyParameters> InternalLayer24<PMP> {
-    #[inline]
-    pub(crate) unsafe fn to_packed_field_array(self) -> [PackedMontyField31AVX2<PMP>; 24] {
-        unsafe { transmute(self) }
-    }
-
-    #[inline]
-    #[must_use]
-    pub(crate) fn from_packed_field_array(vector: [PackedMontyField31AVX2<PMP>; 24]) -> Self {
-        unsafe { transmute(vector) }
-    }
-}
-
 /// Use hard coded methods to compute `x -> x^D` for the even index entries and small `D`.
 /// Inputs should be signed 32-bit integers in `[-P, ..., P]`.
 /// Outputs will also be signed integers in `(-P, ..., P)` stored in the odd indices.
@@ -94,3 +69,16 @@ pub(crate) fn add_rc_and_sbox<PMP: PackedMontyParameters, const D: u64>(
         *val = PackedMontyField31AVX2::<PMP>::from_vector(output);
     }
 }
+
+/// Applies the S-Box `x -> x^D` to a packed vector. Output is in canonical form.
+#[inline(always)]
+#[must_use]
+pub(crate) fn sbox<PMP: PackedMontyParameters, const D: u64>(
+    val: PackedMontyField31AVX2<PMP>,
+) -> PackedMontyField31AVX2<PMP> {
+    unsafe {
+        let vec = val.to_vector();
+        let out = apply_func_to_even_odd::<PMP>(vec, exp_small::<PMP, D>);
+        PackedMontyField31AVX2::<PMP>::from_vector(out)
+    }
+}

crates/backend/koala-bear/src/monty_31/x86_64_avx512/poseidon_helpers.rs

@@ -5,10 +5,8 @@
 use core::arch::x86_64::{self, __m512i};
 use core::mem::transmute;
 
-use crate::{
-    MontyParameters, PackedMontyField31AVX512, PackedMontyParameters, apply_func_to_even_odd, packed_exp_3,
-    packed_exp_5, packed_exp_7,
-};
+use super::{apply_func_to_even_odd, packed_exp_3, packed_exp_5, packed_exp_7};
+use crate::{MontyParameters, PackedMontyField31AVX512, PackedMontyParameters};
 
 /// A specialized representation of the Poseidon state for a width of 16.
 ///
@@ -34,29 +32,6 @@ impl<PMP: PackedMontyParameters> InternalLayer16<PMP> {
     }
 }
 
-/// A specialized representation of the Poseidon state for a width of 24.
-///
-/// Same split as `InternalLayer16` but for width 24.
-#[derive(Clone, Copy)]
-#[repr(C)]
-pub struct InternalLayer24<PMP: PackedMontyParameters> {
-    pub(crate) s0: PackedMontyField31AVX512<PMP>,
-    pub(crate) s_hi: [__m512i; 23],
-}
-
-impl<PMP: PackedMontyParameters> InternalLayer24<PMP> {
-    #[inline]
-    pub(crate) unsafe fn to_packed_field_array(self) -> [PackedMontyField31AVX512<PMP>; 24] {
-        unsafe { transmute(self) }
-    }
-
-    #[inline]
-    #[must_use]
-    pub(crate) fn from_packed_field_array(vector: [PackedMontyField31AVX512<PMP>; 24]) -> Self {
-        unsafe { transmute(vector) }
-    }
-}
-
 /// Use hard coded methods to compute `x -> x^D` for the even index entries and small `D`.
 /// Inputs should be signed 32-bit integers in `[-P, ..., P]`.
 /// Outputs will also be signed integers in `(-P, ..., P)` stored in the odd indices.
@@ -94,3 +69,16 @@ pub(crate) fn add_rc_and_sbox<PMP: PackedMontyParameters, const D: u64>(
         *val = PackedMontyField31AVX512::<PMP>::from_vector(output);
     }
 }
+
+/// Applies the S-Box `x -> x^D` to a packed vector. Output is in canonical form.
+#[inline(always)]
+#[must_use]
+pub(crate) fn sbox<PMP: PackedMontyParameters, const D: u64>(
+    val: PackedMontyField31AVX512<PMP>,
+) -> PackedMontyField31AVX512<PMP> {
+    unsafe {
+        let vec = val.to_vector();
+        let out = apply_func_to_even_odd::<PMP>(vec, exp_small::<PMP, D>);
+        PackedMontyField31AVX512::<PMP>::from_vector(out)
+    }
+}