def test_exclusive_scan():
for target in ["cuda", "rocm"]:
if not tvm.testing.device_enabled(target):
print("Skip because %s is not enabled" % target)
continue
with tvm.target.Target(target + " -libs=thrust") as tgt:
if not thrust_check_func[target](tgt,
"tvm.contrib.thrust.sum_scan"):
print("skip because thrust is not enabled...")
return
for ishape in [(10, ), (10, 10), (10, 10, 10)]:
values = te.placeholder(ishape, name="values", dtype="int32")
scan, reduction = exclusive_scan(values, return_reduction=True)
s = schedule_scan([scan, reduction])
dev = tvm.device(target, 0)
f = tvm.build(s, [values, scan, reduction], target)
values_np = np.random.randint(0, 10,
size=ishape).astype(np.int32)
values_np_out = np.zeros(values_np.shape, np.int32)
if len(ishape) == 1:
reduction_shape = ()
else:
reduction_shape = ishape[:-1]
reduction_np_out = np.zeros(reduction_shape, np.int32)
values_in = tvm.nd.array(values_np, dev)
values_out = tvm.nd.array(values_np_out, dev)
reduction_out = tvm.nd.array(reduction_np_out, dev)
f(values_in, values_out, reduction_out)
ref_values_out = np.cumsum(values_np, axis=-1,
dtype="int32") - values_np
tvm.testing.assert_allclose(values_out.asnumpy(),
ref_values_out,
rtol=1e-5)
ref_reduction_out = np.sum(values_np, axis=-1)
tvm.testing.assert_allclose(reduction_out.asnumpy(),
ref_reduction_out,
rtol=1e-5)
def test_exclusive_scan():
if not is_thrust_available():
print("skip because thrust is not enabled...")
return
for ishape in [(10, ), (10, 10), (10, 10, 10)]:
values = te.placeholder(ishape, name="values", dtype="int32")
with tvm.target.Target("cuda"):
scan, reduction = exclusive_scan(values, return_reduction=True)
s = schedule_scan([scan, reduction])
ctx = tvm.gpu(0)
f = tvm.build(s, [values, scan, reduction], "cuda")
values_np = np.random.randint(0, 10, size=ishape).astype(np.int32)
values_np_out = np.zeros(values_np.shape, np.int32)
if len(ishape) == 1:
reduction_shape = ()
else:
reduction_shape = ishape[:-1]
reduction_np_out = np.zeros(reduction_shape, np.int32)
values_in = tvm.nd.array(values_np, ctx)
values_out = tvm.nd.array(values_np_out, ctx)
reduction_out = tvm.nd.array(reduction_np_out, ctx)
f(values_in, values_out, reduction_out)
ref_values_out = np.cumsum(values_np, axis=-1,
dtype="int32") - values_np
tvm.testing.assert_allclose(values_out.asnumpy(),
ref_values_out,
rtol=1e-5)
ref_reduction_out = np.sum(values_np, axis=-1)
tvm.testing.assert_allclose(reduction_out.asnumpy(),
ref_reduction_out,
rtol=1e-5)