def update_lib(lib): test_dir = os.path.dirname( os.path.realpath(os.path.expanduser(__file__))) source_dir = os.path.join(test_dir, "..", "..", "..") contrib_path = os.path.join(source_dir, "src", "runtime", "contrib") kwargs = {} kwargs["options"] = ["-O2", "-std=c++14", "-I" + contrib_path] tmp_path = util.tempdir() lib_name = 'lib.so' lib_path = tmp_path.relpath(lib_name) lib.export_library(lib_path, fcompile=False, **kwargs) lib = runtime.load_module(lib_path) return lib
def load_tvm(self, export_dir): """Load tvm module from export directory""" self.export_dir = export_dir self.tvm_lib = load_module(os.path.join(export_dir, TVM_ASSETS[0])) with open(os.path.join(export_dir, TVM_ASSETS[1]), "r", encoding="utf8") as f: self.tvm_graph = f.read() with open(os.path.join(export_dir, TVM_ASSETS[2]), "rb") as f: self.tvm_params = relay.load_param_dict(f.read()) self.tvm_module = graph_executor.create(self.tvm_graph, self.tvm_lib, device=self.dev) self.tvm_module.set_input(**self.tvm_params) return self.tvm_module
def build_cutlass_kernels(lib, sm, tmp_dir="./tmp", lib_path="compile.so", threads=-1, use_fast_math=False): """Compile CUTLASS kernels in lib and return the runtime module ready to run. Parameters ---------- lib : GraphExecutorFactoryModule The output from relay.build containing compiled host code and non-cutlass kernels. sm : int An integer specifying the compute capability. For example, 75 for Turing and 80 or 86 for Ampere. tmp_dir : string, optional A temporary directory where intermediate compiled artifacts will be stored. lib_path : string, optional The path to a shared library which will be generated as the result of the build process. threads : int, optional The number of threads to use for compiling generated kernels. Only available for CUDA 11.2 or later. Use all physical cores by default. use_fast_math : bool, optional Whether or not to use faster but less accurate math intrinsics. Returns ------- updated_lib : runtime.Module The updated module with compiled cutlass kernels. """ kwargs = _get_cutlass_compile_options(sm, threads, use_fast_math) lib.export_library(lib_path, workspace_dir=tmp_dir, **kwargs) return runtime.load_module(lib_path)
def update_lib(lib, backend): test_dir = path.dirname(path.realpath(path.expanduser(__file__))) source_dir = path.join(test_dir, "..", "..", "..") contrib_path = path.join(source_dir, "src", "runtime", "contrib") sim_opts = [path.join(test_dir, "lib", "sim_lib.cc")] opts = compile_adder() if backend == "adder" else sim_opts kwargs = {} kwargs["options"] = [ "-O2", "-std=c++14", "-I" + contrib_path, "-I" + path.join(test_dir, "lib"), ] + opts tmp_path = util.tempdir() lib_name = "lib.so" lib_path = tmp_path.relpath(lib_name) lib.export_library(lib_path, fcompile=False, **kwargs) lib = runtime.load_module(lib_path) return lib