chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,245 @@
|
||||
# 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)
|
||||
Reference in New Issue
Block a user