test_codegen_verilog.py 2.23 KB
Newer Older
1
import tvm
2
from tvm.contrib import verilog
3 4 5 6 7 8 9 10
import numpy as np

def lower(s, args, name):
    binds = {}
    arg_list = []

    for x in args:
        assert isinstance(x, tvm.tensor.Tensor)
11
        buf = tvm.decl_buffer(x.shape, dtype=x.dtype, name=x.op.name)
12 13
        binds[x] = buf
        arg_list.append(buf)
14
    s = s.normalize()
15 16
    bounds = tvm.schedule.InferBound(s)
    stmt = tvm.schedule.ScheduleOps(s, bounds)
17
    stmt = tvm.ir_pass.StorageFlatten(stmt, binds, 64)
18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35
    stmt = tvm.ir_pass.CanonicalSimplify(stmt)
    stmt = tvm.ir_pass.Simplify(stmt)
    stmt = tvm.ir_pass.SplitPipeline(stmt, True)
    fapi = tvm.ir_pass.MakeAPI(stmt, name, arg_list, 0)
    return fapi

@tvm.register_func
def tvm_callback_verilog_postproc(code):
    """Hook to inspect the verilog code before actually run it"""
    print(code)
    return code

def test_add_pipeline():
    nn = 128
    n = tvm.convert(nn)
    A = tvm.placeholder((n,), name='A', dtype='int32')
    B = tvm.placeholder((n,), name='B', dtype='int32')
    C = tvm.compute(A.shape, lambda i: A[i] + B[i], name='C')
36
    s = tvm.create_schedule(C.op)
37

38 39
    px, x = s[C].split(C.op.axis[0], nparts=1)
    s[C].bind(px, tvm.thread_axis("pipeline"))
40
    fapi = lower(s, [A, B, C], "myadd")
41
    fsplits = [x for x in tvm.ir_pass.SplitHostDevice(fapi)]
42
    fsplits[0] = tvm.ir_pass.LowerTVMBuiltin(fsplits[0])
43 44 45
    print("------")

    def check_target(device, host="stackvm"):
46
        if not tvm.module.enabled(host):
47
            return
48
        if not tvm.module.enabled(device):
49 50
            return
        ctx = tvm.vpi(0)
51 52
        mhost = tvm.codegen.build_module(fsplits[0], host)
        mdev = tvm.codegen.build_module(fsplits[1:], device)
53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69
        mhost.import_module(mdev)
        code = mdev.get_source()
        f = mhost.entry_func
        # launch the kernel.
        n = nn
        a = tvm.nd.array((np.random.uniform(size=n) * 128).astype(A.dtype), ctx)
        b = tvm.nd.array((np.random.uniform(size=n) * 128).astype(A.dtype), ctx)
        c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
        f(a, b, c)
        print("Check correctness...")
        np.testing.assert_allclose(
            c.asnumpy(), a.asnumpy() + b.asnumpy())
    check_target("verilog")


if __name__ == "__main__":
    test_add_pipeline()