RSTSR：科学计算矩阵与高维张量库

项目地址：

• github：https://github.com/RESTGroup/rstsr
• gitee mirror: https://gitee.com/restgroup/rstsr

RSTSR 的定位与 Python 的 NumPy/SciPy，或 Rust 的 ndarray 类似，是高维度大型稠密张量的处理框架。但 RSTSR 具有与 ndarray、nalgebra，以及机器学习框架 candle、Burn 不同的功能与特性。

RSTSR 的定位是科学计算的基石框架。它是中间层库，下接 Faer 或 OpenBLAS 等后端 (backend)，上可应用于各类科学计算程序。

RSTSR 是下述两个英文名称的简称：

• REST Subproject for Tensor Support：为 Rust 电子结构程序 REST 提供张量功能支撑的子项目；
• An n-Dimensional Rust Tensor Toolkit：高维度张量工具。

RSTSR 在存储结构与算子实现上，与现有库有不同的设计理念，以希望达到 Fast (高运行效率) / Intuitive (更友好的编程体验、更高的代码可读性) / Extensible (可扩展后端、算子)。

RSTSR 文档资料

目前 RSTSR 的文档资料已有雏形。包括

• 用户文档 restgroup.github.io，中英文；
• API 文档 docs.rs/rstsr，英文。

比较重要或典型的文档包括

• NumPy-RSTSR 对照表：习惯 NumPy/SciPy 的用户，可以快速上手的参考文档；
• API Specification of crate rstsr-core：RSTSR 主要核心函数列表；
• Array API standard Fullfillment：RSTSR 对 Python array API standard 即 Python 语言下高维张量基本函数 API 与约定 (broadcasting, indexing, type promotion 等) 的实现情况。

RSTSR 特色

同时参考用户文档：为何选择 RSTSR。

• 便利的使用：与 NumPy 类似的使用方法，与 ndarray 相似的所有权设计。
  • RSTSR 许多函数支持基于 Trait 的重载。典型的例子是基础索引 (tensor.i) 与张量生成 (rt::asarray)。
• 语法糖：矩阵乘法可以用 % 符号实现，即 &a % &b 等价于 a.matmul(&b)。
• 允许多后端：允许用户切换后端，允许库开发者实现自己的后端。
  • 多后端的设计，其本来的目的是打算支持 CUDA、HIP、HDF5 等异构后端的存储或计算。但受限于时间与精力，目前这些后端还尚未实现。我们在不远的未来将尝试实现 CUDA 后端。
• 支持并行 (rayon 框架)，且可以动态地控制线程数量。
• 支持行优先与列优先，且可以通过动态地设置后端以作控制。

目前实现的后端包括

RSTSR 已经实现的重要功能

• 通过泛型，理论上支持所有数据类型。科学计算关心的数据类型 (单双浮点、复数) 在矩阵线性代数上有额外实现或加速。
• reshape：与 ndarray 的 to_shape 有些许不同，RSTSR 仅在必须时作内存复制。RSTSR 基本复刻了 NumPy 的运作模式：这种“必须”的情景并不一定要求张量是 row-major 连续的，一些分段连续的子张量也可能通过 reshape 返回张量视窗，而不是复制了数据的张量。
• tensor.i / tensor.slice：完整的基础索引 (basic-indexing) 功能。其中，
  • tensor.i(a) 可以用于作单元素索引；
  • tensor.i(a..b), tensor.i(..b), tensor.i(a..) 等可以用于缩小张量；
  • tensor.i((.., a, ..b)) 是混合的高维度基础索引；
  • tensor.i((Ellipsis, None)) 可以在最后插入一个维度 (Ellipsis 相当于 Python 的 ...、None 相当于 NumPy 的 np.newaxis)；
  • tensor.i(slice!(-3, 1, -2)) 可以作负步长的索引。
• to_broadcast：实现了行优先与列优先的广播规则 (参考 Python Array API 对 Broadcasting rules 的说明)。
• 基于 generic 的类型系统 (与 ndarray 类似，但与 candle 与 Burn 有差异)。实现了接近于 NumPy 的 Type Promotion。
• 部分实用的线性代数接口 (已实现的一些功能在 API 文档 中能找到)。

示范案例

use rstsr::prelude::*;

// 3x2 matrix with c-contiguous memory layout
let a = rt::asarray((vec![6., 2., 7., 4., 8., 5.], [3, 2].c()));

// 2x4x3 matrix by arange and reshaping
let b = rt::arange(24.);
let b = b.reshape((-1, 4, 3));
// in one line, you can also
// let b = rt::arange(24.).into_shape((-1, 4, 3));

// broadcasted matrix multiplication
let c = &b % &a;

// print the result
println!("{:6.1}", c);
// output:
// [[[   23.0   14.0]
//   [   86.0   47.0]
//   [  149.0   80.0]
//   [  212.0  113.0]]
//
//  [[  275.0  146.0]
//   [  338.0  179.0]
//   [  401.0  212.0]
//   [  464.0  245.0]]]

// print layout of the result
println!("{:?}", c.layout());
// output:
// 3-Dim (dyn), contiguous: Cc
// shape: [2, 4, 2], stride: [8, 2, 1], offset: 0

RSTSR 发展现状

目前 RSTSR 在 CPU 上的功能实现已经到可以实用的阶段。作为实战案例，我们在 Rust 语言下使用 RSTSR，可以实现计算化学精度的“金标准”CCSD(T) 的高效计算 (CPU FLOP 利用率对于 CCSD 在 45%、CCSD(T) 在 35% 左右，单节点性能与目前最好的计算化学程序 Psi4 (CCSD) 与 PySCF (CCSD(T)) 等接近或更好)。

基于 RSTSR 项目，我们也发展了一些与 RSTSR 关联的子项目：

• FFI 库，github: RESTGroup/rstsr-ffi，目前其用途主要是暴露 BLAS FFI (目前已有的 BLAS/LAPACK FFI 未实现 ILP64 与 dynamic-loading 功能)；
• TBLIS 封装，github: RESTGroup/tblis-rs，可以用于实现高效 CPU 张量缩并等计算；
• 张量缩并路径寻找，github: RESTGroup/opt-einsum-path；
• RSTSR 文档，github: RESTGroup/rstsr-book。

RSTSR 目前实现了大多数 Python Array API 对高维度张量库的函数接口要求。

RSTSR 还是一个非常新的项目；其开发是从 2024 年 8 月开始。也有很多没有来得及实现的功能。目前亟需发展或改进的功能包括

• GPU 支持 (优先 CUDA)；
• API 文档、用户文档、开发者文档；
• 更完善的线性代数支持；
• scipy.optimize 的支持。

RSTSR 通过泛型，允许实现任意类型数据的后端。因此 RSTSR 理论上可以通过在底层数据中插入新的 field (如梯度张量、算子 enum 等)，以实现自动导数。但 RSTSR 的设计目的是为科学计算作基础，以发展 REST 电子结构程序为最关键目标，因此实现机器学习框架不在目前 RSTSR 的发展方向中。

如果感兴趣，欢迎在 Github Issue 或 Discussion 中交流！

我们并非计算机科班出生的程序员，而是天坑专业的科研工作者。作为开源项目，我们欢迎大家对项目的发展提出期望和想法，帮助这个项目更好的发展。我们也希望这个项目能帮助到大家，为 Rust 目前的科学计算生态贡献一些力量。

致谢

RSTSR 项目是 REST 电子结构程序的子项目。REST 与 RSTSR 受自然科学基金与专项的资助。