Make align generic across all vector types

Author: novacrazy

crates/thermite-compensated/src/lib.rs

@@ -1424,6 +1424,14 @@ impl<V: CompensatedFloatVector> GenericVector for Compensated<V> {
         }
     }
 
+    #[inline(always)]
+    fn align<const OFFSET: usize>(self, other: Self) -> Self {
+        Self {
+            value: self.value.align::<OFFSET>(other.value),
+            error: self.error.align::<OFFSET>(other.error),
+        }
+    }
+
 
     fn map<F>(mut self, f: F) -> Self
     where

crates/thermite-dual/src/vector.rs

@@ -592,6 +592,14 @@ impl<V: DualFloatVector, const N: usize> GenericVector for Dual<V, N> {
         }
     }
 
+    #[inline(always)]
+    fn align<const OFFSET: usize>(self, other: Self) -> Self {
+        Self {
+            re: self.re.align::<OFFSET>(other.re),
+            dual: array_each!([V::ZERO; N], |j| self.dual[j].align::<OFFSET>(other.dual[j])),
+        }
+    }
+
     #[inline(always)] fn map<F>(mut self, f: F) -> Self
     where
         F: Fn(Self::Element) -> Self::Element,

crates/thermite/src/backend/macros.rs

@@ -99,6 +99,45 @@ macro_rules! impl_byte_align_alignr256 {
     };
 }
 
+/// Whole-register byte-shift `align` override (`Register::align`) for a full,
+/// unpadded 128-bit integer register with native full-width byte shifts
+/// (`bshli`/`bshri` = `pslldq`/`psrldq`), e.g. SSE2 where there is no `palignr`.
+/// Drop into an `impl Register` block.
+///
+/// `align::<OFFSET>(a, b)` is `(a >> ob) | (b << (16 - ob))` in bytes, where
+/// `ob = OFFSET * size_of::<Element>()`. The match is keyed on `ob` so each
+/// arm's shift counts are literals (stable rejects a const expr of `OFFSET` in
+/// const-generic position); `ob` const-folds to one arm. `ob > 16` means
+/// `OFFSET > LANES` (out of range) and falls back to the generic default.
+///
+/// NOTE: 128-bit only - AVX2 `_mm256_bslli/bsrli_epi128` shift per 128-bit lane.
+macro_rules! impl_byteshift_align {
+    () => {
+        fn align<const OFFSET: usize>(a: Storage<Self>, b: Storage<Self>) -> Storage<Self> {
+            match const { OFFSET * core::mem::size_of::<Self::Element>() } {
+                0  => Self::bitor(Self::bshri::<0>(a),  Self::bshli::<16>(b)),
+                1  => Self::bitor(Self::bshri::<1>(a),  Self::bshli::<15>(b)),
+                2  => Self::bitor(Self::bshri::<2>(a),  Self::bshli::<14>(b)),
+                3  => Self::bitor(Self::bshri::<3>(a),  Self::bshli::<13>(b)),
+                4  => Self::bitor(Self::bshri::<4>(a),  Self::bshli::<12>(b)),
+                5  => Self::bitor(Self::bshri::<5>(a),  Self::bshli::<11>(b)),
+                6  => Self::bitor(Self::bshri::<6>(a),  Self::bshli::<10>(b)),
+                7  => Self::bitor(Self::bshri::<7>(a),  Self::bshli::<9>(b)),
+                8  => Self::bitor(Self::bshri::<8>(a),  Self::bshli::<8>(b)),
+                9  => Self::bitor(Self::bshri::<9>(a),  Self::bshli::<7>(b)),
+                10 => Self::bitor(Self::bshri::<10>(a), Self::bshli::<6>(b)),
+                11 => Self::bitor(Self::bshri::<11>(a), Self::bshli::<5>(b)),
+                12 => Self::bitor(Self::bshri::<12>(a), Self::bshli::<4>(b)),
+                13 => Self::bitor(Self::bshri::<13>(a), Self::bshli::<3>(b)),
+                14 => Self::bitor(Self::bshri::<14>(a), Self::bshli::<2>(b)),
+                15 => Self::bitor(Self::bshri::<15>(a), Self::bshli::<1>(b)),
+                16 => Self::bitor(Self::bshri::<16>(a), Self::bshli::<0>(b)),
+                _  => Self::swizzle_const::<$crate::swizzle::AlignIndices<OFFSET, Self::Lanes>>(a, b),
+            }
+        }
+    };
+}
+
 macro_rules! impl_bit_casts {
     ($($from:ty as $to:ty => $conv:ident),* $(,)?) => {
         const _: () = {$(

crates/thermite/src/backend/x86_v1/registers/i16x8.rs

@@ -233,6 +233,8 @@ impl Register for I16x8V1 {
 
     // no `pshufb` on SSE2, so variable permutes fall back to the scalar defaults
     const HAS_PERMUTEV: bool = false;
+
+    impl_byteshift_align!();
 }
 
 #[rustfmt::skip] #[thermite_macros::inline_always]

crates/thermite/src/backend/x86_v1/registers/i32x4.rs

@@ -223,6 +223,8 @@ impl Register for I32x4V1 {
 
     // no `pshufb` on SSE2, so variable permutes fall back to the scalar defaults
     const HAS_PERMUTEV: bool = false;
+
+    impl_byteshift_align!();
 }
 
 #[rustfmt::skip] #[thermite_macros::inline_always]

crates/thermite/src/backend/x86_v1/registers/i64x2.rs

@@ -239,6 +239,8 @@ impl Register for I64x2V1 {
     }
 
     const HAS_PERMUTEV: bool = false;
+
+    impl_byteshift_align!();
 }
 
 #[thermite_macros::inline_always]

crates/thermite/src/backend/x86_v1/registers/i8x16.rs

@@ -235,6 +235,8 @@ impl Register for I8x16V1 {
 
     // no `pshufb` on SSE2, so variable permutes fall back to the scalar defaults
     const HAS_PERMUTEV: bool = false;
+
+    impl_byteshift_align!();
 }
 
 #[rustfmt::skip] #[thermite_macros::inline_always]

crates/thermite/src/backend/x86_v1/registers/u16x8.rs

@@ -198,6 +198,8 @@ impl Register for U16x8V1 {
     }
 
     const HAS_PERMUTEV: bool = false;
+
+    impl_byteshift_align!();
 }
 
 #[rustfmt::skip] #[thermite_macros::inline_always]

crates/thermite/src/backend/x86_v1/registers/u32x4.rs

@@ -209,6 +209,8 @@ impl Register for U32x4V1 {
     }
 
     const HAS_PERMUTEV: bool = false;
+
+    impl_byteshift_align!();
 }
 
 #[thermite_macros::inline_always]

crates/thermite/src/backend/x86_v1/registers/u64x2.rs

@@ -233,6 +233,8 @@ impl Register for U64x2V1 {
     }
 
     const HAS_PERMUTEV: bool = false;
+
+    impl_byteshift_align!();
 }
 
 #[rustfmt::skip] #[thermite_macros::inline_always]