Pilot: ndarray-simd leaky_relu kernel as additional x86_64 candidate

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

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::<f16>(&[1024], 16).unwrap() };
-    let input = input.as_slice_mut::<f16>().unwrap();
+    let input = unsafe { input.as_slice_mut_unchecked::<f16>() };
     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::<f32>(&[1024], 16).unwrap() };
-    let input = input.as_slice_mut::<f32>().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::<f32>(&[1024], 64).unwrap() };
+    let input = unsafe { input.as_slice_mut_unchecked::<f32>() };
     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)]

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