# 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. # ruff: noqa: E722 """External kernel integration fro TIR""" import json import logging import os import tempfile from pathlib import Path from typing import Any import tvm_ffi from tvm import __version__ as tvm_version from tvm import tirx from tvm.ir import Expr, PointerType, is_prim_expr from tvm.runtime import Module, const from tvm.support import nvcc class BaseKernel: # pylint: disable=too-few-public-methods """Base class for external kernels.""" def compile_to_device_module( self, launch_args, *args, **kwargs ) -> tuple[str, Module, list[Any]]: """Compile the kernel to a device module.""" raise NotImplementedError() def _format_tvm_module_metadata(self, kernel_name, arg_types, launch_param_tags): """Format the TVM module metadata.""" tvm_metadata = """{{ "tvm_version": "{version}", "func_info": {{ "{kernel_name}": {{ "name": "", "arg_types": {arg_types}, "launch_param_tags": {launch_param_tags} }} }} }}""".format_map( { "version": tvm_version, "kernel_name": kernel_name, "arg_types": json.dumps(arg_types), "launch_param_tags": json.dumps(launch_param_tags), } ) return tvm_metadata def _create_cuda_module( self, binary_data, kernel_arg_types, launch_param_tags, kernel_name, fmt="ptx" ): """ Create a CUDA module from compiled binary (PTX or cubin) and metadata. Parameters ---------- binary_data : str or bytes The compiled binary data (PTX as str, cubin as bytes). kernel_arg_types : List[str] The types of the kernel arguments. launch_param_tags : List[str] The tags of the launch parameters. kernel_name : str The name of the kernel. fmt : str The format of the binary data: "ptx" or "cubin". Returns ------- kernel_module : Module The CUDA module. """ tvm_metadata = self._format_tvm_module_metadata( kernel_name, kernel_arg_types, launch_param_tags ) # Build the FunctionInfo map in-memory from the JSON metadata, then # construct the CUDA module via the FFI registry without going to # disk. Avoids the load_from_file dispatch path entirely. if isinstance(binary_data, str): binary_bytes = binary_data.encode("utf-8") else: binary_bytes = bytes(binary_data) load_meta = tvm_ffi.get_global_func("runtime.LoadMetaDataFromJSON") fmap = load_meta(tvm_metadata) create_cuda = tvm_ffi.get_global_func("ffi.Module.create.cuda") kernel_module = create_cuda(binary_bytes, fmt, fmap, {}) return kernel_module class SourceKernel(BaseKernel): # pylint: disable=too-few-public-methods """A kernel from source code.""" def __init__(self, source_code: str): self.source_code = source_code def compile_to_device_module( # pylint: disable=arguments-differ self, grid: list[list[int | tirx.Expr]], *args: list[Any], **kwargs: dict[str, Any], ) -> tuple[str, Module, list[Any]]: """Compile the kernel to a device module.""" from tvm.relax.frontend.nn import ( # pylint: disable=import-outside-toplevel SourceModule, ) kernel_name = kwargs["kernel_name"] assert len(grid) == 2, ( "grid should be two list of integers, representing the dimension of " "['blockIdx.x', 'blockIdx.y', 'blockIdx.z'] and " "['threadIdx.x', 'threadIdx.y', 'threadIdx.z']" ) assert isinstance(grid[0], list | tuple) and isinstance(grid[1], list | tuple) launch_param_tags = ["blockIdx.x", "blockIdx.y", "blockIdx.z"][: len(grid[0])] + [ "threadIdx.x", "threadIdx.y", "threadIdx.z", ][: len(grid[1])] runtime_args = [arg if isinstance(arg, Expr) else const(arg) for arg in args] kernel_arg_types = [] for arg in runtime_args: if isinstance(arg.ty, PointerType): kernel_arg_types.append("handle") else: assert is_prim_expr(arg) kernel_arg_types.append(str(arg.ty.dtype)) runtime_args = runtime_args + list(grid[0]) + list(grid[1]) # Reuse compilation path from SourceModule compile_options = SourceModule.get_compile_options("cu") source_code = self.source_code try: source_path = Path(source_code) if source_path.is_file(): with open(source_path) as f: source_code = f.read() except: # pylint: disable=bare-except pass with tempfile.TemporaryDirectory() as temp_dir: # Check if NVSHMEM is used - requires cubin output for device library linking use_nvshmem = ( "#include " in source_code or "#include " in source_code ) target_format = "cubin" if use_nvshmem else "ptx" output_path = f"{temp_dir}/{kernel_name}.{target_format}" compiler = os.environ.get("TVM_CUDA_COMPILE_MODE", "nvrtc") nvcc.compile_cuda( source_code, target_format=target_format, options=compile_options, path_target=output_path, compiler=compiler, ) if target_format == "ptx": with open(output_path) as f: binary_data = f.read() else: with open(output_path, "rb") as f: binary_data = f.read() kernel_module = self._create_cuda_module( binary_data, kernel_arg_types, launch_param_tags, kernel_name, fmt=target_format ) return kernel_name, kernel_module, runtime_args def call_kernel( kernel, launch_args: list[int | tirx.Expr | list[int | tirx.Expr]], *args: list[Any], **kwargs: dict[str, Any], ): """ Call an external kernel. Parameters ---------- kernel : Any The external kernel to call. launch_args : List[Union[int, tirx.Expr, List[Union[int, tirx.Expr]]]] The launch arguments. A list of integers for grid size, block size, and shared memory size. The actual requirements depend on the kernel. args : List[tirx.Expr] The arguments to pass to the kernel. kwargs : Dict[str, Any] Additional keyword arguments to pass to the kernel or compilation. """ from tvm.script.ir_builder.ir import ( # pylint: disable=import-outside-toplevel module_get_attr, module_set_attr, ) from .ir import call_packed # pylint: disable=import-outside-toplevel kernel_type = f"{type(kernel).__module__}.{type(kernel).__qualname__}" if kernel_type == "triton.runtime.jit.JITFunction": from .triton import TritonKernel # pylint: disable=import-outside-toplevel kernel = TritonKernel(kernel) elif kernel_type == "builtins.str": kernel = SourceKernel(kernel) else: raise ValueError(f"Unsupported kernel type {kernel_type}") kernel_name, kernel_module, runtime_args = kernel.compile_to_device_module( launch_args, *args, **kwargs ) # Attach the kernel module to the current IRModule external_mods: list[Module] = module_get_attr("external_mods") or [] kernel_exists = any([mod.implements_function(kernel_name) for mod in external_mods]) if kernel_exists: logging.debug("Kernel %s already exists in the IRModule", kernel_name) else: external_mods.append(kernel_module) module_set_attr("external_mods", external_mods, True) return call_packed(kernel_name, *runtime_args)