Skip to content

v0.2.0

Choose a tag to compare

@ehsanmok ehsanmok released this 19 Sep 03:04
· 175 commits to main since this release

The big things, and why you'd want this release

Dense linear algebra on Tensor is now most of scipy.linalg, not a third of it. 0.1 gave you cholesky, lu_factor, qr_factor and solve, blocked and running on MAX's GEMM. 0.2 adds the eigendecomposition, the SVD, the Schur form, and the matrix functions built on top of it, all with the same shape: device-resident, blocked, cubic term in linalg.matmul. If you were reaching for NumPy or SciPy to diagonalize a matrix, take an SVD, or compute expm/logm/sqrtm, that call now has a Mojo counterpart that runs on your GPU.

You can now borrow instead of copying. Every subpackage moved onto a TensorLike bound this release, so a routine written against it takes an owned Tensor or a borrowed View of a sub-block interchangeably. Factoring a 4x4 corner of a 1000x1000 matrix, or filtering a slice of a recording, no longer costs a copy first.

Thousands of small linear-algebra problems, one GPU thread each. map_blocks is new: hand it a batch of small problems in a structure-of-arrays layout and it runs one whole factorization, solve, or eigenproblem per SIMD lane or per GPU thread. examples/advanced/batched_solve.mojo runs 4096 independent 4x4 Cholesky solves this way, plus the derivative of every one of them at Dual, from one call.

scipy.signal's filter design is complete. butter, cheby1, cheby2, ellip, and the iirfilter front door over all four, plus freqz, are new. If you design filters, this closes the one part of numax.signal that wasn't here yet.

Two real bugs are fixed. maximum/minimum silently dropped NaN instead of propagating it (a NumPy-visible difference, now fixed, with fmax/fmin added as the NaN-skipping alternative). And the new spectral routines' gpu=True path, which answered wrong rather than raising, is now a compile-time error instead, so a silent wrong answer is no longer reachable.

Plus 38 new NumPy/SciPy names picked up in a full audit, and a documented reason for every name still missing (docs/parity.md).

Improvements and the detailed API surface

Spectral decompositions (new), over Tensor:

Function What it gives you
eigh, eigvalsh Symmetric eigendecomposition: ascending eigenvalues, orthonormal eigenvectors
svd, svdvals Singular value decomposition, any m >= n
eigvals General eigenvalues (re, im), Francis double-shift QR
schur Real Schur form A = Z T Z^T, the base every matrix function past expm builds on
sytrd, gebrd, hessenberg The tridiagonal, bidiagonal, and Hessenberg reductions underneath the four above, each its own TensorTridiagonal/TensorSVD-style result with a .q()

Every one takes a block size (block == n recovers the unblocked algorithm, pinned by tests) and a gpu: Bool, currently False-only: gpu=True is now a compile error on these six plus pinv/cond/matrix_rank, guarding the Metal correctness bug mentioned above rather than papering over it.

Matrix functions (new), over Tensor, on the Schur form: expm (scaling and squaring, Padé), logm (inverse scaling and squaring), sqrtm (Björck-Hammarling, handles repeated/defective eigenvalues), funm(a, f) (any FloatLike scalar function via block Parlett), cosm, sinm, fractional_matrix_power.

TensorLike and View (new): the trait bound behind owned/borrowed genericity, in numax/core/tensorlike.mojo. Every one of core, linalg, stats, signal, fft, integrate, interpolate, optimize, special and io was moved onto it this release. dim[T, i] reads a compile-time extent off either type; is_row_major[T] refuses a strided block at compile time where a routine needs to flatten.

map_blocks (new): numax/core/tensor.mojo. One small problem (a factorization, a solve, an eigenproblem) per lane instead of one scalar, structure-of-arrays so a lane's load is contiguous.

IIR filter design (new), over Tensor: butter, cheby1, cheby2, ellip (lowpass/highpass/bandpass/bandstop), iirfilter (dispatch by name), freqz. SciPy's exact route: analog prototype, frequency warping, lp2lp/lp2hp/lp2bp/lp2bs, bilinear transform.

The 38-name surface sweep, grouped by area, all built from routines already in the library:

  • Elementwise: sign, square, reciprocal, rint, degrees, radians, fmax, fmin
  • Array manipulation: flatten, dstack, atleast_3d, rot90, array_split
  • Sorting/selection: compress, partition, argpartition
  • Polynomials: polyval, polyder, polyint, roots, polyfit
  • Linear algebra: orth, null_space, polar, tanm, rq
  • Signal: csd, coherence, istft, decimate, oaconvolve, peak_prominences, zpk2tf
  • Special/integration: betaln, fixed_quad, dblquad
  • Statistics: ks_2samp, wilcoxon

