[Roadmap] Sparse TensorIR

[yzh119]

Design: #368

1. Sparse block lowering. (transformation)
   1. Sparse/Dense coordinates transformation. (@MasterJH5574 WIP)
      1. The same sparse iterator viewed in different sparse buffers.
   2. + block body
   3. tir.reads/writes for sparse buffers.
   4. convert SparseBuffers into normal Buffers
   5. convert SparseBlocks into normal Blocks
   6. ...
2. Python-side frontend.
   1. Python-binding of sparse buffers/blocks/formats. (@yzh119 WIP)
   2. Specialize sparse formats.
   3. ...
3. TIR intrinsics.
   1. lower_bound (@yzh119 WIP)
      1. Return offset or value (prefer offset)
   2. boundary check
4. Examples
   1. End-to-end BSR SpMM/SDDMM example and optimizations (@yzh119 )
   2. Sparse Softmax.

\infy: tensorization.

[tqchen]

Thanks @yzh119 . can you paste a TVMScript example mockup of the sparse TensorIR, and transformation allowed? This would help us greatly understand the relation between the current design rationale and the new one.

[yzh119]

Yes I'll elaborate more here.

[yzh119]

Sparse Formats

In Sparse TIR we have four kinds of axis.

1. Fixed Sparse.
2. Variable Sparse.
3. Fixed Dense
4. Variable Dense.

They can represent both sparse matrices and ragged tensors.
The length of fixed axis is a constant while the length of variable axis depends on its coordinates.
We use format tree to describe the dependency among axes.

For example, if we want to represent an irregular batched matrix multiplication:

For i in range(b):
    For j in range(n[i]):
        For k in range(m[i]):
            For l in range(k[i]):
                 C[i, j, k] = C[i, j, k] + A[i, j, l] * B[i, l, k]

Its dependent tree is:

     i
    /|\
   j k l

The following syntax describe how do we define such structure in Sparse TIR.

i = tir.sp.FixedDenseAxis(b)
j = tir.sp.VariableDenseAxis(N)
k = tir.sp.VariableDenseAxis(M)      (where N, M, K are input 1-dim buffers in this case)
l = tir.sp.VariableDenseAxis(K)
Fmt = tir.sp.format([[i, j], [i, k], [i, l]], (i, j, k, l))
A = tir.sp.match_buffer(A_handle, fmt, (i, j, l))
B = tir.sp.match_buffer(B_handle, fmt, (i, j, k))
C = tir.sp.match_buffer(C_handle, fmt, (i, k, l))

Sparse Blocks

A sparse block indicates

1. the axes involved in the block
2. whether these axes are used as spatial/reduce axis
3. how do we reorder/fuse these axes.

with tir.sp.block([i, j, k], [spatial, spatial, reduce], [[0, 1], [2]]) as [vi, vj, vk]:
    pass

[yzh119]

After discussion with @MasterJH5574 and @tqchen , we decide to update the syntax as following:

def sddmm(a: ty.handle, b: ty.handle, c: ty.handle, fmt: ty.handle):
    N = tir.var('n')
    M = tir.var('m')
    B = tir.var('b')  # block size
    K = tir.var('k')
    i = tir.match_axis(fmt, 'i', N)
    j = tir.match_axis(fmt, 'j', M)
    k = tir.match_axis(fmt, 'k', K)
    bi = tir.match_axis(fmt, 'bi', B)
    bj = tir.match_axis(fmt, 'bj', B)
    A = tir.match_buffer(a, (i, bi, k), 'float32')
    B = tir.match_buffer(b, (tir.to_dense(j), bj, k), 'float32')
    C = tir.match_buffer(c, (i, j, bi, bj), 'float32')
    for vi, vj in tir.fuse(tir.cord(i), tir.cord(j)):
        for vk in tir.cord(k):
            for vbi in tir.cord(bi):
                for vbj in tir.cord(vbj):
                    with tir.block([], 'sddmm'):
                        tir.block_attr({'sparse': True})
                        with tir.init():
                            C[vi, vj, vbi, vbj] = 0.
                        C[vi, vj, vbi, vbj] = C[vi, vj, vbi, vbj] +\
                            A[vi, vbi, vk] * B[vj, vbj, vk]

where tir.match_axis, tir.pos, tir.cord, tir.fuse are newly introduced keywords for sparse support in TIR.

Below we describe detailed syntax of our new design.

Format Definition

We write format definition in Python outside out TIR scripts:

fmt = tir.format(
    {
         "i": (tir.kDenseFixed, None),
         "j": (tir.kSparseVariable, "i")
    }
)

We specify the format via a Python dictionary:

1. key indicates the axis name.
2. value is a tuple where the first element indicate the axis property and the second element refers to its parent axis.

each axis has no more than one parent.

Sparse Tensor Declaration

Sparse Tensor is declared in Python as well.

indptr = [...]
indices = [...]
a = tir.sparse.tensor(
    data,
    indptr,
    indices,
)

where data is a tvm runtime array, indptr and indices are list of tvm runtime arrays, and they both have inherent data types.

Sparse Support in TIR scripts.

match_axis

