Test combine multiple module into one shared lib.
()
| 158 | |
| 159 | @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") |
| 160 | def test_combine_module_llvm(): |
| 161 | """Test combine multiple module into one shared lib.""" |
| 162 | pytest.importorskip("cloudpickle") # needed by popen_pool.PopenWorker |
| 163 | |
| 164 | # graph |
| 165 | nn = 12 |
| 166 | n = tvm.runtime.convert(nn) |
| 167 | A = te.placeholder((n,), name="A") |
| 168 | B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B") |
| 169 | mod1 = tvm.IRModule.from_expr(te.create_prim_func([A, B]).with_attr("global_symbol", "myadd1")) |
| 170 | mod2 = tvm.IRModule.from_expr(te.create_prim_func([A, B]).with_attr("global_symbol", "myadd2")) |
| 171 | |
| 172 | def check_llvm(): |
| 173 | dev = tvm.cpu(0) |
| 174 | temp = utils.tempdir() |
| 175 | fadd1 = tvm.tirx.build(mod1, "llvm") |
| 176 | fadd2 = tvm.tirx.build(mod2, "llvm") |
| 177 | path1 = temp.relpath("myadd1.o") |
| 178 | path2 = temp.relpath("myadd2.o") |
| 179 | path_dso = temp.relpath("mylib.so") |
| 180 | fadd1.write_to_file(path1) |
| 181 | fadd2.write_to_file(path2) |
| 182 | # create shared library with multiple functions |
| 183 | cc.create_shared(path_dso, [path1, path2]) |
| 184 | m = tvm.runtime.load_module(path_dso) |
| 185 | fadd1 = m["myadd1"] |
| 186 | fadd2 = m["myadd2"] |
| 187 | a = tvm.runtime.tensor(np.random.uniform(size=nn).astype(A.dtype), dev) |
| 188 | b = tvm.runtime.tensor(np.zeros(nn, dtype=A.dtype), dev) |
| 189 | fadd1(a, b) |
| 190 | np.testing.assert_equal(b.numpy(), a.numpy() + 1) |
| 191 | fadd2(a, b) |
| 192 | np.testing.assert_equal(b.numpy(), a.numpy() + 1) |
| 193 | |
| 194 | def check_system_lib(): |
| 195 | dev = tvm.cpu(0) |
| 196 | if not tvm.testing.device_enabled("llvm"): |
| 197 | print("Skip because llvm is not enabled") |
| 198 | return |
| 199 | temp = utils.tempdir() |
| 200 | print("Running popen check") |
| 201 | fadd1 = tvm.tirx.build(mod1.with_attr("system_lib_prefix", ""), "llvm") |
| 202 | fadd2 = tvm.tirx.build(mod2.with_attr("system_lib_prefix", ""), "llvm") |
| 203 | path1 = temp.relpath("myadd1.o") |
| 204 | path2 = temp.relpath("myadd2.o") |
| 205 | path_dso = temp.relpath("mylib.so") |
| 206 | fadd1.write_to_file(path1) |
| 207 | fadd2.write_to_file(path2) |
| 208 | cc.create_shared(path_dso, [path1, path2]) |
| 209 | |
| 210 | def popen_check(): |
| 211 | import ctypes |
| 212 | |
| 213 | import tvm.runtime |
| 214 | |
| 215 | # Load dll, will trigger system library registration |
| 216 | ctypes.CDLL(path_dso) |
| 217 | # Load the system wide library |
no test coverage detected
searching dependent graphs…