apache--tvm/python/tvm/tirx/script/builder/external_kernel.py

# 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 <nvshmem.h>" in source_code or "#include <nvshmemx.h>" 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)