wehub-resource-sync
52 lines
2.0 KiB
Python
52 lines
2.0 KiB
Python
"""A compiler pass that dispatches patterns to CUBLAS."""
|
|
|
|
import tvm
|
|
from tvm import IRModule, relax
|
|
from tvm.relax.backend import get_patterns_with_prefix
|
|
|
|
try:
|
|
import tvm.relax.backend.cuda.cublas as _cublas # noqa: F401
|
|
import tvm.relax.backend.rocm.hipblas as _hipblas # noqa: F401
|
|
except ImportError:
|
|
# Note: legacy path of cublas/hipblas for backward compatibility
|
|
pass
|
|
|
|
|
|
@tvm.transform.module_pass(opt_level=0, name="BLASDispatch")
|
|
class BLASDispatch:
|
|
"""A compiler pass that dispatches patterns to cuBLAS/hipBLAS."""
|
|
|
|
def __init__(self, target: tvm.target.Target) -> None:
|
|
if target.kind.name == "cuda":
|
|
self.has_blas = tvm.get_global_func("relax.ext.cublas", True)
|
|
if not self.has_blas:
|
|
raise Exception("cuBLAS is not enabled.")
|
|
self.patterns = get_patterns_with_prefix("cublas")
|
|
elif target.kind.name == "rocm":
|
|
self.has_blas = tvm.get_global_func("relax.ext.hipblas", True)
|
|
if not self.has_blas:
|
|
raise Exception("hipBLAS is not enabled.")
|
|
self.patterns = get_patterns_with_prefix("hipblas")
|
|
else:
|
|
raise Exception(f"Unsupported target {target.kind.name} for BLAS dispatch.")
|
|
|
|
def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
|
|
"""IRModule-level transformation"""
|
|
model_names = [
|
|
gv.name_hint for gv, func in mod.functions.items() if isinstance(func, relax.Function)
|
|
]
|
|
# exclude single batch decode
|
|
model_names = [name for name in model_names if "batch" in name or "decode" not in name]
|
|
mod = tvm.transform.Sequential(
|
|
[
|
|
relax.transform.FuseOpsByPattern(
|
|
self.patterns,
|
|
bind_constants=False,
|
|
annotate_codegen=True,
|
|
entry_functions=model_names,
|
|
),
|
|
relax.transform.RunCodegen({}, entry_functions=model_names),
|
|
]
|
|
)(mod)
|
|
return mod
|