SparseTIR 解读（论文、源码）

`Introduction`

以Halide/TVM为代表的张量运算编译器，引入了计算与调度分离的概念，使得大家可以只用写一套计算描述(Tensor Expression，只与计算的数学形式有关)，用不同的调度原语(Schedule Primitive)来描述如何去优化程序(如何做矩阵分块，绑定线程，利用缓存，设计流水线，利用硬件的加速单元)，这个过程可以手动，也可以利用自动化的调度模板生成（例如 AutoTVM）并搜索，从而为不同的硬件后端生成高效的代码。

但是最初的Halide和TVM都是针对稠密/规整运算设计，因此很难对于深度学习中的稀疏算子做统一的抽象和优化，在之前TVM对于稀疏矩阵需要手写每一个手动优化好的算子，无法利用调度原语设计统一的优化方案。

稀疏运算编译器，随着 TACO (The Tensor Algebra Compiler)等项目的出现重新进入人们视野，其主要思路是计算与数据结构分离，使得大家可以只写一套计算描述和一套数据结构的描述。而在图形学中很多数据有 spatially sparse(全局稀疏，局部稠密)的特性，使用专门设计的的层次化数据结构来存储会有更高的效率。

因此能否把这两类工作结合到一起，从而减少手写稀疏算子的负担，应对深度学习中各种不同算子/格式的组合，是 SparseTIR 这个项目的初衷。

包括TACO在内的稀疏矩阵编译器，利用他们强大的格式抽象可以表示结构化稀疏，但是目前并没有很好的硬件加速指令支持，从而并没能实现对结构化稀疏运算的有效加速。这是因为在张量编译器中，为了使用这些硬件指令，我们将程序中的子结构与硬件指令的描述相匹配，并替换为相应的指令实现，这个操作由特定的调度原语实现。

Halide和TACO这些编译器把所有的调度信息存储在schedule tree(调度树)/provenance graph(来源图)这样的中间数据结构上，这样的设计过于中心化：每个调度原语都会对调度树本身带来一些约束，调度原语之间可能会互相影响，给引入新的硬件指令带来了困难。

TVM中新的算子级别抽象 TensorIR 的提出正是为了解决如何使用硬件加速单元的问题，TensorIR 放弃了 schedule tree这些中间数据结构，而是把所有的信息记录在IR本身，通过这样的设计解耦了调度原语之间的相互制约，调度原语本身只会与IR进行交互，因此加入新的硬件指令不需要考虑其对其它调度原语的影响，可以更有效率地支持大量的硬件指令。

每一个调度原语都被写成了一个带参数的程序转换函数(效果上等价于Compiler Pass)，为了与硬件加速单元适配，TensorIR 提出了 block (块)这样的抽象，每一个 block都是整个运算的一个子结构(可以描述计算或者数据搬运)，调度原语可以操作这些 block及其外部的循环结构，与加速单元的语义描述相匹配的 block可以被替换为相应的硬件加速指令。

相比于原来的Tensor Expression前端，TensorIR的表达力足够强大，有循环和”块”这样的数据结构，已经可以用来手写一些稀疏算子(TVM原有的Tensor Expression只能描述循环体内的操作)，但是相比于TACO，仍然需要手动处理稀疏数据结构的编码和解码过程。

因此，面临的问题几个：

TVM 等传统编译器并不支持稀疏运算的支持
TVM 中初始的 Relay 甚至 TensorIR 都不能很好的描述稀疏数据结构
TACO 没有很好的硬件加速指令支持，运算速度较慢
TACO 中的调度算法的设计过于中心化，每个调度原语都会对调度树本身带来一些约束，调度原语之间可能会互相影响。

因此，这里提出了 SparseTIR 来应对这些问题

系统架构

SparseTIR 整个的编译流程可以总结为上图。

SparseTIR 分为了三个阶段：Stage I(坐标空间计算)→ Stage II(位置空间计算)→ Stage III(TensorIR)，并设计了两个lowering(递降)算法，分别表示I到II阶段，II到III阶段的转换。

程序被变换成TensorIR之后，就可以直接使用TVM已有的pass继续生成不同设备的代码。

在SparseTIR中，我们也希望所有的调度都在IR本身完成，这里将所有的调度原语也分为了三类，分别在IR的三个阶段使用，这样的设计有别于以往的张量编译器中single-shot scheduling(一次性发射所有的调度指令)的设计，有助于我们同时考虑高层级和低层级的优化:

在第一阶段中，我们可以进行稀疏格式感知的优化(例如DecomposeFormat)。
在第二和第三阶段中，我们的IR更接近与底层，可以进行硬件感知的优化(例如使用硬件加速指令等等)。

其中不同的阶段的调度选择，以及格式分解的选择，构成了SparseTIR的整体搜索空间，因此可以进行(格式*调度的整体优化)

SparseTIR没有像TACO一样提出具体的搜索算法，而是提出了一个框架让用户自己指定搜索空间。

需要注意的是 SparseTIR 暂时未支持 Auto Scheduling，也没有任何默认的算法

`SparseTIR` 的抽象设计

首先明确， SparseTIR 是依赖于 TVM 与 TensorIR 的，因此，这个实际上是以 TensorIR 为基础的一个处理稀疏数据拓展版本。

这里作者提出一个观点：

我们观察了非结构化稀疏运算的GPU算子常用优化，主要要点在于load-balance(负载均衡)和访存优化。这两者都属于被研究了几十年的老大难问题，性能高度依赖于非零元素本身的分布。深度学习中很多稀疏矩阵的非零元素分布是固定的(stationary-sparsity)，如果我们能在编译期就利用这些信息提前做好调度，那么可以省掉很多运行时调度的开销。

因此，SparseTIR 设计了可组合式的抽象：Composable Formats(可组合的格式) 和 Composable Transformations(可组合的程序变换)

可组合的格式

Composable Formats 代表，虽然很多矩阵本身很难挖掘出整体结构，但是我们把它分解成很多更小的结构化稀疏矩阵之和，不同的部分用不同的格式来存，其中每一个小矩阵都更加对硬件友好，相比于把整个稀疏矩阵做 padding 或者 block-sparse 化，带来的块内碎片会显著更少。如下图所示：

可组合的程序变换

把程序变换分解到 IR 的各个层级，多级调度，使得我们可以在同一个框架下完成数据感知的高语义层级优化和偏硬件层级的优化。

格式抽象与矩阵存储

格式描述是 SparseTIR 的一等公民，其对于矩阵的描述与存储延续了 TACO 的一些理念：

SparseTIR 创建了一个类似的 axis(sparse iterator) 结构，保留了TACO的 dense/compressed(sparse)属性，并增加了一组属性fixed/variable，表示当前维度是定长还是变长

这有助于我们生成更高效的代码——编译期能得到的信息越多，可能做的优化就越多。

传统的稠密矩阵的每一个维度在 SparseTIR 中都会使用 dense_fixed来标注，代表其连续而且定长；
对于非零元位置不连续(稀疏)的维度，我们使用 sparse来标注，需要指定 indices数组存储其非零元的坐标；
对于长度不一定的维度，我们使用 variable来标注，需要指定 indptr数组，储存每一个fiber(fiber是”行”和”列”在高维稀疏矩阵中的拓展)的初始指针位置。

所有的 sparse和 variable的维度都必须指定一个其依赖的维度，如果把每个axis向其依赖的axis连一条边，那么在一个SparseTIR的程序中，所有的axis会组成一个森林，而森林中每一个树的根节点都是 dense_fixed的。

例子如下：

