wehub-resource-sync
244 lines
9.6 KiB
Python
244 lines
9.6 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
|
|
"""The build utils in python."""
|
|
|
|
import tvm
|
|
from tvm import ir
|
|
from tvm.ir.module import IRModule
|
|
from tvm.target import Target
|
|
from tvm.tirx import PrimFunc
|
|
|
|
|
|
def split_host_device_mods(mod: IRModule) -> tuple[IRModule, dict[Target, IRModule]]:
|
|
"""Split an IRModule into host and device modules.
|
|
|
|
This function takes an IRModule containing functions with different target attributes
|
|
and separates them into host (CPU) and device (GPU/accelerator) modules. Functions
|
|
are categorized based on their target attribute in func_attr.
|
|
|
|
Parameters
|
|
----------
|
|
mod : tvm.IRModule
|
|
The input module to split.
|
|
The module should contain functions with target attributes in their func_attr.
|
|
Functions with "cpu" in their target string are considered host functions,
|
|
while others are considered device functions.
|
|
|
|
Returns
|
|
-------
|
|
host_mod : tvm.IRModule
|
|
The module containing host functions (CPU-targeted functions)
|
|
device_mod_dict : Dict[Target, tvm.IRModule]
|
|
A dict mapping targets to device modules. Each device module contains
|
|
functions targeting the same device (e.g., CUDA GPU, OpenCL, etc.)
|
|
|
|
Examples
|
|
--------
|
|
Given an IRModule with the following functions:
|
|
|
|
.. code-block:: python
|
|
|
|
@I.ir_module
|
|
class Module:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(a: T.int32, b: T.int32) -> T.int32:
|
|
T.func_attr({"target": T.target({"arch": "sm_90", "keys": ["cuda", "gpu"],
|
|
"kind": "cuda", "max_num_threads": 1024}))
|
|
return a + b
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add_host(a: T.int32, b: T.int32) -> T.int32:
|
|
T.func_attr({"target": T.target({"keys": ["cpu"], "kind": "c"}))
|
|
return a + b
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main_kernel(A: T.handle, B: T.handle, C: T.handle, length: T.int32):
|
|
T.func_attr({"target": T.target({"arch": "sm_90", "keys": ["cuda", "gpu"],
|
|
"kind": "cuda"}),
|
|
"calling_conv": 2, # kDeviceKernelLaunch for device kernels
|
|
"tirx.is_global_func": True})
|
|
# ... kernel implementation
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def main(self_handle: T.handle, args: T.handle, num_args: T.int32, result: T.handle):
|
|
T.func_attr({"target": T.target({"keys": ["cpu"], "kind": "c"}),
|
|
"calling_conv": 1, # kCPackedFunc for entry functions
|
|
"tirx.is_entry_func": True})
|
|
# ... main function implementation
|
|
|
|
The function will return:
|
|
- host_mod: Contains `add_host` and `main` functions (CPU targets)
|
|
- device_mod_dict: Contains a CUDA module with `add` and `main_kernel` functions
|
|
|
|
Notes
|
|
-----
|
|
- Functions are categorized based on string matching of their target attribute
|
|
- Functions with "cpu" in the target string are considered host functions
|
|
- Device functions are grouped by their target to create separate modules
|
|
- The function uses string-based target matching due to target hash limitations
|
|
- All functions must have a `calling_conv` attribute in their func_attr:
|
|
- Private helper functions (private=True): use `calling_conv: 0` (kDefault, by default)
|
|
- Public entry functions: use `calling_conv: 1` (kCPackedFunc)
|
|
- Device kernel functions: use `calling_conv: 2` (kDeviceKernelLaunch)
|
|
"""
|
|
|
|
def is_host_func(f):
|
|
target = f.attrs.get("target", tvm.target.Target("llvm"))
|
|
return target.kind.name in ["llvm", "c"]
|
|
|
|
host_mod = tvm.tirx.transform.Filter(is_host_func)(mod)
|
|
device_mod = tvm.tirx.transform.Filter(lambda f: not is_host_func(f))(mod)
|
|
# TODO(syfeng): Here we use str as key since target hash is not correct
|
|
target_str2target = {}
|
|
device_func_dict = {}
|
|
device_mod_dict: dict[Target, IRModule] = {}
|
|
for gv, func in device_mod.functions.items():
|
|
target = func.attrs.get("target", None)
|
|
target_str = str(target) if target is not None else ""
|
|
target_str2target[target_str] = target # This might be overridden by the last one
|
|
device_func_dict.setdefault(target_str, dict()).update({gv: func})
|
|
for target_str in target_str2target.keys():
|
|
target = target_str2target[target_str]
|
|
device_mod_dict[target] = tvm.IRModule(device_func_dict[target_str], attrs=device_mod.attrs)
|
|
return host_mod, device_mod_dict
|
|
|
|
|
|
def codegen_build(mod: IRModule, target: Target) -> tvm.runtime.Module:
|
|
"""Build a runtime module from an IRModule and a Target."""
|
|
if tvm.ir.transform.PassContext.current().config.get("tirx.disable_assert", False):
|
|
mod = tvm.tirx.transform.SkipAssert()(mod)
|
|
build_f_name = "target.build." + target.kind.name
|
|
bf = tvm.get_global_func(build_f_name)
|
|
if bf is None:
|
|
raise ValueError(f"{build_f_name} is not enabled")
|
|
return bf(mod, target)
|
|
|
|
|
|
def tir_to_runtime(
|
|
host_mod: IRModule, device_mod_dict: dict[Target, IRModule], target_host: Target
|
|
):
|
|
"""Convert a collection of TIR IRModules (keyed by Target) into a single runtime Module."""
|
|
|
|
# Get the first module to get the attributes
|
|
# necessary for tests/python/codegen/test_target_codegen_blob.py::test_cuda_multi_lib
|
|
mhost_all = ir.IRModule({}, attrs=host_mod.attrs)
|
|
|
|
mhost_all.update(host_mod)
|
|
device_modules = []
|
|
for target, device_mod in device_mod_dict.items():
|
|
if len(device_mod.functions) != 0:
|
|
device_modules.append(codegen_build(device_mod, target))
|
|
|
|
mhost = codegen_build(mhost_all, target_host)
|
|
for dev_mod in device_modules:
|
|
if dev_mod is not None:
|
|
mhost.import_module(dev_mod)
|
|
return mhost
|
|
|
|
|
|
def build(
|
|
mod: PrimFunc | IRModule,
|
|
target: str | Target | None = None,
|
|
pipeline: None | str | tvm.transform.Pass = "default",
|
|
):
|
|
"""Build a function with a signature, generating code for devices
|
|
coupled with target information.
|
|
|
|
Parameters
|
|
----------
|
|
mod : Union[PrimFunc, IRModule]
|
|
The input to be built.
|
|
target : Optional[Union[str, Target]]
|
|
The target for compilation.
|
|
pipeline : Union[None, str, tvm.transform.Pass]
|
|
The pipeline to use for compilation.
|
|
|
|
Returns
|
|
-------
|
|
tvm.runtime.Module
|
|
A module combining both host and device code.
|
|
"""
|
|
# Convert PrimFunc to IRModule
|
|
if isinstance(mod, PrimFunc):
|
|
mod = tvm.IRModule.from_expr(mod)
|
|
else:
|
|
assert isinstance(mod, tvm.IRModule)
|
|
|
|
# Step 0: Determine the target in environment
|
|
# It's used to bind the PrimFunc without target attr to serve as a default target
|
|
target_to_bind = Target.current() if target is None else target
|
|
if target_to_bind is None:
|
|
target_to_bind = "llvm"
|
|
assert target_to_bind is not None
|
|
target_to_bind = Target(target_to_bind)
|
|
|
|
# Step 1: Determine the target to search for tirx pipeline
|
|
target = Target.current() if target is None else target
|
|
if target is None:
|
|
for func in mod.functions.values():
|
|
f_target = func.attrs.get("target", None)
|
|
if f_target is not None:
|
|
target = f_target
|
|
break
|
|
if target is not None:
|
|
target = Target(target)
|
|
|
|
# Step 2: Determine the host target
|
|
target_host = "llvm" if tvm.runtime.enabled("llvm") else "c"
|
|
if target is not None:
|
|
if target.host is not None:
|
|
target_host = target.host
|
|
elif (
|
|
tvm.device(target.kind.name, 0).dlpack_device_type() == tvm.cpu(0).dlpack_device_type()
|
|
):
|
|
target_host = target
|
|
target_host = Target(target_host)
|
|
target_to_bind = target_to_bind.with_host(target_host)
|
|
|
|
# Step 3: Bind the target to the input module
|
|
mod = tvm.tirx.transform.BindTarget(target_to_bind)(mod)
|
|
|
|
# Step 4: Apply the tirx pipeline
|
|
if pipeline is not None:
|
|
# custom pipeline
|
|
assert isinstance(pipeline, str)
|
|
pipeline, finalize_host_passes, finalize_device_passes = tvm.tirx.get_tir_pipeline(pipeline)
|
|
else:
|
|
# default pipeline depends on the target
|
|
pipeline, finalize_host_passes, finalize_device_passes = tvm.tirx.get_default_tir_pipeline(
|
|
target
|
|
)
|
|
mod = pipeline(mod)
|
|
|
|
# Step 5: Get host and device modules
|
|
host_mod, device_mod_dict = split_host_device_mods(mod)
|
|
|
|
# Step 6: Apply finalization passes
|
|
host_mod = finalize_host_passes()(host_mod)
|
|
device_mod_dict = {
|
|
target: finalize_device_passes()(device_mod)
|
|
for target, device_mod in device_mod_dict.items()
|
|
}
|
|
|
|
# Convert TIR IRModules to runtime Module by calling target.build
|
|
return tir_to_runtime(host_mod, device_mod_dict, target_host)
|
|
|
|
|
|
tvm.register_global_func("tirx.build", build)
|