Add dynamic swizzles for 128-bit u8/i8 vectors

fearless_simd/src/generated/avx2.rs

@@ -967,6 +967,16 @@ impl Simd for Avx2 {
         kernel(self, a)
     }
     #[inline(always)]
+    fn swizzle_dyn_i8x16(self, a: i8x16<Self>, idxs: u8x16<Self>) -> i8x16<Self> {
+        crate::kernel!(
+            #[inline(always)]
+            fn kernel(token: Avx2, a: i8x16<Avx2>, idxs: u8x16<Avx2>) -> i8x16<Avx2> {
+                _mm_shuffle_epi8(a.into(), idxs.into()).simd_into(token)
+            }
+        );
+        kernel(self, a, idxs)
+    }
+    #[inline(always)]
     fn reinterpret_u8_i8x16(self, a: i8x16<Self>) -> u8x16<Self> {
         crate::kernel!(
             #[inline(always)]
@@ -1331,6 +1341,16 @@ impl Simd for Avx2 {
         kernel(self, a, b)
     }
     #[inline(always)]
+    fn swizzle_dyn_u8x16(self, a: u8x16<Self>, idxs: u8x16<Self>) -> u8x16<Self> {
+        crate::kernel!(
+            #[inline(always)]
+            fn kernel(token: Avx2, a: u8x16<Avx2>, idxs: u8x16<Avx2>) -> u8x16<Avx2> {
+                _mm_shuffle_epi8(a.into(), idxs.into()).simd_into(token)
+            }
+        );
+        kernel(self, a, idxs)
+    }
+    #[inline(always)]
     fn widen_u8x16(self, a: u8x16<Self>) -> u16x16<Self> {
         crate::kernel!(
             #[inline(always)]

fearless_simd/src/generated/avx512.rs

@@ -1050,6 +1050,16 @@ impl Simd for Avx512 {
         kernel(self, a)
     }
     #[inline(always)]
+    fn swizzle_dyn_i8x16(self, a: i8x16<Self>, idxs: u8x16<Self>) -> i8x16<Self> {
+        crate::kernel!(
+            #[inline(always)]
+            fn kernel(token: Avx512, a: i8x16<Avx512>, idxs: u8x16<Avx512>) -> i8x16<Avx512> {
+                _mm_shuffle_epi8(a.into(), idxs.into()).simd_into(token)
+            }
+        );
+        kernel(self, a, idxs)
+    }
+    #[inline(always)]
     fn reinterpret_u8_i8x16(self, a: i8x16<Self>) -> u8x16<Self> {
         crate::kernel!(
             #[inline(always)]
@@ -1486,6 +1496,16 @@ impl Simd for Avx512 {
         kernel(self, a, b)
     }
     #[inline(always)]
+    fn swizzle_dyn_u8x16(self, a: u8x16<Self>, idxs: u8x16<Self>) -> u8x16<Self> {
+        crate::kernel!(
+            #[inline(always)]
+            fn kernel(token: Avx512, a: u8x16<Avx512>, idxs: u8x16<Avx512>) -> u8x16<Avx512> {
+                _mm_shuffle_epi8(a.into(), idxs.into()).simd_into(token)
+            }
+        );
+        kernel(self, a, idxs)
+    }
+    #[inline(always)]
     fn widen_u8x16(self, a: u8x16<Self>) -> u16x16<Self> {
         crate::kernel!(
             #[inline(always)]

fearless_simd/src/generated/fallback.rs

@@ -1165,6 +1165,14 @@ impl Simd for Fallback {
         .simd_into(self)
     }
     #[inline(always)]
+    fn swizzle_dyn_i8x16(self, a: i8x16<Self>, idxs: u8x16<Self>) -> i8x16<Self> {
+        let mut dest = [Default::default(); 16];
+        for (i, idx) in idxs.val.0.iter().copied().enumerate() {
+            dest[i] = a.val.0.get(idx as usize).copied().unwrap_or(0);
+        }
+        dest.simd_into(self)
+    }
+    #[inline(always)]
     fn reinterpret_u8_i8x16(self, a: i8x16<Self>) -> u8x16<Self> {
         a.bitcast()
     }
@@ -1768,6 +1776,14 @@ impl Simd for Fallback {
         result.simd_into(self)
     }
     #[inline(always)]
+    fn swizzle_dyn_u8x16(self, a: u8x16<Self>, idxs: u8x16<Self>) -> u8x16<Self> {
+        let mut dest = [Default::default(); 16];
+        for (i, idx) in idxs.val.0.iter().copied().enumerate() {
+            dest[i] = a.val.0.get(idx as usize).copied().unwrap_or(0);
+        }
+        dest.simd_into(self)
+    }
+    #[inline(always)]
     fn widen_u8x16(self, a: u8x16<Self>) -> u16x16<Self> {
         [
             a[0usize] as u16,

fearless_simd/src/generated/neon.rs

@@ -910,6 +910,16 @@ impl Simd for Neon {
         kernel(self, a)
     }
     #[inline(always)]
+    fn swizzle_dyn_i8x16(self, a: i8x16<Self>, idxs: u8x16<Self>) -> i8x16<Self> {
+        crate::kernel!(
+            #[inline(always)]
+            fn kernel(token: Neon, a: i8x16<Neon>, idxs: u8x16<Neon>) -> i8x16<Neon> {
+                vqtbl1q_s8(a.into(), idxs.into()).simd_into(token)
+            }
+        );
+        kernel(self, a, idxs)
+    }
+    #[inline(always)]
     fn reinterpret_u8_i8x16(self, a: i8x16<Self>) -> u8x16<Self> {
         crate::kernel!(
             #[inline(always)]
@@ -1266,6 +1276,16 @@ impl Simd for Neon {
         }
     }
     #[inline(always)]
+    fn swizzle_dyn_u8x16(self, a: u8x16<Self>, idxs: u8x16<Self>) -> u8x16<Self> {
+        crate::kernel!(
+            #[inline(always)]
+            fn kernel(token: Neon, a: u8x16<Neon>, idxs: u8x16<Neon>) -> u8x16<Neon> {
+                vqtbl1q_u8(a.into(), idxs.into()).simd_into(token)
+            }
+        );
+        kernel(self, a, idxs)
+    }
+    #[inline(always)]
     fn widen_u8x16(self, a: u8x16<Self>) -> u16x16<Self> {
         crate::kernel!(
             #[inline(always)]

fearless_simd/src/generated/simd_trait.rs

@@ -306,6 +306,8 @@ pub trait Simd:
     fn combine_i8x16(self, a: i8x16<Self>, b: i8x16<Self>) -> i8x32<Self>;
     #[doc = "Negate each element of the vector, wrapping on overflow."]
     fn neg_i8x16(self, a: i8x16<Self>) -> i8x16<Self>;
+    #[doc = "Swizzle a vector of elements according to an index vector.\n\nThe behavior for out-of-bound indices is unspecified and backend-dependent, but it is guaranteed to not panic."]
+    fn swizzle_dyn_i8x16(self, a: i8x16<Self>, idxs: u8x16<Self>) -> i8x16<Self>;
     #[doc = "Reinterpret the bits of this vector as a vector of `u8` elements.\n\nThe total bit width is preserved; the number of elements changes accordingly."]
     fn reinterpret_u8_i8x16(self, a: i8x16<Self>) -> u8x16<Self>;
     #[doc = "Reinterpret the bits of this vector as a vector of `u32` elements.\n\nThe total bit width is preserved; the number of elements changes accordingly."]
@@ -388,6 +390,8 @@ pub trait Simd:
     fn max_u8x16(self, a: u8x16<Self>, b: u8x16<Self>) -> u8x16<Self>;
     #[doc = "Combine two vectors into a single vector with twice the width.\n\n`a` provides the lower elements and `b` provides the upper elements."]
     fn combine_u8x16(self, a: u8x16<Self>, b: u8x16<Self>) -> u8x32<Self>;
+    #[doc = "Swizzle a vector of elements according to an index vector.\n\nThe behavior for out-of-bound indices is unspecified and backend-dependent, but it is guaranteed to not panic."]
+    fn swizzle_dyn_u8x16(self, a: u8x16<Self>, idxs: u8x16<Self>) -> u8x16<Self>;
     #[doc = "Zero-extend each element to a wider integer type.\n\nThe number of elements in the result is half that of the input."]
     fn widen_u8x16(self, a: u8x16<Self>) -> u16x16<Self>;
     #[doc = "Reinterpret the bits of this vector as a vector of `u32` elements.\n\nThe total bit width is preserved; the number of elements changes accordingly."]

fearless_simd/src/generated/simd_types.rs

@@ -5,7 +5,7 @@
 
 use crate::{
     Bytes, Select, Simd, SimdBase, SimdCvtFloat, SimdCvtTruncate, SimdFrom, SimdInto, SimdMask,
-    seal::Seal,
+    SimdSwizzleDyn, seal::Seal,
 };
 #[doc = "A SIMD vector of 4 [`f32`] elements.\n\nYou may construct this vector type using the [`Self::splat`], [`Self::from_slice`], [`Self::simd_from`], [`Self::from_fn`], and [`Self::block_splat`] methods.\n\n```rust\n# use fearless_simd::{prelude::*, f32x4};\nfn construct_simd<S: Simd>(simd: S) {\n    // From a single scalar value:\n    let a = f32x4::splat(simd, 1.0);\n    let b = f32x4::simd_from(simd, 1.0);\n\n    // From a slice:\n    let c = f32x4::from_slice(simd, &[1.0, 2.0, 3.0, 4.0]);\n\n    // From an array:\n    let d = f32x4::simd_from(simd, [1.0, 2.0, 3.0, 4.0]);\n\n    // From an element-wise function:\n    let e = f32x4::from_fn(simd, |i| i as f32);\n}\n```"]
 #[derive(Clone, Copy)]
@@ -445,6 +445,13 @@ impl<S: Simd> crate::SimdCombine<S> for i8x16<S> {
         self.simd.combine_i8x16(self, rhs.simd_into(self.simd))
     }
 }
+impl<S: Simd> SimdSwizzleDyn<S> for i8x16<S> {
+    #[doc = "Swizzle a vector of elements according to an index vector.\n\nThe behavior for out-of-bound indices is unspecified and backend-dependent, but it is guaranteed to not panic."]
+    #[inline(always)]
+    fn swizzle_dyn(self, idxs: u8x16<S>) -> Self {
+        self.simd.swizzle_dyn_i8x16(self, idxs)
+    }
+}
 #[doc = "A SIMD vector of 16 [`u8`] elements.\n\nYou may construct this vector type using the [`Self::splat`], [`Self::from_slice`], [`Self::simd_from`], [`Self::from_fn`], and [`Self::block_splat`] methods.\n\n```rust\n# use fearless_simd::{prelude::*, u8x16};\nfn construct_simd<S: Simd>(simd: S) {\n    // From a single scalar value:\n    let a = u8x16::splat(simd, 1);\n    let b = u8x16::simd_from(simd, 1);\n\n    // From a slice:\n    let c = u8x16::from_slice(simd, &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]);\n\n    // From an array:\n    let d = u8x16::simd_from(simd, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]);\n\n    // From an element-wise function:\n    let e = u8x16::from_fn(simd, |i| i as u8);\n}\n```"]
 #[derive(Clone, Copy)]
 #[repr(C, align(16))]
@@ -630,6 +637,13 @@ impl<S: Simd> crate::SimdCombine<S> for u8x16<S> {
         self.simd.combine_u8x16(self, rhs.simd_into(self.simd))
     }
 }
+impl<S: Simd> SimdSwizzleDyn<S> for u8x16<S> {
+    #[doc = "Swizzle a vector of elements according to an index vector.\n\nThe behavior for out-of-bound indices is unspecified and backend-dependent, but it is guaranteed to not panic."]
+    #[inline(always)]
+    fn swizzle_dyn(self, idxs: u8x16<S>) -> Self {
+        self.simd.swizzle_dyn_u8x16(self, idxs)
+    }
+}
 #[doc = "A SIMD mask of 16 logical lanes corresponding to 8-bit vector elements.\n\nThe storage representation of this type is intentionally opaque and may vary depending on the SIMD level.\n\nYou can construct this mask type using the [`Self::splat`], [`Self::from_bitmask`], [`Self::from_slice`], and [`Self::simd_from`] methods.\n\n```rust\n# use fearless_simd::{prelude::*, mask8x16};\nfn construct_mask<S: Simd>(simd: S) {\n    // From a single boolean value:\n    let a = mask8x16::splat(simd, true);\n    let b = mask8x16::simd_from(simd, true);\n\n    // From signed integer mask lanes:\n    let c = mask8x16::from_slice(simd, &[-1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]);\n    let d = mask8x16::simd_from(simd, [-1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]);\n\n    // From a compact bitmask (same mask as above, least significant bit maps to lane 0):\n    let e = mask8x16::from_bitmask(simd, 0b0001);\n\n    // By setting individual lanes:\n    let mut f = mask8x16::splat(simd, false);\n    f.set(0, true);\n}\n```"]
 #[derive(Clone, Copy)]
 pub struct mask8x16<S: Simd> {

fearless_simd/src/generated/sse4_2.rs

@@ -975,6 +975,16 @@ impl Simd for Sse4_2 {
         kernel(self, a)
     }
     #[inline(always)]
+    fn swizzle_dyn_i8x16(self, a: i8x16<Self>, idxs: u8x16<Self>) -> i8x16<Self> {
+        crate::kernel!(
+            #[inline(always)]
+            fn kernel(token: Sse4_2, a: i8x16<Sse4_2>, idxs: u8x16<Sse4_2>) -> i8x16<Sse4_2> {
+                _mm_shuffle_epi8(a.into(), idxs.into()).simd_into(token)
+            }
+        );
+        kernel(self, a, idxs)
+    }
+    #[inline(always)]
     fn reinterpret_u8_i8x16(self, a: i8x16<Self>) -> u8x16<Self> {
         crate::kernel!(
             #[inline(always)]
@@ -1336,6 +1346,16 @@ impl Simd for Sse4_2 {
         }
     }
     #[inline(always)]
+    fn swizzle_dyn_u8x16(self, a: u8x16<Self>, idxs: u8x16<Self>) -> u8x16<Self> {
+        crate::kernel!(
+            #[inline(always)]
+            fn kernel(token: Sse4_2, a: u8x16<Sse4_2>, idxs: u8x16<Sse4_2>) -> u8x16<Sse4_2> {
+                _mm_shuffle_epi8(a.into(), idxs.into()).simd_into(token)
+            }
+        );
+        kernel(self, a, idxs)
+    }
+    #[inline(always)]
     fn widen_u8x16(self, a: u8x16<Self>) -> u16x16<Self> {
         crate::kernel!(
             #[inline(always)]

fearless_simd/src/generated/wasm.rs

@@ -585,6 +585,10 @@ impl Simd for WasmSimd128 {
         i8x16_neg(a.into()).simd_into(self)
     }
     #[inline(always)]
+    fn swizzle_dyn_i8x16(self, a: i8x16<Self>, idxs: u8x16<Self>) -> i8x16<Self> {
+        i8x16_swizzle(a.into(), idxs.into()).simd_into(self)
+    }
+    #[inline(always)]
     fn reinterpret_u8_i8x16(self, a: i8x16<Self>) -> u8x16<Self> {
         <v128>::from(a).simd_into(self)
     }
@@ -794,6 +798,10 @@ impl Simd for WasmSimd128 {
         }
     }
     #[inline(always)]
+    fn swizzle_dyn_u8x16(self, a: u8x16<Self>, idxs: u8x16<Self>) -> u8x16<Self> {
+        i8x16_swizzle(a.into(), idxs.into()).simd_into(self)
+    }
+    #[inline(always)]
     fn widen_u8x16(self, a: u8x16<Self>) -> u16x16<Self> {
         let low = u16x8_extend_low_u8x16(a.into());
         let high = u16x8_extend_high_u8x16(a.into());

fearless_simd/src/traits.rs

@@ -156,6 +156,11 @@ pub trait SimdCvtFloat<T: Seal>: Seal {
     fn float_from(x: T) -> Self;
 }
 
+/// Dynamic swizzle for byte vectors.
+pub trait SimdSwizzleDyn<S: Simd>: SimdBase<S> + Seal {
+    fn swizzle_dyn(self, idxs: crate::u8x16<S>) -> Self;
+}
+
 /// Concatenation of two SIMD vectors.
 ///
 /// This is implemented on all vectors 256 bits and lower, producing vectors of up to 512 bits.

fearless_simd_gen/src/generic.rs

@@ -229,7 +229,7 @@ pub(crate) fn generic_op(op: &Op, ty: &VecType) -> TokenStream {
                 }
             }
         }
-        OpSig::StoreInterleaved { .. } => {
+        OpSig::StoreInterleaved { .. } | OpSig::SwizzleDyn => {
             panic!("The generic fallback is not implemented for this operation")
         }
         OpSig::Split { .. }

fearless_simd_gen/src/mk_fallback.rs

@@ -219,6 +219,24 @@ impl Level for Fallback {
                     }
                 }
             }
+            OpSig::SwizzleDyn => {
+                assert!(
+                    matches!(vec_ty.scalar, ScalarType::Unsigned | ScalarType::Int)
+                        && vec_ty.scalar_bits == 8
+                        && vec_ty.len == 16,
+                    "swizzle_dyn is currently only supported for 8-bit 128-bit vectors"
+                );
+
+                quote! {
+                    #method_sig {
+                        let mut dest = [Default::default(); 16];
+                        for (i, idx) in idxs.val.0.iter().copied().enumerate() {
+                            dest[i] = a.val.0.get(idx as usize).copied().unwrap_or(0);
+                        }
+                        dest.simd_into(self)
+                    }
+                }
+            }
             OpSig::Shift => {
                 let items = make_list(
                     (0..vec_ty.len)

fearless_simd_gen/src/mk_neon.rs

@@ -294,6 +294,23 @@ impl Level for Neon {
                     quote! { #vbsl(#reinterpret(a.into()), b.into(), c.into()).simd_into(#token) }
                 })
             }
+            OpSig::SwizzleDyn => {
+                assert!(
+                    matches!(vec_ty.scalar, ScalarType::Unsigned | ScalarType::Int)
+                        && vec_ty.scalar_bits == 8
+                        && vec_ty.len == 16,
+                    "swizzle_dyn is currently only supported for 8-bit 128-bit vectors"
+                );
+
+                let (tbl, idxs) = match vec_ty.scalar {
+                    ScalarType::Unsigned => (quote! { vqtbl1q_u8 }, quote! { idxs.into() }),
+                    ScalarType::Int => (quote! { vqtbl1q_s8 }, quote! { idxs.into() }),
+                    _ => unreachable!(),
+                };
+                self.kernel_method(op, vec_ty, |token| {
+                    quote! { #tbl(a.into(), #idxs).simd_into(#token) }
+                })
+            }
             OpSig::Combine { combined_ty } => {
                 let combined_wrapper = combined_ty.aligned_wrapper();
                 let combined_arch_ty = self.arch_ty(&combined_ty);

fearless_simd_gen/src/mk_simd_types.rs

@@ -8,14 +8,14 @@ use crate::{
     generic::generic_op_name,
     ops::{
         F32_TO_I32, F32_TO_I32_PRECISE, F32_TO_U32, F32_TO_U32_PRECISE, I32_TO_F32, Op, OpSig,
-        TyFlavor, U32_TO_F32, vec_trait_ops_for,
+        SWIZZLE_DYN, TyFlavor, U32_TO_F32, vec_trait_ops_for,
     },
     types::{SIMD_TYPES, ScalarType, VecType},
 };
 
 pub(crate) fn mk_simd_types() -> TokenStream {
     let mut result = quote! {
-        use crate::{Bytes, Select, Simd, SimdBase, SimdFrom, SimdInto, SimdMask, SimdCvtFloat, SimdCvtTruncate, seal::Seal};
+        use crate::{Bytes, Select, Simd, SimdBase, SimdFrom, SimdInto, SimdMask, SimdCvtFloat, SimdCvtTruncate, SimdSwizzleDyn, seal::Seal};
     };
     for ty in SIMD_TYPES {
         let name = ty.rust();
@@ -214,6 +214,23 @@ pub(crate) fn mk_simd_types() -> TokenStream {
                 }
             });
         }
+        if matches!(ty.scalar, ScalarType::Unsigned | ScalarType::Int)
+            && ty.scalar_bits == 8
+            && ty.len == 16
+        {
+            let swizzle_dyn = generic_op_name("swizzle_dyn", ty);
+            let idxs_ty = VecType::new(ScalarType::Unsigned, 8, 16).rust();
+            let doc = SWIZZLE_DYN.format_docstring(TyFlavor::VecImpl);
+            conditional_impls.push(quote! {
+                impl<S: Simd> SimdSwizzleDyn<S> for #name<S> {
+                    #[doc = #doc]
+                    #[inline(always)]
+                    fn swizzle_dyn(self, idxs: #idxs_ty<S>) -> Self {
+                        self.simd.#swizzle_dyn(self, idxs)
+                    }
+                }
+            });
+        }
         result.extend(quote! {
             #[doc = #doc]
             #[derive(Clone, Copy)]

fearless_simd_gen/src/mk_wasm.rs

@@ -330,6 +330,20 @@ impl Level for WasmSimd128 {
                     }
                 }
             }
+            OpSig::SwizzleDyn => {
+                assert!(
+                    matches!(vec_ty.scalar, ScalarType::Unsigned | ScalarType::Int)
+                        && vec_ty.scalar_bits == 8
+                        && vec_ty.len == 16,
+                    "swizzle_dyn is currently only supported for 8-bit 128-bit vectors"
+                );
+
+                quote! {
+                    #method_sig {
+                        i8x16_swizzle(a.into(), idxs.into()).simd_into(self)
+                    }
+                }
+            }
             OpSig::Combine { combined_ty } => generic_block_combine(method_sig, &combined_ty, 128),
             OpSig::Split { half_ty } => generic_block_split(method_sig, &half_ty, 128),
             OpSig::Zip { select_low } => {

fearless_simd_gen/src/mk_x86.rs

@@ -306,6 +306,18 @@ impl Level for X86 {
             OpSig::Zip { select_low } => self.handle_zip(op, vec_ty, select_low),
             OpSig::Unzip { select_even } => self.handle_unzip(op, vec_ty, select_even),
             OpSig::Slide { granularity } => self.handle_slide(method_sig, vec_ty, granularity),
+            OpSig::SwizzleDyn => {
+                assert!(
+                    matches!(vec_ty.scalar, ScalarType::Unsigned | ScalarType::Int)
+                        && vec_ty.scalar_bits == 8
+                        && vec_ty.len == 16,
+                    "swizzle_dyn is currently only supported for 8-bit 128-bit vectors"
+                );
+
+                self.kernel_method(op, vec_ty, |token| {
+                    quote! { _mm_shuffle_epi8(a.into(), idxs.into()).simd_into(#token) }
+                })
+            }
             OpSig::Cvt {
                 target_ty,
                 scalar_bits,