test_scan.py 1.58 KB
Newer Older
1 2 3 4
import tvm
import numpy as np

def test_scan():
5 6
    m = tvm.var("m")
    n = tvm.var("n")
7 8 9
    X = tvm.placeholder((m, n), name="X")
    s_state = tvm.placeholder((m, n))
    s_init = tvm.compute((1, n), lambda _, i: X[0, i])
10 11
    s_update = tvm.compute((m, n), lambda t, i: s_state[t-1, i] + X[t, i])
    res = tvm.scan(s_init, s_update, s_state)
12 13

    # schedule
14
    s = tvm.create_schedule(res.op)
15
    num_thread = 256
16 17
    block_x = tvm.thread_axis(None, "blockIdx.x")
    thread_x = tvm.thread_axis((0, num_thread), "threadIdx.x")
18 19 20 21 22 23
    xo, xi = s[s_init].split(s_init.op.axis[1], factor=num_thread)
    s[s_init].bind(xo, block_x)
    s[s_init].bind(xi, thread_x)
    xo, xi = s[s_update].split(s_update.op.axis[1], factor=num_thread)
    s[s_update].bind(xo, block_x)
    s[s_update].bind(xi, thread_x)
24 25

    # one line to build the function.
26
    def check_device(device):
27 28
        ctx = tvm.context(device, 0)
        if not ctx.exist:
29
            print("skip because %s is not enabled.." % device)
30
            return
31
        fscan = tvm.build(s, [X, res],
32
                          device,
33 34 35 36 37 38 39 40 41 42 43
                          name="myscan")
        # launch the kernel.
        n = 1024
        m = 10
        a_np = np.random.uniform(size=(m, n)).astype(res.dtype)
        a = tvm.nd.array(a_np, ctx)
        b = tvm.nd.array(np.zeros((m, n), dtype=res.dtype), ctx)
        fscan(a, b)
        np.testing.assert_allclose(
            b.asnumpy(), np.cumsum(a_np, axis=0))

44
    check_device("vulkan")
45
    check_device("cuda")
46
    check_device("metal")
47
    check_device("opencl")
48 49 50 51


if __name__ == "__main__":
    test_scan()