philbin 1.0.1

//! Philbin detects the level of SIMD support on the current CPU and dispatches
//! to the best-supported SIMD implementation. This module provides the
//! primitives.
//!
//! The `SimdLevel` enum represents the detected SIMD level. It can only be
//! constructed through the `SimdLevel::best()` function which performs the
//! mentioned runtime CPU detection.
//!
//! Each `SimdLevel` variant stores a struct that implements the corresponding
//! trait. E.g.:
//!
//! ```text
//!      VARIANT       STRUCT    TRAIT
//!   SimdLevel::Neon   Neon    HasNeon
//!   SimdLevel::Sse2   Sse2    HasSse2
//!   SimdLevel::Avx    Avx     HasAvx
//!   SimdLevel::Avx2   Avx2    HasAvx2
//!   ...
//! ```
//!
//! Hopefully you've noticed the naming pattern. The structs are all zero-sized
//! and thus _do not exist at runtime_; they are a compile-time-only construct
//! and implement the corresponding marker trait.
//!
//! A function that wants to use e.g. Avx2 would look something like:
//!
//! ```text
//! fn aes_encrypt_round(state: Self, round_key: Self) -> Self {
//!   with_simd! {
//!     fn helper(
//!       token: impl HasAvx2,
//!       state: __m256i,
//!       round_key: __m256i
//!     ) -> __m256i {
//!       _mm256_aesenc_epi128(state, round_key)
//!     }
//!   }
//!   ... use helper() ...
//! }
//! ```
//!
//! The `with_simd!` macro (defined in this module) takes a function which takes
//! an `impl HasAvx` parameter. The caller can pass the `Avx2` (or `Avx512`)
//! struct which implements that trait. The only way to actually get an instance
//! of that struct is to call `SimdLevel::best()` and then extract the stored
//! struct instance.
//!
//! The `with_simd!` macro encapsulates the necessary `unsafe` machinery.
//!
//! All of this put together gives us a _safe_ API for calling functions
//! annotated with the required `target_features` attribute. Mistakes are caught
//! at compile-time!

use crate::arch::*;
use crate::base::block::Block;
use sealed::SimdBuild;
use std::sync::LazyLock;

// Describes the level of SIMD support.
//
// Keep in mind we "emulate" bigger SIMD vectors with smaller ones when
// necessary. For instance, if the CPU only supports 128-bit SIMD (SSE2),
// we'll use a struct holding 2x 128-bit SIMD vectors to provide a 256-bit SIMD
// abstraction to the rest of the code. This of course doesn't match the perf
// of a native 256-bit SIMD vector, but it can improve perf because of CPU
// instruction pipelining.
//
// NOTE: Everything above `Fallback` requires support for CPU AES instructions,
// because without that SIMD support is irrelevant (for AEGIS).
//
// NOTE: Every `SimdLevel` variant has full support for every _lower_ variant.
#[derive(Ord, PartialOrd, Eq, PartialEq, Copy, Clone)]
#[repr(u8)]
pub enum SimdLevel {
  // See [`highest_level`] for info on which CPU features are required for which
  // `SimdLevel` variant.

  // Apple sillicon always supports NEON + AES so it doesn't need Fallback.
  #[cfg(not(all(
    target_arch = "aarch64",
    any(target_os = "macos", target_os = "ios")
  )))]
  Fallback(Fallback) = 0, // No SIMD or AES support => pure software stack.

  #[cfg(target_arch = "aarch64")]
  Neon(Neon) = 1, // 128 bit SIMD vectors ("neon") + AES support ("aes")

  #[cfg(target_arch = "x86_64")]
  Sse2(Sse2) = 1, // 128 bit SIMD vectors ("sse2") + AES-NI ("aes")

  #[cfg(target_arch = "x86_64")]
  Avx(Avx) = 2, // 128 bit SIMD vectors ("sse2") + AES-NI ("aes") + AVX ("avx")

  #[cfg(target_arch = "x86_64")]
  Avx2(Avx2) = 3, // 256 bit SIMD vectors ("avx2") + VAES ("vaes")

  #[cfg(target_arch = "x86_64")]
  Avx512(Avx512) = 4, // 512 bit SIMD vectors ("avx512f") + VAES ("vaes")
}