1
# I is a dense fixed iterator.
2
I = T.dense_fixed(m)
3
# J1 is a sparse fixed iterator, whose dependent iterator is I
4
# it has maximum length n and number of non-zero elements per row: c,
5
# the indices data are stored in the region started from indices_1 handle,
6
# and the index data type (in indices array) is int32.
7
J1 = T.sparse_fixed(I, (n, c), indices_1, idtype="int32")
8
# J2 is a dense variable iterator, whose dependent iterator is I,
9
# it has a maximum length of n,
10
# the indptr data are stored in the region started from indptr_2 handle,
11
# and the index data type (in indptr array) is int32.
12
J2 = T.dense_variable(I, n, indptr_2, idtype="int32")
13
# J3 is a sparse variable iterator, whose dependent iterator is J1,
14
# it has maximum length of n1, number of elements nnz in the space composed of (I, J1, J3),
15
# the indptr data are stored in the region started from indptr_3 handle,
16
# and the indices data are stored in the region started from indices_3 handle,
17
# the index data type (of indptr and indices array) is "int64")
18
J3 = T.sparse_variable(J1, (n1, nnz), (indptr_3, indices_3), idtype="int64")

其存储的实例如下：

1
I = T.dense_fixed(3)
2

3
J = T.sparse_variable(I, (4, 6), (indices_1=[1, 0, 2, 3, 1, 3], indptr=[0, 1, 4, 6]), idtype="int32")
4

5
A = T.match_sparse_buffer(a, (I, J), dtype="float32", scope="global")

这里的 $A$ 就是一个 CSR 形式表示的矩阵，上图例子可以表示 $C = A B$

采用迭代器的存储设计，使得描述更易于拓展，可以定制更加复杂的矩阵结构，也方便描述一些具有特定分布的稀疏矩阵。

解决了矩阵的存储描述，我们还需要解决矩阵相乘的抽象，换而言之就是如何进行计算：

我们需要一个数据结构来指定程序的迭代空间，因此我们创建了一个叫做sparse iteration(稀疏迭代空间)的数据结构，其语法结构如下：

1
with T.sp_iter([I, J, K], "SSR", "sddmm") as [i, j, k]:
2
    with T.init():
3
        Y[i, j] = 0.0
4
    Y[i, j] = Y[i, j] + A[i, k] * B[j, k] * X[i, j]

其接收一组迭代器( [I, J, K] )，并生成一组循环变量 i, j, k，在这个稀疏迭代空间的主体中我们可以写任意的数学表达式，来描述我们的计算，在稀疏迭代空间的语义中，我们无需考虑 A, B ,X , Y 具体的存储格式，只需当作他们是稠密数据结构来描述数学运算，我们将这个阶段的数学运算称为(坐标空间计算)。

`Stage` 解读

`Stage I`

格式分解
空间坐标计算
调度

在前文中我们提到可以把计算分解到更小的结构化稀疏矩阵中，在SparseTIR中，这通过一个pass: DecomposeFormat来实现:

用户需要提供一组FormatRewriteRule(使用TVMScript)来指定新的子矩阵格式，作为参数传给DecomposeFormat这个pass，随后这个pass将会生成新的子矩阵所需要的axis和sparse buffer，并把原来在原有稀疏迭代空间中的矩阵运算改写成在新的结构化稀疏迭代空间上的运算

如下图例子所示，将CSR格式的计算分解为BSR格式的计算后生成的稀疏矩阵-矩阵乘法（SpMM）的IR，块大小为 2，每行有 2 个非零列的 ELL 格式。除了对新格式的稀疏计算外，还产生了另外两个将数据从原始格式复制到新格式的稀疏迭代。当要分解的稀疏矩阵是静止的，我们可以在预处理步骤中进行数据复制，以避免运行时格式转换的开销。

这个 pass 会为新建的矩阵分配 axis 与缓冲区，并且会生成迭代空间，如：若定义了 $C = A \times B$ ，代码如下：

1
with T.sp_iter([I, J, K], "SSR", "sddmm") as [i, j, k]:
2
    with T.init():
3
        C[i, j] = 0.0
4
    C[i, j] = Y[i, j] + A[i, k] * B[j, k]

在分解矩阵后，pass 会重写计算规则如下：

这一步就被称之为 空间坐标计算

注意这个pass本身对数据是不可知的，如何生成新的子矩阵每个维度的indptr/indices等信息，需要用户在运行时完成

SparseTIR为一些论文中所使用的格式分解模式提供了一些运行时函数，对于通用的格式分解，如何自动化地推断出indptr/indices等信息，仍然是一个有趣的研究问题

关于第一阶段的调度，定义了两个调度源语，sparse_reorder 和 sparse_fuse 。

sparse_reorder 稀疏迭代中的稀疏轴的顺序影响第二阶段生成的循环的顺序。这个基元可以对稀疏轴的顺序进行操作
sparse_fuse 将一个给定的稀疏迭代中的几个迭代器融合成一个。例如想要一个单一的循环而不是两个在所有非零元素上迭代的嵌套循环。

我们可以从下图理解这两个源语：

`Stage II`

区分 Stage-I 和 Stage-II 的一个重要标志是坐标空间(coordinate space)和位置空间(position space)的区分

其含义是逻辑上的坐标和在压缩存储格式下的物理位置，例如下图中，假设我们以 1-base 计数，元素-1的坐标为4，而位置为2

对应到 SparseTIR 中，我们在1阶段的迭代变量和稀疏矩阵的访问语义均是在坐标空间中，到了2阶段，我们希望把所有的稀疏迭代数据结构改写成TensorIR 中的嵌套循环和block ，并把稀疏矩阵访问和循环变量的语义改写到坐标空间，使得2阶段的IR与TensorIR的调度原语兼容。

我们设计了一个pass: Sparse Iteration Lowering，用于从1阶段到2阶段IR的转换，坐标空间到位置空间的转换算法涉及稀疏结构的压缩/解压过程。

这个 pass 分为四步：

Auxiliary buffer materialization 在创建 axis 时，indptr和 indices 的指针被指定为参数。在第二阶段，我们需要明确声明这些辅助缓冲区，以便在确定循环范围和翻译坐标时访问它们的值。除了辅助缓冲区外，我们还创建了指示缓冲区值域的提示，这些提示在第二阶段执行调度时用于整数集分析。下图展示了其工作方式：
Nested loop generation 这一步将第一阶段的稀疏迭代重组为第二阶段的嵌套循环：我们在稀疏迭代中的每个 axis 发出一个循环。产生的循环从 $0$ 开始。它们被 TensorIR 的块结构分开，建立边界以防止跨块循环的重新排序。此外，我们在最内层生成的循环内添加一个块，并将原始稀疏迭代的主体置于其中。下图显示了不同稀疏迭代的嵌套循环结构：
Coordinate translation 这一步将用于访问稀疏缓冲区的指数从坐标空间改写为位置空间，以弥补第一阶段和第二阶段之间的语义差距。具体的做法可以参照下图：

这里使用了在第一步中创建的辅助缓冲区 J_indptr 与 J_indices
Read/Write Region Analysis 缓冲区的读/写区域信息对于 TensorIR 的块构造是必要的。我们执行一个缓冲区区域分析来收集缓冲区的访问信息，并将每个块内访问的所有读/写区域联合起来，并将它们注释为块属性

关于这阶段的调度，其负责操作循环（fuse/reorder/split），在内存层次中移动数据（cache_read/cache_write），将循环与物理/逻辑线程绑定以使其并行化，并在硬件中使用矢量/张量指令（vectorize/tensorize）。由于在第二阶段已经完成了 SparseTIR 与 TensorIR 中调度源语的兼容，因此，我们在第二阶段完全支持 TVM 调度源语。

