# 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)