master ← tqchen:master
opened 05:10AM - 25 Jul 18 UTC
Ref #1276
## Usage
We can have a pragma attribute pragma_import_llvm, which… can be used in either ir builder or schedule
## Example
```python
def test_llvm_import():
# extern "C" is necessary to get the correct signature
# we can put inline asm here
cc_code = """
extern "C" float my_add(float x, float y) {
return x + y;
}
"""
n = 10
A = tvm.placeholder((n,), name='A')
B = tvm.compute((n,), lambda *i:
tvm.call_pure_extern("float32", "my_add", A(*i), 1.0),
name='B')
def check_llvm(use_file):
if not tvm.module.enabled("llvm"):
return
if not clang.find_clang(required=False):
print("skip because clang is not available")
return
temp = util.tempdir()
ll_path = temp.relpath("temp.ll")
# Create LLVM ir from c source code
ll_code = clang.create_llvm(cc_code, output=ll_path)
s = tvm.create_schedule(B.op)
if use_file:
s[B].pragma(s[B].op.axis[0], "import_llvm", ll_path)
else:
s[B].pragma(s[B].op.axis[0], "import_llvm", ll_code)
# BUILD and invoke the kernel.
f = tvm.build(s, [A, B], "llvm")
ctx = tvm.cpu(0)
# launch the kernel.
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx)
f(a, b)
np.testing.assert_allclose(
b.asnumpy(), a.asnumpy() + 1.0)
check_llvm(use_file=True)
```