diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py index d36554bbd516..8775dafa17d1 100644 --- a/python/tvm/runtime/module.py +++ b/python/tvm/runtime/module.py @@ -444,6 +444,10 @@ def load_module(path, fmt=""): This function will automatically call cc.create_shared if the path is in format .o or .tar """ + if os.path.isfile(path): + path = os.path.realpath(path) + else: + raise ValueError("cannot find file %s" % path) # c++ compiler/linker cc = os.environ.get("CXX", "g++") diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc index 5eaf1e890d8a..b643459be4b6 100644 --- a/src/target/source/codegen_c.cc +++ b/src/target/source/codegen_c.cc @@ -960,11 +960,19 @@ void CodeGenC::VisitStmt_(const EvaluateNode* op) { return; } else if (call->op.same_as(builtin::tvm_struct_set())) { ICHECK_EQ(call->args.size(), 4); + int kind = call->args[2].as()->value; + std::string ref = GetStructRef(call->args[3].dtype(), call->args[0], call->args[1], kind); std::string value = PrintExpr(call->args[3]); - std::string ref = GetStructRef(call->args[3].dtype(), call->args[0], call->args[1], - call->args[2].as()->value); + std::string cast; + if (kind == builtin::kArrStrides) { + // cast void* to int64_t* + cast = call->args[3]->dtype.is_handle() ? "(int64_t*)" : ""; + } else if (kind == builtin::kArrDeviceType) { + // cast int to enum + cast = "(DLDeviceType)"; + } this->PrintIndent(); - this->stream << ref << " = " << value << ";\n"; + this->stream << ref << " = " << cast << value << ";\n"; return; } } diff --git a/src/target/source/codegen_c_host.cc b/src/target/source/codegen_c_host.cc index 3ec64ed2ace9..af4bb48d1d73 100644 --- a/src/target/source/codegen_c_host.cc +++ b/src/target/source/codegen_c_host.cc @@ -268,14 +268,17 @@ void CodeGenCHost::VisitExpr_(const CallNode* op, std::ostream& os) { // NOLINT // NOTE: cannot rely on GetUnique for global decl_stream declarations // because it is reset between AddFunction(). std::string packed_func_name = func_name + "_packed"; - if (declared_globals_.insert(packed_func_name).second) { - // Still reserve the name among unique names. - ICHECK(GetUniqueName(packed_func_name) == packed_func_name) - << "Expected name " << packed_func_name << " to not be taken"; - decl_stream << "static void* " << packed_func_name << " = NULL;\n"; + std::string unique_name; + auto it = declared_globals_.find(packed_func_name); + if (it != declared_globals_.end()) { + unique_name = it->second; + } else { + unique_name = GetUniqueName(packed_func_name); + declared_globals_[packed_func_name] = unique_name; + decl_stream << "static void* " << unique_name << " = NULL;\n"; } - this->PrintGetFuncFromBackend(func_name, packed_func_name); - this->PrintFuncCall(packed_func_name, num_args); + this->PrintGetFuncFromBackend(func_name, unique_name); + this->PrintFuncCall(unique_name, num_args); } else if (op->op.same_as(builtin::tvm_throw_last_error())) { this->PrintIndent(); this->stream << "return -1;\n"; diff --git a/src/target/source/codegen_c_host.h b/src/target/source/codegen_c_host.h index 97fe7ab39efa..eace09f13a07 100644 --- a/src/target/source/codegen_c_host.h +++ b/src/target/source/codegen_c_host.h @@ -24,8 +24,8 @@ #ifndef TVM_TARGET_SOURCE_CODEGEN_C_HOST_H_ #define TVM_TARGET_SOURCE_CODEGEN_C_HOST_H_ -#include #include +#include #include #include "codegen_c.h" @@ -63,8 +63,8 @@ class CodeGenCHost final : public CodeGenC { private: std::string module_name_; - /* \brief tracks declared global variables which live despite GetUniqueName */ - std::set declared_globals_; + /* \brief mapping global packed func to the unique name */ + std::unordered_map declared_globals_; /* \brief names of the functions declared in this module */ Array function_names_; /*! \brief whether to emit asserts in the resulting C code */ diff --git a/tests/python/contrib/test_cblas.py b/tests/python/contrib/test_cblas.py index 946d93385e6f..2245bb6c1527 100644 --- a/tests/python/contrib/test_cblas.py +++ b/tests/python/contrib/test_cblas.py @@ -42,6 +42,13 @@ def get_numpy(a, b, bb, transa, transb): b = b.transpose() return np.dot(a, b) + bb + def compile(f, name="test_matmul_add", ext=".so"): + path = name + ext + f.export_library(path) + mod = tvm.runtime.load_module(path) + f = mod[name] + return f + def verify(target="llvm"): if not tvm.testing.device_enabled(target): print("skip because %s is not enabled..." % target) @@ -50,7 +57,10 @@ def verify(target="llvm"): print("skip because extern function is not available") return dev = tvm.cpu(0) - f = tvm.build(s, [A, B, D, bias], target) + name = "test_matmul_add" + f = tvm.build(s, [A, B, D, bias], target, name=name) + if target == "c": + f = compile(f, name) a = tvm.nd.array(np.random.uniform(size=ashape).astype(A.dtype), dev) b = tvm.nd.array(np.random.uniform(size=bshape).astype(B.dtype), dev) d = tvm.nd.array(np.zeros((n, m), dtype=D.dtype), dev) @@ -60,7 +70,8 @@ def verify(target="llvm"): d.asnumpy(), get_numpy(a.asnumpy(), b.asnumpy(), bb, transa, transb), rtol=1e-5 ) - verify() + verify("llvm") + verify("c") def test_matmul_add(): @@ -164,6 +175,13 @@ def get_numpy(a, b, transa, transb): b = b.transpose(0, 2, 1) return tvm.topi.testing.batch_matmul(a, b) + def compile(f, name="test_batch_matmul", ext=".so"): + path = name + ext + f.export_library(path) + mod = tvm.runtime.load_module(path) + f = mod[name] + return f + def verify(target="llvm"): if not tvm.testing.device_enabled(target): print("skip because %s is not enabled..." % target) @@ -172,7 +190,10 @@ def verify(target="llvm"): print("skip because extern function is not available") return dev = tvm.cpu(0) - f = tvm.build(s, [A, B, D], target) + name = "test_batch_matmul" + f = tvm.build(s, [A, B, D], target, name=name) + if target == "c": + f = compile(f, name) a = tvm.nd.array(np.random.uniform(size=ashape).astype(A.dtype), dev) b = tvm.nd.array(np.random.uniform(size=bshape).astype(B.dtype), dev) d = tvm.nd.array(np.zeros((batch, n, m), dtype=D.dtype), dev) @@ -181,7 +202,8 @@ def verify(target="llvm"): d.asnumpy(), get_numpy(a.asnumpy(), b.asnumpy(), transa, transb), rtol=1e-5 ) - verify() + verify("llvm") + verify("c") def test_batch_matmul(): diff --git a/tests/python/unittest/test_target_codegen_c_host.py b/tests/python/unittest/test_target_codegen_c_host.py index 6102beba0d25..a61bb39bd211 100644 --- a/tests/python/unittest/test_target_codegen_c_host.py +++ b/tests/python/unittest/test_target_codegen_c_host.py @@ -191,6 +191,42 @@ def check_c(): check_c() +def test_call_packed(): + def fake_func(fname="fake.func"): + ib = tvm.tir.ir_builder.create() + A = ib.pointer("float32", name="A") + fake_func1 = tvm.tir.call_packed(fname, A[0]) + + ib.emit(fake_func1) + body = ib.get() + return A, body + + def check_global_packed_func(): + fname = "fake.func" + A, body = fake_func(fname) + func1 = tvm.tir.PrimFunc([A], body).with_attr("global_symbol", "func1") + B, body = fake_func() + func2 = tvm.tir.PrimFunc([B], body).with_attr("global_symbol", "func2") + mod = tvm.IRModule({"fake_func1": func1, "fake_func2": func2}) + fcode = tvm.build(mod, None, "c") + src = fcode.get_source() + + # there are two locations calling the packed func + assert src.count(fname) == 2 + + suffix = "_packed" + packed_func_name = fname + suffix + # func name will be standardized by GetUniqueName and not exists anymore + assert src.find(packed_func_name) == -1 + + packed_func_real_name = "_".join(fname.split(".")) + suffix + func_declaration = "static void* %s = NULL;" % packed_func_real_name + # src only has 1 valid declaration + assert src.count(func_declaration) == 1 + + check_global_packed_func() + + if __name__ == "__main__": test_add() test_add_pipeline() @@ -198,3 +234,4 @@ def check_c(): test_ceil() test_floor() test_round() + test_call_packed()