[APPS] add an external dll call example (#2156)

apps/extension/python/tvm_ext/__init__.py

@@ -8,7 +8,9 @@ import tvm
 def load_lib():
     """Load library, the functions will be registered into TVM"""
     curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
-    lib = ctypes.CDLL(os.path.join(curr_path, "../../lib/libtvm_ext.so"))
+    # load in as global so the global extern symbol is visible to other dll.
+    lib = ctypes.CDLL(
+        os.path.join(curr_path, "../../lib/libtvm_ext.so"), ctypes.RTLD_GLOBAL)
     return lib
 
 _LIB = load_lib()

apps/extension/src/tvm_ext.cc

@@ -66,6 +66,11 @@ TVM_REGISTER_GLOBAL("device_api.ext_dev")
   });
 }  // namespace tvm_ext
 
+// External function exposed to runtime.
+extern "C" float TVMTestAddOne(float y) {
+  return y + 1;
+}
+
 // This callback approach allows extension allows tvm to extract
 // This way can be helpful when we want to use a header only
 // minimum version of TVM Runtime.

apps/extension/tests/test_ext.py

@@ -49,7 +49,27 @@ def test_extract_ext():
     assert fdict["mul"](3, 4) == 12
 
 
+def test_extern_call():
+    n = 10
+    A = tvm.placeholder((n,), name='A')
+    B = tvm.compute((n,), lambda *i: tvm.call_extern("float32", "TVMTestAddOne", A(*i)), name='B')
+    s = tvm.create_schedule(B.op)
+
+    def check_llvm():
+        if not tvm.module.enabled("llvm"):
+            return
+        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.zeros(n, dtype=B.dtype), ctx)
+        f(a, b)
+        tvm.testing.assert_allclose(b.asnumpy(), a.asnumpy() + 1)
+    check_llvm()
+
+
 if __name__ == "__main__":
+    test_extern_call()
     test_ext_dev()
     test_ext_vec()
     test_bind_add()