Skip to content

plan-ndarray.md

Latest commit

 

History

History
181 lines (166 loc) · 11 KB

plan-ndarray.md

File metadata and controls

181 lines (166 loc) · 11 KB
Raw

Implementation plan — N-dimensional arrays in pure Go (go-ndarray/ndarray)

Goal: a pure-Go (CGO=0) NumPy-style N-dimensional array library, as a standalone, reusable module with correct scalar kernels now and a clean path to SIMD acceleration across all six 64-bit Go targets.

1. Why this module?

The Go scientific stack has a gap. gonum is excellent but its hand-written assembly kernels are amd64-only (other arches fall back to generic Go), and it is matrix/mat-centric rather than a general N-d array. The wider opportunity is cross-language: Ruby has no cgo-free ndarray — Numo::NArray, NMatrix and friends are all C extensions. A pure-Go core with SIMD kernels generated for every arch is therefore a durable moat, and a natural backend for go-embedded-ruby.

So the design separates structure (shape, strides, broadcasting, views) from numeric inner loops (the kernels), and keeps the kernels behind a narrow contiguous-slice API so accelerated variants can drop in unchanged.

2. Architecture

ndarray/
  ndarray.go            public API: Array (data/shape/strides/offset),
                        constructors, introspection, reshape/transpose/copy,
                        broadcasting binary ops, map/unary, whole-array and
                        per-axis reductions
  slicing.go            NumPy basic-indexing Slice(...Index) -> strided views
  creation.go           Linspace, Eye/Identity (Reshape -1 lives in ndarray.go)
  ufunc.go              unary math ufuncs + broadcasting comparison/min-max ops
  manipulation.go       Flatten/Squeeze/ExpandDims/Concatenate/Stack/V+HStack
  linalg.go             MatMul/Dot/Inner/Outer
  internal/kernels/     portable scalar inner loops over flat []float64
                        (Add/Sub/Mul/Div, Map, comparisons, Maximum/Minimum,
                        Sum/Prod/Max/Min, Abs, {Sum,Prod,Max,Min}Axis over
                        [outer][axisLen][inner], MatMul GEMM, Dot1D)
  docs/plan-ndarray.md  this roadmap

Array is row-major (C-order): a flat data []float64 described by shape, per-axis strides (in elements) and a base offset. Reshape and Transpose are zero-copy views where possible; broadcasting materialises operands to the broadcast shape before invoking a kernel.

3. Kernels and SIMD (the accelerator path)

Every numeric loop lives in internal/kernels as a pure-Go scalar function over contiguous []float64. These are the reference and the fallback.

Phase 1 replaces them with kernels generated by go-asmgen — Plan 9 assembly emitted for all six 64-bit Go SIMD targets: amd64, arm64, riscv64, loong64, ppc64le (VSX), s390x (vector, big-endian) — selected at runtime via build tags / CPU-feature detection, behind the same kernel signatures and the same tests. Because the public API and the test suite address the kernels through that narrow interface, no caller changes and correctness stays pinned by the existing 100%-coverage suite (plus the per-arch CI jobs already wired in .github/workflows/ci.yml).

