Optimér ytelse: Axiom brings Numpy API to Apple Silicon

Show HN: Axiom – C++ tensor-bibliotek med NumPy-API, optimeret til Apple Silicon

Den klassiske AI-udvikler-tragedie: Du prototype’r i NumPy eller PyTorch, alt er smidigt, og så kommer “den rigtige verden” med C++ til edge deployment. Pludselig bruger du længere tid på at oversætte API’er end på at løse problemet. Projektet Axiom (vist på Hacker News) går direkte efter den smerte: et C++ tensor-bibliotek, der prøver at ligne Python-oplevelsen så meget, at omskrivningen bliver… mekanisk.

For folk på macOS er ambitionen ekstra spids: GPU på Metal via MPSGraph, så det ikke kun er matmul, der får lov at leve et lykkeligt liv, mens alt andet falder tilbage på CPU’en. Og ja, det er målrettet Apple Silicon, med den praktiske bonus at unified memory på M-serien kan gøre CPU↔GPU-flyt mere “flyt” end “kopiér”.

NumPy/PyTorch-API i C++ – uden mental oversættelse

Axioms kerneidé er enkel på papiret og svær i praksis: samme metode-navne, samme broadcasting-regler, operator overloading, dynamiske shapes og runtime dtypes. Det betyder, at et velkendt mønster fra transformer-verdenen:

scores = Q.matmul(K.transpose(-2, -1)) / math.sqrt(64)
output = scores.softmax(-1).matmul(V)

…bliver til C++ uden at du skal “oversætte din hjerne” til et andet bibliotek:

auto scores = Q.matmul(K.transpose(-2, -1)) / std::sqrt(64.0f);
auto output = scores.softmax(-1).matmul(V);

Det lyder som en lille ting, men hvis du nogensinde har forsøgt at få Eigen til at føles som tensor-arbejde (spoiler: det føler ikke det samme tilbage), så forstår du pointen. Axiom positionerer sig også imod xtensor: mindre template-eksplosion, mere runtime-typer og fejlbeskeder, der kan læses uden at tilkalde en arkæolog.

Hvad får du i “boxen”?

Projektet beskrives som cirka 28k LOC og leverer over 100 operationer: aritmetik, reduktioner, aktiveringer (relu/gelu/silu/softmax), pooling, FFT og “rigtig” lineær algebra via LAPACK (SVD, QR, Cholesky, eigendecomposition og solvere). Der er også indbygget einops-style rearrange, så du kan skrive noget á la tensor.rearrange("b h w c -> b c h w") uden at blive tvunget ud i en weekend med dimensionel origami.

På performance-siden nævner Axiom Highway SIMD på tværs af arkitekturer (NEON, AVX2/AVX-512, SSE, WASM, RISC-V). På Apple-platforme betyder det i praksis, at NEON er en borger i første klasse, hvilket klæder alt fra en MacBook Pro til en Mac mini samt nyere iPad-klasser, hvor CPU-side tensor-arbejde stadig kan være relevant.

Metal via MPSGraph: GPU til mere end bare matmul

Den mest Apple-specifikke del er GPU-stakken: MPSGraph betyder, at man lægger sig tæt op ad Apples Metal-optimerede ML-infrastruktur, som også spiller sammen med performance-logikken i nyere M4 Pro-maskiner. Axiom påstår eksplicit, at “alle ops” kan køre på GPU—ikke kun matrixmultiplikation—og at kompilerede grafer caches for at undgå genkompilering.

Tallene: M4 Pro, Eigen, PyTorch og NumPy

Axiom deler benchmarks på en M4 Pro mod Eigen (med OpenBLAS), PyTorch og NumPy. Udvalgte highlights:

• Matmul 2048×2048: 3.196 GFLOPS (Eigen 2.911 / PyTorch 2.433)
• ReLU 4096×4096: 123 GB/s (Eigen 117 / PyTorch 70)
• FFT2 2048×2048: 14,9 ms (PyTorch 27,6 ms / NumPy 63,5 ms)

Det er lovende, men også den type tal, der altid bør læses med den sunde skepsis, man normalt reserverer til “op til”-formuleringer på batteritid. Data-layout (row-major vs column-major), dtype, trådning, cache-adfærd og om du rammer et “godt” kernel-path i Metal, kan flytte nålen meget. Axiom understøtter begge memory layouts (row-major default, column-major via as_f_contiguous()), hvilket i praksis er en subtil men vigtig detalje, hvis du vil matche forventninger fra eksisterende Python-kode.

Det mest interessante her er, at Axiom prøver at gøre API-kompatibilitet til et performance-våben på Apple-platformen: Når du kan skrive “PyTorch-agtig” kode i C++ og samtidig udnytte MPSGraph på macOS, får du en ret direkte vej fra prototypen til en shipping build—uden at hoppe mellem tre biblioteker og fem forskellige semantikker. Det er sjældent, at “mindre mental overhead” faktisk også kan betyde “mindre runtime overhead”, men på Apple Silicon med unified memory er det i det mindste plausibelt.

Hvordan prøver man det?

Koden ligger på GitHub, og du kan bygge med en simpel make release eller hente det ind via CMake FetchContent. Hvis du vil nørde videre i konteksten af Apple-platforme og performance, kan du også følge emnet hos We❤️Apple.

git clone https://github.com/frikallo/axiom.git
cd axiom && make release

Min vurdering er, at Axiom rammer et reelt hul for udviklere, der elsker Python-ergonomi men hader produktions-virkeligheden: C++ på edge, stramme latenskrav og platformsspecifik acceleration. Hvis projektet kan holde API-løftet stabilt og samtidig følge med Apple’s evigt skiftende GPU/driver-virkelighed i macOS, kan det blive et af de mere pragmatiske bud på “PyTorch-feel, C++-deploy” — især hvis du arbejder tæt på Core ML-økosystemet og vil have noget, der føles hjemmevant ved siden af Xcode.

Hvis du vil læse den oprindelige tråd og reaktionerne, ligger den på Hacker News, og selve projektet på GitHub. For relaterede nyheder om GPU-accelereret udvikling på macOS kan du også holde øje med Metal-dækningen hos We❤️Apple.

Kilde: github.com/Frikallo/axiom · Diskussion: news.ycombinator.com