def test_const_saveload_json():
# save load json
x = tvm.const(1, "int32")
y = tvm.const(10, "int32")
z = x + y
z = z + z
json_str = tvm.save_json(z)
zz = tvm.load_json(json_str)
assert tvm.save_json(zz) == tvm.save_json(z)
def test_with():
n = tvm.var('n')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((n, l), name='A')
B = tvm.placeholder((m, l), name='B')
with tvm.tag_scope(tag="gemm"):
k = tvm.reduce_axis((0, l), name='k')
C = tvm.compute((n, m),
lambda i, j: tvm.sum(A[i, k] * B[j, k], axis=k),
attrs={
"hello": 1,
"arr": [10, 12]
})
assert C.op.tag == 'gemm'
assert "hello" in C.op.attrs
assert "xx" not in C.op.attrs
assert C.op.attrs["hello"].value == 1
CC = tvm.load_json(tvm.save_json(C))
assert CC.op.attrs["hello"].value == 1
assert CC.op.attrs["arr"][0].value == 10
# str format happened to be json compatible
assert json.loads(str(CC.op.attrs))["arr"][1] == 12
def test_function_attrs():
param_names = ['a', 'b', 'c', 'd']
params = tvm.convert([relay.var(n, shape=(5, 2)) for n in param_names])
ret_type = relay.TupleType(tvm.convert([]))
body = relay.Tuple(tvm.convert([]))
type_params = tvm.convert([])
fn = relay.Function(params, body, ret_type, type_params)
model_params = {}
for param in params[:1]:
cty = param.type_annotation
tensor = np.random.rand(*[int(sh) for sh in cty.shape]).astype(cty.dtype)
model_params[param] = tvm.nd.array(tensor)
fn = fn.set_params(model_params)
assert fn.params == params
assert fn.body == body
assert fn.type_params == type_params
assert fn.span == None
str(fn)
check_json_roundtrip(fn)
json_str = tvm.save_json(fn)
fn_after = tvm.load_json(json_str)
model_params_after = fn_after.get_params()
after_keys = [item[0] for item in model_params_after.items()]
for key1, key2 in zip(model_params, after_keys):
assert key1.name_hint == key2.name_hint
p1 = model_params[key1]
p2 = model_params_after[key2]
np.testing.assert_allclose(p1.data.asnumpy(), p2.data.asnumpy())
def test_schedule_create():
m = tvm.var('m')
n = tvm.var('n')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
B = tvm.placeholder((n, l), name='B')
AA = tvm.compute((m, l), lambda i, j: A[i, j])
T = tvm.compute((m, n, l), lambda i, j, k: AA(i, k) * B(j, k))
s = tvm.create_schedule(T.op)
s[AA].set_scope("shared")
xo, xi = s[T].split(T.op.axis[0], factor=10)
xi1, xi2 = s[T].split(xi, factor=2)
s[AA].compute_at(s[T], xi1)
xo, xi = s[AA].split(AA.op.axis[0], factor=10)
s[T].reorder(xi2, xi1)
assert T.op.axis[1] in s[T].leaf_iter_vars
# save load json
json_str = tvm.save_json(s)
s_loaded = tvm.load_json(json_str)
assert isinstance(s_loaded, tvm.schedule.Schedule)
assert(str(s_loaded.outputs[0].body) == str(s.outputs[0].body))
# pickle unpickle
dump = pkl.dumps(s)
s_loaded = pkl.loads(dump)
assert isinstance(s_loaded, tvm.schedule.Schedule)
assert(str(s_loaded.outputs[0].body) == str(s.outputs[0].body))
def test_schedule_create():
m = tvm.var('m')
n = tvm.var('n')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
B = tvm.placeholder((n, l), name='B')
AA = tvm.compute((m, l), lambda i, j: A[i, j])
T = tvm.compute((m, n, l), lambda i, j, k: AA(i, k) * B(j, k))
s = tvm.create_schedule(T.op)
s[AA].set_scope("shared")
xo, xi = s[T].split(T.op.axis[0], factor=10)
xi1, xi2 = s[T].split(xi, factor=2)
s[AA].compute_at(s[T], xi1)
xo, xi = s[AA].split(AA.op.axis[0], factor=10)
s[T].reorder(xi2, xi1)
assert T.op.axis[1] in s[T].leaf_iter_vars
# save load json
json_str = tvm.save_json(s)
s_loaded = tvm.load_json(json_str)
assert isinstance(s_loaded, tvm.schedule.Schedule)
assert(str(s_loaded.outputs[0].body) == str(s.outputs[0].body))
# pickle unpickle
dump = pkl.dumps(s)
s_loaded = pkl.loads(dump)
assert isinstance(s_loaded, tvm.schedule.Schedule)
assert(str(s_loaded.outputs[0].body) == str(s.outputs[0].body))
def test_make_sum():
A = tvm.placeholder((2, 10), name='A')
k = tvm.reduce_axis((0,10), "k")
B = tvm.compute((2,), lambda i: tvm.sum(A[i, k], axis=k), name="B")
json_str = tvm.save_json(B)
BB = tvm.load_json(json_str)
assert B.op.body[0].combiner is not None
assert BB.op.body[0].combiner is not None
def test_map_save_load_json():
a = tvm.var('a')
b = tvm.var('b')
amap = tvm.convert({a: 2, b: 3})
json_str = tvm.save_json(amap)
amap = tvm.load_json(json_str)
assert len(amap) == 2
dd = {kv[0].name: kv[1].value for kv in amap.items()}
assert (dd == {"a": 2, "b": 3})
def test_make_smap():
# save load json
x = tvm.const(1)
y = tvm.const(10)
z = x + y
smap = tvm.convert({"z": z, "x": x})
json_str = tvm.save_json(tvm.convert([smap]))
arr = tvm.load_json(json_str)
assert len(arr) == 1
assert arr[0]["z"].a == arr[0]["x"]
def test_make_smap():
# save load json
x = tvm.const(1, "int32")
y = tvm.const(10, "int32")
z = tvm.expr.Add(x, y)
smap = tvm.convert({"z": z, "x": x})
json_str = tvm.save_json(tvm.convert([smap]))
arr = tvm.load_json(json_str)
assert len(arr) == 1
assert arr[0]["z"].a == arr[0]["x"]
def test_map_save_load_json():
a = tvm.var('a')
b = tvm.var('b')
amap = tvm.convert({a: 2,
b: 3})
json_str = tvm.save_json(amap)
amap = tvm.load_json(json_str)
assert len(amap) == 2
dd = {kv[0].name : kv[1].value for kv in amap.items()}
assert(dd == {"a": 2, "b": 3})
def test_env_func():
@tvm.register_func("test.env_func")
def test(x):
return x + 1
f = tvm.get_global_func("test.env_func")
x = tvm.get_env_func("test.env_func")
assert x.name == "test.env_func"
json_str = tvm.save_json([x])
y = tvm.load_json(json_str)[0]
assert y.name == x.name
assert y(1) == 2
assert y.func(1) == 2
x = tvm.make.node("attrs.TestAttrs", name="xx", padding=(3,4), func=y)
assert x.name == "xx"
assert x.padding[0].value == 3
assert x.padding[1].value == 4
assert x.axis == 10
x = tvm.load_json(tvm.save_json(x))
assert isinstance(x.func, tvm.container.EnvFunc)
assert x.func(10) == 11
def test_tensor_reduce():
m = tvm.var('m')
n = tvm.var('n')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
B = tvm.placeholder((n, l), name='B')
T = tvm.compute((m, n, l), lambda i, j, k: A[i, k] * B[j, k])
rv = tvm.reduce_axis((0, A.shape[1]), "k")
C = tvm.compute((m, n), lambda i, j: tvm.sum(T(i, j, rv+1), axis=rv))
# json load save
C_json = tvm.save_json(C)
C_loaded = tvm.load_json(C_json)
assert(isinstance(C_loaded, tvm.tensor.Tensor))
assert(str(C_loaded) == str(C))
def test_ndarray_reflection():
x = np.random.uniform(size=(10, 2)).astype("float32")
xx = tvm.nd.array(x)
xnode = tvm.make.node("NDArrayWrapper", name="xx", array=xx)
xnode2 = tvm.make.node("NDArrayWrapper", name="x2", array=xx)
assert xnode.array.same_as(xx)
json_str = tvm.save_json([xnode, xnode2])
json_dict = json.loads(json_str)
b64_str = json_dict["b64ndarrays"][0]
decoded = py_str(base64.b64encode(base64.b64decode(b64_str)))
assert b64_str == decoded
xlist = tvm.load_json(json_str)
np.testing.assert_equal(xlist[0].array.asnumpy(), xx.asnumpy())
assert xlist[1].array == xlist[0].array
def test_ndarray_reflection():
x = np.random.uniform(size=(10, 2)).astype("float32")
xx = tvm.nd.array(x)
xnode = tvm.make.node("NDArrayWrapper", name="xx", array=xx)
xnode2 = tvm.make.node("NDArrayWrapper", name="x2", array=xx)
assert xnode.array.same_as(xx)
json_str = tvm.save_json([xnode, xnode2])
json_dict = json.loads(json_str)
b64_str = json_dict["b64ndarrays"][0]
decoded = py_str(base64.b64encode(base64.b64decode(b64_str)))
assert b64_str == decoded
xlist = tvm.load_json(json_str)
np.testing.assert_equal(xlist[0].array.asnumpy(), xx.asnumpy())
assert xlist[1].array == xlist[0].array
def test_tensor_reduce():
m = tvm.var('m')
n = tvm.var('n')
l = tvm.var('l')
A = tvm.placeholder((m, l), name='A')
B = tvm.placeholder((n, l), name='B')
T = tvm.compute((m, n, l), lambda i, j, k: A[i, k] * B[j, k])
rv = tvm.reduce_axis((0, A.shape[1]), "k")
C = tvm.compute((m, n), lambda i, j: tvm.sum(T(i, j, rv+1), axis=rv))
# json load save
C_json = tvm.save_json(C)
C_loaded = tvm.load_json(C_json)
assert(isinstance(C_loaded, tvm.tensor.Tensor))
assert(str(C_loaded) == str(C))
def test_span():
span = relay.Span(None, 1, 1)
assert span.source == None
assert span.lineno == 1
assert span.col_offset == 1
assert span.same_as(span)
assert span == span
assert isinstance(span, relay.base.Span)
str(span)
# span is not a node so we can't use graph_equal
# to test the round trip
back = tvm.load_json(tvm.save_json(span))
assert back.source == span.source
assert back.lineno == span.lineno
assert back.col_offset == span.col_offset
def test_span():
span = relay.Span(None, 1, 1)
assert span.source == None
assert span.lineno == 1
assert span.col_offset == 1
assert span.same_as(span)
assert span == span
assert isinstance(span, relay.base.Span)
str(span)
# span is not a node so we can't use graph_equal
# to test the round trip
back = tvm.load_json(tvm.save_json(span))
assert back.source == span.source
assert back.lineno == span.lineno
assert back.col_offset == span.col_offset
def test_scan_multi_out():
m = tvm.size_var("m")
n = tvm.size_var("n")
x1 = tvm.placeholder((m, n))
s1 = tvm.placeholder((m, n))
x2 = tvm.placeholder((m, n))
s2 = tvm.placeholder((m, n))
s1_init = tvm.compute((1, n), lambda _, i: x1[0, i])
s2_init = tvm.compute((1, n), lambda _, i: x2[0, i])
s1_update = tvm.compute(
(m, n), lambda t, i: s1[t - 1, i] + s2[t - 1, i] + x1[t, i])
s2_update = tvm.compute((m, n), lambda t, i: x2[t, i] + s2[t - 1, i])
r0, r1 = tvm.scan([s1_init, s2_init], [s1_update, s2_update], [s1, s2])
assert (r0.value_index == 0)
assert (r1.value_index == 1)
json_str = tvm.save_json(r0.op)
zz = tvm.load_json(json_str)
assert isinstance(zz, tvm.tensor.ScanOp)
def test_scan_multi_out():
m = tvm.var("m")
n = tvm.var("n")
x1 = tvm.placeholder((m, n))
s1 = tvm.placeholder((m, n))
x2 = tvm.placeholder((m, n))
s2 = tvm.placeholder((m, n))
s1_init = tvm.compute((1, n), lambda _, i: x1[0, i])
s2_init = tvm.compute((1, n), lambda _, i: x2[0, i])
s1_update = tvm.compute((m, n), lambda t, i: s1[t-1, i] + s2[t-1, i] + x1[t, i])
s2_update = tvm.compute((m, n), lambda t, i: x2[t, i] + s2[t-1,i])
r0, r1 = tvm.scan([s1_init, s2_init],
[s1_update, s2_update],
[s1, s2])
assert(r0.value_index == 0)
assert(r1.value_index == 1)
json_str = tvm.save_json(r0.op)
zz = tvm.load_json(json_str)
assert isinstance(zz, tvm.tensor.ScanOp)
def test_with():
n = tvm.var('n')
m = tvm.var('m')
l = tvm.var('l')
A = tvm.placeholder((n, l), name='A')
B = tvm.placeholder((m, l), name='B')
with tvm.tag_scope(tag="gemm"):
k = tvm.reduce_axis((0, l), name='k')
C = tvm.compute((n, m), lambda i, j: tvm.sum(A[i, k] * B[j, k], axis=k),
attrs={"hello" : 1, "arr": [10, 12]})
assert C.op.tag == 'gemm'
assert "hello" in C.op.attrs
assert "xx" not in C.op.attrs
assert C.op.attrs["hello"].value == 1
CC = tvm.load_json(tvm.save_json(C))
assert CC.op.attrs["hello"].value == 1
assert CC.op.attrs["arr"][0].value == 10
# str format happened to be json compatible
assert json.loads(str(CC.op.attrs))["arr"][1] == 12
def test_make_attrs():
try:
x = tvm.make.node("attrs.TestAttrs", unknown_key=1, name="xx")
assert False
except tvm.TVMError as e:
assert str(e).find("unknown_key") != -1
try:
x = tvm.make.node("attrs.TestAttrs", axis=100, name="xx")
assert False
except tvm.TVMError as e:
assert str(e).find("upper bound") != -1
x = tvm.make.node("attrs.TestAttrs", name="xx", padding=(3,4))
assert x.name == "xx"
assert x.padding[0].value == 3
assert x.padding[1].value == 4
assert x.axis == 10
dattr = tvm.make.node("DictAttrs", x=1, y=10, name="xyz", padding=(0,0))
assert dattr.x.value == 1
datrr = tvm.load_json(tvm.save_json(dattr))
assert dattr.name.value == "xyz"
def test_make_attrs():
try:
x = tvm.make.node("attrs.TestAttrs", unknown_key=1, name="xx")
assert False
except tvm.TVMError as e:
assert str(e).find("unknown_key") != -1
try:
x = tvm.make.node("attrs.TestAttrs", axis=100, name="xx")
assert False
except tvm.TVMError as e:
assert str(e).find("upper bound") != -1
x = tvm.make.node("attrs.TestAttrs", name="xx", padding=(3,4))
assert x.name == "xx"
assert x.padding[0].value == 3
assert x.padding[1].value == 4
assert x.axis == 10
dattr = tvm.make.node("DictAttrs", x=1, y=10, name="xyz", padding=(0,0))
assert dattr.x.value == 1
datrr = tvm.load_json(tvm.save_json(dattr))
assert dattr.name.value == "xyz"
print(tvm.lower(s, [A, B], simple_mode=True))
s = tvm.create_schedule(B.op)
ko, ki = s[B].split(B.op.reduce_axis[0], factor=16)
BF = s.rfactor(B, ki)
print(tvm.lower(s, [A, B], simple_mode=True))
xo, xi = s[B].split(s[B].op.axis[0], factor=32)
s[B].bind(xo, tvm.thread_axis("blockIdx.x"))
s[B].bind(xi, tvm.thread_axis("threadIdx.y"))
tx = tvm.thread_axis("threadIdx.x")
s[B].bind(s[B].op.reduce_axis[0], tx)
s[BF].compute_at(s[B], s[B].op.reduce_axis[0])
s[B].set_store_predicate(tx.var.equal(0))
fcuda = tvm.build(s, [A, B], "cuda")
print(fcuda.imported_modules[0].get_source())
print(tvm.save_json(s))
nn = 128
ctx = tvm.gpu(0)
a = tvm.nd.array(np.random.uniform(size=(nn, nn)).astype(A.dtype), ctx)
b = tvm.nd.array(np.zeros(nn, dtype=B.dtype), ctx)
fcuda(a, b)
tvm.testing.assert_allclose(
b.asnumpy(), np.sum(a.asnumpy(), axis=1), rtol=1e-4)
evaluator = fcuda.time_evaluator(fcuda.entry_name, ctx, number=400, repeat=100)
times = evaluator(a, b)
print('%f' % times.mean)
def test_array_save_load_json():
a = tvm.convert([1,2,3])
json_str = tvm.save_json(a)
a_loaded = tvm.load_json(json_str)
assert(a[1].value == 2)
def check_json_roundtrip(node):
json_str = tvm.save_json(node)
back = tvm.load_json(json_str)
assert graph_equal(back, node)
def check_json_roundtrip(node):
json_str = tvm.save_json(node)
back = tvm.load_json(json_str)
assert graph_equal(back, node)
def test_array_save_load_json():
a = tvm.convert([1, 2, 3])
json_str = tvm.save_json(a)
a_loaded = tvm.load_json(json_str)
assert (a[1].value == 2)