cfg_select! {
  target_arch = "x86_64" => {
    fn highest_level() -> SimdLevel {
      // ["->" means "transitively implies presence of"]
      // ["..." means "and other features we don't care about"]
      //
      // aes -> sse2
      // avx -> sse2
      // vaes -> (avx2, aes) -> (sse2, ...)
      // avx512f -> (avx2, ...) -> (sse2, ...)
      //
      // Source:
      //   https://doc.rust-lang.org/reference/attributes/codegen.html#x86-or-x86_64

      use std::arch::is_x86_feature_detected;
      let has_avx512f = is_x86_feature_detected!("avx512f");

      // There are Intel CPUs that support `vaes`, but do not support `avx512f`.
      // This is confusing since the VAES full name is
      // "AVX-512 Vector AES instructions".
      // Intel's consumer CPUs with P and E cores fall into this weird bucket.
      //
      // Implementing a special AVX2 + AES (so without VAES) SIMD level wouldn't
      // bring any meaningful perf benefits over SSE2 + AES + AVX.
      let has_vaes = is_x86_feature_detected!("vaes");

      // Having a separate AVX level might seem odd at first: AEGIS uses only
      // integer operations so the 256-bit floating point operations AVX brings
      // aren't useful. So why not skip plain AVX entirely? AVX2 is necessary
      // for 256-bit integer operations.
      //
      // Yet AVX brings "non-destructive" VEX-encoded operations (instructions
      // with three operands instead of just two) which are highly beneficial to
      // AEGIS. Even though it's still 128-bit SIMD, the perf difference vs SSE2
      // is _significant_: +50% on the x2 AEGIS variants and +150-200% on the x4
      // AEGIS variants.
      let has_avx = is_x86_feature_detected!("avx");
      let has_aes = is_x86_feature_detected!("aes");

      if has_avx512f && has_vaes {
        SimdLevel::Avx512(Avx512::new())
      } else if has_vaes {
        SimdLevel::Avx2(Avx2::new())
      } else if has_avx && has_aes {
        SimdLevel::Avx(Avx::new())
      } else if has_aes {
        SimdLevel::Sse2(Sse2::new())
      } else {
        SimdLevel::Fallback(Fallback::new())
      }
    }
  }

  all(target_arch = "aarch64", any(target_os = "macos", target_os = "ios")) => {
    fn highest_level() -> SimdLevel {
      // The `is_aarch64_feature_detected` macro is documented to only work on
      // Linux and always returns `false` on other platforms. Well that just
      // won't do.
      //
      // Meanwhile, *every* Apple A-series and M-series chip (their AArch64
      // chips) supports the `aes` feature, so if the crate is being compiled
      // for `aarch64` AND (`macos` OR `ios`), we're good.
      SimdLevel::Neon(Neon::new())
    }
  }

  target_arch = "aarch64" => {
    fn highest_level() -> SimdLevel {
      // The Linux path is covered by `is_aarch64_feature_detected`.
      //
      // For reasons beyond my comprehension, Rust docs imply that `neon` must
      // be runtime detected on AArch64. This is silly since ARM *requires*
      // NEON for AArch64.[^1][^2]
      //
      // BUT the `aes` feature "implicitly enables" the `neon` feature[^3] in
      // rustc-speak, so we don't need to check for `neon`.
      //
      // [^1]: https://en.wikipedia.org/wiki/AArch64
      // [^2]: https://chromium.googlesource.com/libyuv/libyuv/+/HEAD/docs/feature_detection.md
      // [^3]: https://doc.rust-lang.org/reference/attributes/codegen.html#aarch64
      //
      // P.S. That one guy using Windows on AArch64 gets the fallback `soft`
      // path.

      use std::arch::is_aarch64_feature_detected;
      if is_aarch64_feature_detected!("aes") {
        SimdLevel::Neon(Neon::new())
      } else {
        SimdLevel::Fallback(Fallback::new())
      }
    }
  }

  _ => {
    fn highest_level() -> SimdLevel {
      SimdLevel::Fallback(Fallback::new())
    }
  }
}

impl SimdLevel {
  /// Returns the highest supported (best) SIMD level for the current CPU.
  #[inline(always)]
  pub fn best() -> SimdLevel {
    // Cache the result of our CPU inspection to avoid needless recomputation.
    static CPU_INFO: LazyLock<SimdLevel> =
      LazyLock::new(|| highest_level().env_bounded());
    *CPU_INFO
  }

