diff --git a/python/tvm/_ffi/libinfo.py b/python/tvm/_ffi/libinfo.py index 02f038cb06ad..c885c0b2cc79 100644 --- a/python/tvm/_ffi/libinfo.py +++ b/python/tvm/_ffi/libinfo.py @@ -70,7 +70,7 @@ def find_lib_path(name=None, search_path=None, optional=False): if os.environ.get("TVM_LIBRARY_PATH", None): dll_path.append(os.environ["TVM_LIBRARY_PATH"]) - if sys.platform.startswith("linux"): + if sys.platform.startswith("linux") or sys.platform.startswith("freebsd"): dll_path.extend(split_env_var("LD_LIBRARY_PATH", ":")) dll_path.extend(split_env_var("PATH", ":")) elif sys.platform.startswith("darwin"): diff --git a/python/tvm/contrib/cc.py b/python/tvm/contrib/cc.py index 1b6a62fd9834..9643d9b650fd 100644 --- a/python/tvm/contrib/cc.py +++ b/python/tvm/contrib/cc.py @@ -17,6 +17,7 @@ """Util to invoke C/C++ compilers in the system.""" # pylint: disable=invalid-name import sys +import os import subprocess from .._ffi.base import py_str @@ -39,7 +40,11 @@ def create_shared(output, objects, options=None, cc="g++"): cc : Optional[str] The compiler command. """ - if sys.platform == "darwin" or sys.platform.startswith("linux"): + if ( + sys.platform == "darwin" + or sys.platform.startswith("linux") + or sys.platform.startswith("freebsd") + ): _linux_compile(output, objects, options, cc, compile_shared=True) elif sys.platform == "win32": _windows_shared(output, objects, options) @@ -103,7 +108,9 @@ def get_target_triple(): # assign so as default output format create_shared.output_format = "so" if sys.platform != "win32" else "dll" create_shared.get_target_triple = get_target_by_dump_machine( - "g++" if sys.platform == "darwin" or sys.platform.startswith("linux") else None + os.environ.get( + "CXX", "g++" if sys.platform == "darwin" or sys.platform.startswith("linux") else None + ) ) diff --git a/python/tvm/rpc/server.py b/python/tvm/rpc/server.py index 03a124b4686b..b25ed4675641 100644 --- a/python/tvm/rpc/server.py +++ b/python/tvm/rpc/server.py @@ -73,6 +73,9 @@ def load_module(file_name): @tvm._ffi.register_func("tvm.rpc.server.download_linked_module", override=True) def download_linked_module(file_name): """Load module from remote side.""" + # c++ compiler/linker + cc = os.environ.get("CXX", "g++") + # pylint: disable=import-outside-toplevel path = temp.relpath(file_name) @@ -80,7 +83,7 @@ def download_linked_module(file_name): # Extra dependencies during runtime. from tvm.contrib import cc as _cc - _cc.create_shared(path + ".so", path) + _cc.create_shared(path + ".so", path, cc=cc) path += ".so" elif path.endswith(".tar"): # Extra dependencies during runtime. @@ -89,7 +92,7 @@ def download_linked_module(file_name): tar_temp = util.tempdir(custom_path=path.replace(".tar", "")) _tar.untar(path, tar_temp.temp_dir) files = [tar_temp.relpath(x) for x in tar_temp.listdir()] - _cc.create_shared(path + ".so", files) + _cc.create_shared(path + ".so", files, cc=cc) path += ".so" elif path.endswith(".dylib") or path.endswith(".so"): pass diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py index d9166b5f4417..0559a016a9c2 100644 --- a/python/tvm/runtime/module.py +++ b/python/tvm/runtime/module.py @@ -17,6 +17,7 @@ # pylint: disable=invalid-name, unused-import, import-outside-toplevel """Runtime Module namespace.""" +import os import ctypes import struct from collections import namedtuple @@ -393,13 +394,17 @@ def load_module(path, fmt=""): This function will automatically call cc.create_shared if the path is in format .o or .tar """ + + # c++ compiler/linker + cc = os.environ.get("CXX", "g++") + # High level handling for .o and .tar file. # We support this to be consistent with RPC module load. if path.endswith(".o"): # Extra dependencies during runtime. from tvm.contrib import cc as _cc - _cc.create_shared(path + ".so", path) + _cc.create_shared(path + ".so", path, cc=cc) path += ".so" elif path.endswith(".tar"): # Extra dependencies during runtime. @@ -408,7 +413,7 @@ def load_module(path, fmt=""): tar_temp = _util.tempdir(custom_path=path.replace(".tar", "")) _tar.untar(path, tar_temp.temp_dir) files = [tar_temp.relpath(x) for x in tar_temp.listdir()] - _cc.create_shared(path + ".so", files) + _cc.create_shared(path + ".so", files, cc=cc) path += ".so" # TODO(weberlo): we should probably use a more distinctive suffix for uTVM object files elif path.endswith(".obj"):