Feature request: shuffle methods

[RReverser]

Looks like LLVM autovectorizer struggles with turning array operations into shuffles, so direct control would be very welcome. This would be particularly useful for things like bitonic sorting.

[Lokathor]

I'd love to have this. I think the shuffle pattern has to be a const value on x86, which might make it awkward, but even so it's a good idea.

If someone does the PR I'll try to review it as quickly as possible.

[RReverser]

I think the shuffle pattern has to be a const value on x86

Yeah, IIRC stdlib intrinsics solve this by taking such params (in shuffle and other ops with similar restriction) via const generics.

Same should work here.

[mcroomp]

What stable intrinsics would this translate to? Sadly most of the < 32 bit shuffle intrinsics are kind of special case

[Lokathor]

presumably these ones: https://docs.rs/safe_arch/latest/safe_arch/?search=shuffle

[ErikMcClure]

I attempted to implement this, but unless I am mistaken, a unified API is currently completely impossible without using a macro until some form of const expressions is stabilized. This is because the safe_arch API and the u32x4_shuffle API are incompatible. You must either specify all lanes separately and combine them at compile time:

if #[cfg(target_feature="sse")] {
  Self { sse: shuffle_abi_f32_all_m128::<{I0 | I1 << 2 | I2 << 4 | I3 << 6}>(self.sse, b.sse) }
} else if #[cfg(target_feature="simd128")] {
  Self {
    simd: u32x4_shuffle::<I0, I1, I2, I3>(self.simd, b.simd)
  }
}

Or specify them all at once and separate them at compile time:

if #[cfg(target_feature="sse")] {
  Self { sse: shuffle_abi_f32_all_m128::<I>(self.sse, b.sse) }
} else if #[cfg(target_feature="simd128")] {
  Self {
    simd: u32x4_shuffle::<{I0 & 0b11}, {(I1 >> 2) & 0b11}, {(I2 >> 4) & 0b11} {(I3 >> 6) & 0b11}>(self.simd, b.simd)
  }
} 

Both of which require operations on const generics, and because rust does not have bitfields, there is no way to "sneak around" the problem, as crates like bitfield-struct generate bitwise operations in the accessors.

If this API were to be supported now, before some form of const expressions is stabilized, it would require either a divergent API (change the function signature based on the target architecture, which kind of ruins the point of the abstraction), or a macro would have to be involved. I considered one option where a macro was only used to generate the indices, but because of an edge-case in the macro rules, this isn't possible, because you must wrap the macro in {} in one case but not the other. Even if you avoid wrapping it in {}, generating a blank list of elements currently isn't considered a valid syntactic result from a macro:

#[cfg(target_feature = "sse")]
pub fn shuffle<const I: i32>(self, b: Self) -> Self { }

#[cfg(not(target_feature = "sse"))]
pub fn shuffle<const I0: usize, const I1: usize, const I2: usize, const I3: usize>(self, b: Self) { }

#[cfg(target_feature = "sse")]
macro_rules! shuffle_idx {
    ($($x:literal),+) => {{ $crate::calc_shuffle_idx(&[$($x),*]) as i32 }};
}

#[cfg(not(target_feature = "sse"))]
macro_rules! shuffle_idx {
    ($($idx:literal),+) => { $(idx),+ }; // not valid
}

fn impl_f32x4_shuffle() {
  let a = f32x4::from([1.0, 2.0, 3.0, 4.0]);
  let b = f32x4::from([5.0, 6.0, 7.0, 8.0]);
  // Still requires divergent implementations, failing to unify the API
  #[cfg(target_feature = "sse")]
  a.shuffle::<{ shuffle_idx!(0, 0, 0, 0) }>(b);
  #[cfg(not(target_feature = "sse"))]
  // Doesn't even compile because a bare list of literals is not a valid macro result
  a.shuffle::<shuffle_idx!(0, 0, 0, 0)>(b);
}

This means you have to use a macro to call the function:

#[cfg(target_feature = "sse")]
pub const fn calc_shuffle_idx(idx: &[usize]) -> usize {
  calc_shuffle_idx_inner(idx.len().ilog2(), 0, idx)
}

#[macro_export]
#[cfg(target_feature = "sse")]
macro_rules! f32x4_shuffle {
    ([$($x:literal),+], $a:expr, $b:expr) => { $a.shuffle::<{$crate::calc_shuffle_idx(&[$($x),*]) as i32 }>($b) };
}

#[macro_export]
#[cfg(not(target_feature = "sse"))]
macro_rules! f32x4_shuffle {
    ([$($idx:literal),+], $a:expr, $b:expr) => {
      $a.shuffle::<$($idx),*>($b)
    };
}

impl f32x4 {
  #[must_use]
  #[inline]
  #[cfg(target_feature = "sse")]
  pub fn shuffle<const I: i32>(self, b: Self) -> Self {
    Self { sse: shuffle_abi_f32_all_m128::<I>(self.sse, b.sse) }
  }
  
  #[must_use]
  #[inline]
  #[cfg(not(target_feature = "sse"))]
  pub fn shuffle<
    const I0: usize,
    const I1: usize,
    const I2: usize,
    const I3: usize,
  >(
    self,
    b: Self,
  ) -> Self {
    pick! {
      if #[cfg(target_feature="simd128")] {
        Self {
          simd: u32x4_shuffle::<I0, I1, I2, I3>(self.simd, b.simd)
        }
      } else {
        Self { arr: [
          self.arr[I0],
          self.arr[I1],
          b.arr[I2],
          b.arr[I3],
        ]}
      }
    }
  }
}

Note that I didn't attempt implementing ARM NEON shuffles because it requires specializations to determine which instruction to use.

It might seem as though you could avoid the macro if we could simply have a constant function that returned a tuple for #[cfg(not(target_feature="sse"))]. Unfortunately, this isn't possible, because in order to extract an element out of a tuple, you need an expression. Because the function only returns a single value as far as the parser is concerned, there is no way to switch between the single-parameter generic signature and the multi-parameter generic without the macro.

The version presented here leaves the shuffle method as a publicly accessible function on f32x4 in case the user wants to interact with it, but you could also put it in a separate module. I tested this example here with f32x4 for SSE and WASM, and it seems to work:

#[test]
fn impl_f32x4_shuffle() {
  let a = f32x4::from([1.0, 2.0, 3.0, 4.0]);
  let b = f32x4::from([5.0, 6.0, 7.0, 8.0]);
  assert_eq!(
    f32x4_shuffle!([0, 0, 0, 0], a, b),
    f32x4::new([1.0, 1.0, 5.0, 5.0])
  );
  assert_eq!(
    f32x4_shuffle!([0, 1, 2, 3], a, b),
    f32x4::new([1.0, 2.0, 7.0, 8.0])
  );
  assert_eq!(
    f32x4_shuffle!([1, 2, 2, 0], a, b),
    f32x4::new([2.0, 3.0, 7.0, 5.0])
  );
}

Because adding this to all the different types would be a pain in the ass, I will only create a PR using this approach if you approve of this direction ahead of time. Maybe someone else has a better idea, but I couldn't find one.

[Lokathor]

(sorry about the delay in reply) I think it's better we save ourselves the pain in the ass and just wait for Rust to add the necessary language feature.