  #[cfg(all(test, target_arch = "x86_64"))]
  pub fn supported<T: Simd>() -> Option<T> {
    (SimdLevel::best() >= T::LEVEL).then(T::new)
  }

  // Returns the LOWER ("worse") SimdLevel between `self` and the level named by
  // the `PHILBIN_BACKEND` env var. Note that for security reasons we only read
  // this env var when the `internal_test_only_eats_babies` feature is enabled,
  // which is off by default.
  fn env_bounded(self) -> SimdLevel {
    if cfg!(not(feature = "internal_test_only_eats_babies")) {
      // SECURITY: Prod path does NOT read any env vars.
      return self;
    }

    let backend_str = match std::env::var("PHILBIN_BACKEND") {
      Err(_) => return self,
      Ok(val) => val.to_lowercase(),
    };

    let env_level = match backend_str.as_str() {
      #[cfg(not(all(
        target_arch = "aarch64",
        any(target_os = "macos", target_os = "ios")
      )))]
      "fallback" => SimdLevel::Fallback(Fallback::new()),

      #[cfg(target_arch = "aarch64")]
      "neon" => SimdLevel::Neon(Neon::new()),

      #[cfg(target_arch = "x86_64")]
      "sse2" => SimdLevel::Sse2(Sse2::new()),

      #[cfg(target_arch = "x86_64")]
      "avx" => SimdLevel::Avx(Avx::new()),

      #[cfg(target_arch = "x86_64")]
      "avx2" => SimdLevel::Avx2(Avx2::new()),

      #[cfg(target_arch = "x86_64")]
      "avx512" => SimdLevel::Avx512(Avx512::new()),

      unknown => panic!(
        "Value for PHILBIN_BACKEND is unknown or not supported on this \
        platform: {unknown}"
      ),
    };

    // SAFETY: We must use min(runtime, env var) to prevent selecting a SIMD
    // level through the env var that the CPU doesn't actually support.
    // The only way to construct a `SimdLevel` remains `SimdLevel::best()` and
    // this function can only _reduce_ the runtime-detected SIMD support which
    // is always safe.
    self.min(env_level)
  }
}

#[doc = include_str!("../rustdoc/is_hw_accelerated.md")]
#[must_use]
pub fn is_hw_accelerated() -> bool {
  cfg_select! {
    all(target_arch = "aarch64",
        any(target_os = "macos", target_os = "ios")) => true,
    _ => SimdLevel::best() != SimdLevel::Fallback(Fallback::new())
  }
}

// SIMD structs implement this trait and thus define which Block implementations
// to use for that SIMD level.
pub trait Simd: sealed::SimdBuild + Copy {
  type Block128: Block<SelfArray = [u8; 16], Simd = Self>;
  type Block256: Block<SelfArray = [u8; 32], Simd = Self>;
  type Block512: Block<SelfArray = [u8; 64], Simd = Self>;
  // Needed to map `T: Simd` back to a `SimdLevel`.
  #[cfg(all(test, target_arch = "x86_64"))]
  const LEVEL: SimdLevel;

  // Calls the provided closure inside a function annotated
  // `#[target_features(enable = "...")]`. The closure can use intrinsics for
  // that SIMD level.
  //
  // The (internal) function wrapping the closure is guaranteed to NOT use any
  // `#[inline]` attributes. (The compiler is still free to inline the function
  // if it so chooses.) This can be used to provide the compiler with an
  // inlining "barrier function", AKA a function into which other functions are
  // inlined, but which itself is not inlined into callers because it is
  // (transitively) very large.
  fn with_target_features<F: FnOnce(Self) -> R, R>(self, func: F) -> R;
}

mod sealed {
  // Trait uses the "sealed trait" pattern to prevent wider access to `new()`
  // for structs that impl `Simd` for safety reasons. (A struct that impl's
  // `Simd` is proof of a runtime CPU feature check.)
  pub trait SimdBuild {
    fn new() -> Self;
  }
}

// Takes a string and a list of item definitions. Expands to the same item
// definitions plus the string as the item's rustdoc.
macro_rules! with_doc {
  ($doc:literal $( $it:item )+ ) => {
    $(
      #[doc = $doc]
      $it
    )+
  };
}

