prepare_test_libs.py

import tvm
import os

def prepare_test_libs(base_path):
    n = tvm.var("n")
    A = tvm.placeholder((n,), name='A')
    B = tvm.compute(A.shape, lambda *i: A(*i) + 1.0, name='B')
    s = tvm.create_schedule(B.op)

    s[B].bind(B.op.axis[0], tvm.thread_axis("blockIdx.x"))

    fadd_syslib = tvm.build(s, [A, B], "opencl", name="addone")
    syslib_path = os.path.join(base_path, "libtest_addone.so")
    fadd_syslib.export_library(syslib_path)

if __name__ == "__main__":
    curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
    prepare_test_libs(os.path.join(curr_path, "./lib"))