Correctness fixes: maximum/minimum now propagate NaN on float dtypes (previously followed SIMD's maxNum/minNum, which silently skip it); gpu=True on any spectral routine or pinv/cond/matrix_rank is a compile error, not a wrong answer.

Faster, measured on the same machine before and after: quantile/percentile by three-way quickselect instead of a sort (1,096 ms to 151 ms at 2^24 elements); cov/corrcoef by Welford means plus one GEMM (125 ms to 4.1 ms); the FFT engine fused into one bit-reversal-plus-six-stages kernel (3x behind pocketfft on a 2^20 transform, down from 14x); lfilter/filtfilt reading through the tensor's own host mapping instead of copying into a List (filtfilt now ahead of SciPy).

Platform: bumped to max-core ==26.6 and mojo >=1.0.0,<2.0.0 (from 26.5/1.0).

Performance: NVIDIA A10G and Apple M3 Pro, and where numax wins

Every number below is measured on one processor at a time and never mixed with the other (docs/performance.md has the full harness for each).

Where numax wins outright:

  • The GEMM every factorization runs on beats the vendor path on both platforms. On the A10G (Ampere), MAX's matmul reaches 20,459 GFLOP/s against cuBLAS's FP32 path at 15,342 (PyTorch) and 14,995 (CuPy), a 33-36% lead. On the CPU side (EPYC 7R32), it now reaches 93% of OpenBLAS, up from 79% before this release's MAX bump.
  • The elementwise kernel matches hand-written CUDA and beats every ML framework's dispatch. On the A10G, a FloatLike Gaussian kernel with no device name and no dtype in its signature hits 61,013 M elem/s, within 1% of a hand-written cupy.ElementwiseKernel (60,352) and 10% ahead of torch.compile (53,670), at 83% of the card's 600 GB/s bandwidth spec. Same story on the M3 Pro's Metal GPU: 14,465 against torch.compile's 13,831, and about 3x ahead of eager PyTorch and MLX (4,807 and 4,866).
  • solve and qr_factor overtake PyTorch's MPS backend as matrices get bigger. On the M3 Pro, solve passes PyTorch MPS at n = 2048 (59.7 against 50.4 GFLOP/s) and qr_factor is 1.48x ahead at n = 512.
  • BLAS-1 beats the vendor CPU libraries outright. dot and nrm2 on Tensor beat OpenBLAS's own sdot/snrm2 by 2.5-5x on the EPYC and Accelerate's by 1.7-3.2x on the M3 Pro, because routing them through MAX's rowwise scaffolder threads them where the vendor routines don't.
  • A run of NumPy/SciPy-shaped CPU wins on the M3 Pro (against Accelerate): cov/corrcoef 3.3x ahead, welch 8.9x, norm.cdf 6.3x, interp 9.4x, filtfilt now ahead of SciPy outright.

Where numax is still behind, stated rather than hidden:

  • The new spectral decompositions ship CPU-only this release. eigh reaches 0.45 of LAPACK at n = 1024, svd 0.056; gpu=True is a compile error rather than a live path (see the correctness fix above). Closing the device path is the headline item for 0.3.
  • The factorizations' device solvers still trail cuSOLVER's fused kernels, even though the GEMM underneath wins: on the A10G, cholesky is 62.5 GFLOP/s against cuSOLVER's 595.0, lu_factor 30.6 against 282.9. The gap is a single-block, effectively serial panel kernel and its launch count, not the multiply. cuSOLVER's advantage is a parallel panel; that's the next lever.
  • quantile is 2.85x behind NumPy, because it still downloads the whole tensor before selecting on it.

What 0.2.0 represents

0.1 proved the two-axis idea (one kernel, several conformers; one tensor, any device) on a narrow slice of dense linear algebra. 0.2 is where the Tensor tier stops being three factorizations and a solve and becomes close to the whole of scipy.linalg, with genuine wins on the primitives underneath it (the GEMM, the elementwise dispatch, BLAS-1) and an honestly stated gap on the parts still catching up (the device spectral path, the factorization panels).

Requires max-core ==26.6 and mojo >=1.0.0,<2.0.0.

[dependencies]
numax = { git = "https://github.com/ehsanmok/numax.git", tag = "v0.2.0" }

Full diff: v0.1.0...v0.2.0