值得一提的是，SparseTIR 目前只支持生成嵌套循环，尚不支持生成 co-iterations

co-iteration结构指在两个稀疏向量的交集/并集做迭代，但具体的介绍在论文中并没有提及，暂时不清楚这应该是一个什么样的问题

`Stage III`

此阶段是我们的目标IR: TensorIR，在2阶段中我们依然保留了axes 和 sparse buffers这些数据结构，而3阶段中我们需要把它们全部改写成普通的 buffer，删去所有的稀疏抽象。

这里我们通过 Sparse Buffer Lowering 方法来解决：

利用 indptr 信息，计算出在一维连续储存下的offset，举例而言， $A [i, j]$ 会被翻译为 $A [J_in d pt r [i] + j]$ ，于是，所有的 sparse buffers 都能被扁平化为一维向量的形式，如下图所示：

比如，二阶段中的

1
with T.init():
2
  C[vi, vbi, vf] = T.float32(0)

会被翻译为：

1
with T.init():
2
  C_data[vi * (blk * feat_size) + vbi * feat_size + vf] = T.float32(0)

C 被改写为了一维向量，而对这些 buffer 的访问规则也进行了相应的改写，第3阶段的IR是完全的TensorIR，可以直接使用TVM的编译栈生成后端代码。

`TVM` 后端上的变化

由于在格式分解之后我们会生成许多在这些小矩阵上的算子，对于 CUDA 后端，启动这些 kernel 会有一定的开销，因此我们在 TVM 的加入了Horizontal Fusion作为后处理，把这些小矩阵上的算子融合到同一个算子中，例如下图所示：

图源自论文原作者

源码解析

在之前提到，SparseTIR 是 TIR 的一个稀疏化版本，以 TIR 为基础发展而来的，因此我们只需要在 TIR 版本上看 SparseTIR 的改进即可。

并且 SparseTIR 的工作重点在数据的表示与处理上，所以这里就从data structure → pass → schedule 来解读源码

在 sparsetir/python/tvm/script/tir/__init__.pyi 中列出了 sparse 增加的数据结构与函数

1
"""
2
sparse
3
"""
4

5
class Axis: ...
6

7
def dense_fixed(length: PrimExpr, idtype: str = "int32"): ...
8
def dense_variable(
9
    parent_axis: Axis,
10
    sizes: Tuple[Union[PrimExpr, int], Union[PrimExpr, int]],
11
    data: Var,
12
    idtype: str = "int32",
13
): ...
14
def sparse_fixed(
15
    parent_axis: Axis,
16
    sizes: Tuple[Union[PrimExpr, int], Union[PrimExpr, int]],
17
    data: Var,
18
    idtype: str = "int32",
19
): ...
20
def sparse_variable(
21
    parent_axis: Axis,
22
    sizes: Tuple[Union[PrimExpr, int], Union[PrimExpr, int]],
23
    data: Tuple[Var, Var],
24
    idtype: str = "int32",
25
): ...
26
def match_sparse_buffer(data: Var, axes: Sequence[Axis], dtype: str = "float32"): ...
27
def alloc_sparse_buffer(axes: Sequence[Axis], dtype: str, scope: str = "float32"): ...
28
def assume_buffer_domain(buf: Buffer, dom: Tuple[PrimExpr, PrimExpr]): ...
29

30
class iter(ContextManager):
31
    def __init__(
32
        self, axes: Sequence[Axis], iter_types: str, name_hint: str = ""
33
    ) -> None: ...
34
    def __enter__(self) -> Sequence[IterVar]: ...

数据结构

`Axis`

在之前提到过，这是一个迭代器，其实现如下：

1
@tvm._ffi.register_object("tir.sparse.Axis")
2
class Axis(Object):
3
    name: str
4
    parent: Optional["Axis"]
5
    length: PrimExpr
6
    nnz: PrimExpr
7
    nnz_cols: Optional[PrimExpr]
8
    indptr: Optional[Var]
9
    indices: Optional[Var]
10
    idtype: str
11
    sorted: bool
12

13
    def __init__(
14
        self, name, parent, length, nnz, nnz_cols, indptr, indices, idtype, sorted
15
    ) -> None:
16
        self.__init_handle_by_constructor__(_ffi_api.Axis, name, parent, length, nnz, nnz_cols, indptr, indices, idtype, sorted)  # type: ignore

这里向 tvm 框架中注册了一个 Object 类，其名称为 Axis

注意这里的 PrimExpr 实际上是一个 Node 类，称为源语表达式，用于调度时的分析等，而 Node 的定义如下：

1
class Node(Object):
2
    """Base class of all IR Nodes, implements astext function."""
3

4
    def astext(self, show_meta_data=True, annotate=None):
5
        return _ffi_api.AsText(self, show_meta_data, annotate)
6

7
    def __str__(self):
8
        return _ffi_api.PrettyPrint(self)

而 Var 是可以看作是一个迭代器类，定义如下

1
@tvm._ffi.register_object("tir.Var")
2
class Var(PrimExprWithOp):
3
    def __init__(self, name: str, dtype: Union[str, ir.Type], span: Optional[Span] = None):
4
        self.__init_handle_by_constructor__(_ffi_api.Var, name, dtype, span)  # type: ignore

Axis 具体的参数解释如下：

name : Axis 的名称
parent : Axis 的父节点
length : 当前 Axis 的长度上界
nnz : 从根节点（最初的 Axis ）到当前 Axis 的非零元素的累计数量
nnz_cols : 当前行非零列数，仅对固定轴有效
indptr : 索引指针
indices : 非零元素的索引
idtype : 索引的数据类型（可以是 int32 ， float16 等）
sorted : 索引是否能被排序

矩阵声明

矩阵分为四种：

稠密定长
稠密变长
稀疏定长
稀疏变长

dense_fixed_axis 定义如下：

1
def dense_fixed_axis(name: str, length: PrimExpr, idtype: str) -> Axis:
2
    return Axis(name, None, length, length, length, None, None, idtype, True)

实际在使用时，我们可以通过如下方式使用：

1
I = T.dense_fixed(m)

其中 m 是这里要求的 length

构造的映射通过下面来实现：

1
@register
2
class DenseFixed(SpecialStmt):
3
    """Special Stmt for creating dense fixed axis."""
4

5
    def __init__(self):
6
        def dense_fixed(length: PrimExpr, idtype: str = "int32", span: Optional[Span] = None):
7
            names = [x.id.name for x in self.node.lhs]
8
            if len(names) != 1:
9
                self.context.report_error(
10
                    f"`dense_fixed` expected assign to only one var, but got {names}", span
11
                )
12

13
            axis = dense_fixed_axis(names[0], length, idtype)
14
            self.context.func_sp_axes.append(axis)
15
            self.context.update_symbol(names[0], axis, self.node)
16

17
        super().__init__(dense_fixed, def_symbol=True)

dense_variable_axis 定义如下

1
def dense_variable_axis(
2
    name: str, parent: Axis, length: PrimExpr, nnz: PrimExpr, indptr: Var, idtype: str
3
) -> Axis:
4
    return Axis(name, parent, length, nnz, None, indptr, None, idtype, True)

应用时如下：

1
J2 = T.dense_variable(I, n, indptr_2, idtype="int32")

这里 (I, J2)构成了第一维连续，第二维也连续但是不定长的迭代空间，第一维的长度为 n ，第二位的最大长度也为 n