with_doc! { r#"
  Structs implementing this trait are considered "proof" that the current CPU
  supports that instruction set.

  Functions that want the caller to prove that the instruction set is available
  should take an `impl HasXXX` param.

  # Safety

  Should only be implemented for types whose creation is guarded by the
  corresponding runtime CPU feature check.
  "#
  #[allow(unused, reason = "We don't actually use `HasFallback`, but various
    code is simpler if it exists (avoids a special case).")]
  pub unsafe trait HasFallback {}
  #[cfg(target_arch = "x86_64")]
  pub unsafe trait HasSse2 {}
  #[cfg(target_arch = "x86_64")]
  pub unsafe trait HasAvx : HasSse2 {}
  #[cfg(target_arch = "x86_64")]
  pub unsafe trait HasAvx2 : HasAvx {}
  #[cfg(target_arch = "x86_64")]
  pub unsafe trait HasAvx512 : HasAvx2 {}
  #[cfg(target_arch = "aarch64")]
  pub unsafe trait HasNeon {}
}

// Takes a SIMD struct name and a list of marker traits it should implement,
// then produces a definition of that struct.
//
// Call example:
//
//   gen_simd_struct{ Avx2 => HasSse2, HasAvx2 }
macro_rules! gen_simd_struct {
  (
    $name:ident =>
      $block128:ident, $block256:ident, $block512:ident;
      $( $trait:ident ),+ )
  => {
    #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
    // We use a private (unnamed) field to prevent struct creation without
    // going through `new()`.
    pub struct $name(());

    impl sealed::SimdBuild for $name {
      // Only this module is allowed to create the struct.
      fn new() -> $name {
        $name(())
      }
    }

    impl Simd for $name {
      type Block128 = $block128::<Self>;
      type Block256 = $block256::<Self>;
      type Block512 = $block512::<Self>;
      #[cfg(all(test, target_arch = "x86_64"))]
      const LEVEL: SimdLevel = SimdLevel::$name($name(()));

      pastey::paste! {
        #[inline(always)]
        fn with_target_features<F: FnOnce(Self) -> R, R>(self, func: F) -> R {
          annotate_with_target_features! { $name =>
            // No #[inline] because we want to support _not_ inlining!
            fn [<helper_ $name:lower>]<F: FnOnce($name) -> R, R>(
              token: $name,
              func: F
            ) -> R
            {
              func(token)
            }
          }

          #[allow(unused_unsafe, reason = "HasFallback doesn't need unsafe")]
          // SAFETY: Safe because `helper_*` requires a $token argument of
          // $token_type and `#[target_features]` are selected by $token_type.
          // Getting a $token_type instance requires calling `SimdLevel::best()`
          // which performs a runtime CPU feature check.
          unsafe { [<helper_ $name:lower>](self, func) }
        }
      }
    }

    $(
      unsafe impl $trait for $name {}
    )+
  };
}

cfg_select! {
  target_arch = "x86_64" => {
    gen_simd_struct!{
      Sse2 => Sse1x128, Sse2x128, Sse4x128; HasSse2
    }
    gen_simd_struct!{
      // AVX is still just 128 bit (integer) SIMD so uses SSE2 Block impls. The
      // value comes from the top-level `#[target_feature(enable = "aes,avx")]`
      // because the `avx` CPU feature gives us VEX-encoded operations.
      Avx => Sse1x128, Sse2x128, Sse4x128; HasSse2, HasAvx
    }
    gen_simd_struct!{
      Avx2 => Sse1x128, Avx1x256, Avx2x256; HasSse2, HasAvx, HasAvx2
    }
    gen_simd_struct!{
      Avx512 => Sse1x128, Avx1x256, Avx1x512;
                HasSse2, HasAvx, HasAvx2, HasAvx512
    }
    gen_simd_struct!{
      Fallback => Fallback128, Fallback256, Fallback512; HasFallback
    }
  }

  all(target_arch = "aarch64", any(target_os = "macos", target_os = "ios")) => {
    gen_simd_struct!{
      Neon => Neon1x128, Neon2x128, Neon4x128; HasNeon
    }
  }

  target_arch = "aarch64" => {
    gen_simd_struct!{
      Neon => Neon1x128, Neon2x128, Neon4x128; HasNeon
    }
    gen_simd_struct!{
      Fallback => Fallback128, Fallback256, Fallback512; HasFallback
    }
  }

  _ => {
    gen_simd_struct!{
      Fallback => Fallback128, Fallback256, Fallback512; HasFallback
    }
  }
}

