Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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"))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.