同样，映射的类实现如下：

1
@register
2
class DenseVariable(SpecialStmt):
3
    """Special Stmt for creating dense variable axis."""
4

5
    def __init__(self):
6
        def dense_variable(
7
            parent_axis: Axis,
8
            shape: Tuple[PrimExpr, PrimExpr],
9
            indptr_var: Var,
10
            idtype: str = "int32",
11
            span: Optional[Span] = None,
12
        ):
13
            names = [x.id.name for x in self.node.lhs]
14
            if len(names) != 1:
15
                self.context.report_error(
16
                    f"`dense_variable` expected assign to only one var, but got {names}", span
17
                )
18

19
            length, nnz = shape
20
            axis = dense_variable_axis(names[0], parent_axis, length, nnz, indptr_var, idtype)
21
            self.context.func_sp_axes.append(axis)
22
            self.context.update_symbol(names[0], axis, self.node)
23

24
        super().__init__(dense_variable, def_symbol=True)

这里的 nnz 可以为 None

sparse_fixed_axis

1
def sparse_fixed_axis(
2
    name: str,
3
    parent: Axis,
4
    length: PrimExpr,
5
    nnz_cols: PrimExpr,
6
    indices: Var,
7
    idtype: str,
8
    sorted: bool = True,
9
) -> Axis:
10
    return Axis(
11
        name, parent, length, parent.nnz * nnz_cols, nnz_cols, None, indices, idtype, sorted
12
    )
13

14

15
@register
16
class SparseFixed(SpecialStmt):
17
    """Special Stmt for creating sparse fixed axis."""
18

19
    def __init__(self):
20
        def sparse_fixed(
21
            parent_axis: Axis,
22
            shape: Tuple[PrimExpr, PrimExpr],
23
            indices_var: Var,
24
            idtype: str = "int32",
25
            sorted: bool = True,
26
            span: Optional[Span] = None,
27
        ):
28
            names = [x.id.name for x in self.node.lhs]
29
            if len(names) != 1:
30
                self.context.report_error(
31
                    f"`sparse_fixed` expected assign to only one var, but got {names}", span
32
                )
33

34
            length, nnz_cols = shape
35
            axis = sparse_fixed_axis(
36
                names[0], parent_axis, length, nnz_cols, indices_var, idtype, sorted=sorted
37
            )
38
            self.context.func_sp_axes.append(axis)
39
            self.context.update_symbol(names[0], axis, self.node)
40

41
        super().__init__(sparse_fixed, def_symbol=True)

具体使用如下：

1
J1 = T.sparse_fixed(I, (n, c), indices_1, idtype="int32")

注意这里，我们同时指定了 length 与 nnz 的值

sparse_variable_axis

1
def sparse_variable_axis(
2
    name: str,
3
    parent: Axis,
4
    length: PrimExpr,
5
    nnz: PrimExpr,
6
    indptr: Var,
7
    indices: Var,
8
    idtype: str,
9
    sorted: bool = True,
10
) -> Axis:
11
    return Axis(name, parent, length, nnz, None, indptr, indices, idtype, sorted)
12

13
@register
14
class SparseVariable(SpecialStmt):
15

16
    def __init__(self):
17
        def sparse_variable(
18
            parent_axis: Axis,
19
            shape: Tuple[PrimExpr, PrimExpr],
20
            data: Tuple[Var, Var],
21
            idtype: str = "int32",
22
            sorted: bool = True,
23
            span: Optional[Span] = None,
24
        ):
25
            names = [x.id.name for x in self.node.lhs]
26
            if len(names) != 1:
27
                self.context.report_error(
28
                    f"`sparse_variable` expected assign to only one var, but got {names}", span
29
                )
30

31
            length, nnz = shape
32
            indptr_var, indices_var = data
33
            axis = sparse_variable_axis(
34
                names[0], parent_axis, length, nnz, indptr_var, indices_var, idtype, sorted=sorted
35
            )
36
            self.context.func_sp_axes.append(axis)
37
            self.context.update_symbol(names[0], axis, self.node)
38

39
        super().__init__(sparse_variable, def_symbol=True)

具体使用如下：

1
J3 = T.sparse_variable(J1, (n1, nnz), (indptr_3, indices_3), idtype="int64")

在这里我们还指定了 indices 的值

以上就是基础的数据结构的定义，需要注意的是，这里实际上还是抽象，真正底层的实现是使用 tvm 中的 Array

对于 Buffer 部分 python 部分只有定义，完整的实现在 cpp 中 src/tir/transforms/lower_sparse_buffer.cc，在此实现中，会将声明和 match_sparse_buffer 降低到 match_buffers，然后通过 lower_match_buffer.cc 中的实现，将指针与 buffer 绑定在一起

当然，如果作用域不是全局，那么其储存位置在 shared memory 上，属于临时分配的 buffer，不需要与指针进行绑定

编译

在之前提到，SparseTIR 的目标是转化为 TensorIR ，但是为实现可组合格式与可组合变换，他加入了两个 pass 来完成这一点：

DecomposeFormat 目的是把计算分解到更小的结构化稀疏矩阵，最后进行空间坐标计算
Sparse Iteration Lowering 目的是将空间坐标转换为位置坐标，进行稀疏结构的压缩/解压

下面对三个阶段进行阐述，详细解读这两个 pass 的做法

`Stage I`

考虑 DecomposeFormat，其需要传递一个 FormatRewriteRule 类的参数，此类的定义如下：

1
@tvm._ffi.register_object("tir.sparse.FormatRewriteRule")
2
class FormatRewriteRule(Object):
3

4
    def __init__(
5
        self,
6
        name: str,
7
        new_format_desc: tvm.tir.PrimFunc,
8
        buffers_to_rewrite: List[str],
9
        axes_before_rewrite: List[str],
10
        axes_after_rewrite: List[str],
11
        axis_map: Dict[str, List[str]],
12
        idx_map: Union[Callable, IndexMap],
13
        inv_idx_map: Union[Callable, IndexMap],
14
    ) -> None:
15
        if isinstance(idx_map, Callable):
16
            idx_map = IndexMap.from_func(idx_map)
17
        if isinstance(inv_idx_map, Callable):
18
            inv_idx_map = IndexMap.from_func(inv_idx_map)
19
        self.__init_handle_by_constructor__(
20
            _ffi_api.FormatRewriteRule,
21
            name,
22
            new_format_desc,
23
            buffers_to_rewrite,
24
            axes_before_rewrite,
25
            axes_after_rewrite,
26
            axis_map,
27
            idx_map,
28
            inv_idx_map,
29
        )  # type: ignore

实际上需要传递的就是：

重写方法的名称
新矩阵的存储格式
需要被重写的缓冲区
重写前的迭代器
重写后生成的迭代器
重写前后迭代器的映射方式
格式的映射方式

示例如下：

1
FormatRewriteRule(
2
    0 + "_" + 1,
3
    ell.specialize({nnz_cols_symbol: bucket_size}),
4
    ["A"],
5
    ["I", "J"],
6
    ["O", "I", "J"],
7
    {"I": ["O", "I"], "J": ["J"]},
8
    csr2ell_index_map,
9
    csr2ell_inv_index_map,
10
)

其中 ell.specialize 是一个脚本序列化参数的方法，通过将缺少的参数填入来完成一个脚本函数的生成，比如上面，就会将 bucket_size 的值填入所有 nnz_cols_symbol 占用的位置

最后的是两个 csr2ell 的索引映射方式（索引包括 indices 与 indptr）