// Takes a SIMD level literal and a function definition.
// Expands to the provided function annotated with the correct
// `#[target_features]` annotation for that SIMD level.
//
// WARNING: Calling the annotated function *requires* a runtime check for that
// SIMD level. In this crate that means the function should take a `HasXXX`
// parameter, e.g. `HasAvx2`.
//
// Call examples:
//
//   with_target_features! {
//     Avx512 => fn foo()
//   }
//
//   with_target_features! {
//     Sse2 => fn foo()
//   }
//
// See `simd_level.rs` for the reasoning behind the feature strings.
macro_rules! annotate_with_target_features {
  (Avx512 => $func:item) => {
    #[target_feature(enable = "avx512f,vaes")]
    $func
  };

  (Avx2 => $func:item) => {
    #[target_feature(enable = "vaes")]
    $func
  };

  (Avx => $func:item) => {
    #[target_feature(enable = "aes,avx")]
    $func
  };

  (Sse2 => $func:item) => {
    #[target_feature(enable = "aes")]
    $func
  };

  (Neon => $func:item) => {
    #[target_feature(enable = "aes")]
    $func
  };

  // This will be used by the software-only fallback path. Unlikely to ever
  // be taken on x86-64, except by ancient CPUs.
  (Fallback => $func:item) => {
    #[cold] // Marks the whole call site basic block as cold
    #[inline(never)]
    $func
  };

  (HasAvx512 => $func:item) => {
    crate::arch::annotate_with_target_features! { Avx512 => $func }
  };

  (HasAvx2 => $func:item) => {
    crate::arch::annotate_with_target_features! { Avx2 => $func }
  };

  (HasAvx => $func:item) => {
    crate::arch::annotate_with_target_features! { Avx => $func }
  };

  (HasSse2 => $func:item) => {
    crate::arch::annotate_with_target_features! { Sse2 => $func }
  };

  (HasNeon => $func:item) => {
    crate::arch::annotate_with_target_features! { Neon => $func }
  };

  (HasFallback => $func:item) => {
    crate::arch::annotate_with_target_features! { Fallback => $func }
  };
}
pub(crate) use annotate_with_target_features;

// Takes a function declaration and expands to a version of the function with
// the correct `#[target_features]`. Also puts that impl behind a wrapper
// function that abstracts away the `unsafe` block, BUT the function _requires_
// that the caller provides a value implementing the correct SIMD trait. The
// first parameter of the provided function must be a by-value token.
//
// The trait must be one of: HasSse2, HasAvx2, HasAvx512, HasNeon, HasFallback.
//
// Since:
//
// - the token is used to select the `#[target_features]` attribute to emit,
// - the generated function requires passing the token,
// - creating the token is guarded by `SimdLevel`,
//
// this macro provides a _safe_ abstraction for calling a function annotated
// with `#[target_features]` from a non-target-features context.
//
// Call example:
//
//    with_simd! {
//      fn xor(
//        token: impl HasSse2,
//        first: __m128i,
//        second: __m128i
//      ) -> __m128i {
//        _mm_xor_si128(first, second)
//      }
//    }
//
// This idea is inspired by the work in the `fearless_simd` crate. Huge thanks
// to everyone working on that crate!
macro_rules! with_simd {
  (
    fn $name:ident(
      $token:ident : impl $token_type:ident,
      $( $arg:ident : $arg_type:ty ),* $(,)?
    ) -> $ret:ty
    $body:block
  ) => {
    #[inline(always)]
    fn $name(
      $token: impl $token_type,
      $($arg: $arg_type),*
    ) -> $ret {
      crate::arch::annotate_with_target_features! { $token_type =>
        #[inline] // #[inline(always)] not supported with #[target_features]
        fn __real_impl(
          $token: impl $token_type,
          $($arg: $arg_type),*
        ) -> $ret {
          let _ = $token;
          $body
        }
      }

      // SAFETY: Safe because `__real_impl` requires a $token argument of
      // $token_type and `#[target_features]` are selected by $token_type.
      // Getting a $token_type instance requires calling `SimdLevel::best()`
      // which performs a runtime CPU feature check.
      unsafe { __real_impl( $token, $($arg),* ) }
    }
  };
}
pub(crate) use with_simd;

#[cfg(test)]
mod tests {
  use super::*;

  #[test]
  fn best_simd_basic() {
    let _ = SimdLevel::best();
  }
}