wehub-resource-sync
294 lines
9.3 KiB
Python
294 lines
9.3 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=invalid-name, no-member
|
|
"""VM build logics"""
|
|
|
|
import tvm
|
|
from tvm import relax
|
|
from tvm.ir.module import IRModule
|
|
from tvm.runtime import Executable
|
|
from tvm.tirx.function import PrimFunc
|
|
|
|
from . import _ffi_api
|
|
|
|
|
|
class VMExecutable(Executable):
|
|
"""The virtual machine executable object emitted by the VM compiler or the ExecBuilder."""
|
|
|
|
def __init__(self, mod: tvm.runtime.Module):
|
|
super().__init__(mod)
|
|
self._stats = self.mod["stats"]
|
|
self._as_text = self.mod["as_text"]
|
|
self._as_python = self.mod["as_python"]
|
|
|
|
def stats(self) -> str:
|
|
"""print the detailed statistics of the executable."""
|
|
return self._stats()
|
|
|
|
def as_text(self) -> str:
|
|
"""print the instructions as text format."""
|
|
return self._as_text()
|
|
|
|
def as_python(self) -> str:
|
|
"""print the instructions as python program."""
|
|
return self._as_python()
|
|
|
|
|
|
def _vmcodegen(
|
|
builder: "relax.ExecBuilder",
|
|
mod: tvm.IRModule,
|
|
exec_mode: str = "bytecode",
|
|
) -> tvm.IRModule:
|
|
"""Running VM codegen.
|
|
|
|
Parameters
|
|
----------
|
|
builder: relax.ExecBuilder
|
|
ExecBuilder to collect the vm executable.
|
|
|
|
mod: IRModule
|
|
The input IRModule to be built.
|
|
|
|
exec_mode: {"bytecode", "compiled"}
|
|
The execution mode.
|
|
|
|
Return
|
|
------
|
|
leftover: IRModule
|
|
Left over IRModule that may contain extra functions.
|
|
"""
|
|
|
|
if exec_mode == "bytecode":
|
|
return _ffi_api.VMCodeGen(builder, mod) # type:ignore
|
|
if exec_mode == "compiled":
|
|
return _ffi_api.VMTIRCodeGen(builder, mod) # type: ignore
|
|
raise ValueError(f"Unknown exec_mode {exec_mode}")
|
|
|
|
|
|
def _auto_attach_system_lib_prefix(
|
|
tir_mod: tvm.IRModule,
|
|
target: tvm.target.Target | None = None,
|
|
system_lib: bool | None = None,
|
|
):
|
|
"""Automatically detect system lib req and attach prefix attr"""
|
|
if target is not None:
|
|
host = target if target.host is None else target.host
|
|
if system_lib is None:
|
|
system_lib = False
|
|
if "wasm" in host.attrs.get("mtriple", ""):
|
|
system_lib = True
|
|
|
|
if system_lib:
|
|
if tir_mod.get_attr("system_lib_prefix") is None:
|
|
return tir_mod.with_attr("system_lib_prefix", "")
|
|
return tir_mod
|
|
|
|
|
|
def _is_device_module(mod: tvm.runtime.Module) -> bool:
|
|
return mod.kind in ["cuda", "opencl", "metal", "hip", "vulkan", "webgpu"]
|
|
|
|
|
|
def _vmlink(
|
|
builder: "relax.ExecBuilder",
|
|
target: str | tvm.target.Target | None,
|
|
tir_mod: tvm.IRModule | None = None,
|
|
tir_pipeline: str | tvm.transform.Pass | None = "default",
|
|
ext_libs: list[tvm.runtime.Module] | None = None,
|
|
params: dict[str, list] | None = None,
|
|
*,
|
|
system_lib: bool | None = None,
|
|
):
|
|
"""
|
|
Internal codegen function to make executable.
|
|
|
|
This function is only used for unit-testing purpoes.
|
|
|
|
Use build instead.
|
|
|
|
Parameters
|
|
----------
|
|
builder: relax.ExecBuilder
|
|
Builder used to collect executables.
|
|
|
|
target : Optional[Union[str, tvm.target.Target]]
|
|
A build target which can have optional host side compilation target.
|
|
If the target is not specified, the target in the vdevice list will be used.
|
|
For multi-target compilation, the vdevice should be annotated.
|
|
|
|
tir_mod: IRModule
|
|
The input TIR IRModule to be linked together.
|
|
|
|
ext_libs: List[tvm.runtime.Module]
|
|
List of compiled external modules.
|
|
|
|
params: Optional[Dict[str, list]]
|
|
Extra parameter mappings.
|
|
|
|
Returns
|
|
-------
|
|
ex: tvm.relax.Executable
|
|
An executable that can be loaded by virtual machine.
|
|
"""
|
|
if isinstance(target, str):
|
|
target = tvm.target.Target(target)
|
|
if params is None:
|
|
params = {}
|
|
if ext_libs is None:
|
|
ext_libs = []
|
|
lib = None
|
|
relax_ext_libs = []
|
|
tir_ext_libs = []
|
|
if tir_mod is not None and len(tir_mod.get_global_vars()) > 0:
|
|
tir_mod = _auto_attach_system_lib_prefix(tir_mod, target, system_lib)
|
|
lib = tvm.tirx.build(tir_mod, target=target, pipeline=tir_pipeline)
|
|
for ext_mod in ext_libs:
|
|
if _is_device_module(ext_mod):
|
|
tir_ext_libs.append(ext_mod)
|
|
else:
|
|
relax_ext_libs.append(ext_mod)
|
|
if lib is not None:
|
|
for mod in tir_ext_libs:
|
|
lib.import_module(mod)
|
|
elif len(tir_ext_libs) > 0:
|
|
print("Warning: No TIR module is found, but external modules for TIR are provided.")
|
|
lib = _ffi_api.VMLink(builder, target, lib, relax_ext_libs, params) # type: ignore
|
|
return VMExecutable(lib)
|
|
|
|
|
|
def build(
|
|
mod: tvm.IRModule,
|
|
target: str | tvm.target.Target | None = None,
|
|
params: dict[str, list] | None = None,
|
|
relax_pipeline: None | str | tvm.transform.Pass = "default",
|
|
tir_pipeline: None | str | tvm.transform.Pass = "default",
|
|
exec_mode: str = "bytecode",
|
|
*,
|
|
system_lib: bool | None = None,
|
|
) -> Executable:
|
|
"""
|
|
Build an IRModule to VM executable.
|
|
|
|
Parameters
|
|
----------
|
|
mod: IRModule
|
|
The input IRModule to be built.
|
|
|
|
target : Optional[Union[str, tvm.target.Target]]
|
|
A build target which can have optional host side compilation target.
|
|
|
|
When TVM compiles device specific program such as CUDA,
|
|
we also need host(CPU) side code to interact with the driver
|
|
to setup the dimensions and parameters correctly.
|
|
host is used to specify the host side codegen target.
|
|
By default, llvm is used if it is enabled,
|
|
otherwise a c backend is used.
|
|
|
|
params: Optional[Dict[str, list]]
|
|
Parameters for the input IRModule that will be bound.
|
|
|
|
relax_pipeline : str = "default"
|
|
The Relax compilation pipeline to use.
|
|
|
|
tir_pipelinie : str = "default"
|
|
The TIR compilation pipeline to use.
|
|
|
|
exec_mode: {"bytecode", "compiled"}
|
|
The execution mode.
|
|
|
|
system_lib: Optional[bool]
|
|
Whether to build system lib that is being packed statically and
|
|
auto registers generated functions to the system.
|
|
By default auto detects based on the target.
|
|
|
|
Returns
|
|
-------
|
|
ex: tvm.relax.Executable
|
|
An executable that can be loaded by virtual machine.
|
|
|
|
Example
|
|
-------
|
|
|
|
.. code-block:: python
|
|
|
|
class InputModule:
|
|
@R.function
|
|
def foo(x: Tensor((3, 4), "float32"), y: Tensor((3, 4), "float32")):
|
|
z = R.add(x, y)
|
|
return z
|
|
|
|
mod = InputModule
|
|
target = tvm.target.Target("llvm", host="llvm")
|
|
ex = tvm.compile(mod, target)
|
|
"""
|
|
|
|
def _extract_attrs(mod: tvm.IRModule):
|
|
attrs = dict(mod.attrs) if mod.attrs else {}
|
|
ext_libs = attrs.get("external_mods", [])
|
|
constants = attrs.get("const_name_to_constant", {})
|
|
return ext_libs, constants
|
|
|
|
if isinstance(target, str):
|
|
target = tvm.target.Target(target)
|
|
if not params:
|
|
params = {}
|
|
|
|
if relax_pipeline is not None:
|
|
if isinstance(relax_pipeline, str):
|
|
# For GPU targets, prefer the target-specific pipeline which
|
|
# includes DLight scheduling. Without it, TIR functions generated
|
|
# from ops like Clip/ReLU6 lack thread bindings and fail
|
|
# VerifyMemory. CPU targets continue to use the generic pipeline
|
|
# since the CPU-specific pipeline applies fusion passes that can
|
|
# incorrectly remove call_pure_packed calls whose results are
|
|
# unused but whose side effects are relied upon.
|
|
_is_gpu = target is not None and "gpu" in target.keys
|
|
if relax_pipeline == "default" and _is_gpu:
|
|
try:
|
|
relax_pipeline = relax.get_default_pipeline(target)
|
|
except (ValueError, AttributeError):
|
|
relax_pipeline = relax.get_pipeline(relax_pipeline)
|
|
else:
|
|
relax_pipeline = relax.get_pipeline(relax_pipeline)
|
|
if target is None:
|
|
mod = relax_pipeline(mod)
|
|
else:
|
|
with target:
|
|
mod = relax_pipeline(mod)
|
|
|
|
ext_libs, constants = _extract_attrs(mod)
|
|
params.update(dict(constants))
|
|
builder = relax.ExecBuilder()
|
|
mod = _vmcodegen(builder, mod, exec_mode)
|
|
return _vmlink(
|
|
builder=builder,
|
|
target=target,
|
|
tir_mod=_filter_tir(mod),
|
|
tir_pipeline=tir_pipeline,
|
|
ext_libs=ext_libs,
|
|
params=params,
|
|
system_lib=system_lib,
|
|
)
|
|
|
|
|
|
def _filter_tir(mod: tvm.IRModule) -> tvm.IRModule | None:
|
|
tir_mod = {gvar: func for gvar, func in mod.functions.items() if isinstance(func, PrimFunc)}
|
|
|
|
if tir_mod:
|
|
return IRModule(tir_mod, attrs=mod.attrs)
|
|
else:
|
|
return None
|