中文内容
在上一篇文章中，我们介绍了 Universal Sparse Tensor（UST），它使开发者能够将张量的稀疏性与其内存布局解耦，从而获得更大的灵活性和更高的性能。我们很高兴地宣布，UST 已集成到 nvmath-python v0.9.0 中，以加速稀疏科学计算和深度学习应用。
本文将逐步介绍 UST 的关键特性、实现细节和性能概览，包括：
- 零成本互操作性：与 PyTorch、SciPy 和 CuPy 进行无需数据移动的转换。
- 自定义格式：定义新的稀疏性方案。
- 多态操作：稀疏性无关函数会自动使用优化内核或生成自定义稀疏代码，从而无需为新格式手动编码。
- PyTorch 注入：轻松将 UST 的性能优势注入现有 PyTorch 模型。
- 透明缓存：避免 JIT/LTO 重新编译和重新规划——将开销摊销到后续对同一操作的重复执行中。
张量格式 DSL
UST 使用一种简单的领域特定语言（DSL）来描述常见（例如 COO、CSR、CSC 和 BSR）以及较少见的稀疏张量存储格式。例如，CSC 格式描述如下。
(i, j) -> (j : dense, i : compressed) # CSC
将这种 DSL 集成到编程语言中需要做出具有各种权衡的设计决策。
例如，在 C++ 库 MatX 中，该 DSL 纯粹通过类型和模板实现，这使得针对特定存储格式进行有趣的编译期优化成为可能。例如，在下面的分派方法 foo() 中，格式特定的测试可以在编译期求值，这意味着只会编译并交付相关分支，以略高的编译时间换取更小的二进制文件大小。
template<TensorType> foo(TensorType a) {
if constexpr (TensorType::Format::isCOO()) {
... dispatch to COO code ...
} else if constexpr (TensorType::Format::isCSR()) {
... dispatch to CSR code ...
} else ...
}
在 nvmath-python 中，我们采用了一种更符合 Python 风格的方式来呈现该 DSL，在检查格式对象时以部分性能换取运行时灵活性。例如，上面所示的 CSC 格式表示如下。
i, j = ( Dimension(dimension_name="i"), Dimension(dimension_name="j") )
CSC = TensorFormat( [i, j], {j: LevelFormat.DENSE, i: LevelFormat.COMPRESSED} )
这些对象的一个主要优势是，一切都可以在运行时动态构建（包括从字符串解析）。然而，执行特定于格式的任务需要在运行时检查实际内容。由于此类决策通常发生在性能关键路径之外，因此为通用性做出取舍似乎是可以接受的选择。
与 PyTorch、SciPy、CuPy 的互操作性
nvmath-python UST 实现提供了与 PyTorch、SciPy、CuPy 和 NumPy 张量的互操作性。转换基本上是零成本的，这意味着将 COO、CSR、CSC、BSR、BSC 和 DIA 等稠密或稀疏格式转换为 UST 对象，或从 UST 对象转换回来，都无需数据移动或复制。相反，UST 对象会引用原始数据结构的存储缓冲区。
例如，以下从 COO 格式的 SciPy 稀疏矩阵到 nvmath-python UST 对象的转换，会使用行数组、列数组和值数组来构成 UST 的位置、坐标和值数组。
row = np.array([0, 0, 1, 1, 2, 3], dtype=np.int32) col = np.array([0, 1, 1, 3, 2, 3], dtype=np.int32) val = np.array([1, 2, 3, 4, 5, 6], dtype=np.float32) coo = sps.coo_array((val, (row, col)), shape=(4, 8)) # SciPy ust = Tensor.from_package(coo)
此转换会自动生成 COO 的 DSL，如下所示。转换其他格式将生成这些格式对应的 DSL。
TensorFormat( [i, j] -> {i: (<LevelFormat.COMPRESSED>, <LevelProperty.NONUNIQUE>),
j: <LevelFormat.SINGLETON>} )
UST 对象会按如下方式转换回原始包中的稀疏张量。这同样是零成本转换，只需传递对组成缓冲区的引用即可。
coo = ust.to_package() # this yields a SciPy coo format again
较不常见的格式
开发者还可以使用 DSL 定义自己的新型稀疏格式。例如，以下代码将一个密集 PyTorch 张量转换为 2 位 delta 压缩格式（类似于 CSR，但使用坐标的运行时差值），这是 PyTorch 原生不支持的格式。
A = torch.tensor([[1, 0, 0, 0, 0, 0, 0, 2],
[0, 0, 3, 4, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 5]], dtype=torch.float64)
U = Tensor.from_package(A) # dense UST
delta_format = TensorFormat([i, j], {i: LevelFormat.DENSE,
j: (LevelFormat.DELTA, 2)}) # use 2-bit
S = U.convert(tensor_format=delta_format) # sparse UST
这会得到图 1 所示的稀疏 UST 存储，其中每个坐标使用到下一个已存储元素的距离，而不是直接的列索引。请注意，在该距离会超过 3（使用 2 位时的最大值）的情况下，需要进行填充（红色显示的零）。
打印和绘图实用工具
提供了各种实用工具，帮助用户调试、测试并熟悉 nvmath-python UST 实现。首先，双下划线方法 __repr__ 提供了一种明确的字符串表示，可用于打印张量存储内容。
A = torch.arange(4 * 5).reshape(4, 5).cuda().to_sparse_csr() U = Tensor.from_package(A) print(U)
此代码片段会产生以下输出，其中显示了用于数值的类型（int64）、UST 存储中的位置和坐标数组（同样为 int64）、设备（cuda）、维度和层级范围、存储的数据、用于存储的字节数摘要（5*8+19*8+19*8），以及稀疏度（(1 – 19 / 20) x 100%）。
---- Sparse Tensor<VAL=int64,POS=int64,CRD=int64,DIM=2,LVL=2> format : [i, j] -> (i: <LevelFormat.DENSE>, j: <LevelFormat.COMPRESSED>) device : cuda dim : [4, 5] lvl : [4, 5] nse : 19 pos[1] : [0, 4, 9, 14, 19] #5 crd[1] : [1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4] #19 values : [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19] #19 data : 344 bytes sparsity : 5.00% ----
draw_storage() 和 draw() 方法会生成张量存储（任意维度）或张量内容（最高支持 3D）的图像。前一种方法用于生成图 1。对前面的 UST 示例调用后一种方法会得到图 2 所示的图像，其中灰色表示已存储的非零元素（在此情况下几乎为所有元素）。
这两种方法会显示相对较小示例的 UST，因为它们无法扩展到非常大的规模。对于更大的张量，可以使用 draw_raw() 方法（2D 或 3D）来获得仅显示张量非零结构的可视化效果，如图 3 所示，该图使用的是取自 SuiteSparse 库的一个矩阵。
多态操作
UST 的张量格式 DSL 使开发者能够专注于张量稀疏性。然而，定义张量的存储只是需求的一部分。开发者还可以指定要执行的操作。
nvmath-python UST 实现使开发者能够使用 matmul() 和 solve() 等与稀疏性无关的操作来定义计算（不同于张量索引记法等替代方案）。这种方法很强大，因为系统会检查操作数格式以确定执行路径：对于常见格式，分派到优化的手写库或内核；当不存在优化方案时，则依赖自动稀疏代码生成。
这种设计能够使用新颖的稀疏格式，而无需显式手动编码，同时仍可通过使用现有库解决方案处理标准格式来实现高性能。用于规划和执行 matmul() 操作的语法如下所示。按照 nvmath-python 的约定，该操作可分为规划阶段（在分派或 JIT LTO 代码生成之间做出决定，并缓存以供未来复用，同时设置所有必需参数）和执行阶段（执行实际操作）。
# c := a @ b + c
with Matmul(a, b, c, beta=1.0) as mm:
mm.plan()
mm.execute()
这种方法可将初始规划成本分摊到同一 matmul() 调用的多次执行中，这在例如迭代式稀疏求解器，或深度学习中的（剪枝）线性层中非常有价值。
将 UST 注入 PyTorch
PyTorch 稀疏张量与 nvmath-python 之间的互操作性使开发者能够在 PyTorch 模型中尝试新颖的稀疏存储方案。尽管功能强大，但 AI 研究人员并不愿意重写模型代码，将线性层内部执行的 GEMM 操作替换为上述 UST 多态操作，即使这样能带来好得多的性能。
UST 的 nvmath-python 实现提供了一种无需重写原始模型代码即可将 UST“注入”现有模型的方法。这是通过将 UST 包装到 torch.Tensor 的子类中，并为 UST 的多态操作提供包装器来实现的。通过为常见的 PyTorch 张量操作（例如 dot、mv、mm 和 linear）提供包装器，研究人员无需深入了解其内部工作机制即可使用 UST。
这些包装器还添加了另一层缓存（位于已缓存的 JIT/LTO 查找之上），用于复用针对类似操作已规划的 MatMul 实例，从而将后续运行时间缩短到仅包含执行阶段。
还提供了 reformat_model() 方法，可遍历所有线性层的权重，并允许用户定义的函数检查每个权重矩阵，并可能对其进行剪枝和稀疏化（或跳过）。
weights = torchvision.models.get_model_weights(model_name).DEFAULT
model = torchvision.models.get_model(model_name, weights=weights)
model.to(device)
model.eval()
...
reformat_model(model, func=reformat)
...
with torch.inference_mode():
prediction = model(batch)
用户提供的重格式化方法的一般形式如下所示。返回 None 表示保持权重矩阵不变。此后，推理将在所有已被修改的线性层应用中使用 UST。
def reformat(weight):
... inspect (possibly even prune weight), then decide on storage ...
if (...)
return TorchUST.from_torch(weight.to_sparse_csr())
return None
}
性能示例
首先，我们比较在原生支持的格式（COO、CSR 和 DIA）下使用 CuPy 的 @ 操作的 SpMV 性能，以及针对 COO 和 CSR 使用 PyTorch mv() 的性能，并将其与 UST 在来自 Matrix Market 的 1,489,752 x 1,48,9752 矩阵 atmosmodl 上的 matmul() 进行对比。对于后者，我们仅测量 execute() 运行时间；对于需要重复乘法的应用而言，这是公平的比较（因为 CuPy 和 PyTorch 都不提供规划设置）。所得运行时间如图 4 所示（纵轴采用对数刻度）。
在此场景中，UST 实现了 1.1 到 444 倍不等的加速。CuPy 和 PyTorch 的 COO 实现性能较差，这一点值得注意。对于 CSR 格式，所有版本都使用 cuSPARSE 实现，而 UST 受益于规划阶段的复用。DIA 格式下观察到的加速源于 UST 使用了专用内核，相比之下，CuPy 的方法是先转换为 CSR（而 PyTorch 不支持 DIA）。
第二个实验基于论文 MACKO: Sparse Matrix-Vector Multiplication for Low Sparsity，该论文提供了 SpMV 操作的一种高效实现，而 SpMV 是深度学习中单 token 推理的重要步骤。作者提出了 MACKO 格式，本质上就是图 1 中的增量压缩格式，并提供了令人印象深刻的 NVIDIA CUDA 实现，该实现在 50% 及以上的非结构化稀疏度下已经相较于稠密实现带来了加速。
由于 nvmath-python UST 提供了一种简便方式，可将新的手写内核纳入查找机制中，我们尝试将其实现用作增量压缩格式的后端内核实现。图 5 比较了针对 8192×8192 均匀随机稀疏矩阵的稠密实现（GEMV）、cuSPARSE 的 SPMV 实现、MACKO UST 以及原始 MACKO 实现的运行时间，稀疏度从 0% 稀疏（稠密）变化到 100% 稀疏（零矩阵）。MACKO UST 的表现略好，因为填充会在每一行末尾停止（当然，原始 MACKO 也很容易纳入这一优化）。
了解更多
如需更深入了解 nvmath-python 的 UST 实现，请参阅 nvmath-python 在线文档。
标签
