DragonArrayLoongson optimization lab

DragonArray docs

DragonArray NumPy patch set and wheels

DragonArray is an unofficial, experimental LoongArch LASX optimization patch set for NumPy on Loongson 3A6000 systems.

In this patch I completed a LASX ufunc enablement and optimization pass for NumPy 2.5.0-dev on the Loongson 3A6000. The patch wires more ufunc dispatch paths through NumPy's existing 256-bit LASX universal-SIMD backend, adds a small set of missing 256-bit primitives, fixes one backend primitive, and adds new vectorized math kernels for cases that still fell back to scalar code on this CPU.

This is an independent community patch set, not an official NumPy release.

DragonArray builds on NumPy's own LoongArch work. The bulk of the LASX npyv_* backend — arithmetic, load/store, compare, select, and most other universal-SIMD primitives — is upstream NumPy work contributed by the LoongArch / Loongson community. DragonArray does not reimplement that backend. It uses it.

Headline

1.81× faster, raw geometric mean across 135 operation × dtype pairs vs. stock NumPy 2.5.0-dev on the Loongson 3A6000 @ 2.0 GHz. Both builds use the same upstream NumPy 2.5.0-dev base revision and the same GCC 15.2.0 toolchain: one stock build, one with the DragonArray LASX patch set. Each build runs its best SIMD on the contiguous path (stock: LSX; DragonArray: LASX). This is the real speedup from installing this wheel instead of stock NumPy on the same chip — not a scalar-fallback comparison.

dtype family ops geomean speedup
float16 (LASX bridge) 20 2.77×
bool 2 8.31×
float32 27 1.61×
float64 27 1.25×
integer 59 1.86×
all 135 1.81×

See BENCHMARKS.md for the full per-op tables and OPTIMIZATIONS.md for the chronological technical log.

What DragonArray added

DragonArray's work falls into four layers:

  1. a small number of missing LASX backend primitives;
  2. one LASX backend bug fix;
  3. math and integer kernels composed on top of NPYV primitives;
  4. dispatch, build, and runtime feature-detection wiring.

New LASX primitives

DragonArray adds 8 genuinely new 256-bit primitives. Everything else at the primitive level comes from NumPy's existing LASX backend.

float16 bridge primitives

NumPy's LASX backend does not expose float16 NPYV support, so DragonArray adds a small f16 bridge in numpy/_core/src/umath/npyv_f16_bridge.h:

primitive intrinsic does
npyv_load_f16 __lasx_xvld load 16× f16, 256-bit
npyv_store_f16 __lasx_xvst store 16× f16
npyv_f16_to_f32_lo __lasx_xvfcvtl_s_h low 8 f16 → 8 f32
npyv_f16_to_f32_hi __lasx_xvfcvth_s_h high 8 f16 → 8 f32
npyv_f32_to_f16 __lasx_xvfcvt_h_s 8+8 f32 → 16 f16

These are used by the float16 transcendental bridge.

int↔float cast helpers

DragonArray also adds 3 internal cast helpers in numpy/_core/src/umath/npyv_exp_log.h, because the portable NPYV API does not expose lane-wise int↔float casts in the form these kernels need. Each helper also has a 128-bit LSX twin behind the same macro.

primitive LASX intrinsic does
npyv__cvt_f32_s32 __lasx_xvffint_s_w s32 → f32
npyv__cvt_f64_s64 __lasx_xvffint_d_l s64 → f64
npyv__cvt_s64_f64 __lasx_xvftintrne_l_d f64 → s64, round-nearest

Backend fix

DragonArray rewrites npyv_tobits_b16 in lasx/conversion.h.

The upstream version only observed 8 of the 16 lanes. The DragonArray version uses xvmskltz_h instead of the broken xvsat_hu / xvpickev_b sequence. This is a bug fix to an existing primitive, not a new primitive.

Math kernels composed on NPYV

Most of the patch is not new primitive work. It is vectorized math kernel work built on NumPy's existing NPYV backend plus the small primitive additions above.

DragonArray adds or rewrites NPYV paths for:

The transcendental kernels are compositions built on NPYV primitives. They are not themselves backend primitives.

Payne-Hanek range reduction

DragonArray adds Payne-Hanek support for full-range f64 trigonometric kernels:

float16 transcendental bridge

DragonArray adds a float16 transcendental bridge in npyv_f16_kernels.h and the loops_half driver.

The bridge loads 16 f16 lanes, converts them to two f32 vectors, runs the transcendental at f32 SIMD width, then packs the result back to f16.

Integer LASX kernels

DragonArray adds vector-by-vector LASX kernels for integer floor_divide / remainder on int8, int16, and int32.

These use LASX xvdiv / xvmod and cover cases where stock NumPy had no SIMD integer-divide path on LoongArch. This is a LoongArch-specific advantage: mainstream x86 SIMD does not provide equivalent packed integer division instructions.

Dispatch and build wiring

DragonArray wires LASX into more ufunc dispatch units, including:

The patch also keeps some paths deliberately on LSX where measurements showed LASX was not a win, including comparison-related bool packing paths and selected little-endian unary-fp loops.

For tanh, DragonArray keeps f32 on the Highway SVML port and uses the NPYV path for f64, with gating to avoid duplicate symbols.

CPU feature detection

DragonArray adds runtime LSX / LASX feature detection for LoongArch, including:

What DragonArray did not do

DragonArray did not create NumPy's LoongArch SIMD backend from scratch.

It did not create the bulk of the LASX npyv_* primitive layer. Existing upstream NumPy already provided the main 256-bit LASX universal-SIMD backend: add, multiply, divide, load, store, load2, compare, select, and the ordinary machinery needed by NPYV kernels.

DragonArray also did not make the transcendental kernels into new backend primitives. The exp, log, sin, cos, cbrt, tanh, power, and related ufunc implementations are vectorized kernels composed on top of NPYV primitives.

The accurate primitive-level accounting is:

Accuracy

Accuracy is tested against libm with randomized, edge-case, and special-value test sets. Current results are libm-comparable for the tested ranges; see OPTIMIZATIONS.md for max-ULP tables and known caveats.

Install

The experimental wheels are CPU-specific. Choose the wheel matching both your processor and Python version. Do not ask pip to select from the whole wheel directory: LA464 and LA664 use the same linux_loongarch64 platform tag, so pip cannot distinguish them.

Use this inside a virtual environment. It replaces the numpy package in that environment with the DragonArray build.

Loongson 3C5000 (LA464)

Loongson 3A6000 (LA664)

Pass the selected link directly to pip:

python -m pip install --upgrade --force-reinstall --no-deps \
    '<paste the matching wheel URL from above>'

Wheels are built for Python 3.12, 3.13, and 3.14 (linux_loongarch64). Do not install an LA664 wheel on a 3C5000. 3.9–3.11 are not available — this is NumPy 2.5.0.dev, whose pyproject.toml sets requires-python = ">=3.12". libstdc++ and libgcc are statically linked, built with /opt/loongson-gcc-15.2.0, so the wheels run even on systems with an older GLIBCXX than the build toolchain.

Verify

import numpy as np
print(np.__version__)        # 2.5.0.dev0+dragon.unofficial.1
np.show_runtime()            # LASX / LSX appear in the CPU dispatch info

Build from source

If you'd rather build it yourself, the changes are also available as a patch against NumPy 2.5.0.dev0, and as a git branch:

Apply the patch to a NumPy 2.5.0.dev0 checkout (or check out that branch), then build per the provenance below.

Build provenance

License

These modified, unofficial wheels are distributed under the applicable BSD-3-Clause license terms. NumPy is BSD-3-Clause. The Payne-Hanek reduction table generator is ported from SLEEF (Boost Software License 1.0, © Naoki Shibata). Polynomial coefficients and algorithms are derived from FDLIBM / SunPro permissively licensed sources, with notices preserved, and from Cephes where noted. Existing copyright notices, license texts, and third-party attributions are preserved in the source and wheel metadata.

Disclaimer

DragonArray is an independent project and is not affiliated with, endorsed by, sponsored by, or approved by NumFOCUS, the NumPy project, Loongson Technology, or any hardware vendor. NumPy is a trademark of NumFOCUS, Inc. Loongson and LoongArch are trademarks or registered trademarks of their respective owners. All third-party names and marks are used solely to identify compatibility, architecture support, or upstream software.