这里提供了很强的自定义性，用户可以设计自己的矩阵存储格式，只需要重写索引的映射规则，那么就可以进行转化

传入这个自定义的参数后，执行

1
mod = tvm.IRModule.from_expr(csrmm)
2
mod = format_decompose(mod, rewrites)

注意这里的 mod 是一个 Tensor IR 模块，通过传入这个参数，进行矩阵格式的重组。

此 pass 的实现在 src/tir/transforms/sparse_format_decompose.cc 中

定义如下：

1
Pass SparseFormatDecompose(Array<FormatRewriteRule> composable_formats,
2
                           bool include_format_rewrite_blks) {
3
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
4
    return SparseFormatDecompose(std::move(composable_formats), std::move(f),
5
                                 include_format_rewrite_blks);
6
  };
7
  return CreatePrimFuncPass(pass_func, 0, "tir.SparseFormatDecompose", {});
8
}

对照 python 中的接口：

1
def format_decompose(
2
    mod: IRModule,
3
    composable_formats: List["FormatRewriteRule"],
4
    include_format_rewrite_blks: bool = True,
5
):
6
    """Rewrite the sparse format of sparse buffers in the TIR scripts.
7

8
    Parameters
9
    ----------
10
    mod : IRModule
11
        The IRModule to lower.
12
    composable_formats : List[FormatRewriteRule]
13
        Composable formats is a list of rewrite rules.
14
    include_format_rewrite_blks : bool
15
        Whether to include format rewrite blocks in the output.
16
    """
17
    if not isinstance(mod, IRModule):
18
        raise TypeError("Expected IRModule, but got {}".format(type(mod)))
19
    return SparseFormatDecompose(composable_formats, include_format_rewrite_blks)(mod)

实际上就是返回生成的 IR ，这里我们考虑 lambda 函数 pass_func 中的 SparseFormatDecompose :

1
PrimFunc SparseFormxatDecompose(Array<FormatRewriteRule> composable_formats, PrimFunc f,
2
                                bool include_format_rewrite_blks = true) {
3
  CHECK(composable_formats.size() >= 1)
4
      << "The given composable formats length should be greater than or equal to 1.";
5
  // Only apply this pass to TIR that is not from TE schedules
6
  if (!IsFromLegacyTESchedule(f) && SparseTIRLevel(f) == 2) {
7
    // SparseFormatDecomposer rewriter(composable_formats);
8
    PrimFuncNode* fptr = f.CopyOnWrite();
9
    Array<PrimFunc> format_descs;
10
    Array<Axis> old_sp_axes = f->sp_axes;
11
    Array<SparseBuffer> old_buffers;
12
    for (const auto& kv : f->buffer_map) {
13
      old_buffers.push_back(Downcast<SparseBuffer>(kv.second));
14
    }
15
    for (const FormatRewriteRule& rule : composable_formats) {
16
      format_descs.push_back(AddSuffix(rule->new_format_desc, "_" + rule->name));
17
    }
18
    fptr->params = UpdateParams(format_descs, f->params);
19
    fptr->buffer_map = UpdateBufferMap(format_descs, f->buffer_map);
20
    fptr->sp_axes = UpdateSparseAxes(format_descs, f->sp_axes);
21
    Array<Stmt> format_rewrite_blks, compute_blks;
22
    // generate format rewrite blocks and compute blocks for each rule
23
    for (size_t i = 0; i < composable_formats.size(); ++i) {
24
      SparseFormatDecomposer rewriter(composable_formats[i], format_descs[i], old_sp_axes,
25
                                      old_buffers);
26
      rewriter(f->body);
27
      for (const Stmt& sp_iter : rewriter.format_rewrites_blks) {
28
        format_rewrite_blks.push_back(sp_iter);
29
      }
30
      for (const Stmt& sp_iter : rewriter.compute_blks) {
31
        compute_blks.push_back(sp_iter);
32
      }
33
    }
34
    // merge format rewrite and compute blocks.
35
    Array<Stmt> all_blks;
36
    if (include_format_rewrite_blks) {
37
      for (const Stmt& sp_iter : format_rewrite_blks) {
38
        all_blks.push_back(sp_iter);
39
      }
40
    }
41
    for (const Stmt& sp_iter : compute_blks) {
42
      all_blks.push_back(sp_iter);
43
    }
44
    Stmt new_body = all_blks.size() == 1 ? all_blks[0] : SeqStmt(all_blks);
45
    fptr->body = BlockRealize({}, const_true(), Block({}, {}, {}, "root", new_body, NullOpt, {}));
46
    // add composable flag.
47
    if (composable_formats.size() > 1) {
48
      Map<String, ObjectRef> new_attr_dict = fptr->attrs->dict;
49
      new_attr_dict.Set("composable", Integer(1));
50
      fptr->attrs = DictAttrs(new_attr_dict);
51
    }
52
    return f;
53
  } else {
54
    return f;
55
  }
56
}

注意到这里需要判断 SparseTIRLevel 是否为 2，防止 lowering 错误

主要是因为 SparseTIR 提供了与 TIR 相互转化的功能

这里返回的 f 就是分解完矩阵后的 IR了，在这之后，消除多余的缓冲区即可：

1
def RemovePreprocess():
2
    # Remove the preprocess blocks/sparse iterations in the module.
3

4
    return _ffi_api.RemovePreprocess()  # type: ignore

其 C++ 实现为：

1
PrimFunc RemovePreprocess(PrimFunc f) {
2
  if (!IsFromLegacyTESchedule(f)) {
3
    PrimFuncNode* fptr = f.CopyOnWrite();
4
    PreprocessRemover remover;
5
    fptr->body = remover(fptr->body);
6
    // insert extra parameters
7
    for (const Var& var : remover.extra_buffer_vars) {
8
      fptr->params.push_back(var);
9
      ICHECK(remover.extra_buffer_map.count(var))
10
          << "Internal error, extra_buffer_map do not have key " << var;
11
      fptr->buffer_map.Set(var, remover.extra_buffer_map.Get(var).value());
12
    }
13
    return f;
14
  } else {
15
    return f;
16
  }
17
}
18

19
Pass RemovePreprocess() {
20
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
21
    return RemovePreprocess(std::move(f));
22
  };
23
  return CreatePrimFuncPass(pass_func, 0, "tir.RemovePreprocess", {});
24
}

这里就完成了分解后的迭代与计算规则重写，最后返回 IR

接着，我们需要进行调度，在这一阶段支持的调度只有 reorder 与 fuse ，其接口如下：

1
    def sparse_reorder(self, block: SparseIterationRV, new_order: List[SpIterVar]) -> None:
2
        _ffi_api.ScheduleSparseReorder(  # type: ignore # pylint: disable=no-member
3
            self,
4
            block,
5
            new_order,
6
        )
7

8
    def sparse_fuse(self, block: SparseIterationRV, iters_to_fuse: List[SpIterVar]) -> None:
9
        _ffi_api.ScheduleSparseFuse(  # type: ignore # pylint: disable=no-member
10
            self,
11
            block,
12
            iters_to_fuse,
13
        )

我们可以通过下面的方法来使用(mod为IR)：

1
sch = tir.Schedule(mod)
2
sp_iteration = sch.get_sparse_iteration("sddmm")
3
i, j, k = sch.get_sp_iters(sp_iteration)
4
sch.sparse_reorder(sp_iteration, [k, i, j])
5
sch.sparse_fuse(sp_iteration, [i, j])

而其具体实现为：

