这个仓库是什么？

这是一个简约的、自包含的稀疏Cholesky求解器，支持在CPU和GPU上解决，易于集成到您的张量管道中。

当我们正在撰写我们的“几何逆渲染大步”论文[1]时，我们发现将现有的稀疏线性求解器连接到我们的管道相当具有挑战性，我们通过添加对大型项目的依赖（即cusparse和scikit-sparse），仅使用其功能的一小部分来实现这一点。因此，我们决定实现自己的库，该库只有一个目的：使用Cholesky分解高效地在GPU或CPU上解决稀疏线性系统。

底层，它依赖于CHOLMOD进行稀疏矩阵分解。在求解阶段，它使用CHOLMOD的CPU版本，并使用构建求解器时运行的一次分析步骤的结果，在GPU上进行快速求解[2]。

它实现了与其他框架相当的性能，并附带漂亮的依赖项。

在Linux Ryzen 3990X工作站上使用TITAN RTX运行的基准测试

Python绑定使用nanobind生成，这使得它容易与大多数张量框架（Numpy、PyTorch、JAX...）进行互操作

安装

使用PyPI（推荐）

pip install cholespy

从源码

git clone --recursive https://github.com/rgl-epfl/cholespy
pip install ./cholespy

文档

该模块中只有一个类，包含两个变体：CholeskySolverF 和 CholeskySolverD。它们之间的唯一区别是 CholeskySolverF 使用单精度解算系统，而 CholeskySolverD 使用双精度。这在GPU上解算时非常有用，因为CPU版本依赖于CHOLMOD，而CHOLMOD只支持双精度。

默认支持最常用的张量框架（PyTorch、NumPy、TensorFlow...）。您可以直接将它们传递到模块中，无需手动转换。

由于两种变体具有相同的签名，以下仅详细说明 CholeskySolverF：

cholespy.CholeskySolverF(n_rows, ii, jj, x, type)

参数

• n_rows - 稀疏矩阵中的行数。
• ii - 稀疏矩阵表示中的第一个索引数组。如果 type 是 COO，则此数组为行索引数组。如果是 CSC（或 CSR），则此数组为列（或行）索引数组，其中列（或行）k 的行（或列）索引存储在 jj[ii[k]:ii[k+1]] 中，相应的条目存储在 x[ii[k]:ii[k+1]] 中。
• jj - 稀疏矩阵表示中的第二个索引数组。如果 type 是 COO，则此数组为列索引数组。如果是 CSC（或 CSR），则此数组为行（或列）索引数组。
• x - 非零条目的数组。
• type - 矩阵表示类型，类型为 MatrixType。可用类型包括 MatrixType.COO、MatrixType.CSC 和 MatrixType.CSR。

cholespy.CholeskySolverF.solve(b, x)

参数

• b - 要解算的方程的右端。可以是向量或矩阵。如果是矩阵，则其形状必须为 (n_rows, n_rhs)。它必须在与传递给求解器构造函数的张量相同的设备上。如果使用CUDA数组，则 n_rhs 的最大支持值为 128。
• x - 解的占位符。它必须在同一设备上，并且具有与 b 相同的形状。

x 和 b 必须 与所使用的求解器具有相同的数据类型，即 CholeskySolverF 为 float32 或 CholeskySolverD 为 float64。由于 x 将就地修改，因此不支持隐式类型转换。

示例用法

from cholespy import CholeskySolverF, MatrixType
import torch

# Identity matrix
n_rows = 20
rows = torch.arange(n_rows, device='cuda')
cols = torch.arange(n_rows, device='cuda')
data = torch.ones(n_rows, device='cuda')

solver = CholeskySolverF(n_rows, rows, cols, data, MatrixType.COO)

b = torch.ones(n_rows, device='cuda')
x = torch.zeros_like(b)

solver.solve(b, x)
# b = [1, ..., 1]

参考

[1] Nicolet, B.，Jacobson, A.，& Jakob, W. (2021). 大步推进几何逆渲染。ACM Transactions on Graphics (TOG)，40(6)，1-13。

[2] Naumov, M. (2011). 在GPU上的预条件迭代方法中并行求解稀疏三角线性系统。NVIDIA Corp.，Westford，MA，USA，技术报告NVR-2011，1。