mirror of
https://github.com/oxc-project/oxc.git
synced 2026-11-04 10:29:58 +00:00
refactor(lexer): make Vec128 and Vec256 exist for all CPU types on all platforms (#28770)
Pure refactor - no change to codegen. Preparatory step for compiling all of the lexer's code on all platforms (#28772). Previously, the `simd/avx2` module was only compiled on x86_64, and `simd/neon` only on aarch64. So `Vec128<Avx2>` and `Vec256<Avx2>` only existed on x86_64, and `Vec128<Neon>` and `Vec256<Neon>` only on aarch64. Other platforms had no vector types at all. Compile both modules on all platforms. On other platforms, a vector wraps a `VecData128` / `VecData256` instead of the platform's SIMD type, which has the same size and alignment: ```rust type Vec128Inner = cfg_select! { target_arch = "x86_64" => __m128i, _ => VecData128, }; ``` There, the CPU type is uninhabited, so the vector can never be created. Add a `#[cpu_impl]` attribute macro, for each `impl` block of a CPU type's vector methods. It wraps each method's body in `cfg_select!`, so the body is only compiled on that CPU's platform, and is `unreachable!()` elsewhere. Signatures don't change, so the methods exist on every platform, and code which calls them compiles everywhere. ```rust #[cpu_impl(Avx2)] impl Vec256<Avx2> { pub fn movemask(self) -> u32 { // SAFETY: `Vec256` type guarantees platform supports AVX2 unsafe { _mm256_movemask_epi8(self.inner).cast_unsigned() } } } ``` expands to: ```rust #[cfg_attr( not(target_arch = "x86_64"), allow(dead_code, unused_variables, clippy::unused_self, clippy::needless_pass_by_value) )] impl Vec256<Avx2> { pub fn movemask(self) -> u32 { cfg_select! { target_arch = "x86_64" => { // SAFETY: `Vec256<Avx2>` type guarantees platform supports AVX2 unsafe { _mm256_movemask_epi8(self.inner).cast_unsigned() } } _ => unreachable!(), } } } ``` The `simd` tests are compiled on all platforms too.
This commit is contained in:
overlookmotel committed
1 parent
6d5be86ea1
commit
e920957d12
13 files changed
+359
-63
No files matched your search
@@ -16,6 +16,7 @@
|
||||
//! and uses [`cpu_select!`] to pick an implementation.
|
||||
//! [`#[cpu_version]`] marks each implementation with the CPU feature set it's for.
|
||||
//! [`#[cpu_multiversion]`] turns a function into one with a separate entry point for each CPU feature set.
|
||||
//! [`#[cpu_impl]`] marks methods which only work on one CPU feature set's platform.
|
||||
//!
|
||||
//! The attribute macros are in the `oxc_lexer_macros` crate.
|
||||
//!
|
||||
@@ -59,10 +60,11 @@
|
||||
//!
|
||||
//! [`#[cpu_version]`]: macro@cpu_version
|
||||
//! [`#[cpu_multiversion]`]: macro@cpu_multiversion
|
||||
//! [`#[cpu_impl]`]: macro@cpu_impl
|
||||
|
||||
use std::{convert::Infallible as Never, mem};
|
||||
|
||||
pub(crate) use oxc_lexer_macros::{cpu_multiversion, cpu_version};
|
||||
pub(crate) use oxc_lexer_macros::{cpu_impl, cpu_multiversion, cpu_version};
|
||||
|
||||
mod select_macro;
|
||||
|
||||
|
||||
@@ -2,6 +2,10 @@
|
||||
//!
|
||||
//! `Vec256<Avx2>` wraps an [`__m256i`], and `Vec128<Avx2>` wraps an [`__m128i`].
|
||||
//!
|
||||
//! On other platforms, they can never be created, and wrap a [`VecData256`] / [`VecData128`],
|
||||
//! which have the same sizes and alignments. Their method bodies are only compiled on x86_64,
|
||||
//! but code using the types can be compiled and linted on any platform.
|
||||
//!
|
||||
//! Each vector holds an [`Avx2`], so a vector can only exist if the CPU supports AVX2.
|
||||
//! That removes the need for the caller of a vector's methods to promise it does.
|
||||
//!
|
||||
@@ -22,8 +26,10 @@
|
||||
//! but every vector operation is a function call, passing the vectors via memory.
|
||||
//!
|
||||
//! [`Avx2`]: crate::cpu::Avx2
|
||||
//! [`__m128i`]: std::arch::x86_64::__m128i
|
||||
//! [`__m256i`]: std::arch::x86_64::__m256i
|
||||
//! [`__m128i`]: https://doc.rust-lang.org/std/arch/x86_64/struct.__m128i.html
|
||||
//! [`__m256i`]: https://doc.rust-lang.org/std/arch/x86_64/struct.__m256i.html
|
||||
//! [`VecData128`]: super::VecData128
|
||||
//! [`VecData256`]: super::VecData256
|
||||
|
||||
mod vec128;
|
||||
mod vec256;
|
||||
|
||||
@@ -5,17 +5,18 @@
|
||||
//! These are 128-bit operations, so the intrinsics are from SSE2 to SSE4.1,
|
||||
//! which every CPU with AVX2 supports.
|
||||
//!
|
||||
//! [`__m128i`]: std::arch::x86_64::__m128i
|
||||
//! [`__m128i`]: https://doc.rust-lang.org/std/arch/x86_64/struct.__m128i.html
|
||||
|
||||
#![expect(
|
||||
clippy::inline_always,
|
||||
reason = "wrappers must inline to get the `target_feature`s of the function calling them"
|
||||
)]
|
||||
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
use std::arch::x86_64::*;
|
||||
|
||||
use crate::{
|
||||
cpu::Avx2,
|
||||
cpu::{Avx2, cpu_impl},
|
||||
simd::{
|
||||
Vec128, Vec128Impl, VecData128,
|
||||
arrays::{IntArray64, IntArray128},
|
||||
@@ -23,19 +24,32 @@ use crate::{
|
||||
},
|
||||
};
|
||||
|
||||
/// Type which `Vec128<Avx2>` wraps.
|
||||
///
|
||||
/// On platforms other than x86_64, `Vec128<Avx2>` can never be created,
|
||||
/// and it wraps a [`VecData128`], for same size and alignment as `__m128i`.
|
||||
type Vec128Inner = cfg_select! {
|
||||
target_arch = "x86_64" => __m128i,
|
||||
_ => VecData128,
|
||||
};
|
||||
|
||||
#[cpu_impl(Avx2)]
|
||||
impl Vec128<Avx2> {
|
||||
/// Create a [`Vec128`] wrapping an [`__m128i`].
|
||||
///
|
||||
/// Internal method only. `pub(super)` so `Vec256` methods can use it.
|
||||
///
|
||||
/// [`__m128i`]: https://doc.rust-lang.org/std/arch/x86_64/struct.__m128i.html
|
||||
#[inline(always)]
|
||||
pub(super) fn new(inner: __m128i, cpu: Avx2) -> Self {
|
||||
pub(super) fn new(inner: Vec128Inner, cpu: Avx2) -> Self {
|
||||
Self { inner, cpu }
|
||||
}
|
||||
}
|
||||
|
||||
#[cpu_impl(Avx2)]
|
||||
impl Vec128Impl for Vec128<Avx2> {
|
||||
type Cpu = Avx2;
|
||||
type Inner = __m128i;
|
||||
type Inner = Vec128Inner;
|
||||
|
||||
// ----------------------------------------
|
||||
// Creation methods
|
||||
@@ -380,6 +394,7 @@ impl Vec128Impl for Vec128<Avx2> {
|
||||
}
|
||||
|
||||
/// Methods which are specific to this architecture.
|
||||
#[cpu_impl(Avx2)]
|
||||
impl Vec128<Avx2> {
|
||||
/// Gather the top bit of each byte of this [`Vec128`] into a 16-bit mask.
|
||||
///
|
||||
|
||||
@@ -5,33 +5,47 @@
|
||||
//! Some AVX2 instructions treat a 256-bit vector as 2 separate 128-bit halves
|
||||
//! (e.g. `shuffle_u8s`, `unpack_lo_u32`). The docs of those methods say so.
|
||||
//!
|
||||
//! [`__m256i`]: std::arch::x86_64::__m256i
|
||||
//! [`__m256i`]: https://doc.rust-lang.org/std/arch/x86_64/struct.__m256i.html
|
||||
|
||||
#![expect(
|
||||
clippy::inline_always,
|
||||
reason = "wrappers must inline to get the `target_feature`s of the function calling them"
|
||||
)]
|
||||
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
use std::arch::x86_64::*;
|
||||
|
||||
use crate::{
|
||||
cpu::Avx2,
|
||||
cpu::{Avx2, cpu_impl},
|
||||
simd::{Vec128, Vec256, Vec256Impl, VecData256, arrays::IntArray256, operand::Operand256},
|
||||
};
|
||||
|
||||
/// Type which `Vec256<Avx2>` wraps.
|
||||
///
|
||||
/// On platforms other than x86_64, `Vec256<Avx2>` can never be created,
|
||||
/// and it wraps a [`VecData256`], for same size and alignment as `__m256i`.
|
||||
type Vec256Inner = cfg_select! {
|
||||
target_arch = "x86_64" => __m256i,
|
||||
_ => VecData256,
|
||||
};
|
||||
|
||||
#[cpu_impl(Avx2)]
|
||||
impl Vec256<Avx2> {
|
||||
/// Create a [`Vec256`] wrapping an [`__m256i`].
|
||||
///
|
||||
/// Internal method only.
|
||||
///
|
||||
/// [`__m256i`]: https://doc.rust-lang.org/std/arch/x86_64/struct.__m256i.html
|
||||
#[inline(always)]
|
||||
fn new(inner: __m256i, cpu: Avx2) -> Self {
|
||||
fn new(inner: Vec256Inner, cpu: Avx2) -> Self {
|
||||
Self { inner, cpu }
|
||||
}
|
||||
}
|
||||
|
||||
#[cpu_impl(Avx2)]
|
||||
impl Vec256Impl for Vec256<Avx2> {
|
||||
type Cpu = Avx2;
|
||||
type Inner = __m256i;
|
||||
type Inner = Vec256Inner;
|
||||
|
||||
// ----------------------------------------
|
||||
// Creation methods
|
||||
@@ -401,6 +415,7 @@ impl Vec256Impl for Vec256<Avx2> {
|
||||
}
|
||||
|
||||
/// Methods which are specific to this architecture.
|
||||
#[cpu_impl(Avx2)]
|
||||
impl Vec256<Avx2> {
|
||||
/// Gather the top bit of each byte of this [`Vec256`] into a 32-bit mask.
|
||||
///
|
||||
|
||||
@@ -23,7 +23,6 @@ use std::{
|
||||
ptr,
|
||||
};
|
||||
|
||||
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
use super::{
|
||||
CpuSimd, Vec128Impl, Vec256Impl,
|
||||
arrays::{IntArray128, IntArray256},
|
||||
@@ -112,7 +111,6 @@ impl VecData128 {
|
||||
///
|
||||
/// [`Vec128`]: super::Vec128
|
||||
/// [`Vec128::read`]: super::Vec128::read
|
||||
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
#[inline(always)]
|
||||
pub fn read<C: CpuSimd>(&self, cpu: C) -> C::Vec128 {
|
||||
C::Vec128::read(self, cpu)
|
||||
@@ -240,7 +238,6 @@ impl VecData256 {
|
||||
///
|
||||
/// [`Vec256`]: super::Vec256
|
||||
/// [`Vec256::read`]: super::Vec256::read
|
||||
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
#[inline(always)]
|
||||
pub fn read<C: CpuSimd>(&self, cpu: C) -> C::Vec256 {
|
||||
C::Vec256::read(self, cpu)
|
||||
|
||||
@@ -8,48 +8,44 @@
|
||||
//! [`Vec128`] and [`Vec256`] are in the `vec128` and `vec256` modules, along with the
|
||||
//! [`Vec128Impl`] and [`Vec256Impl`] traits, which hold the methods which every feature set has,
|
||||
//! and the `!`, `&`, `|` and `^` operators.
|
||||
//! Each feature set implements the traits in its own module,
|
||||
//! `avx2` on x86_64 and `neon` on aarch64, along with the methods which only it has.
|
||||
//! Other platforms have no vector types.
|
||||
//! Each feature set implements the traits in its own module, `avx2` or `neon`,
|
||||
//! along with the methods which only it has.
|
||||
//!
|
||||
//! The vector types and all their methods exist on every platform, so code which uses them
|
||||
//! compiles everywhere. Each feature set's method bodies are only compiled on its own platform
|
||||
//! (`avx2` on x86_64, `neon` on aarch64). Elsewhere, the feature set's type is uninhabited,
|
||||
//! so its vectors can never be created, and its methods are `unreachable!()`.
|
||||
//! See [`cpu_impl`].
|
||||
//!
|
||||
//! The other modules:
|
||||
//!
|
||||
//! - `data` has [`VecData128`] and [`VecData256`], aligned byte arrays to read into vectors.
|
||||
//! - `arrays` has traits for arrays of integers which vectors can read without alignment.
|
||||
//! - `operand` has traits which let an operation take either a vector or a scalar.
|
||||
//!
|
||||
//! [`cpu_impl`]: crate::cpu::cpu_impl
|
||||
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
mod avx2;
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
mod neon;
|
||||
use crate::cpu::{Avx2, Neon};
|
||||
|
||||
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
mod arrays;
|
||||
#[cfg(any(test, target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
mod avx2;
|
||||
mod data;
|
||||
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
mod neon;
|
||||
mod operand;
|
||||
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
mod vec128;
|
||||
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
mod vec256;
|
||||
|
||||
#[cfg(any(test, target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
pub use data::{VecData128, VecData256};
|
||||
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
pub use vec128::Vec128;
|
||||
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
use vec128::Vec128Impl;
|
||||
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
pub use vec256::Vec256;
|
||||
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
|
||||
use vec128::Vec128Impl;
|
||||
use vec256::Vec256Impl;
|
||||
|
||||
#[cfg(all(test, any(target_arch = "x86_64", target_arch = "aarch64")))]
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
/// CPU feature set which has SIMD vector types.
|
||||
#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
|
||||
pub trait CpuSimd: Copy {
|
||||
/// [`Vec128`] for this CPU feature set.
|
||||
type Vec128: Vec128Impl<Cpu = Self>;
|
||||
@@ -57,14 +53,12 @@ pub trait CpuSimd: Copy {
|
||||
type Vec256: Vec256Impl<Cpu = Self>;
|
||||
}
|
||||
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
impl CpuSimd for crate::cpu::Avx2 {
|
||||
type Vec128 = Vec128<Self>;
|
||||
type Vec256 = Vec256<Self>;
|
||||
impl CpuSimd for Avx2 {
|
||||
type Vec128 = Vec128<Avx2>;
|
||||
type Vec256 = Vec256<Avx2>;
|
||||
}
|
||||
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
impl CpuSimd for crate::cpu::Neon {
|
||||
type Vec128 = Vec128<Self>;
|
||||
type Vec256 = Vec256<Self>;
|
||||
impl CpuSimd for Neon {
|
||||
type Vec128 = Vec128<Neon>;
|
||||
type Vec256 = Vec256<Neon>;
|
||||
}
|
||||
@@ -3,6 +3,10 @@
|
||||
//! `Vec128<Neon>` wraps a [`uint8x16_t`], and `Vec256<Neon>` wraps a [`uint8x16x2_t`].
|
||||
//! Neon registers are 128 bits, so a `Vec256<Neon>` is a pair of registers.
|
||||
//!
|
||||
//! On other platforms, they can never be created, and wrap a [`VecData128`] / pair of
|
||||
//! [`VecData128`]s, which have the same sizes and alignments. Their method bodies are only
|
||||
//! compiled on aarch64, but code using the types can be compiled and linted on any platform.
|
||||
//!
|
||||
//! Each vector holds a [`Neon`], so a vector can only exist if the CPU supports Neon.
|
||||
//! That removes the need for the caller of a vector's methods to promise it does.
|
||||
//!
|
||||
@@ -40,8 +44,9 @@
|
||||
//! [`Neon`]: crate::cpu::Neon
|
||||
//! [`Vec256::shuffle_u8s`]: super::Vec256::shuffle_u8s
|
||||
//! [`shuffle_u8s`]: super::Vec128::shuffle_u8s
|
||||
//! [`uint8x16_t`]: std::arch::aarch64::uint8x16_t
|
||||
//! [`uint8x16x2_t`]: std::arch::aarch64::uint8x16x2_t
|
||||
//! [`uint8x16_t`]: https://doc.rust-lang.org/std/arch/aarch64/struct.uint8x16_t.html
|
||||
//! [`uint8x16x2_t`]: https://doc.rust-lang.org/std/arch/aarch64/struct.uint8x16x2_t.html
|
||||
//! [`VecData128`]: super::VecData128
|
||||
|
||||
mod vec128;
|
||||
mod vec256;
|
||||
|
||||
@@ -3,17 +3,18 @@
|
||||
//! Implements [`Vec128Impl`], and has the methods which only Neon has.
|
||||
//! The `neon` module docs list the methods which behave differently from AVX2.
|
||||
//!
|
||||
//! [`uint8x16_t`]: std::arch::aarch64::uint8x16_t
|
||||
//! [`uint8x16_t`]: https://doc.rust-lang.org/std/arch/aarch64/struct.uint8x16_t.html
|
||||
|
||||
#![expect(
|
||||
clippy::inline_always,
|
||||
reason = "thin wrappers, which should compile down to just the intrinsics they call"
|
||||
)]
|
||||
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
use std::arch::aarch64::*;
|
||||
|
||||
use crate::{
|
||||
cpu::Neon,
|
||||
cpu::{Neon, cpu_impl},
|
||||
simd::{
|
||||
Vec128, Vec128Impl, VecData128,
|
||||
arrays::{IntArray64, IntArray128},
|
||||
@@ -21,19 +22,32 @@ use crate::{
|
||||
},
|
||||
};
|
||||
|
||||
/// Type which `Vec128<Neon>` wraps.
|
||||
///
|
||||
/// On platforms other than aarch64, `Vec128<Neon>` can never be created,
|
||||
/// and it wraps a [`VecData128`], for same size and alignment as `uint8x16_t`.
|
||||
type Vec128Inner = cfg_select! {
|
||||
target_arch = "aarch64" => uint8x16_t,
|
||||
_ => VecData128,
|
||||
};
|
||||
|
||||
#[cpu_impl(Neon)]
|
||||
impl Vec128<Neon> {
|
||||
/// Create a [`Vec128`] wrapping a [`uint8x16_t`].
|
||||
///
|
||||
/// Internal method only. `pub(super)` so `Vec256` methods can use it.
|
||||
///
|
||||
/// [`uint8x16_t`]: https://doc.rust-lang.org/std/arch/aarch64/struct.uint8x16_t.html
|
||||
#[inline(always)]
|
||||
pub(super) fn new(inner: uint8x16_t, cpu: Neon) -> Self {
|
||||
pub(super) fn new(inner: Vec128Inner, cpu: Neon) -> Self {
|
||||
Self { inner, cpu }
|
||||
}
|
||||
}
|
||||
|
||||
#[cpu_impl(Neon)]
|
||||
impl Vec128Impl for Vec128<Neon> {
|
||||
type Cpu = Neon;
|
||||
type Inner = uint8x16_t;
|
||||
type Inner = Vec128Inner;
|
||||
|
||||
// ----------------------------------------
|
||||
// Creation methods
|
||||
@@ -369,6 +383,7 @@ impl Vec128Impl for Vec128<Neon> {
|
||||
}
|
||||
|
||||
/// Methods which are specific to this architecture.
|
||||
#[cpu_impl(Neon)]
|
||||
impl Vec128<Neon> {
|
||||
/// Shift each of the 16 bytes of this [`Vec128`] left by `BITS` bits.
|
||||
#[cfg_attr(not(test), expect(dead_code))]
|
||||
|
||||
@@ -6,33 +6,49 @@
|
||||
//! and its bitwise methods work on each half separately.
|
||||
//! It only has the few methods which Neon code needs, far fewer than `Vec256<Avx2>`.
|
||||
//!
|
||||
//! [`uint8x16x2_t`]: std::arch::aarch64::uint8x16x2_t
|
||||
//! [`uint8x16x2_t`]: https://doc.rust-lang.org/std/arch/aarch64/struct.uint8x16x2_t.html
|
||||
|
||||
#![expect(
|
||||
clippy::inline_always,
|
||||
reason = "thin wrappers, which should compile down to just the intrinsics they call"
|
||||
)]
|
||||
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
use std::arch::aarch64::*;
|
||||
|
||||
use crate::{
|
||||
cpu::Neon,
|
||||
cpu::{Neon, cpu_impl},
|
||||
simd::{Vec128, Vec256, Vec256Impl, VecData256, arrays::IntArray256, operand::Operand256},
|
||||
};
|
||||
|
||||
/// Type which `Vec256<Neon>` wraps.
|
||||
///
|
||||
/// On platforms other than aarch64, `Vec256<Neon>` can never be created, and it wraps
|
||||
/// a pair of [`VecData128`]s, for same size and alignment as `uint8x16x2_t`.
|
||||
///
|
||||
/// [`VecData128`]: crate::simd::VecData128
|
||||
type Vec256Inner = cfg_select! {
|
||||
target_arch = "aarch64" => uint8x16x2_t,
|
||||
_ => [crate::simd::VecData128; 2],
|
||||
};
|
||||
|
||||
#[cpu_impl(Neon)]
|
||||
impl Vec256<Neon> {
|
||||
/// Create a [`Vec256`] wrapping a [`uint8x16x2_t`].
|
||||
///
|
||||
/// Internal method only.
|
||||
///
|
||||
/// [`uint8x16x2_t`]: https://doc.rust-lang.org/std/arch/aarch64/struct.uint8x16x2_t.html
|
||||
#[inline(always)]
|
||||
fn new(inner: uint8x16x2_t, cpu: Neon) -> Self {
|
||||
fn new(inner: Vec256Inner, cpu: Neon) -> Self {
|
||||
Self { inner, cpu }
|
||||
}
|
||||
}
|
||||
|
||||
#[cpu_impl(Neon)]
|
||||
impl Vec256Impl for Vec256<Neon> {
|
||||
type Cpu = Neon;
|
||||
type Inner = uint8x16x2_t;
|
||||
type Inner = Vec256Inner;
|
||||
|
||||
// ----------------------------------------
|
||||
// Creation methods
|
||||
@@ -424,6 +440,7 @@ impl Vec256Impl for Vec256<Neon> {
|
||||
}
|
||||
|
||||
/// Methods which are specific to this architecture.
|
||||
#[cpu_impl(Neon)]
|
||||
impl Vec256<Neon> {
|
||||
/// Rearrange the bytes of this [`Vec256`] into a [`Vec128`].
|
||||
///
|
||||
|
||||
@@ -10,23 +10,19 @@
|
||||
use std::{array, mem::transmute};
|
||||
|
||||
use crate::{
|
||||
cpu::DynCpu,
|
||||
cpu::{Avx2, DynCpu, Neon},
|
||||
simd::{
|
||||
CpuSimd, Vec128, Vec256, VecData128, VecData256, operand::Operand128, vec128::Vec128Impl,
|
||||
vec256::Vec256Impl,
|
||||
},
|
||||
};
|
||||
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
use crate::cpu::Avx2;
|
||||
#[cfg(target_arch = "aarch64")]
|
||||
use crate::cpu::Neon;
|
||||
|
||||
/// Get this platform's CPU type with SIMD vectors, if the CPU supports it.
|
||||
pub(super) fn simd_cpu() -> Option<
|
||||
cfg_select! {
|
||||
target_arch = "x86_64" => Avx2,
|
||||
target_arch = "aarch64" => Neon,
|
||||
_ => Avx2,
|
||||
},
|
||||
> {
|
||||
match DynCpu::detect() {
|
||||
@@ -91,6 +87,16 @@ pub(super) fn map_u16(bytes: [u8; 32], f: impl Fn(u16) -> u16) -> [u8; 32] {
|
||||
out
|
||||
}
|
||||
|
||||
/// Each CPU type's vectors have the same size and alignment on every platform,
|
||||
/// including those where they wrap a [`VecData128`] or [`VecData256`] instead of a SIMD type.
|
||||
#[test]
|
||||
fn size_and_alignment() {
|
||||
assert_eq!((size_of::<Vec128<Avx2>>(), align_of::<Vec128<Avx2>>()), (16, 16));
|
||||
assert_eq!((size_of::<Vec256<Avx2>>(), align_of::<Vec256<Avx2>>()), (32, 32));
|
||||
assert_eq!((size_of::<Vec128<Neon>>(), align_of::<Vec128<Neon>>()), (16, 16));
|
||||
assert_eq!((size_of::<Vec256<Neon>>(), align_of::<Vec256<Neon>>()), (32, 16));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vec128_create() {
|
||||
let Some(cpu) = simd_cpu() else { return };
|
||||
|
||||
@@ -62,14 +62,20 @@ impl Cpu {
|
||||
quote!(crate::cpu::DynCpu::#name)
|
||||
}
|
||||
|
||||
/// The `target_arch` of the platform which supports the feature set.
|
||||
/// `None` for `Generic`, which every platform supports.
|
||||
pub fn target_arch(self) -> Option<&'static str> {
|
||||
match self {
|
||||
Cpu::Avx2 => Some("x86_64"),
|
||||
Cpu::Neon => Some("aarch64"),
|
||||
Cpu::Generic => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// `#[cfg(...)]` attribute for the platform which supports the feature set.
|
||||
/// `None` for `Generic`, which every platform supports.
|
||||
pub fn cfg_attr(self) -> Option<TokenStream> {
|
||||
let arch = match self {
|
||||
Cpu::Avx2 => "x86_64",
|
||||
Cpu::Neon => "aarch64",
|
||||
Cpu::Generic => return None,
|
||||
};
|
||||
let arch = self.target_arch()?;
|
||||
Some(quote!(#[cfg(target_arch = #arch)]))
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
//! `#[cpu_impl(Avx2)]`: an `impl` block whose methods only work on one CPU feature set's platform.
|
||||
//!
|
||||
//! See the docs on [`cpu_impl`] for what the output looks like.
|
||||
//! The examples in this file are the output for this `impl` block:
|
||||
//!
|
||||
//! ```ignore
|
||||
//! #[cpu_impl(Avx2)]
|
||||
//! impl Vec256<Avx2> {
|
||||
//! /// Gather the top bit of each byte of this [`Vec256`] into a 32-bit mask.
|
||||
//! #[inline(always)]
|
||||
//! pub fn movemask(self) -> u32 {
|
||||
//! // SAFETY: `Vec256` type guarantees platform supports AVX2
|
||||
//! unsafe { _mm256_movemask_epi8(self.inner).cast_unsigned() }
|
||||
//! }
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! [`cpu_impl`]: crate::cpu_impl
|
||||
|
||||
use proc_macro2::{Delimiter, Group, Span, TokenStream, TokenTree};
|
||||
use quote::quote;
|
||||
use syn::{Error, Ident, Result, parse2};
|
||||
|
||||
use crate::cpu::Cpu;
|
||||
|
||||
/// The expansion of `#[cpu_impl(Avx2)]` on an `impl` block.
|
||||
///
|
||||
/// `args` is the attribute's argument: `Avx2` or `Neon`. `item` is the `impl` block.
|
||||
///
|
||||
/// ```ignore
|
||||
/// #[cfg_attr(
|
||||
/// not(target_arch = "x86_64"),
|
||||
/// allow(dead_code, unused_variables, clippy::unused_self, clippy::needless_pass_by_value)
|
||||
/// )]
|
||||
/// impl Vec256<Avx2> {
|
||||
/// /// Gather the top bit of each byte of this [`Vec256`] into a 32-bit mask.
|
||||
/// #[inline(always)]
|
||||
/// pub fn movemask(self) -> u32 {
|
||||
/// cfg_select! {
|
||||
/// target_arch = "x86_64" => {
|
||||
/// // SAFETY: `Vec256` type guarantees platform supports AVX2
|
||||
/// unsafe { _mm256_movemask_epi8(self.inner).cast_unsigned() }
|
||||
/// }
|
||||
/// _ => unreachable!(),
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
pub fn expand(args: TokenStream, item: TokenStream) -> Result<TokenStream> {
|
||||
let arch = parse_arch(args)?;
|
||||
|
||||
let mut head = item.into_iter().collect::<Vec<_>>();
|
||||
let body = match head.pop() {
|
||||
Some(TokenTree::Group(body)) if body.delimiter() == Delimiter::Brace && is_impl(&head) => {
|
||||
body
|
||||
}
|
||||
_ => return Err(Error::new(Span::call_site(), "`cpu_impl` must be on an `impl` block")),
|
||||
};
|
||||
|
||||
let mut items = Group::new(Delimiter::Brace, items(body.stream(), arch));
|
||||
items.set_span(body.span());
|
||||
|
||||
Ok(quote! {
|
||||
#[cfg_attr(
|
||||
not(target_arch = #arch),
|
||||
allow(dead_code, unused_variables, clippy::unused_self, clippy::needless_pass_by_value)
|
||||
)]
|
||||
#(#head)* #items
|
||||
})
|
||||
}
|
||||
|
||||
/// The `target_arch` of the feature set in the macro's arguments,
|
||||
/// e.g. `"x86_64"` for `#[cpu_impl(Avx2)]`.
|
||||
fn parse_arch(args: TokenStream) -> Result<&'static str> {
|
||||
let cpu = parse2::<Ident>(args.clone()).ok().and_then(|ident| Cpu::from_ident(&ident));
|
||||
cpu.and_then(Cpu::target_arch)
|
||||
.ok_or_else(|| Error::new_spanned(args, "`cpu_impl` takes `Avx2` or `Neon`"))
|
||||
}
|
||||
|
||||
/// Check if `head`, the tokens before an item's `{ ... }`, start an `impl` block.
|
||||
/// That's the `impl` keyword, after any attributes.
|
||||
fn is_impl(head: &[TokenTree]) -> bool {
|
||||
let mut tokens = head.iter();
|
||||
loop {
|
||||
match tokens.next() {
|
||||
// An attribute is `#` followed by `[...]`
|
||||
Some(TokenTree::Punct(punct)) if punct.as_char() == '#' => {
|
||||
tokens.next();
|
||||
}
|
||||
Some(TokenTree::Ident(ident)) => return ident == "impl",
|
||||
_ => return false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The items of the `impl` block, with each method's body replaced by [`arch_only_body`].
|
||||
///
|
||||
/// The tokens are scanned, not parsed. A method's body is the first `{ ... }` after `fn`.
|
||||
/// A `;` ends an item which has no body, e.g. `type Inner = __m256i;`.
|
||||
/// All other tokens are output unchanged.
|
||||
fn items(tokens: TokenStream, arch: &str) -> TokenStream {
|
||||
let mut in_fn = false;
|
||||
tokens
|
||||
.into_iter()
|
||||
.map(|token| {
|
||||
match &token {
|
||||
TokenTree::Ident(ident) if ident == "fn" => {
|
||||
in_fn = true;
|
||||
}
|
||||
TokenTree::Punct(punct) if punct.as_char() == ';' => {
|
||||
in_fn = false;
|
||||
}
|
||||
TokenTree::Group(body) if in_fn && body.delimiter() == Delimiter::Brace => {
|
||||
in_fn = false;
|
||||
return TokenTree::Group(arch_only_body(body, arch));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
token
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// A method body which only compiles `body` on `arch`, and is `unreachable!()` elsewhere.
|
||||
///
|
||||
/// ```ignore
|
||||
/// {
|
||||
/// cfg_select! {
|
||||
/// target_arch = "x86_64" => {
|
||||
/// // SAFETY: `Vec256` type guarantees platform supports AVX2
|
||||
/// unsafe { _mm256_movemask_epi8(self.inner).cast_unsigned() }
|
||||
/// }
|
||||
/// _ => unreachable!(),
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
fn arch_only_body(body: &Group, arch: &str) -> Group {
|
||||
let tokens = quote! {
|
||||
cfg_select! {
|
||||
target_arch = #arch => #body
|
||||
_ => unreachable!(),
|
||||
}
|
||||
};
|
||||
let mut new_body = Group::new(Delimiter::Brace, tokens);
|
||||
new_body.set_span(body.span());
|
||||
new_body
|
||||
}
|
||||
@@ -8,6 +8,8 @@
|
||||
//!
|
||||
//! - [`cpu_multiversion`]: One function, compiled once for each feature set.
|
||||
//! - [`cpu_version`]: One version of some code, for one feature set.
|
||||
//! - [`cpu_impl`]: Methods of one feature set's types, which exist on every platform,
|
||||
//! but only work on the platform which has the feature set.
|
||||
//!
|
||||
//! `oxc_lexer`'s `cpu_select!` macro chooses between versions.
|
||||
//!
|
||||
@@ -22,6 +24,7 @@ use syn::{Error, Result};
|
||||
|
||||
mod cpu;
|
||||
mod func;
|
||||
mod impl_block;
|
||||
mod multiversion;
|
||||
mod version;
|
||||
|
||||
@@ -324,6 +327,74 @@ pub fn cpu_version(args: TokenStream, item: TokenStream) -> TokenStream {
|
||||
output(version::expand(args.into(), item.into()))
|
||||
}
|
||||
|
||||
/// Marks an `impl` block as methods of one CPU feature set's types,
|
||||
/// which only work on the platform which has that feature set.
|
||||
///
|
||||
/// Applied to an `impl` block with one of:
|
||||
///
|
||||
/// - `#[cpu_impl(Avx2)]`: The method bodies are only compiled on x86_64.
|
||||
/// - `#[cpu_impl(Neon)]`: The method bodies are only compiled on aarch64.
|
||||
///
|
||||
/// On other platforms, each method's body is replaced with `unreachable!()`.
|
||||
/// Its signature is unchanged, so the method exists on every platform,
|
||||
/// and code which calls it compiles everywhere.
|
||||
///
|
||||
/// Every method must take a value which proves that the CPU supports the feature set,
|
||||
/// e.g. `self` of a `Vec256<Avx2>`, or a `cpu: Avx2` parameter. On other platforms, the feature
|
||||
/// set's type is uninhabited, so that value can't exist, and the method can never be called.
|
||||
/// The macro doesn't check this. A method which doesn't take one would panic if called
|
||||
/// on another platform.
|
||||
///
|
||||
/// ```ignore
|
||||
/// #[cpu_impl(Avx2)]
|
||||
/// impl Vec256<Avx2> {
|
||||
/// /// Gather the top bit of each byte of this [`Vec256`] into a 32-bit mask.
|
||||
/// #[inline(always)]
|
||||
/// pub fn movemask(self) -> u32 {
|
||||
/// // SAFETY: `Vec256` type guarantees platform supports AVX2
|
||||
/// unsafe { _mm256_movemask_epi8(self.inner).cast_unsigned() }
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// expands to:
|
||||
///
|
||||
/// ```ignore
|
||||
/// #[cfg_attr(
|
||||
/// not(target_arch = "x86_64"),
|
||||
/// allow(dead_code, unused_variables, clippy::unused_self, clippy::needless_pass_by_value)
|
||||
/// )]
|
||||
/// impl Vec256<Avx2> {
|
||||
/// /// Gather the top bit of each byte of this [`Vec256`] into a 32-bit mask.
|
||||
/// #[inline(always)]
|
||||
/// pub fn movemask(self) -> u32 {
|
||||
/// cfg_select! {
|
||||
/// target_arch = "x86_64" => {
|
||||
/// // SAFETY: `Vec256` type guarantees platform supports AVX2
|
||||
/// unsafe { _mm256_movemask_epi8(self.inner).cast_unsigned() }
|
||||
/// }
|
||||
/// _ => unreachable!(),
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// On other platforms, nothing calls the methods, and their bodies don't use their parameters.
|
||||
/// The `allow` silences the lints for that. It's `allow`, not `expect`, because which of them fire
|
||||
/// depends on the methods. e.g. methods of a trait implementation are never dead code,
|
||||
/// and clippy doesn't lint their parameters.
|
||||
///
|
||||
/// Only method bodies change. Everything else is output unchanged, so it must compile on every
|
||||
/// platform. That includes the types in the methods' signatures, and other items in the `impl`
|
||||
/// block (e.g. `type Inner = __m256i;`). `__m256i` only exists on x86_64.
|
||||
///
|
||||
/// The tokens of the `impl` block are scanned, not parsed, so the macro is quick to compile.
|
||||
/// A method's body is the first `{ ... }` after `fn`.
|
||||
#[proc_macro_attribute]
|
||||
pub fn cpu_impl(args: TokenStream, item: TokenStream) -> TokenStream {
|
||||
output(impl_block::expand(args.into(), item.into()))
|
||||
}
|
||||
|
||||
/// The macro's output, or the error as a `compile_error!`.
|
||||
fn output(result: Result<TokenStream2>) -> TokenStream {
|
||||
result.unwrap_or_else(Error::into_compile_error).into()
|
||||
|
||||
Reference in new issue
Block a user