From 36bebeacda0b68c24c35a981a4d34bf0f99757d1 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 4 Sep 2026 08:29:24 +0000 Subject: [PATCH] Add an ndarray-simd-backed leaky_relu kernel as an additional x86_64 candidate linalg had no element-wise kernel built on the ndarray fork's simd::F32x16 polyfill instead of hand intrinsics/asm, so there was no real measurement of that path against the existing hand-tuned AVX-512 leaky_relu kernel. Adds x86_64_ndarray_leaky_relu_f32_16n via the same routine_ew_rust! registration the other x86_64 activation kernels use, declared at X86_64Avx2 (the tier ndarray's F32x16 actually compiles to in a non-avx512f-target build), wires it into the existing leaky_relu_frame_tests! correctness harness, and adds it to the leaky_relu bench alongside the existing AVX-512 asm kernel and the generic scalar reference. The existing kernel's ISA level (avx512f) keeps it the selected default: nothing about kernel selection changes. The bench file itself had bit-rotted (Tensor::as_slice_mut no longer exists, an aarch64-only #[target_feature(enable = "fp16")] variant, and 16-byte tensor alignment too narrow for the AVX-512 kernels' vmovaps) and needed a minimal repair to build and run at all; fixed alongside adding the new kernel calls. --- Cargo.lock | 1 + linalg/Cargo.toml | 6 ++++++ linalg/benches/leaky_relu.rs | 28 ++++++++++++++++------------ linalg/src/x86_64/act.rs | 36 ++++++++++++++++++++++++++++++++++++ 4 files changed, 59 insertions(+), 12 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0b85cbeb3c..abde11c3ee 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5005,6 +5005,7 @@ dependencies = [ "libc", "log", "minijinja", + "ndarray", "nu-ansi-term", "num-traits", "pastey 0.2.3", diff --git a/linalg/Cargo.toml b/linalg/Cargo.toml index ef4f3d92b3..810dc9343d 100644 --- a/linalg/Cargo.toml +++ b/linalg/Cargo.toml @@ -33,6 +33,12 @@ tract-data.workspace = true [target.'cfg(target_arch = "riscv64")'.dependencies] libc.workspace = true +# Pilot: one x86_64 element-wise kernel (leaky_relu) calls the AdaWorldAPI ndarray +# fork's `simd::F32x16` polyfill instead of hand-rolled intrinsics/asm, as an +# additional candidate alongside the existing hand-tuned AVX-512 kernel. +[target.'cfg(target_arch = "x86_64")'.dependencies] +ndarray.workspace = true + [build-dependencies] cc.workspace = true half.workspace = true diff --git a/linalg/benches/leaky_relu.rs b/linalg/benches/leaky_relu.rs index f0f0d37632..7093f68450 100644 --- a/linalg/benches/leaky_relu.rs +++ b/linalg/benches/leaky_relu.rs @@ -1,17 +1,18 @@ use criterion::*; use tract_data::prelude::*; +#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))] use tract_linalg::element_wise::ElementWiseKer; fn leaky_relu_f16(c: &mut Criterion) { let mut group = c.benchmark_group("leaky_relu_f16"); group.throughput(Throughput::Elements(1024)); let mut input = unsafe { Tensor::uninitialized_aligned::(&[1024], 16).unwrap() }; - let input = input.as_slice_mut::().unwrap(); + let input = unsafe { input.as_slice_mut_unchecked::() }; let alpha = f16::from_f32(0.1); group.bench_function("rust", |b| b.iter(|| rust_fp16(input, alpha))); - group.bench_function("rust_with_f16", |b| b.iter(|| unsafe { rust_with_fp16(input, alpha) })); group.bench_function("linalg", |b| b.iter(|| linalg16(input, alpha))); + #[cfg(target_arch = "aarch64")] group.bench_function("linalg-asm", |b| { b.iter(|| tract_linalg::arm64::arm64fp16_leaky_relu_f16_16n::run(input, alpha)) }); @@ -24,14 +25,6 @@ fn rust_fp16(input: &mut [f16], alpha: f16) { } } -#[target_feature(enable = "fp16")] -#[inline(never)] -unsafe fn rust_with_fp16(input: &mut [f16], alpha: f16) { - for x in input { - *x = if *x > f16::ZERO { *x } else { *x * alpha } - } -} - #[inline(never)] fn linalg16(input: &mut [f16], alpha: f16) { tract_linalg::routines::Func::LeakyRelu @@ -44,14 +37,25 @@ fn linalg16(input: &mut [f16], alpha: f16) { fn leaky_relu_f32(c: &mut Criterion) { let mut group = c.benchmark_group("leaky_relu_f32"); group.throughput(Throughput::Elements(1024)); - let mut input = unsafe { Tensor::uninitialized_aligned::(&[1024], 16).unwrap() }; - let input = input.as_slice_mut::().unwrap(); + // 64-byte aligned: the AVX-512 kernels below load/store with `vmovaps`, which + // requires zmm-width (64-byte) alignment. + let mut input = unsafe { Tensor::uninitialized_aligned::(&[1024], 64).unwrap() }; + let input = unsafe { input.as_slice_mut_unchecked::() }; let alpha = 0.1f32; group.bench_function("rust", |b| b.iter(|| rust_fp32(input, alpha))); group.bench_function("linalg", |b| b.iter(|| linalg32(input, alpha))); + #[cfg(target_arch = "aarch64")] group.bench_function("linalg-asm", |b| { b.iter(|| tract_linalg::arm64::arm64simd_leaky_relu_f32_8n::run(input, alpha)) }); + #[cfg(target_arch = "x86_64")] + group.bench_function("linalg-avx512-asm", |b| { + b.iter(|| tract_linalg::x86_64::act::x86_64_avx512_leaky_relu_f32_64n::run(input, alpha)) + }); + #[cfg(target_arch = "x86_64")] + group.bench_function("linalg-ndarray-simd", |b| { + b.iter(|| tract_linalg::x86_64::act::x86_64_ndarray_leaky_relu_f32_16n::run(input, alpha)) + }); } #[inline(never)] diff --git a/linalg/src/x86_64/act.rs b/linalg/src/x86_64/act.rs index 018ba49042..2ab83a0ee7 100644 --- a/linalg/src/x86_64/act.rs +++ b/linalg/src/x86_64/act.rs @@ -160,6 +160,42 @@ unsafe fn x86_64_avx512_leaky_relu_f32_64n_run(buf: &mut [f32], alpha: f32) { } } +// leaky_relu(x) = x > 0 ? x : alpha * x, built on `ndarray::simd::F32x16` instead of +// hand-rolled intrinsics/asm. Additional candidate alongside +// `x86_64_avx512_leaky_relu_f32_64n` above: `F32x16` resolves at ndarray's own compile +// time to a native `__m512`-backed type under `avx512f`, or to a pair of `__m256` +// halves otherwise, so this single body serves both without tract choosing between +// them itself. Declared at `X86_64Avx2` because that is what a tract build (no global +// `avx512f` target feature) actually links against. +routine_ew_rust!(x86_64; + f32, + x86_64_ndarray_leaky_relu_f32_16n, + 16, + 16, + #[inline(never)] + fn run(buf: &mut [f32], alpha: f32) { + debug_assert!(buf.len() % Self::nr() == 0); + debug_assert!(buf.as_ptr() as usize % Self::alignment_bytes() == 0); + x86_64_ndarray_leaky_relu_f32_16n_run(buf, alpha) + }, + func(LeakyRelu), + param, + isa(X86_64Avx2) +); + +#[cfg(target_arch = "x86_64")] +fn x86_64_ndarray_leaky_relu_f32_16n_run(buf: &mut [f32], alpha: f32) { + use ndarray::simd::F32x16; + let alpha_v = F32x16::splat(alpha); + let zero = F32x16::splat(0.0); + for chunk in buf.as_chunks_mut::<{ F32x16::LANES }>().0 { + let x = F32x16::from_slice(chunk.as_slice()); + let scaled = x * alpha_v; + let y = x.simd_gt(zero).select(x, scaled); + y.copy_to_slice(chunk.as_mut_slice()); + } +} + // Tanh-form GELU (pow=3) matching tract's GeluApproximate: // gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) // Composed at the kernel level (mirrors arm64): save the original x, compute