1
SparseIteration SparseReorder(ScheduleState self, const SparseIteration& sp_iteration,
2
                              const Array<SpIterVar>& new_order) {
3
  // Step 1. Check whether the iterators in `new_order` are the same as `sp_iteration`'s iterators.
4
  CheckValidInputIterators(self, new_order, sp_iteration->sp_iter_vars);
5

6
  // Step 2. Check whether the new order does not break the iterator dependency.
7
  CheckDependency(self, new_order);
8

9
  // Step 3. Create the new SparseIteration.
10
  ObjectPtr<SparseIterationNode> p_new_sp_iteration =
11
      make_object<SparseIterationNode>(*sp_iteration.get());
12
  p_new_sp_iteration->sp_iter_vars = new_order;
13
  SparseIteration new_sp_iteration(p_new_sp_iteration);
14

15
  UpdateIRModule(self, sp_iteration, new_sp_iteration);
16
  return new_sp_iteration;
17
}
18

19
SparseIteration SparseFuse(ScheduleState self, const SparseIteration& sp_iteration,
20
                           const Array<SpIterVar>& iters_to_fuse) {
21
  // Step 1. Check match or not.
22
  int match_pos = CheckFuseMatch(self, iters_to_fuse, sp_iteration->sp_iter_vars);
23

24
  ObjectPtr<SparseIterationNode> p_new_sp_iteration =
25
      make_object<SparseIterationNode>(*sp_iteration.get());
26
  Array<SpIterVar> new_sp_iters;
27
  for (int i = 0; i < match_pos; ++i) {
28
    new_sp_iters.push_back(sp_iteration->sp_iter_vars[i]);
29
  }
30
  Array<Axis> axis_group;
31
  for (const SpIterVar& sp_iter_var : iters_to_fuse) {
32
    axis_group.push_back(sp_iter_var->axis);
33
  }
34
  for (size_t i = 0; i < iters_to_fuse.size(); ++i) {
35
    const SpIterVar& sp_iter_var = iters_to_fuse[i];
36
    Axis new_axis = FusedAxis(axis_group, i);
37
    new_sp_iters.push_back(SpIterVar(sp_iter_var->var, sp_iter_var->is_reduction, new_axis));
38
  }
39
  for (size_t i = match_pos + iters_to_fuse.size(); i < sp_iteration->sp_iter_vars.size(); ++i) {
40
    new_sp_iters.push_back(sp_iteration->sp_iter_vars[i]);
41
  }
42
  p_new_sp_iteration->sp_iter_vars = new_sp_iters;
43
  SparseIteration new_sp_iteration(p_new_sp_iteration);
44

45
  UpdateIRModule(self, sp_iteration, new_sp_iteration);
46
  return new_sp_iteration;
47
}

较为复杂的是 fuse 的实现，这里根据 fuse 规则，重写了循环迭代，并且生成了新的 Axis 来适应新的循环迭代器

`Stage II`

在这一步，我们需要将 SparseTIR 的迭代重写为 TIR 的迭代形式

因此，在这里加入了一次 lowering：

1
def lower_sparse_iter(mod: IRModule, check_invalid_binary_search: bool = False):
2
    """Lower sparse iterators in Sparse TIR.
3
    """
4

5
    if not isinstance(mod, IRModule):
6
        raise TypeError("Expected IRModule, but got {}".format(type(mod)))
7
    return LowerSparseIter(check_invalid_binary_search)(mod)

其 C++ 实现为：

1
PrimFunc LowerSparseIter(PrimFunc f, bool check_invalid_binary_search) {
2
  // Only apply this pass to TIR that is not from TE schedules
3
  if (!IsFromLegacyTESchedule(f) && SparseTIRLevel(f) == 2) {
4
    PrimFuncNode* fptr = f.CopyOnWrite();
5
    // Step 1. Update the PrimFunc's buffer map.
6
    Map<Axis, SparseBuffer> axis_indptr_map, axis_indices_map;
7
    Array<BufferDomain> buf_doms;
8
    std::tie(axis_indptr_map, axis_indices_map, fptr->buffer_map, fptr->sp_axes, buf_doms) =
9
        UpdateMetadata(f);
10
    // Step 2. Lower iterations.
11
    IterTransformer lower_sparse(axis_indptr_map, axis_indices_map, fptr->sp_axes,
12
                                 check_invalid_binary_search);
13
    Stmt body = lower_sparse(std::move(fptr->body));
14
    // Step 3. Wrap with root block, insert bsearch blocks and allocated buffers.
15
    if (!lower_sparse.bsearch_structures.empty()) {
16
      Array<Stmt> seq;
17
      for (const auto& bsearch_struct : lower_sparse.bsearch_structures) {
18
        seq.push_back(bsearch_struct.body);
19
      }
20
      seq.push_back(body);
21
      body = SeqStmt(seq);
22
    }
23
    buf_doms = Concat(buf_doms, lower_sparse.alloc_buf_doms);
24
    Block root_block({}, {}, {}, "root", body, NullOpt, lower_sparse.root_alloc_buffers, {},
25
                     buf_doms);
26
    fptr->body = BlockRealize({}, const_true(), std::move(root_block));
27
    // Step 4. Lower sparse tir level.
28
    Map<String, ObjectRef> new_attr_dict = fptr->attrs->dict;
29
    new_attr_dict.Set("sparse_tir_level", Integer(1));
30
    fptr->attrs = DictAttrs(new_attr_dict);
31
    // Step 5. postprocess bufferload with possible invalid indices
32
    if (check_invalid_binary_search) {
33
      fptr->body = InvalidIndicesPostProcess()(std::move(fptr->body));
34
    }
35
    return f;
36
  } else {
37
    return f;
38
  }
39
}
40

41
Pass LowerSparseIter(bool check_invalid_binary_search) {
42
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
43
    return LowerSparseIter(std::move(f), check_invalid_binary_search);
44
  };
45
  return CreatePrimFuncPass(pass_func, 0, "tir.LowerSparseIter", {});
46
}

