Skip to content

Instantly share code, notes, and snippets.

@ruslo
Created December 20, 2018 00:46
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • Save ruslo/4fa15e70910abd371515a5c4890950af to your computer and use it in GitHub Desktop.
Save ruslo/4fa15e70910abd371515a5c4890950af to your computer and use it in GitHub Desktop.
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, "test_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