wehub-resource-sync
555 lines
18 KiB
Python
555 lines
18 KiB
Python
"""Helper functions for target auto-detection."""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Callable, List, Optional, Tuple # noqa: UP035
|
|
|
|
from tvm import IRModule, relax
|
|
from tvm.ir.transform import Pass
|
|
from tvm.support import ndk, tar, xcode
|
|
from tvm.target import Target
|
|
from tvm_ffi import get_global_func, register_global_func
|
|
|
|
from . import logging
|
|
from .auto_device import AUTO_DETECT_DEVICES, detect_device, device2str
|
|
from .constants import MLC_MULTI_ARCH
|
|
from .style import bold, green, red
|
|
|
|
if TYPE_CHECKING:
|
|
from mlc_llm.compiler.compile import CompileArgs
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# TODO: add help message on how to specify the target manually
|
|
HELP_MSG = """TBD"""
|
|
FOUND = green("Found")
|
|
NOT_FOUND = red("Not found")
|
|
BuildFunc = Callable[[IRModule, "CompileArgs", Pass], None]
|
|
|
|
|
|
def detect_target_and_host(
|
|
target_hint: str,
|
|
host_hint: str = "auto",
|
|
enable_subgroups: Optional[bool] = None,
|
|
) -> Tuple[Target, BuildFunc]: # noqa: UP006
|
|
"""Detect the configuration for the target device and its host, for example, target GPU and
|
|
the host CPU.
|
|
|
|
Parameters
|
|
----------
|
|
target_hint : str
|
|
The hint for the target device.
|
|
|
|
host_hint : str
|
|
The hint for the host CPU, default is "auto".
|
|
"""
|
|
target, build_func = _detect_target_gpu(target_hint)
|
|
if target.host is None:
|
|
target = Target(target, host=_detect_target_host(host_hint))
|
|
target = _apply_webgpu_subgroups(target, enable_subgroups)
|
|
if target.kind.name == "cuda":
|
|
# Enable thrust for CUDA
|
|
target_dict = dict(target.export())
|
|
target_dict["libs"] = (
|
|
(target_dict["libs"] + ["thrust"]) if "libs" in target_dict else ["thrust"]
|
|
)
|
|
target = Target(target_dict)
|
|
_register_cuda_hook(target)
|
|
elif target.kind.name == "rocm":
|
|
target_dict = dict(target.export())
|
|
extra_libs = ["thrust", "rocblas", "miopen", "hipblas"]
|
|
target_dict["libs"] = (
|
|
(target_dict["libs"] + extra_libs) if "libs" in target_dict else extra_libs
|
|
)
|
|
target = Target(target_dict)
|
|
return target, build_func
|
|
|
|
|
|
def _apply_webgpu_subgroups(target: Target, enable_subgroups: Optional[bool]) -> Target:
|
|
if not enable_subgroups:
|
|
return target
|
|
if target.kind.name != "webgpu":
|
|
logger.warning(
|
|
"--enable-subgroups is only supported for WebGPU targets; ignoring for %s",
|
|
target.kind.name,
|
|
)
|
|
return target
|
|
target_dict = dict(target.export())
|
|
target_dict["supports_subgroups"] = True
|
|
return Target(target_dict)
|
|
|
|
|
|
def _detect_target_gpu(hint: str) -> Tuple[Target, BuildFunc]: # noqa: UP006
|
|
if hint in ["iphone", "macabi", "android", "webgpu", "mali", "opencl"]:
|
|
hint += ":generic"
|
|
if hint == "auto" or hint in AUTO_DETECT_DEVICES:
|
|
target: Optional[Target] = None
|
|
device = detect_device(hint)
|
|
if device is not None:
|
|
device_str = device2str(device)
|
|
try:
|
|
target = Target.from_device(device)
|
|
except ValueError:
|
|
logger.info("%s: Cannot detect target from device: %s", NOT_FOUND, device_str)
|
|
if target is None:
|
|
raise ValueError(f"No target detected from device: {hint}. Please specify explicitly")
|
|
logger.info(
|
|
'%s configuration of target device "%s": %s',
|
|
FOUND,
|
|
bold(device_str),
|
|
target.export(),
|
|
)
|
|
return target, _build_default()
|
|
if hint in PRESET:
|
|
preset = PRESET[hint]
|
|
target = Target(preset["target"])
|
|
build = preset.get("build", _build_default)
|
|
return target, build()
|
|
if _is_device(hint):
|
|
logger.info("Detecting target device: %s", hint)
|
|
target = Target.from_device(hint)
|
|
logger.info("%s target: %s", FOUND, target.export())
|
|
return target, _build_default()
|
|
try:
|
|
logger.info("Try creating device target from string: %s", hint)
|
|
target = Target(hint)
|
|
logger.info("%s target: %s", FOUND, target.export())
|
|
return target, _build_default()
|
|
except Exception as err:
|
|
logger.info("%s: Failed to create target", NOT_FOUND)
|
|
raise ValueError(f"Invalid target: {hint}") from err
|
|
|
|
|
|
def _detect_target_host(hint: str) -> Target:
|
|
"""Detect the host CPU architecture."""
|
|
if hint == "auto":
|
|
target_triple = get_global_func("tvm.codegen.llvm.GetDefaultTargetTriple")()
|
|
target = Target.from_device("cpu")
|
|
logger.info("%s host LLVM triple: %s", FOUND, bold(target.attrs["mtriple"]))
|
|
logger.info("%s host LLVM CPU: %s", FOUND, bold(target.attrs["mcpu"]))
|
|
return target
|
|
target_triple = hint
|
|
logger.info("Using LLVM triple specified by --host: %s", bold(target_triple))
|
|
return Target({"kind": "llvm", "mtriple": target_triple})
|
|
|
|
|
|
def _is_device(device: str):
|
|
if " " in device:
|
|
return False
|
|
if device.count(":") != 1:
|
|
return False
|
|
return True
|
|
|
|
|
|
def _add_system_lib_prefix(mod: IRModule, prefix: str, is_system_lib: bool) -> IRModule:
|
|
if is_system_lib and prefix:
|
|
mod = mod.with_attrs({"system_lib_prefix": prefix})
|
|
elif is_system_lib:
|
|
logger.warning(
|
|
"%s is not specified when building a static library",
|
|
bold("--system-lib-prefix"),
|
|
)
|
|
elif prefix:
|
|
logger.warning(
|
|
"--system-lib-prefix is specified, but it will not take any effect "
|
|
"when building the shared library"
|
|
)
|
|
return mod
|
|
|
|
|
|
def _build_metal_x86_64():
|
|
def build(mod: IRModule, args: "CompileArgs", pipeline=None):
|
|
output = args.output
|
|
mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=False)
|
|
assert output.suffix == ".dylib"
|
|
relax.build(
|
|
mod,
|
|
target=args.target,
|
|
relax_pipeline=pipeline,
|
|
).export_library(
|
|
str(output),
|
|
fcompile=xcode.create_dylib,
|
|
sdk="macosx",
|
|
arch="x86_64",
|
|
)
|
|
|
|
return build
|
|
|
|
|
|
def _build_iphone():
|
|
@register_global_func("tvm_callback_metal_compile", override=True)
|
|
def compile_metal(src, target):
|
|
libs = target.attrs.get("libs", None)
|
|
if libs:
|
|
return xcode.compile_metal(src, sdk=libs[0])
|
|
return xcode.compile_metal(src)
|
|
|
|
def build(mod: IRModule, args: "CompileArgs", pipeline=None):
|
|
output = args.output
|
|
mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=True)
|
|
assert output.suffix == ".tar"
|
|
relax.build(
|
|
mod,
|
|
target=args.target,
|
|
relax_pipeline=pipeline,
|
|
system_lib=True,
|
|
).export_library(
|
|
str(output),
|
|
fcompile=tar.tar,
|
|
)
|
|
|
|
return build
|
|
|
|
|
|
def _build_android():
|
|
def build(mod: IRModule, args: "CompileArgs", pipeline=None):
|
|
output = args.output
|
|
mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=True)
|
|
assert output.suffix == ".tar"
|
|
ex = relax.build(
|
|
mod,
|
|
target=args.target,
|
|
relax_pipeline=pipeline,
|
|
system_lib=True,
|
|
)
|
|
ex.export_library(
|
|
str(output),
|
|
fcompile=tar.tar,
|
|
)
|
|
if args.debug_dump is not None:
|
|
source = ex.mod.imports[0].imports[0].inspect_source()
|
|
with open(args.debug_dump / "kernel.cl", "w", encoding="utf-8") as f:
|
|
f.write(source)
|
|
|
|
return build
|
|
|
|
|
|
def _build_android_so():
|
|
def build(mod: IRModule, args: "CompileArgs", pipeline=None):
|
|
output = args.output
|
|
mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=False)
|
|
assert output.suffix == ".so"
|
|
ex = relax.build(
|
|
mod,
|
|
target=args.target,
|
|
relax_pipeline=pipeline,
|
|
system_lib=False,
|
|
)
|
|
ex.export_library(
|
|
str(output),
|
|
fcompile=ndk.create_shared,
|
|
)
|
|
if args.debug_dump is not None:
|
|
source = ex.mod.imports[0].imports[0].inspect_source()
|
|
with open(args.debug_dump / "kernel.cl", "w", encoding="utf-8") as f:
|
|
f.write(source)
|
|
|
|
return build
|
|
|
|
|
|
def _build_webgpu():
|
|
def build(mod: IRModule, args: "CompileArgs", pipeline=None):
|
|
output = args.output
|
|
mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=True)
|
|
assert output.suffix == ".wasm"
|
|
|
|
# Try to locate `mlc_wasm_runtime.bc`
|
|
bc_path = None
|
|
bc_candidates = ["web/dist/wasm/mlc_wasm_runtime.bc"]
|
|
if os.environ.get("MLC_LLM_SOURCE_DIR", None):
|
|
mlc_source_home_dir = os.environ["MLC_LLM_SOURCE_DIR"]
|
|
bc_candidates.append(
|
|
os.path.join(mlc_source_home_dir, "web", "dist", "wasm", "mlc_wasm_runtime.bc")
|
|
)
|
|
error_info = (
|
|
"Cannot find library: mlc_wasm_runtime.bc\n"
|
|
+ "Make sure you have run `./web/prep_emcc_deps.sh` and "
|
|
+ "`export MLC_LLM_SOURCE_DIR=/path/to/mlc-llm` so that we can locate the file. "
|
|
+ "We tried to look at candidate paths:\n"
|
|
)
|
|
for candidate in bc_candidates:
|
|
error_info += candidate + "\n"
|
|
if Path(candidate).exists():
|
|
bc_path = candidate
|
|
if not bc_path:
|
|
raise RuntimeError(error_info)
|
|
|
|
relax.build(
|
|
mod,
|
|
target=args.target,
|
|
relax_pipeline=pipeline,
|
|
system_lib=True,
|
|
).export_library(
|
|
str(output),
|
|
libs=[bc_path],
|
|
)
|
|
|
|
return build
|
|
|
|
|
|
def _build_mali():
|
|
def build(mod: IRModule, args: "CompileArgs", pipeline=None):
|
|
output = args.output
|
|
mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=True)
|
|
assert output.suffix == ".so"
|
|
mod = relax.build(
|
|
mod,
|
|
target=args.target,
|
|
relax_pipeline=pipeline,
|
|
system_lib=True,
|
|
)
|
|
if "TVM_NDK_CC" in os.environ:
|
|
mod.export_library(str(output), fcompile=ndk.create_shared)
|
|
else:
|
|
mod.export_library(str(output))
|
|
|
|
return build
|
|
|
|
|
|
def _build_default():
|
|
def build(mod: IRModule, args: "CompileArgs", pipeline=None):
|
|
output = args.output
|
|
if output.suffix in [".tar", ".lib"]:
|
|
system_lib = True
|
|
elif output.suffix in [".so", ".dylib", ".dll"]:
|
|
system_lib = False
|
|
else:
|
|
logger.warning("Unknown output suffix: %s. Assuming shared library.", output.suffix)
|
|
system_lib = False
|
|
mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=system_lib)
|
|
relax.build(
|
|
mod,
|
|
target=args.target,
|
|
relax_pipeline=pipeline,
|
|
system_lib=system_lib,
|
|
).export_library(
|
|
str(output),
|
|
)
|
|
|
|
return build
|
|
|
|
|
|
def detect_cuda_arch_list(target: Target) -> List[int]: # noqa: UP006
|
|
"""Detect the CUDA architecture list from the target."""
|
|
|
|
def convert_to_num(arch_str):
|
|
arch_num_str = "".join(filter(str.isdigit, arch_str))
|
|
assert arch_num_str, f"'{arch_str}' does not contain any digits"
|
|
return int(arch_num_str)
|
|
|
|
assert target.kind.name == "cuda", f"Expect target to be CUDA, but got {target}"
|
|
if MLC_MULTI_ARCH is not None:
|
|
multi_arch = [convert_to_num(x) for x in MLC_MULTI_ARCH.split(",")]
|
|
else:
|
|
assert target.attrs.get("arch", "").startswith("sm_")
|
|
multi_arch = [convert_to_num(target.attrs.get("arch")[3:])]
|
|
multi_arch = list(set(multi_arch))
|
|
return multi_arch
|
|
|
|
|
|
def _register_cuda_hook(target: Target):
|
|
if MLC_MULTI_ARCH is None:
|
|
default_arch = target.attrs.get("arch", None)
|
|
logger.info("Generating code for CUDA architecture: %s", bold(default_arch))
|
|
logger.info(
|
|
"To produce multi-arch fatbin, set environment variable %s. "
|
|
"Example: MLC_MULTI_ARCH=70,72,75,80,86,87,89,90a",
|
|
bold("MLC_MULTI_ARCH"),
|
|
)
|
|
multi_arch = None
|
|
else:
|
|
logger.info("%s %s: %s", FOUND, bold("MLC_MULTI_ARCH"), MLC_MULTI_ARCH)
|
|
multi_arch = [x.strip() for x in MLC_MULTI_ARCH.split(",")]
|
|
logger.info("Generating code for CUDA architecture: %s", multi_arch)
|
|
|
|
@register_global_func("tvm_callback_cuda_compile", override=True)
|
|
def tvm_callback_cuda_compile(code):
|
|
"""use nvcc to generate fatbin code for better optimization"""
|
|
from tvm.support import nvcc
|
|
|
|
if multi_arch is None:
|
|
ptx = nvcc.compile_cuda(code, target_format="fatbin", compiler="nvcc")
|
|
else:
|
|
arch = []
|
|
for compute_version in multi_arch:
|
|
arch += [
|
|
"-gencode",
|
|
f"arch=compute_{compute_version},code=sm_{compute_version}",
|
|
]
|
|
ptx = nvcc.compile_cuda(code, target_format="fatbin", arch=arch, compiler="nvcc")
|
|
return ptx
|
|
|
|
|
|
def detect_system_lib_prefix(
|
|
target_hint: str, prefix_hint: str, model_name: str, quantization: str
|
|
) -> str:
|
|
"""Detect the iOS / Android system lib prefix to identify the library needed to load the app.
|
|
|
|
Parameters
|
|
----------
|
|
target_hint : str
|
|
The hint for the target device.
|
|
|
|
prefix_hint : str
|
|
The hint for the system lib prefix.
|
|
"""
|
|
if prefix_hint == "auto" and (
|
|
target_hint.startswith("iphone")
|
|
or target_hint.startswith("macabi")
|
|
or target_hint.startswith("android")
|
|
):
|
|
prefix = f"{model_name}_{quantization}_".replace("-", "_")
|
|
logger.warning(
|
|
"%s is automatically picked from the filename, %s, this allows us to use the filename "
|
|
"as the model_lib in android/iOS builds. Please avoid renaming the .tar file when "
|
|
"uploading the prebuilt.",
|
|
bold("--system-lib-prefix"),
|
|
bold(prefix),
|
|
)
|
|
return prefix
|
|
if target_hint not in ["iphone", "macabi", "android"]:
|
|
return ""
|
|
return prefix_hint
|
|
|
|
|
|
_MACABI_ARCH = os.environ.get("MLC_MACABI_ARCH", "").strip() or "arm64"
|
|
if _MACABI_ARCH not in ["arm64", "x86_64"]:
|
|
_MACABI_ARCH = "arm64"
|
|
_MACABI_MTRIPLE = (
|
|
"x86_64-apple-ios18.0-macabi" if _MACABI_ARCH == "x86_64" else "arm64-apple-ios18.0-macabi"
|
|
)
|
|
|
|
PRESET = {
|
|
"iphone:generic": {
|
|
"target": {
|
|
"kind": "metal",
|
|
"max_threads_per_block": 256,
|
|
"max_shared_memory_per_block": 32768,
|
|
"thread_warp_size": 1,
|
|
"libs": ["iphoneos"],
|
|
"host": {
|
|
"kind": "llvm",
|
|
"mtriple": "arm64-apple-darwin",
|
|
},
|
|
},
|
|
"build": _build_iphone,
|
|
},
|
|
"macabi:generic": {
|
|
"target": {
|
|
"kind": "metal",
|
|
"max_threads_per_block": 256,
|
|
"max_shared_memory_per_block": 32768,
|
|
"thread_warp_size": 1,
|
|
"libs": ["macosx"],
|
|
"host": {
|
|
"kind": "llvm",
|
|
"mtriple": _MACABI_MTRIPLE,
|
|
},
|
|
},
|
|
"build": _build_iphone,
|
|
},
|
|
"android:generic": {
|
|
"target": {
|
|
"kind": "opencl",
|
|
"host": {
|
|
"kind": "llvm",
|
|
"mtriple": "aarch64-linux-android",
|
|
},
|
|
},
|
|
"build": _build_android,
|
|
},
|
|
"android:adreno": {
|
|
"target": {
|
|
"kind": "opencl",
|
|
"device": "adreno",
|
|
"max_threads_per_block": 512,
|
|
"host": {
|
|
"kind": "llvm",
|
|
"mtriple": "aarch64-linux-android",
|
|
},
|
|
},
|
|
"build": _build_android,
|
|
},
|
|
"android:adreno-so": {
|
|
"target": {
|
|
"kind": "opencl",
|
|
"device": "adreno",
|
|
"max_threads_per_block": 512,
|
|
"host": {
|
|
"kind": "llvm",
|
|
"mtriple": "aarch64-linux-android",
|
|
},
|
|
},
|
|
"build": _build_android_so,
|
|
},
|
|
"windows:adreno_x86": {
|
|
"target": {
|
|
"kind": "opencl",
|
|
"device": "adreno",
|
|
"max_threads_per_block": 512,
|
|
"host": {
|
|
"kind": "llvm",
|
|
"mtriple": "x86_64-pc-windows-msvc",
|
|
},
|
|
},
|
|
},
|
|
"metal:x86-64": {
|
|
"target": {
|
|
"kind": "metal",
|
|
"max_threads_per_block": 256,
|
|
"max_shared_memory_per_block": 32768,
|
|
"thread_warp_size": 1,
|
|
},
|
|
"build": _build_metal_x86_64,
|
|
},
|
|
"webgpu:generic": {
|
|
"target": {
|
|
"kind": "webgpu",
|
|
"host": {
|
|
"kind": "llvm",
|
|
"mtriple": "wasm32-unknown-unknown-wasm",
|
|
},
|
|
},
|
|
"build": _build_webgpu,
|
|
},
|
|
"opencl:generic": {
|
|
"target": {
|
|
"kind": "opencl",
|
|
},
|
|
},
|
|
"mali:generic": {
|
|
"target": {
|
|
"kind": "opencl",
|
|
"host": {
|
|
"kind": "llvm",
|
|
"mtriple": "aarch64-linux-gnu",
|
|
},
|
|
},
|
|
"build": _build_mali,
|
|
},
|
|
"metal:generic": {
|
|
"target": {
|
|
"kind": "metal",
|
|
"max_threads_per_block": 256,
|
|
"max_shared_memory_per_block": 32768,
|
|
"thread_warp_size": 1,
|
|
},
|
|
},
|
|
"vulkan:generic": {
|
|
"target": {
|
|
"kind": "vulkan",
|
|
"max_threads_per_block": 256,
|
|
"max_shared_memory_per_block": 32768,
|
|
"thread_warp_size": 1,
|
|
"supports_float16": 1,
|
|
"supports_int64": 1,
|
|
"supports_int16": 1,
|
|
"supports_int8": 1,
|
|
"supports_8bit_buffer": 1,
|
|
"supports_16bit_buffer": 1,
|
|
"supports_storage_buffer_storage_class": 1,
|
|
},
|
|
},
|
|
}
|