4. Phasing

  • Phase 0 — DONE. float64 dtype. Array with data/shape/strides/offset, row-major. Constructors (New, Zeros, Ones, Full, Arange, FromData). Introspection (Shape, Ndim, Size, At, Set, String). Reshape, Ravel, Transpose, Copy. Elementwise Add/Sub/Mul/Div with full NumPy broadcasting, plus *Scalar variants. Map, Neg, Abs. Whole-array reductions Sum/Prod/Max/Min/Mean. Scalar pure-Go kernels in internal/kernels. 100% statement coverage, CGO=0, six-arch CI.
  • Axis reductions — DONE. Per-axis SumAxis/ProdAxis/MaxAxis/MinAxis/ MeanAxis(axis, keepdims) with NumPy semantics (negative axis from the end, keepdims retains the reduced axis as length 1). The reduced data is viewed as [outer][axisLen][inner] and reduced by {Sum,Prod,Max,Min}Axis kernels whose innermost loop is unit-stride contiguous (SIMD-ready). Differentially checked against numpy 2.4.6; committed tests carry hardcoded expectations.
  • Slicing / views — DONE. Slice(...Index) with NumPy basic-indexing semantics: integer indices (A) drop their axis; range indices (All/R/Rng/From/To/Step) keep the axis as a strided view that shares the backing data (write-through both ways). Negative bounds count from the end, unset bounds default to the natural extreme for the step sign, bounds clamp into range, negative steps reverse. Matches numpy 2.2.4.
  • More creation — DONE. Linspace(start,stop,num) (endpoint inclusive, final sample pinned to stop), Eye(n,m,k) / Identity(n) (k-th diagonal, non-positive m defaults to square), and Reshape -1 dimension inference.
  • Ufuncs — DONE. Unary math (Sqrt/Exp/Log/Log2/Log10/Sin/Cos/ Tan/Floor/Ceil/Round/Square/Power): Sqrt has a SIMD kernel, Exp/Log/Log10 run ports of Arm's optimized-routines (v0.2.0), and the rest go through the Map seam over Go's math; and broadcasting comparisons returning 0/1 float masks (Equal/NotEqual/ Greater/GreaterEqual/Less/LessEqual) plus pairwise Maximum/Minimum. (Round is math.Round / half-away-from-zero, unlike numpy's banker's rounding — documented.) A bool dtype for true masks is deferred to Phase 2.
  • Manipulation — DONE. Flatten (always copy), ExpandDims/Squeeze (length-1 axis views), Concatenate (existing axis), Stack (new axis), and the VStack/HStack conveniences with 1-D promotion. Negative axes throughout. Matches numpy 2.2.4.
  • Linear algebra — DONE (scalar). MatMul (2-D GEMM, ikj order, unit-stride inner loop), Dot (numpy.dot across 1-D/2-D: inner product, matmul, matrix-vector, vector-matrix), Inner (sum over last axis) and Outer. The GEMM/dot kernels live in internal/kernels for the Phase 1 SIMD path. Matches numpy 2.2.4. Blocked/transpose-aware and decomposition variants remain in Phase 4.
  • Performance — DONE (beats NumPy on its core ops; matmul at tuned-BLAS parity). Several levers behind the kernel seam, measured honestly vs NumPy 2.2.4 (see docs/perf.md and the repo's BENCHMARKS.md):
    • Multicore. Add/Sub/Mul/Div/Map, Sum/Max/Min, the axis reductions, Concatenate/Stack, and the GEMM fan out across GOMAXPROCS above a size threshold (small arrays stay serial). Plus an allocation fix: the elementwise/reduce fast path no longer copies broadcast operands for the same-shape contiguous case. Result on large arrays: Add/Mul ~2×, Sum ~2.4×, and the packed GEMM reaches parity with single-threaded vecLib at 1024² (~373 GFLOP/s on an Apple M4 Max) while beating the pure-Go gonum 4–10× at every size. (A many-core join panic in the parallel GEMM was fixed in the 2026-06-23 pass.)
    • SIMD via go-asmgen — DONE for amd64 (SSE2) + arm64 (NEON). Hand-vectorized sum (4-accumulator), sqrt (packed SQRTPD / NEON FSQRT V.2D), max/min (MAXPD/MINPD + NaN scan / NEON FMAX/FMIN since v0.5.2), the elementwise add/sub/mul/div, and the GEMM micro-kernel (NEON 4×8 by-lane FMLA on arm64, AVX2/FMA 6×8 with an SSE2 fallback on amd64), generated by go-asmgen and validated per-arch in CI (bit-identical to the scalar oracle for the elementwise/sqrt/max ops; tight-tolerance for the sum reduction, whose lane-parallel grouping is a valid reordering — the same trade-off as NumPy's pairwise sum). Since v0.2.3 ppc64le has VSX kernels too (sum, dot, sqrt, elementwise, an 8×8 GEMM micro-kernel), through go-asmgen's VSX encoders, and since v0.2.5 loong64 has LASX kernels (used when the CPU reports LASX), and since v0.7.0 riscv64 has RVV kernels (sum, dot, sqrt, elementwise, and since v0.7.1 a 6×8 GEMM micro-kernel; used when the CPU reports V), and since v0.8.0 s390x has vector-facility kernels (sum, dot, sqrt, elementwise, since v0.9.0 a 4×8 GEMM micro-kernel and since v0.9.1 max/min on z14 and later; measured on a real IBM z17). The 32-bit targets keep the validated scalar oracles plus a scalar 4×4 GEMM micro-kernel over the packed panels, with packing + cache-blocking + multicore; they have not been measured against NumPy. Split-CI: the pure-Go core + multicore are held to 100% statement coverage; the generated .s is validated by the per-arch native/qemu execution jobs.
    • SIMD Sqrt/Max and the packed, cache-blocked GEMM — DONE (no longer TODO). Sqrt/Max/Min now WIN vs single-threaded NumPy, and the panel-packed, cache-blocked GEMM with its SIMD-FMA micro-kernel reaches parity with tuned BLAS (single-threaded vecLib at 1024² on M4 Max, ≈0.99× vs multi-threaded OpenBLAS at n=1024 on an arm64 VM) and beats gonum 4–10× everywhere. Remaining headroom is small-n per-call overhead and an AVX2/FMA amd64 micro-kernel.
  • More reductions — DONE. ArgMax/ArgMin (flat and per-axis with keepdims), CumSum/CumProd (per-axis scans + flat), Clip(lo,hi), and the three-operand broadcasting Where(cond,t,f). Match numpy 2.2.4.
  • Boolean / fancy indexing — DONE (on 0/1 float masks). MaskSelect (a[mask]), Nonzero (np.flatnonzero), and Take (integer fancy index into the flattened array, negative indices from the end). These give NumPy boolean indexing today using the 0/1 masks the comparison ufuncs return; a first-class bool dtype (so masks are real booleans) remains Phase 2.
  • Phase 2 — dtypes — DONE. Thirteen dtypes (bool, signed and unsigned integers, float32/64, complex64/128) as a runtime property of Array, with NumPy 2.x promotion; see design-dtypes.md. Next: SIMD kernels for float32, and the fft bridge.
  • Phase 3 — broadcasting ufuncs. A general ufunc framework over views (the axis-reduction primitives above are the first step); more elementwise and reduction ufuncs.
  • Phase 4 — linear algebra. MatMul/Dot, transpose-aware GEMM, decompositions; SIMD/blocked kernels.
  • Phase 5 — Ruby binding. Expose the array type to Ruby through go-embedded-ruby as a cgo-free alternative to Numo::NArray.

5. Decisions

  1. Row-major, strided views. Mirrors NumPy/C-order; enables zero-copy reshape/transpose. Settled.
  2. Kernels behind a contiguous-slice API. Structure and numerics are decoupled so go-asmgen SIMD can replace scalar loops without API churn. Settled.
  3. Errors, not panics, for shape/broadcast problems; panics reserved for programmer index errors (matching Go slice semantics). Settled.
  4. float64 first, dtype abstraction deferred to Phase 2 but kept in mind so the kernel interface generalises cleanly. Superseded: Phase 2 is now designed in design-dtypes.md — the dtype is a runtime property of the one Array type, with NumPy 2.x promotion.

BSD-3-Clause.