Matches axis defined in format and bind it with attributes such as length and number of columns (for ellpack-like formats).

match_buffer

Note that we need to support match buffer given

pos/cord

pos iterate our non-zero elements in the given axis, and cord iterate our all elements in a given axis (from 0 to length).

fuse

fuse several sparse/dense axes, note that the indices conversion is not trivial affine maps so we use the keyword fuse to indicate such fusion.

to_dense

Convert sparse axis to dense ones.

block with sparse attribute

block we sparse attribute need to be lowered via a schedule primitive called lower_sparse.
No dense schedules are allowed to apply inside a block if there is a sparse attribute.

Schedules

lower_sparse(loops)

lower sparse do coordinate conversion according to sparse formats for a specific loop/loops.
Insert proper block isolation is there are dependency between two axis.

reorder(loops)

Reorder loops and check whether they violate format constraints.

fuse(loops)

Fuse loops and check whether they violate format constraints.

[MasterJH5574]

@yzh119 Thanks for the update! I'd like propose another design of formats as follows:

## Format proposal F2
fmt = tir.format({
  "name1": tir.DenseFixedAxis(),
  "name2": tir.DenseVariableAxis(name),  # `name` is the name of the axis it depends on
  "name3": tir.SparseFixedAxis(n_col),
  "name4": tir.SparseVariableAxis(),
})

---

The main difference between this design and the above format design is that in this design, we treat different axis kinds in different ways. It's reasonable because:

• Only a dense-variable axis depends on some other axis, as the length of the axis is determined by the iterator value of the axis it depends on. Sparse-variable and sparse-fixed axes don't depend on other axes, because their lengths are known at compile-time. To be clear, sparse-variable iterators depend on other iterators, but sparse-variable axes do not depend on other axes. I will explain this point in the proposal of SparseIterator.
• To define a sparse-fixed axis, we should provide it's n_col field.

[yzh119]

Could you elaborate more on sparse-variable axis do not depend on other axes? I don't get the point.

n_col could be specified in match_axis IMO.

[yzh119]

@MasterJH5574 A second thought on our proposal of tir.cord(axis, reduction=True), I think it better to bind reduction information to block instead of loops.

Because several blocks might share the same loop iterator, but one might view it as parallel axis and another one view it as reduction axis.

[MasterJH5574]

@yzh119 Here are more thoughts, which were send to the slack channel before 👀.

1. The sddmm example is incomplete - we should provide the format so that the example is more understandable.
2. In tir.match_buffer(...) we sometimes use tir.to_dense(axis) to convert a sparse-variable axis to the dense one. IMO we should have tir.to_sparse_variable(axis) and tir.to_sparse_fixed(axis, n_col) as well.
3. We should provide some examples which contain dense-variable axis and sparse-fixed axis.
4. As we use tir.cord(axis) and tir.pos(axis) to define iterators, I may need to tweak the detailed design of SparseIterator so that each iterator stores only the necessary information, leaving out some redundant information. Not a big deal. I'm going to do it today.

[MasterJH5574]

I think it better to bind reduction information to block instead of loops.
Because several blocks might share the same loop iterator, but one might view it as parallel axis and another one view it as reduction axis.

Yes you're right. It's possible that different blocks view a loop var differently. But in the above design we use opaque blocks (blocks that have no block iter) for sparse block. Therefore it's not very convenient to represent the reduction information in block signatures.

One possible design is to add some block iters, and one block iter in a sparse block is required to be bound to only one sparse iterator. An example might be like

for vi, vj in T.fuse(tir.cord(i), tir.cord(j)):
    for vk in T.cord(k):
        for vbi in T.cord(bi):
            for vbj in T.cord(vbj):
                with T.block('sddmm'):
                    T.block_attr({'sparse': True})
                    vi_ = T.sparse_axis(vi)
                    vj_ = T.sparse_axis(vj)
                    vk_ = T.sparse_axis(vk, reduction=True)
                    vbi_ = T.sparse_axis(vbi)
                    vbj_ = T.sparse_axis(vbj)
                    with T.init():
                        C[vi_, vj_, vbi_, vbj_] = 0.
                    C[vi_, vj_, vbi_, vbj_] = C[vi_, vj_, vbi_, vbj_] + A[vi_, vbi_, vk_] * B[vj_, vbj_, vk_]

(Just an example. The API of "sparse_axis" can change.)

Perhaps we should wait for the block iter/binding refactor in TVM script before converging to a detailed design.

[MasterJH5574]

n_col could be specified in match_axis IMO.

Oh yes I agree. So match_axis might have variable numbers of arguments.

---

Could you elaborate more on sparse-variable axis do not depend on other axes? I don't get the point.

To be more specific:

• What we really concern is the sparse iterator dependency, as some schedule primitives like Fuse and Reorder rely on the iterator dependency - if the iterator dependency requirement isn't satisfied, the primitive cannot be applied.
• Therefore, the key problem is to know that for a given iterator, what iterators it depends on.
• As discussed in a meeting, dense-fixed and sparse-fixed iterators don't depend on any other iterator. A dense-variable iterator depends on a specific iterator, and a sparse-variable iterator depends on a bunch of other iterators. Specifically,

  • for dense-variable axes, consider the batch matrix multiplication workload. In BMM, the second and third axes of A and B are all dense-variable axes, and iterator vi, vj, vk all depend on iterator vb. Why we know that vi, vj, vk depend on iterator vb? Because we defined that relationship in the format as "the second and third axes of A and B depend on the first axis". Note that this relationship can be stored nowhere but in the format. This is why we must store the axis being depended when define a dense-variable axis in format.

    ## BMM example (not written in sparse TIR)
    for vb in range(0, B):
      for vi in range(0, M[vb]):
        for vj in range(0, N[vb]):
          C[vb, vi, vj] = 0.0
          for vk in range(0, K[vb]):
            C[vb, vi, vj] += A[vb, vi, vk] * B[vb, vk, bi]

  • for sparse-variable axes, suppose A is a buffer, and its i-th axis is a sparse-variable one. Now if a sparse-variable iterator vi goes over this axes, its extent relies on the value of the iterators that go over the first i - 1 axes. For example, in the code below we say vi depends on v1, ..., vi-1. How do we know this? By seeing the BufferLoad A[v1, v2, ..., vi-1, vi], we can immediately know the iterators depended by vi. As a result, we don't need to store the relationship information in format. This is why I say a sparse-variable axis never depend on other axes.

    ## A example that is not rigorous
    for v1 in range(0, l1):
      for v2 in range(0, l2):
        ...
        for vi-1 in range(0, li-1):
          for vi in range(0, li):
            ...
            A[v1, v2, ..., vi-1, vi, ...] = ...

I don't know whether my explanation could convince you. I'll post the detailed proposal for sparse iterators and the iterator dependency soon.

[yzh119]

@MasterJH5574 , I'm thinking of the sparse softmax example:

if we write it in the opaque block, the program would look like this (it uses the minus max element trick to avoid overflow):

for vi in tir.cord(i):
    for vj in tir.pos(j, reduction=True):
        with tir.block('A_max'):
            tir.block_attr({'sparse': True})
            with tir.init():
                A_max[vi] = tir.const(-float("inf"), "float32")
            A_max[vi] = tir.max(A_max[vi], A[vi, vj])
    for vj in tir.pos(j):
        with tir.block('A_minus_exp'):
            tir.block_attr({'sparse': True})
            A_exp[vi, vj] = tir.exp(A[vi, vj] - A_max[vi])
    for vj in tir.pos(j, reduction=True):
        with tir.block():
            tir.block_attr({'sparse': True})
            with tir.init('Z'):
                Z[vi] = 0.
            Z[vi] = Z[vi] + A_exp[vi. vj]
    for vj in tir.pos(j):
        with tir.block('out'):
            tir.block_attr({'sparse': True})
            out[vi, vj] = A_exp[vi, vj] / Z[vi]

However, we can fuse Z block and A_minus_exp block via compute_inline.

It would look more natural if we move the reduction attribute to blocks:

for vi in tir.cord(i):
    for vj in tir.pos(j):
        with tir.block(name='A_max', sparse_axes=[vi, tir.reduce_axis(vj)]):
            with tir.init():
                A_max[vi] = tir.const(-float("inf"), "float32")
            A_max[vi] = tir.max(A_max[vi], A[vi, vj])
    for vj in tir.pos(j):
        with tir.block(name='A_minus_exp', sparse_axes=[vi, vj]):
            A_exp[vi, vj] = tir.exp(A[vi, vj] - A_max[vi])
    for vj in tir.pos(j):
        with tir.block(name='Z', sparse_iters=[vi, tir.reduce_axis(vj)]):
            with tir.init():
                Z[vi] = 0.
            Z[vi] = Z[vi] + A_exp[vi. vj]
    for vj in tir.pos(j):
        with tir.block(name='out', sparser_iters=[vi, vj]):
            out[vi, vj] = A_exp[vi, vj] / Z[vi]

Then the block Z and A_max would be under the same loops, we can use compute_at:

for vi in tir.cord(i):
    for vj in tir.pos(j):
        with tir.block(name='A_max', sparse_iters=[vi, tir.reduce_axis(vj)]):
            with tir.init():
                A_max[vi] = tir.const(-float("inf"), "float32")
            A_max[vi] = tir.max(A_max[vi], A[vi, vj])
    for vj in tir.pos(j):
        with tir.block(name='Z_A_minus_exp', sparse_iters=[vi, tir.reduce_axis(vj)]):
            with tir.init():
                Z[vi] = 0.
            Z[vi] = Z[vi] + tir.exp(A[vi, vj] - A_max[vi])
    for vj in tir.pos(j):
        with tir.block(name='out', sparse_iters=[vi, vj]):
            out[vi, vj] = tir.exp(A[vi, vj] - A_max[vi]) / Z[vi]

[yzh119]

For sparse variable axis:

I think you assume that we record axis dependency information in sparse tensors. But what about we don't do so?
By that I mean we don't store axis dependency tree in Sparse Tensors but only stores the list of indptrs and indices, and only match the buffer in TIR scripts.