Pass sparse tensor to tvm.build

From the discussion about running sparse CNNs, I have implemented prototypes of a dense NCHW GEMM convolution, and what I think is a working CSR NCHW GEMM convolution.

I will share the code once it’s a bit more mature.

I’m doing sparse GEMM convolution first, rather than sparse spatial pack, since there are already sparse GEMM methods in TVM I can leverage, whereas a sparse spatial pack would require some cleverness.

However, when building my standalone function as a module, I am getting the following error:

from tvm.contrib import sparse


# create placeholder tensors
...
n = out_c
k = kdim_h * kdim_w * in_c
sparse_weights = sparse.placeholder((n,k), nonzeros=(1-sparsity)*n*k, name='W') # weights 
...


# create output using compositions of te.compute 
conv = gemm_conv_sparse(sparse_weights, ...)


# build function
s = te.create_schedule(conv.op)
dtype = 'float32'

func = tvm.build(s, [data, sparse_weights, conv], target=target, name='gemm_conv2d')

# ^ returns ValueError: args must be Tensor, Buffer or Var

Any ideas on how I’d pass my sparse.placeholder to tvm.build?

1 Like

I have also tried bypassing this issue by passing the three tensor objects inside a sparse array.

from tvm.contrib import sparse


# create placeholder tensors
...
n = out_c
k = kdim_h * kdim_w * in_c
sparse_weights = sparse.placeholder((n,k), nonzeros=(1-sparsity)*n*k, name='W') # weights 
...

# create output using compositions of te.compute 
conv = gemm_conv_sparse_alt(sparse_weights.data, sparse_weights.indices, sparse_weights.indptr, 
                            ...)

# build function
s = te.create_schedule(conv.op)
dtype = 'float32'

func = tvm.build(s, [data, sparse_weights.data, sparse_weights.indices, 
                sparse_weights.indptr, conv], target=target, name='gemm_conv2d')

# ^ works
X = tvm.nd.array(inputs_np.astype(dtype), ctx)
W = tvm.contrib.sparse.array(weights_np.astype(dtype), shape=kernel_shape, ctx=ctx)
c = tvm.nd.array(np.zeros(oshape, dtype=dtype), ctx)

func(X, W.data, W.indices, W.indptr, c)

The final line of this fails with:

 File "../src/runtime/library_module.cc", line 78
TVMError: Check failed: ret == 0 (-1 vs. 0) : Assert fail: (1.6f == float32(arg1.shape[0])), Argument arg1.shape[0] has an unsatisfied constraint

My thought here is that we are passing a placeholder sparse tensor, with only an estimate of our expected sparsity.

When we pass an actual sparse tensor to the function it will have patterns different from what the module was built to expect.

It’s not like we can pass the contents of the actual data W to the function build, as they are NDarrays, rather than the Tensor type the function building expects.

Another work around is to get rid off the sparse placeholder completely. Instead use three standard Tensors for each element in the sparse tensor (i.e. data, indices and index pointers). Then feed those (plus the X matrix) to topi.nn.sparse_dense and things seem to work.

There’s a working implementation in this repo:

@autotvm.template("benchmark/block_sparse")
def block_sparse_template(W_sp_np_data_shape, W_sp_np_indices_shape, W_sp_np_indptr_shape, X_np_shape):
    W_data = te.placeholder(shape=W_sp_np_data_shape, dtype='float32', name='W_data')
    W_indices = te.placeholder(shape=W_sp_np_indices_shape, dtype='int32', name='W_indices')
    W_indptr = te.placeholder(shape=W_sp_np_indptr_shape, dtype='int32', name='W_indptr')
    X = te.placeholder(shape=X_np_shape, dtype='float32', name='X')
    Y = topi.nn.sparse_dense(X, W_data, W_indices, W_indptr)
    ...

Useful resource, thanks. I ended up fixing the approach in my 2nd example (using three ndarrays. Basically one can only pass sparse tensors that have the same sparsity pattern, not just the same level of sparsity as the placeholders you pass.

Thus, when constructing the placeholder tensors for compilation, one needs to give them sizes from the sparse data you will pass to the compiled function.

E.g. if your sparse data is in a SciPy CSR object called W_sp_np, your TVM placeholders would be constructed with:

W_data = te.placeholder(shape=W_sp_np.data.shape, dtype=str(W_sp_np.data.dtype), name='W_data')
W_indices = te.placeholder(shape=W_sp_np.indices.shape, dtype=str(W_sp_np.indices.dtype), name='W_indices')
W_indptr = te.placeholder(shape=W_sp_np.indptr.shape, dtype=str(W_sp_np.indptr.dtype), name='W_indptr')

My (unoptimised) sparse NCHW GEMM convolution is now working, I’ll see if I can draft a tutorial about what I’ve learned once I’ve completed some other things.

Hi I am struggling with sparse cmm usage and could you please help to take a look at this post ?

thanks

Looks like it’s solved, but ping me if you have other issues with sparse stuff. I’m not as well versed as some other developers, but I have been working on it on-and-off the past couple of months.

1 Like

yes, thanks for the hint from this post.