DragonArray 文档
DragonArray NumPy 补丁集与 wheel
DragonArray 是一个非官方、实验性的 LoongArch LASX 优化补丁集,面向 Loongson 3A6000 系统上的 NumPy。
在这个补丁中,我完成了 NumPy 2.5.0-dev 在 Loongson 3A6000 上的 LASX ufunc 启用与优化工作。这个补丁把更多 ufunc dispatch 路径接入 NumPy 现有的 256-bit LASX universal-SIMD backend,补上少量缺失的 256-bit primitive,修复 一个 backend primitive,并为这颗 CPU 上仍然回退到 scalar code 的情况加入新的 vectorized math kernel。
这是一个独立社区补丁集,不是官方 NumPy release。
DragonArray 建立在 NumPy 自己的 LoongArch 工作之上。LASX npyv_* backend
的大部分内容,包括 arithmetic、load/store、compare、select,以及多数其他
universal-SIMD primitive,都是 LoongArch / Loongson 社区贡献到上游 NumPy 的
工作。DragonArray 没有重新实现这个 backend。它是在这个 backend 上继续工作。
总览
在 Loongson 3A6000 @ 2.0 GHz 上,相比原版 NumPy 2.5.0-dev,135 个 operation × dtype 组合的原始几何平均加速为 1.81×。两个构建都使用同一个上游 NumPy 2.5.0-dev 基础版本和同一套 GCC 15.2.0 工具链:一个是原版构建,一个 带 DragonArray LASX 补丁集。每个构建都在 contiguous path 上运行自己最好的 SIMD 路径(原版:LSX;DragonArray:LASX)。这是在同一颗芯片上安装这个 wheel,而不是安装原版 NumPy,所得到的真实加速;不是拿 scalar fallback 做对比。
| 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× |
完整的逐 op 表格见 BENCHMARKS.zh-CN.md。按时间顺序整理的 技术日志见 OPTIMIZATIONS.zh-CN.md。
DragonArray 增加了什么
DragonArray 的工作分为四层:
- 少量缺失的 LASX backend primitive;
- 一个 LASX backend bug fix;
- 构建在 NPYV primitive 之上的 math 和 integer kernel;
- dispatch、build 和 runtime feature-detection wiring。
新增 LASX primitive
DragonArray 增加了 8 个真正新的 256-bit primitive。primitive 层面的其他内容 来自 NumPy 现有的 LASX backend。
float16 bridge primitive
NumPy 的 LASX backend 没有暴露 float16 NPYV 支持,所以 DragonArray 在
numpy/_core/src/umath/npyv_f16_bridge.h 中加入了一个小的 f16 bridge:
| primitive | intrinsic | 作用 |
|---|---|---|
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 |
低 8 个 f16 → 8 个 f32 |
npyv_f16_to_f32_hi |
__lasx_xvfcvth_s_h |
高 8 个 f16 → 8 个 f32 |
npyv_f32_to_f16 |
__lasx_xvfcvt_h_s |
8+8 个 f32 → 16 个 f16 |
这些 primitive 用于 float16 transcendental bridge。
int↔float cast helper
DragonArray 还在 numpy/_core/src/umath/npyv_exp_log.h 中加入了 3 个内部
cast helper,因为 portable NPYV API 没有以这些 kernel 所需的形式暴露
lane-wise int↔float cast。每个 helper 在同一个 macro 后面也有 128-bit LSX
版本。
| primitive | LASX intrinsic | 作用 |
|---|---|---|
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 重写了 lasx/conversion.h 中的 npyv_tobits_b16。
上游版本只观察到了 16 个 lane 中的 8 个。DragonArray 版本使用 xvmskltz_h,
替代有问题的 xvsat_hu / xvpickev_b 序列。这是对现有 primitive 的 bug fix,
不是一个新的 primitive。
基于 NPYV 组合而成的 math kernel
这个补丁的大部分工作不是新的 primitive 工作,而是构建在 NumPy 现有 NPYV backend 和上面少量 primitive 增补之上的 vectorized math kernel 工作。
DragonArray 增加或重写了这些 NPYV 路径:
loops_exponent_log中的 f32 和 f64exp/log,并对 non-AVX2 / non-AVX512 target 启用,所以这条路径也可以让 NEON / ASIMD / VSX 受益;exp2、log2、log10、expm1和log1p;sinh、cosh和 f64tanh;asinh、acosh和atanh;atan、asin、acos、atan2和tan;- full-range f64
sin/cos,使用 SIMD Cody-Waite reduction,加上针对巨大 参数的 Payne-Hanek reduction; - f64
cbrt,使用 FDLIBM 风格的 bit setup 加 Newton refinement; - f32
cbrt,使用 bit setup 加两轮 Halley refinement; power,使用exp(y * log(x))、per-block fast path 和 IEEE-754 special-case handling。
这些 transcendental kernel 是构建在 NPYV primitive 上的组合。它们本身不是 backend primitive。
Payne-Hanek range reduction
DragonArray 为 full-range f64 trigonometric kernel 加入了 Payne-Hanek 支持:
payne_hanek_f64.h中的 scalar Payne-Hanek reduction,基于 FDLIBMk_rem_pio2方法;payne_hanek_simd_table.h中的 SIMD Payne-Hanek table;- table generator
payne_hanek_table.py,移植自 SLEEF 的mkrempitab.c方法,并使用mpmath。
float16 transcendental bridge
DragonArray 在 npyv_f16_kernels.h 和 loops_half driver 中加入了 float16
transcendental bridge。
这个 bridge 会加载 16 个 f16 lane,把它们转换成两个 f32 vector,用 f32 SIMD 宽度运行 transcendental,然后把结果 pack 回 f16。
- float16 transcendentals:一个轻量级 LASX f16↔f32 bridge 可以在 f32 SIMD 宽度上运行 16 个 transcendental,两端 pack/unpack 约为 3-cycle;20 个实测 operation 组合相对原版 NumPy 2.5.0 的原始几何平均加速为 2.77×,其中 9 个组合 的加速至少为 1.10×。原版在 LoongArch 上没有 f16 SIMD 路径。
Integer LASX kernel
DragonArray 为 int8、int16 和 int32 的 integer floor_divide / remainder
加入了 vector-by-vector LASX kernel。
这些 kernel 使用 LASX xvdiv / xvmod,覆盖了原版 NumPy 在 LoongArch 上没有
SIMD integer-divide 路径的情况。这是 LoongArch 特有的优势:主流 x86 SIMD 没有
等价的 packed integer division 指令。
Dispatch 和 build wiring
DragonArray 把 LASX 接入了更多 ufunc dispatch unit,包括:
- arithmetic floating-point loops;
- unary loops;
- unary floating-point loops;
- unary complex loops;
- autovec loops;
- half loops;
- modulo loops;
- arithmetic loops;
- general umath floating-point loops;
- exponent / log loops;
- trigonometric loops。
这个补丁也有意把一些路径保留在 LSX 上,因为实测显示 LASX 在这些地方不是收益, 包括 comparison 相关的 bool packing 路径,以及部分 little-endian unary-fp loop。
对于 tanh,DragonArray 让 f32 保持在 Highway SVML port 上,让 f64 使用 NPYV
路径,并通过 gating 避免 duplicate symbol。
CPU feature detection
DragonArray 为 LoongArch 加入 runtime LSX / LASX feature detection,包括:
npy_cpu_features.{c,h}中基于AT_HWCAP的 LSX / LASX detection;cpu_lasx.c中的 LASX probe;- Meson CPU feature machinery 中的 LoongArch64 feature definitions。
DragonArray 没有做什么
DragonArray 没有从零开始创建 NumPy 的 LoongArch SIMD backend。
它没有创建 LASX npyv_* primitive layer 的大部分内容。现有上游 NumPy 已经提供
主要的 256-bit LASX universal-SIMD backend:add、multiply、divide、load、
store、load2、compare、select,以及 NPYV kernel 所需的普通机制。
DragonArray 也没有把 transcendental kernel 做成新的 backend primitive。
exp、log、sin、cos、cbrt、tanh、power 以及相关 ufunc 实现,
都是构建在 NPYV primitive 之上的 vectorized kernel。
准确的 primitive-level accounting 是:
- 8 个新的 256-bit primitive;
- 1 个现有 LASX backend primitive 修复;
- 许多新的或重写的 ufunc kernel 和 dispatch path,它们都构建在上游 NumPy 的 NPYV LSX/LASX 工作之上。
精度
精度会用 randomized、edge-case 和 special-value test set 与 libm 对比测试。 当前结果在已测试范围内与 libm 的 ULP 表现相当;max-ULP 表格和已知 caveat 见 OPTIMIZATIONS.zh-CN.md。
安装
实验版 wheel 与 CPU 型号绑定。请选择同时匹配处理器和 Python 版本的 wheel。
不要让 pip 从整个 wheel 目录中自动选择:LA464 和 LA664 使用相同的
linux_loongarch64 platform tag,因此 pip 无法区分二者。
建议在 virtual environment 中使用。它会把该环境中的 numpy package 替换为
DragonArray 构建。
龙芯 3C5000(LA464)
- Python 3.12 — numpy-2.5.0.dev0+dragon.unofficial.1-1la464-cp312-cp312-linux_loongarch64.whl
- Python 3.13 — numpy-2.5.0.dev0+dragon.unofficial.1-1la464-cp313-cp313-linux_loongarch64.whl
- Python 3.14 — numpy-2.5.0.dev0+dragon.unofficial.1-1la464-cp314-cp314-linux_loongarch64.whl
龙芯 3A6000(LA664)
- Python 3.12 — numpy-2.5.0.dev0+dragon.unofficial.1-1la664-cp312-cp312-linux_loongarch64.whl
- Python 3.13 — numpy-2.5.0.dev0+dragon.unofficial.1-1la664-cp313-cp313-linux_loongarch64.whl
- Python 3.14 — numpy-2.5.0.dev0+dragon.unofficial.1-1la664-cp314-cp314-linux_loongarch64.whl
把上面匹配的 URL 直接传给 pip:
python -m pip install --upgrade --force-reinstall --no-deps \
'<粘贴上面匹配的 wheel URL>'
wheel 面向 Python 3.12、3.13 和 3.14(linux_loongarch64)构建。没有 3.9–3.11
版本。不要在 3C5000 上安装 LA664 wheel。
这是 NumPy 2.5.0.dev,它的 pyproject.toml 设置了
requires-python = ">=3.12"。libstdc++ 和 libgcc 静态链接,构建使用
/opt/loongson-gcc-15.2.0,所以即使系统上的 GLIBCXX 比构建工具链旧,wheel
也能运行。
验证
import numpy as np
print(np.__version__) # 2.5.0.dev0+dragon.unofficial.1
np.show_runtime() # CPU dispatch 信息中应出现 LASX / LSX
从源码构建
如果你更愿意自行构建,这些改动也以针对 NumPy 2.5.0.dev0 的补丁,以及 git 分支的形式提供:
将补丁应用到 NumPy 2.5.0.dev0 的 checkout(或直接 checkout 该分支),然后按下方 的构建来源说明编译。
构建来源
- Source:NumPy 的
loongson-experimentalbranch (commit2f3ca9a2c3,"ENH: LoongArch 3A6000 LSX/LASX optimization pass")。 - Toolchain:
/opt/loongson-gcc-15.2.0(GCC 15.2.0)。系统 GCC 8.3 太旧, NumPy 要求 ≥ 9.3。 - Build:
python -m build --wheel --no-isolation,并使用-Dcpp_link_args=-static-libstdc++和LDFLAGS="-static-libstdc++ -static-libgcc"。
许可证
这些非官方修改版 wheel 按适用的 BSD-3-Clause 许可证条款分发。NumPy 使用 BSD-3-Clause。Payne-Hanek reduction table generator 移植自 SLEEF(Boost Software License 1.0,© Naoki Shibata)。Polynomial coefficient 和 algorithm 来自 FDLIBM / SunPro 的 permissively licensed source,并保留 notice;另有部分来自 Cephes,具体见相关文件。源码与 wheel metadata 保留现有版权声明、许可证文本和 第三方 attribution。
免责声明
DragonArray 是一个独立项目,与 NumFOCUS、NumPy 项目、龙芯中科或任何硬件厂商均无 隶属关系,亦未获得其认可、赞助或批准。NumPy 是 NumFOCUS, Inc. 的商标。Loongson 和 LoongArch 是其各自所有者的商标或注册商标。所有第三方名称和标识仅用于说明兼容 性、架构支持或上游软件。
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.
