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 的工作分为四层:

  1. 少量缺失的 LASX backend primitive;
  2. 一个 LASX backend bug fix;
  3. 构建在 NPYV primitive 之上的 math 和 integer kernel;
  4. 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 路径:

这些 transcendental kernel 是构建在 NPYV primitive 上的组合。它们本身不是 backend primitive。

Payne-Hanek range reduction

DragonArray 为 full-range f64 trigonometric kernel 加入了 Payne-Hanek 支持:

float16 transcendental bridge

DragonArray 在 npyv_f16_kernels.hloops_half driver 中加入了 float16 transcendental bridge。

这个 bridge 会加载 16 个 f16 lane,把它们转换成两个 f32 vector,用 f32 SIMD 宽度运行 transcendental,然后把结果 pack 回 f16。

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,包括:

这个补丁也有意把一些路径保留在 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,包括:

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。 explogsincoscbrttanhpower 以及相关 ufunc 实现, 都是构建在 NPYV primitive 之上的 vectorized kernel。

准确的 primitive-level accounting 是:

精度

精度会用 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)

龙芯 3A6000(LA664)

把上面匹配的 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 该分支),然后按下方 的构建来源说明编译。

构建来源

许可证

这些非官方修改版 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.