此函数的流程与 [Stage II](#Stage II) 中描述的一致：

更新缓冲区的映射，包括建立辅助缓冲区等
将循环进行展开
在展开的循环中插入 block 进行包装
完成空间坐标到位置坐标的计算映射

生成符合 TIR 一样拥有块结构的迭代后，我们进入这一阶段的调度，而这里的调度我们可以直接使用 TIR 中的调度源语，而不需要自己再去实现 SparseTIR 的调度源语了（因为迭代结构一致）

举例如下：

1
sch = tir.Schedule(mod_preprocess)
2
blk = sch.get_block("binary_search_block_0_0")
3
(i,) = sch.get_loops(blk)
4
io, ii = sch.split(i, [None, 32])
5
sch.bind(ii, "threadIdx.x")
6
sch.bind(io, "blockIdx.x")

显然，这里的调度源语都是 TIR 中已经实现的内容

`Stage III`

在这一步，所做的就是把 buffer 与 axis lowering 到 TIR 的级别上，我们具体的使用如下：

1
mod = lower_sparse_buffer(sch.mod)

这样，得到的 mod 就是一个 TIR 级别的 IR 了

其具体做法如下：

1
PrimFunc LowerSparseBuffer(PrimFunc f) {
2
  // Only apply this pass to TIR that is not from TE schedules
3
  if (!IsFromLegacyTESchedule(f) && SparseTIRLevel(f) == 1) {
4
    bool is_horizontal_fuse = f->HasNonzeroAttr("horizontal_fuse");
5
    PrimFuncNode* fptr = f.CopyOnWrite();
6
    // Step 1. Update the PrimFunc's buffer map.
7
    fptr->buffer_map = std::move(UpdateBufferMap(f));
8
    // Step 2. Lower sparse buffers.
9
    fptr->body = BufferTransformer(fptr->sp_axes, fptr->buffer_map,
10
                                   is_horizontal_fuse)(std::move(fptr->body));
11
    // Step 3. Remove sparse axes
12
    fptr->sp_axes.clear();
13
    // Step 4. Lower sparse tir level
14
    Map<String, ObjectRef> new_attr_dict = fptr->attrs->dict;
15
    new_attr_dict.Set("sparse_tir_level", Integer(0));
16
    fptr->attrs = DictAttrs(new_attr_dict);
17
    return f;
18
  } else {
19
    return f;
20
  }
21
}
22

23
namespace transform {
24

25
/*!
26
 * \brief The lowering pass from TIR to Sparse TIR.
27
 */
28
Pass LowerSparseBuffer() {
29
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
30
    return LowerSparseBuffer(std::move(f));
31
  };
32
  return CreatePrimFuncPass(pass_func, 0, "tir.LowerSparseBuffer", {});
33
}

注意到这里的 SparseTIRlevel 已经降为 1 ，这是因为在 Stage II 中，我们降低了其调度这一级别到 TIR 抽象

所以在这里，我们需要做的就是把 Sparse Buffer 与 Axis 抽象全都删去，将矩阵乘转变为通过 offset 完成的向量乘（也就是说通过偏移量来计算最后的结果），因此这一步实际上就是在计算 offset

完成这一步后，我们就到了 TIR 级别的抽象，相当于我们已经进入了 TVM ，直接使用 TVM 的 build 即可：

1
f = tvm.build(mod, target="cuda")

然后调用函数运行：

1
evaluator = f.time_evaluator(f.entry_name, tvm.cuda(0), number=100)

复现

经作者指出，AE 的所有代码与实验（baseline 等）均在已开源中。此部分在论文的附录 B.3 Description 中提及

在 example 中给出了如何使用 SparseTIR 的示例，如 spmm :

1
import dgl
2
import tvm
3
import tvm.testing
4
import tvm.tir as tir
5
import scipy.sparse as sp
6
import argparse
7
import numpy as np
8
import torch as th
9
from tvm.script import tir as T
10
from tvm.sparse import (
11
    FormatRewriteRule,
12
    lower_sparse_buffer,
13
    lower_sparse_iter,
14
    column_part_hyb,
15
    format_decompose,
16
)
17
import tvm.sparse
18
from utils import get_dataset, ell
19

20

21
@T.prim_func
22
def csrmm(
23
    a: T.handle,
24
    b: T.handle,
25
    c: T.handle,
26
    indptr: T.handle,
27
    indices: T.handle,
28
    m: T.int32,
29
    n: T.int32,
30
    num_tiles: T.int32,
31
    nnz: T.int32,
32
    cwm: T.int32,
33
) -> None:
34
    T.func_attr({"global_symbol": "main", "tir.noalias": True, "sparse_tir_level": 2})
35
    I = T.dense_fixed(m)
36
    J = T.sparse_variable(I, (n, nnz), (indptr, indices), "int32")
37
    J_detach = T.dense_fixed(n)
38
    K1 = T.dense_fixed(num_tiles)
39
    K2 = T.dense_fixed(cwm)
40
    K3 = T.dense_fixed(32)
41
    A = T.match_sparse_buffer(a, (I, J), "float32")
42
    B = T.match_sparse_buffer(b, (J_detach, K1, K2, K3), "float32")
43
    C = T.match_sparse_buffer(c, (I, K1, K2, K3), "float32")
44
    with T.sp_iter([I, J, K1, K2, K3], "SRSSS", "csrmm") as [i, j, k1, k2, k3]:
45
        with T.init():
46
            C[i, k1, k2, k3] = 0.0
47
        C[i, k1, k2, k3] = C[i, k1, k2, k3] + A[i, j] * B[j, k1, k2, k3]
48

49

50
class TorchOpTimer(object):
51
    def __enter__(self):
52
        self.start_event = th.cuda.Event(enable_timing=True)
53
        self.end_event = th.cuda.Event(enable_timing=True)
54
        self.start_event.record()
55
        return self
56

57
    def __exit__(self, type, value, traceback):
58
        self.end_event.record()
59
        th.cuda.synchronize()  # Wait for the events to be recorded!
60
        self.time = self.start_event.elapsed_time(self.end_event) / 1e3
61

62

63
def csr2ell_inv_index_map(o, i, j):
64
    return i, j
65

66

67
def csr2ell_index_map(i, j):
68
    return 0, i, j
69

70

71
cached_bucketing_format = None
72

73

74
def bench_hyb(
75
    g,
76
    x,
77
    y_golden,
78
    feat_size=128,
79
    bucket_sizes=[],
80
    coersening_factor=2,
81
    num_col_parts=1,
82
    use_implicit_unroll=False,
83
):
84
    num_buckets = len(bucket_sizes)
85
    coersening_factor = min(coersening_factor, feat_size // 32)
86
    indptr, indices, _ = g.adj_sparse("csc")
87
    m = g.num_dst_nodes()
88
    n = g.num_src_nodes()
89
    nnz = g.num_edges()
90
    global cached_bucketing_format
91
    if cached_bucketing_format is None:
92
        indptr_nd = tvm.nd.array(indptr.numpy(), device=tvm.cpu())
93
        indices_nd = tvm.nd.array(indices.numpy(), device=tvm.cpu())
94
        cached_bucketing_format = column_part_hyb(
95
            m, n, indptr_nd, indices_nd, num_col_parts, bucket_sizes
96
        )
97
    row_indices, col_indices, mask = cached_bucketing_format
98

99
    # rewrite csrmm
100
    nnz_cols_symbol = ell.params[-1]
101
    rewrites = []
102
    for part_id in range(num_col_parts):
103
        for bucket_id, bucket_size in enumerate(bucket_sizes):
104
            rewrites.append(
105
                FormatRewriteRule(
106
                    str(part_id) + "_" + str(bucket_id),
107
                    ell.specialize({nnz_cols_symbol: bucket_size}),
108
                    ["A"],
109
                    ["I", "J"],
110
                    ["O", "I", "J"],
111
                    {"I": ["O", "I"], "J": ["J"]},
112
                    csr2ell_index_map,
113
                    csr2ell_inv_index_map,
114
                )
115
            )
116
    mod = tvm.IRModule.from_expr(csrmm)
117
    mod = format_decompose(mod, rewrites)
118
    mod = tvm.tir.transform.RemovePreprocess()(mod)
119

120
    # specialize
121
    params = mod["main"].params
122
    param_map = {
123
        params[5]: m,  # m
124
        params[6]: n,  # n
125
        params[7]: feat_size // coersening_factor // 32,  # num_tiles,
126
        params[8]: nnz,  # nnz
127
        params[9]: coersening_factor,  # coersening_factor
128
    }
129
    for part_id in range(num_col_parts):
130
        for bucket_id in range(num_buckets):
131
            param_map[params[10 + 7 * (part_id * num_buckets + bucket_id) + 4]] = m
132
            param_map[params[10 + 7 * (part_id * num_buckets + bucket_id) + 5]] = n
133
            param_map[params[10 + 7 * (part_id * num_buckets + bucket_id) + 6]] = row_indices[
134
                part_id
135
            ][bucket_id].shape[0]
136

137
    mod["main"] = mod["main"].specialize(param_map).with_attr("horizontal_fuse", True)
138

139
    # schedule
140
    sch = tvm.tir.Schedule(mod)
141
    for sp_iter_name in [
142
        "csrmm_{}_{}".format(i, j) for j in range(num_buckets) for i in range(num_col_parts)
143
    ]:
144
        sp_iteration = sch.get_sparse_iteration(sp_iter_name)
145
        o, i, j, k1, k2, k3 = sch.get_sp_iters(sp_iteration)
146
        sch.sparse_fuse(sp_iteration, [o, i])
147

148
    mod = sch.mod
149
    mod = tvm.sparse.lower_sparse_iter(mod)
150
    sch = tvm.tir.Schedule(mod)
151
    for part_id in range(num_col_parts):
152
        for bucket_id, bucket_size in enumerate(bucket_sizes):
153
            is_atomic = num_col_parts > 1 or bucket_id + 1 == num_buckets
154
            blk = sch.get_block("csrmm_{}_{}0".format(part_id, bucket_id))
155
            i, j, foo, foi, fi = sch.get_loops(blk)
156
            sch.reorder(foo, fi, j, foi)
157
            if is_atomic:
158
                sch.annotate(blk, "atomic", True)
159
                write_blk = sch.reverse_cache_write(blk, 0, "local")
160
                sch.reverse_compute_at(write_blk, fi, True)
161
                # sch.unroll(sch.get_loops(write_blk)[-2])
162
            sch.bind(fi, "threadIdx.x")
163
            sch.bind(foo, "blockIdx.y")
164
            sch.unroll(foi)
165
            if use_implicit_unroll:
166
                sch.annotate(foi, "pragma_unroll_explicit", 0)
167
            sch.unroll(j)
168
            if use_implicit_unroll:
169
                sch.annotate(j, "pragma_unroll_explicit", 0)
170
            io, ioi, ii = sch.split(i, [None, bucket_sizes[-1] // bucket_size, 8])
171
            sch.bind(io, "blockIdx.x")
172
            sch.bind(ii, "threadIdx.y")
173
            init_blk = sch.decompose_reduction(blk, fi)
174
            ax0, ax1 = sch.get_loops(init_blk)[-2:]
175
            sch.bind(ax0, "threadIdx.x")
176
            sch.unroll(ax1)
177
            if use_implicit_unroll:
178
                sch.annotate(ax1, "pragma_unroll_explicit", 0)
179

180
    mod = tvm.sparse.lower_sparse_buffer(sch.mod)
181
    mod = tvm.tir.transform.RemoveUnusedArgs()(mod)
182
    f = tvm.build(mod, target="cuda")
183

184
    # prepare nd array
185
    b_nd = tvm.nd.array(
186
        x.numpy().reshape(-1).astype("float32"),
187
        device=tvm.cuda(0),
188
    )
189
    c_nd = tvm.nd.array(np.zeros((n * feat_size,)).astype("float32"), device=tvm.cuda(0))
190
    # prepare args
191
    args = [b_nd, c_nd]
192

193
    for part_id in range(num_col_parts):
194
        for bucket_id, _ in enumerate(bucket_sizes):
195
            weight = tvm.nd.array(
196
                mask[part_id][bucket_id].numpy().reshape(-1).astype("float32"), device=tvm.cuda(0)
197
            )
198
            rows = tvm.nd.array(
199
                row_indices[part_id][bucket_id].numpy().astype("int32"), device=tvm.cuda(0)
200
            )
201
            cols = tvm.nd.array(
202
                col_indices[part_id][bucket_id].numpy().reshape(-1).astype("int32"),
203
                device=tvm.cuda(0),
204
            )
205
            args += [weight, rows, cols]
206

207
    # test accuracy
208
    f(*args)
209
    tvm.testing.assert_allclose(c_nd.numpy().reshape(-1, feat_size), y_golden.numpy(), rtol=1e-4)
210

211
    # evaluate time
212
    evaluator = f.time_evaluator(f.entry_name, tvm.cuda(0), number=100)
213
    print("tir hyb time: {:.5f}ms".format(evaluator(*args).mean * 1000))
214

215

216
col_part_config = {
217
    "arxiv": 1,
218
    "proteins": 8,
219
    "pubmed": 1,
220
    "citeseer": 1,
221
    "cora": 1,
222
    "ppi": 16,
223
    "reddit": 8,
224
    "products": 16,
225
}
226

227
bucketing_config = {
228
    "arxiv": [1, 2, 4, 8, 16, 32],
229
    "proteins": [1, 2, 4, 8, 16, 32, 64, 128, 256],
230
    "pubmed": [1, 2, 4, 8, 16, 32],
231
    "citeseer": [1, 2, 4],
232
    "cora": [1, 2, 4],
233
    "ppi": [1, 2, 4, 8, 16, 32],
234
    "products": [1, 2, 4, 8, 16, 32],
235
    "reddit": [1, 2, 4, 8, 16, 32, 64, 128, 256, 512],
236
}
237

238

239
if __name__ == "__main__":
240
    parser = argparse.ArgumentParser("hybrid format spmm in sparse-tir")
241
    parser.add_argument("--dataset", "-d", type=str, default="arxiv", help="dataset name")
242
    parser.add_argument("--implicit-unroll", "-i", action="store_true", help="use implicit unroll")
243
    args = parser.parse_args()
244
    name = args.dataset
245
    g = get_dataset(name)
246

247
    for feat_size in [32, 64, 128, 256, 512]:
248
        print("feat_size =", feat_size)
249
        x = th.rand((g.num_src_nodes(), feat_size))
250
        y_golden = dgl.ops.copy_u_sum(g, x)
251
        bench_hyb(
252
            g,
253
            x,
254
            y_golden,
255
            feat_size=feat_size,
256
            bucket_sizes=bucketing_config[name],
257
            coersening_factor=2,
258
            num_col_parts=col_part_config[name],
259
            use_implicit_unroll=args.implicit_unroll,
260
        )

`SpMM`

在 RTX 3090Ti 上复现 SparseTIR(no-hyb) 与 SparseTIR(hyb) 的结果如图：

`SDDMM`

在 RTX 3090Ti 上复现 dgl 与 SparseTIR(hyb) 的结果如图：

また夏を追う

最近的笔记

GTest 单元测试模板

如何知道实验什么时候结束

探索

SparseTIR 解读

`Introduction`

系统架构

`SparseTIR` 的抽象设计

可组合的格式

可组合的程序变换

格式抽象与矩阵存储

`Stage` 解读

`Stage I`

`Stage II`

`Stage III`

`TVM` 后端上的变化

源码解析

数据结构

`Axis`

矩阵声明

编译

`Stage I`

`Stage II`

`Stage III`

复现

`SpMM`

`SDDMM`

🕸️ 关系图谱

目录

反向链接

また夏を追う

最近的笔记

GTest 单元测试模板

如何知道实验什么时候结束

探索

SparseTIR 解读

Introduction

系统架构

SparseTIR 的抽象设计

可组合的格式

可组合的程序变换

格式抽象与矩阵存储

Stage 解读

Stage I

Stage II

Stage III

TVM 后端上的变化

源码解析

数据结构

Axis

矩阵声明

编译

Stage I

Stage II

Stage III

复现

SpMM

SDDMM

🕸️ 关系图谱

目录

反向链接

`Introduction`

`SparseTIR` 的抽象设计

`Stage` 解读

`Stage I`

`Stage II`

`Stage III`

`TVM` 后端上的变化

`Axis`

`Stage I`

`Stage II`

`Stage III`

`SpMM`

`SDDMM`