Skip to content

Pilot: ndarray-simd leaky_relu kernel as additional x86_64 candidate - #3

Merged
AdaWorldAPI merged 1 commit into
mainfrom
claude/ndarray-simd-pilot-85232d2e
Sep 4, 2026
Merged

AdaWorldAPI merged 1 commit into
mainfrom
claude/ndarray-simd-pilot-85232d2e

Conversation

@AdaWorldAPI

Copy link
Copy Markdown
Owner

Summary

Staged pilot for migrating tract-linalg's intrinsics/asm element-wise kernels onto the AdaWorldAPI ndarray fork's simd::F32x16 polyfill (see #2). This PR touches exactly one op, additively.

Op picked: leaky_relu (f32). It was chosen over erf.rs (explicitly out of scope — hand-unrolled 4x zmm FMA chains, a bad first pilot) because its existing AVX-512 kernel (x86_64_avx512_leaky_relu_f32_64n in linalg/src/x86_64/act.rs) is short, branch-free (a compare mask + blend), and has no GEMM/microkernel entanglement — a clean one-op comparison.

What's here

  • New kernel: x86_64_ndarray_leaky_relu_f32_16n, built on ndarray::simd::F32x16 (splat/from_slice/simd_gt/select/copy_to_slice) instead of the existing kernel's raw std::arch::asm! AVX-512 body. Registered via the same routine_ew_rust! macro every neighboring x86_64 activation kernel uses, declared isa(X86_64Avx2) — the tier ndarray::simd::F32x16 actually compiles to in a build without a global avx512f target feature (tract doesn't set one).
  • Purely additive: the existing x86_64_avx512_leaky_relu_f32_64n kernel is untouched. Kernel selection is ISA-level-based (isa.rs), and X86_64Avx512f outranks X86_64Avx2, so the old kernel remains the winner on any AVX-512 machine — confirmed by running the dispatcher directly (see below).
  • Correctness: wired into the existing leaky_relu_frame_tests! macro (property + trivial cases against the generic scalar reference) the same way every other kernel in this family is — no new test harness.
  • Benchmark: added the new kernel and the existing AVX-512 asm kernel to benches/leaky_relu.rs, following the existing criterion pattern in that file (which also needed a small pre-existing bit-rot fix — see below).
  • Added ndarray as an x86_64-only dependency of tract-linalg (it wasn't a dependency of this crate before).

Benchmark: real numbers, run on this sandbox (x86_64, AVX-512F present)

leaky_relu_f32/rust                (naive scalar loop)       5.22 µs   196 Melem/s
leaky_relu_f32/linalg              (auto-dispatch)           2.73 µs   375 Melem/s
leaky_relu_f32/linalg-avx512-asm   (existing kernel, direct) 2.73 µs   375 Melem/s
leaky_relu_f32/linalg-ndarray-simd (this PR's kernel, direct) 10.5 µs   97 Melem/s

linalg (auto-dispatch) matching linalg-avx512-asm exactly confirms the existing hand-tuned kernel is still the one selected by default — this PR changes nothing about that.

The ndarray-polyfill kernel is ~3.8x slower than the hand-tuned AVX-512 asm here, and slower than even the naive scalar loop. Plausible reasons (not chased further — this is a pilot, not an optimization pass): this build has no global avx512f target feature, so ndarray::simd::F32x16 compiles to its AVX2-composed-of-two-__m256 fallback rather than a native __m512; the new kernel processes 16 lanes/iteration via unaligned from_slice/copy_to_slice with no unrolling, vs. the existing kernel's 4x-unrolled 64-wide masked-store body.

This is a staged pilot, not a migration

  • Scope is one op, on purpose.
  • The new kernel is additive — nothing is removed, replaced, or repointed. The old kernel remains the default/winning selection.
  • No cutover of default dispatch preference should happen on this evidence. The numbers above say the opposite: on this hardware, the hand-tuned kernel wins by a wide margin. Any future decision to prefer the ndarray-backed path needs real benchmark numbers on real target hardware, reviewed deliberately — not implied by this PR landing.

🍍

Co-Authored-By: Claude Sonnet 5 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_012wrzeZAdwGYTCKoxamwQht

🤖 Generated with Claude Code


Generated by Claude Code

…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.
@cursor

cursor Bot commented Sep 4, 2026

Copy link
Copy Markdown

Bugbot couldn't run - usage limit reached

Bugbot is counted against Cursor usage for this user or team, and this run hit a usage or spend limit.

A user or team admin can review and increase usage limits in the Cursor dashboard.

(requestId: serverGenReqId_61eff607-44cc-44c6-8703-71d8a9b36cf2)

Copy link
Copy Markdown
Owner Author

Correction to the benchmark in the PR description. It was run without any target-cpu flag — this sandbox has no .cargo/config.toml of its own (the v3-baseline config I was reading from was actually the ndarray fork's, from a stale cd earlier in the session). Rustc's plain default baseline has neither AVX2 nor AVX-512 enabled, so ndarray::simd::F32x16 (which selects its backend via compile-time cfg(target_feature = "avx512f"), not a runtime probe) never got a fair shot — it fell all the way back to a lane-by-lane path.

This box is Sapphire-Rapids-class (avx512f/bw/cd/dq/vl/vnni/bf16/fp16, plus amx_tile/int8/bf16), matching the ndarray fork's own .cargo/config-avx512.toml target. Re-running with RUSTFLAGS="-C target-cpu=native":

leaky_relu_f32/linalg-avx512-asm    (existing kernel, direct)  2.45 µs   417 Melem/s
leaky_relu_f32/linalg-ndarray-simd  (this PR's kernel, direct) 2.43 µs   422 Melem/s

With the correct build flags, the ndarray-simd kernel is essentially on par with the hand-tuned AVX-512 asm (very slightly ahead, within noise) — not 3.8x slower as the original PR description said. That number was a build-configuration artifact, not a fair result.

I'll leave the PR description's numbers uncorrected in place (for the record) and treat this comment as the accurate result. Given this, the earlier "no cutover justified" verdict doesn't hold as strongly — the ndarray-backed kernel is a legitimate, competitive candidate on real target hardware. Still not proposing a cutover from this one pilot alone (one op, one shape, one machine) — but the evidence now points toward "worth expanding the pilot to more ops" rather than "the polyfill underperforms."


Generated by Claude Code

@AdaWorldAPI
AdaWorldAPI marked this pull request as ready for review September 4, 2026 12:43
@AdaWorldAPI
AdaWorldAPI merged commit e16e8a4 into main Sep 4, 2026
55 of 66 checks passed
@cursor

cursor Bot commented Sep 4, 2026

Copy link
Copy Markdown

Bugbot couldn't run - usage limit reached

Bugbot is counted against Cursor usage for this user or team, and this run hit a usage or spend limit.

A user or team admin can review and increase usage limits in the Cursor dashboard.

(requestId: serverGenReqId_fb9494e7-5637-4a52-94b1-e5d9a2fa1a47)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants