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
