def test_dynamic_tensor():
dtype = 'float32'
stype = 'csr'
target = 'llvm'
ctx = tvm.context(target, 0)
nr, nc, n = tvm.var('nr'), tvm.var('nc'), tvm.var('n')
A = tvmsp.placeholder(shape=(nr, nc), nonzeros=n, name='A', dtype=dtype)
assert(A.stype == 'csr')
C = tvm.compute(A.data.shape, lambda i: A.data[i] * 2., tag='cs_scatter')
s = tvm.create_schedule(C.op)
_nr, _nc = 3, 5
a = np.maximum(np.random.uniform(size=(_nr, _nc)).astype(dtype)-.6, 0.)
a = tvmsp.array(a, ctx)
assert a.data.dtype == a.dtype
Ab = namedtuple('CSRBuffer', ['data', 'indices', 'indptr'])
Ab.data = tvm.decl_buffer(a.data.shape, a.data.dtype, name='A_data')
Ab.indices = tvm.decl_buffer(a.data.shape, a.data.dtype, name='A_indices')
binds = {A.data: Ab.data, A.indices: Ab.indices}
f = tvm.build(s, [nr, A.data, C], target, binds=binds)
c = tvmsp.array(np.zeros((_nr, _nc), dtype), ctx)
c.data = tvm.nd.empty(a.data.shape, dtype)
c.indices = a.indices
c.indptr = a.indptr
f(a.data.shape[0], a.data, c.data)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() * 2., rtol=1e-5)
def verify_dense_si(batch, in_dim, out_dim, use_bias=True, dtype="float32"):
nonzeros = te.var("nonzeros")
A = tvmsp.placeholder(shape=(batch, in_dim),
nonzeros=nonzeros,
dtype=dtype,
name="A")
B = te.placeholder((out_dim, in_dim), dtype=dtype, name="B")
C = te.placeholder((out_dim, ), dtype=dtype, name="C")
D = topi.sparse.dense(A, B, C if use_bias else None)
s = te.create_schedule(D.op)
# get the test data
def get_ref_data():
mag = 10.0
a_np = np.maximum(
mag *
(np.random.uniform(size=(batch, in_dim)).astype("float32") - 0.5),
0.0).astype(dtype)
b_np = (mag *
(np.random.uniform(size=(out_dim, in_dim)).astype("float32") -
0.5)).astype(dtype)
c_np = (mag * (np.random.uniform(size=(out_dim, )).astype("float32") -
0.5)).astype(dtype)
if use_bias:
d_np = np.dot(a_np, b_np.T) + c_np
else:
d_np = np.dot(a_np, b_np.T)
return (a_np, b_np, c_np, d_np)
a_np, b_np, c_np, d_np = get_ref_data()
def check_device(device):
dev = tvm.device(device, 0)
if not tvm.testing.device_enabled(device):
print("Skip because %s is not enabled" % device)
return
print("Running on target: %s" % device)
a = tvmsp.array(a_np, dev)
b = tvm.nd.array(b_np, dev)
c = tvm.nd.array(c_np, dev)
d = tvm.nd.array(np.zeros(get_const_tuple(D.shape), dtype=dtype), dev)
f = tvm.build(s, [A.data, A.indices, A.indptr, B, C, D],
device,
name="dense")
f(a.data, a.indices, a.indptr, b, c, d)
tvm.testing.assert_allclose(d.asnumpy(), d_np, rtol=1e-4, atol=1e-4)
check_device("llvm")
def verify_dynamic_csrmm(batch, in_dim, out_dim, dtype, use_bias=True):
nr, nc, n = te.var("nr"), te.var("nc"), te.var("n")
A = tvmsp.placeholder(shape=(nr, nc), nonzeros=n, dtype=dtype, name="A")
B = te.placeholder((in_dim, out_dim), dtype=dtype, name="B")
C = te.placeholder((nr, ), dtype=dtype, name="C")
D = topi.sparse.csrmm(A, B, C if use_bias else None)
s = te.create_schedule(D.op)
dtype = A.dtype
# get the test data
def get_ref_data():
a_np = np.random.uniform(size=(batch, in_dim), high=100).astype(dtype)
b_np = np.random.uniform(size=(in_dim, out_dim),
high=100).astype(dtype)
c_np = np.random.uniform(size=(batch, ), high=100).astype(dtype)
if use_bias:
d_np = np.dot(a_np, b_np) + c_np.reshape((batch, 1))
else:
d_np = np.dot(a_np, b_np)
return (a_np, b_np, c_np, d_np)
a_np, b_np, c_np, d_np = get_ref_data()
def check_device(device):
dev = tvm.device(device, 0)
if not tvm.testing.device_enabled(device):
print("Skip because %s is not enabled" % device)
return
print("Running on target: %s" % device)
a = tvmsp.array(a_np, dev)
_nr, _nc, _n = a.shape[0], a.shape[1], a.data.shape[0]
assert a.shape[0] == a.indptr.shape[0] - 1
b = tvm.nd.array(b_np, dev)
c = tvm.nd.array(c_np, dev)
d = tvm.nd.array(np.zeros((_nr, out_dim), dtype=dtype), dev)
f = tvm.build(s, [nr, A.data, A.indices, A.indptr, B, C, D],
device,
name="csrmm")
f(_nr, a.data, a.indices, a.indptr, b, c, d)
tvm.testing.assert_allclose(d.numpy(), d_np, rtol=1e-2, atol=1e-2)
for device in ["llvm"]:
check_device(device)
def test_sparse_array_tuple():
dtype, itype = 'float32', 'int32'
stype = 'csr'
target = 'llvm'
ctx = tvm.context(target, 0)
nr, nc, n = te.size_var('nr'), te.size_var('nc'), te.size_var('n')
A = tvmsp.placeholder(shape=(nr, nc), nonzeros=n, name='A', dtype=dtype)
assert (A.stype == 'csr')
C = te.compute(A.data.shape, lambda i: A.data[i] * 2., tag='cs_scatter')
s = te.create_schedule(C.op)
_nr, _nc = 3, 5
a = np.maximum(np.random.uniform(size=(_nr, _nc)).astype(dtype) - .6, 0.)
# convert to sparse array tuple
source_array = a
ridx, cidx = np.nonzero(source_array)
data = source_array[ridx, cidx]
a_data = _nd.array(data, ctx)
indices = np.nonzero(source_array)[1].astype(itype)
a_indices = _nd.array(indices, ctx)
indptr = [0] + np.apply_along_axis(
np.count_nonzero, axis=1, arr=source_array).tolist()
indptr = np.cumsum(np.array(indptr, itype)).astype(itype)
a_indptr = _nd.array(indptr, ctx)
a_init = (a_data, a_indices, a_indptr)
# construct tvm sparse array with tuple
a = tvmsp.array(a_init, shape=source_array.shape, ctx=ctx)
assert a.data.dtype == a.dtype
Ab = namedtuple('CSRBuffer', ['data', 'indices', 'indptr'])
Ab.data = tvm.tir.decl_buffer(a.data.shape, a.data.dtype, name='A_data')
Ab.indices = tvm.tir.decl_buffer(a.data.shape,
a.data.dtype,
name='A_indices')
binds = {A.data: Ab.data, A.indices: Ab.indices}
f = tvm.build(s, [nr, A.data, C], target, binds=binds)
c = tvmsp.array(np.zeros((_nr, _nc), dtype), ctx)
c.data = tvm.nd.empty(a.data.shape, dtype)
c.indices = a.indices
c.indptr = a.indptr
f(a.data.shape[0], a.data, c.data)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() * 2., rtol=1e-5)
def test_sparse_array_tuple():
dtype, itype = 'float32', 'int32'
stype = 'csr'
target = 'llvm'
ctx = tvm.context(target, 0)
nr, nc, n = tvm.var('nr'), tvm.var('nc'), tvm.var('n')
A = tvmsp.placeholder(shape=(nr, nc), nonzeros=n, name='A', dtype=dtype)
assert(A.stype == 'csr')
C = tvm.compute(A.data.shape, lambda i: A.data[i] * 2., tag='cs_scatter')
s = tvm.create_schedule(C.op)
_nr, _nc = 3, 5
a = np.maximum(np.random.uniform(size=(_nr, _nc)).astype(dtype)-.6, 0.)
# convert to sparse array tuple
source_array = a
ridx, cidx = np.nonzero(source_array)
data = source_array[ridx, cidx]
a_data = _nd.array(data, ctx)
indices = np.nonzero(source_array)[1].astype(itype)
a_indices = _nd.array(indices, ctx)
indptr = [0]+np.apply_along_axis(np.count_nonzero, axis=1, arr=source_array).tolist()
indptr = np.cumsum(np.array(indptr, itype)).astype(itype)
a_indptr = _nd.array(indptr, ctx)
a_init = (a_data, a_indices, a_indptr)
# construct tvm sparse array with tuple
a = tvmsp.array(a_init, shape=source_array.shape, ctx=ctx)
assert a.data.dtype == a.dtype
Ab = namedtuple('CSRBuffer', ['data', 'indices', 'indptr'])
Ab.data = tvm.decl_buffer(a.data.shape, a.data.dtype, name='A_data')
Ab.indices = tvm.decl_buffer(a.data.shape, a.data.dtype, name='A_indices')
binds = {A.data: Ab.data, A.indices: Ab.indices}
f = tvm.build(s, [nr, A.data, C], target, binds=binds)
c = tvmsp.array(np.zeros((_nr, _nc), dtype), ctx)
c.data = tvm.nd.empty(a.data.shape, dtype)
c.indices = a.indices
c.indptr = a.indptr
f(a.data.shape[0], a.data, c.data)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() * 2., rtol=1e-5)
def test_static_tensor():
dtype = 'float32'
stype = 'csr'
target = 'llvm'
ctx = tvm.context(target, 0)
m = tvm.var('m')
n = tvm.var('n')
A = tvmsp.placeholder(shape=(m, n), name='A', dtype=dtype)
assert(A.stype == 'csr')
n = 3
a = np.maximum(np.random.uniform(size=(n,n)).astype(dtype)-.6, 0.)
a = tvmsp.array(a, ctx)
A.data = tvm.placeholder(a.data.shape, dtype, name='A_data')
Ab = tvm.decl_buffer(a.data.shape, dtype, name='A_data')
binds = {A.data: Ab}
C = tvm.compute(A.data.shape, lambda i: A.data[i] * 2., tag='cs_scatter')
s = tvm.create_schedule(C.op)
f = tvm.build(s, [A.data, C], target, binds=binds)
c = tvmsp.array(np.zeros((n,n), dtype), ctx)
c.data = tvm.nd.empty(a.data.shape, dtype)
c.indices = a.indices
c.indptr = a.indptr
f(a.data, c.data)
tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() * 2., rtol=1e-5)
