chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,26 @@
|
||||
# isort: skip_file
|
||||
# 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.
|
||||
"""Package tvm.script.ir_builder.tirx"""
|
||||
|
||||
from .ir import * # pylint: disable=wildcard-import,redefined-builtin
|
||||
from .ir import boolean as bool # pylint: disable=redefined-builtin
|
||||
from .ir import buffer as Buffer
|
||||
from .utils import buffer_proxy, frame_scope, seq_scope
|
||||
from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool, TMEMStages
|
||||
from . import tirx as tile
|
||||
from .tirx import cluster, cta, thread, warp, warpgroup, wg
|
||||
@@ -0,0 +1,21 @@
|
||||
# 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.
|
||||
"""FFI APIs"""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
tvm_ffi.init_ffi_api("script.ir_builder.tirx", __name__) # pylint: disable=protected-access
|
||||
@@ -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)
|
||||
@@ -0,0 +1,110 @@
|
||||
# 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.
|
||||
"""IRBuilder for TIR"""
|
||||
|
||||
from tvm_ffi import register_object as _register_object
|
||||
|
||||
from tvm.script.ir_builder.base import IRBuilderFrame
|
||||
from tvm.tirx import Buffer, Var
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.TIRFrame")
|
||||
class TIRFrame(IRBuilderFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.PrimFuncFrame")
|
||||
class PrimFuncFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.SSBlockFrame")
|
||||
class SBlockFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.SBlockInitFrame")
|
||||
class BlockInitFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.ForFrame")
|
||||
class ForFrame(TIRFrame):
|
||||
def __enter__(self) -> Var | list[Var]: # type: ignore[override]
|
||||
super().__enter__()
|
||||
return self.vars if len(self.vars) > 1 else self.vars[0]
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.AssertFrame")
|
||||
class AssertFrame(TIRFrame): ...
|
||||
|
||||
|
||||
class LetFrame(TIRFrame):
|
||||
def __enter__(self) -> Var:
|
||||
super().__enter__()
|
||||
return self.var
|
||||
|
||||
|
||||
class AllocateFrame(TIRFrame):
|
||||
def __enter__(self) -> Buffer:
|
||||
super().__enter__()
|
||||
return self.buffer_var
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.AttrFrame")
|
||||
class AttrFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.WhileFrame")
|
||||
class WhileFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.IfFrame")
|
||||
class IfFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.ThenFrame")
|
||||
class ThenFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.ElseFrame")
|
||||
class ElseFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.DeclBufferFrame")
|
||||
class DeclBufferFrame(TIRFrame):
|
||||
def __enter__(self) -> Buffer:
|
||||
super().__enter__()
|
||||
return self.buffer
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.LaunchThreadFrame")
|
||||
class LaunchThreadFrame(TIRFrame):
|
||||
def __enter__(self) -> Var:
|
||||
super().__enter__()
|
||||
return self.iter_var.var
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.ComposeOpFrame")
|
||||
class ComposeOpFrame(TIRFrame): ...
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.AllocBufferFrame")
|
||||
class AllocBufferFrame(TIRFrame):
|
||||
def __enter__(self) -> Buffer:
|
||||
super().__enter__()
|
||||
return self.buffer
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.tirx.HintFrame")
|
||||
class HintFrame(TIRFrame): ...
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,19 @@
|
||||
# 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.
|
||||
"""Re-export from canonical location."""
|
||||
|
||||
from tvm.tirx.lang.alloc_pool import TMEMPool, TMEMStages # noqa: F401
|
||||
@@ -0,0 +1,137 @@
|
||||
# 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: RUF005
|
||||
"""Triton kernel integration with TIR"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import triton
|
||||
from packaging import version
|
||||
from triton.runtime.jit import type_canonicalisation_dict
|
||||
|
||||
from tvm import tirx
|
||||
from tvm.ir import PointerType, PrimType, is_prim_expr
|
||||
from tvm.runtime import Module
|
||||
from tvm.topi.utils import get_const_int
|
||||
|
||||
from .external_kernel import BaseKernel
|
||||
|
||||
if version.parse(triton.__version__) < version.parse("3.3.0"):
|
||||
raise ImportError(
|
||||
f"TIR Triton integration requires Triton >= 3.3.0, but found Triton {triton.__version__}"
|
||||
)
|
||||
|
||||
|
||||
class TritonKernel(BaseKernel):
|
||||
"""A kernel from Triton JIT function.
|
||||
|
||||
This class bridges the Triton kernel with TVM runtime. The compilation includes the following
|
||||
steps:
|
||||
- Deduce the kernel signature and generate the Triton kernel
|
||||
- Embed the compiled kernel into the current IRModule as an external module
|
||||
- Generate a call to the Triton kernel following its calling convention via call_packed.
|
||||
"""
|
||||
|
||||
def __init__(self, func):
|
||||
self.func = func
|
||||
|
||||
def compile_to_device_module(
|
||||
self,
|
||||
launch_args: list[int | tirx.Expr],
|
||||
*args: list[Any],
|
||||
**kwargs: dict[str, Any],
|
||||
) -> tuple[str, Module, list[Any]]:
|
||||
"""Compile the kernel to a device module.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
launch_args : List[int]
|
||||
The grid size of the kernel. A list of one to three expressions, representing the number
|
||||
of
|
||||
"blockIdx.x", "blockIdx.y", and "blockIdx.z" respectively.
|
||||
|
||||
args : List[Any]
|
||||
Arguments to the kernel function.
|
||||
|
||||
kwargs : Dict[str, Any]
|
||||
Additional options for the kernel compilation.
|
||||
"""
|
||||
triton_kernel, kernel_args = self._generate_triton_kernel(self.func, *args, **kwargs)
|
||||
kernel_metadata = triton_kernel.metadata
|
||||
ptx = triton_kernel.asm["ptx"]
|
||||
assert kernel_metadata.num_ctas == 1, "Cluster is not supported"
|
||||
num_warps = kernel_metadata.num_warps
|
||||
grid = launch_args
|
||||
launch_param_tags = ["threadIdx.x"] + ["blockIdx.x", "blockIdx.y", "blockIdx.z"][
|
||||
: len(grid)
|
||||
]
|
||||
launch_args = [num_warps * 32] + list(grid)
|
||||
kernel_arg_types = []
|
||||
for arg in kernel_args:
|
||||
if isinstance(arg, int):
|
||||
kernel_arg_types.append("int64")
|
||||
elif isinstance(arg.ty, PointerType):
|
||||
kernel_arg_types.append("handle")
|
||||
else:
|
||||
assert is_prim_expr(arg)
|
||||
kernel_arg_types.append(str(arg.ty.dtype))
|
||||
if triton_kernel.metadata.shared > 0:
|
||||
# Add shared memory size to the launch arguments
|
||||
launch_param_tags.append("tirx.use_dyn_shared_memory")
|
||||
launch_args.append(triton_kernel.metadata.shared)
|
||||
|
||||
kernel_module = self._create_cuda_module(
|
||||
ptx, kernel_arg_types, launch_param_tags, triton_kernel.name
|
||||
)
|
||||
|
||||
return triton_kernel.name, kernel_module, kernel_args + launch_args
|
||||
|
||||
def _generate_triton_kernel(
|
||||
self, func, *args, **kwargs
|
||||
) -> tuple["triton.compiler.CompiledKernel", list[tirx.Expr]]:
|
||||
"""Deduce the kernel signature and generate the Triton kernel"""
|
||||
|
||||
kernel_params = func.params
|
||||
assert len(kernel_params) == len(args), (
|
||||
f"Number of arguments does not match, expected {len(kernel_params)}, got {len(args)}"
|
||||
)
|
||||
|
||||
signature = {}
|
||||
constants = {}
|
||||
kernel_args = [] # Arguments to invoke the kernel
|
||||
for i, arg in enumerate(args):
|
||||
if kernel_params[i].is_constexpr:
|
||||
constants[kernel_params[i].name] = get_const_int(arg)
|
||||
signature[kernel_params[i].name] = "constexpr"
|
||||
kernel_args.append(arg)
|
||||
continue
|
||||
if isinstance(arg.ty, PointerType):
|
||||
assert isinstance(arg, tirx.Var)
|
||||
assert isinstance(arg.ty.element_type, PrimType)
|
||||
elem_type = arg.ty.element_type.dtype
|
||||
pointer_type = "*" + type_canonicalisation_dict[elem_type]
|
||||
signature[kernel_params[i].name] = pointer_type
|
||||
else:
|
||||
assert is_prim_expr(arg)
|
||||
signature[kernel_params[i].name] = type_canonicalisation_dict[arg.ty.dtype]
|
||||
kernel_args.append(arg)
|
||||
|
||||
# TODO: Support default argument in the kernel
|
||||
# TODO: Add specialization for aligned buffer pointers
|
||||
source = triton.compiler.ASTSource(fn=func, signature=signature, constexprs=constants)
|
||||
compiled = triton.compiler.compile(source, options=kwargs)
|
||||
return compiled, kernel_args
|
||||
@@ -0,0 +1,226 @@
|
||||
# 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.
|
||||
"""Utility helpers for TIR IRBuilder."""
|
||||
|
||||
import contextlib
|
||||
|
||||
from tvm import tirx
|
||||
from tvm.tirx import Buffer
|
||||
|
||||
from . import frame
|
||||
from . import ir as T
|
||||
|
||||
|
||||
class _FrameScope:
|
||||
"""Context manager to enter multiple IRBuilder frames without deep nesting.
|
||||
|
||||
This class allows entering multiple frames in a single `with` statement,
|
||||
avoiding the pyramid of nested context managers.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
frames : List[IRBuilderFrame]
|
||||
The list of frames to enter.
|
||||
"""
|
||||
|
||||
def __init__(self, frames):
|
||||
self.frames = frames if isinstance(frames, list | tuple) else [frames]
|
||||
self._stack = None
|
||||
|
||||
def __enter__(self):
|
||||
self._stack = contextlib.ExitStack()
|
||||
self._stack.__enter__()
|
||||
results = [self._stack.enter_context(f) for f in self.frames]
|
||||
return tuple(results) if len(results) > 1 else results[0]
|
||||
|
||||
def __exit__(self, *args):
|
||||
return self._stack.__exit__(*args)
|
||||
|
||||
|
||||
def frame_scope(frames: list[frame.TIRFrame]) -> _FrameScope:
|
||||
"""Enter multiple IRBuilder frames without deep nesting.
|
||||
|
||||
This function provides a way to enter multiple frames in a single `with`
|
||||
statement, which is particularly useful when migrating from cases where
|
||||
allocations don't require nested scopes.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
frames : List[frame.TIRFrame]
|
||||
The list of frames to enter. Each frame's `__enter__` return value
|
||||
will be collected and returned as a tuple.
|
||||
|
||||
Returns
|
||||
-------
|
||||
_FrameScope
|
||||
A context manager that enters all frames and returns their values.
|
||||
"""
|
||||
return _FrameScope(frames)
|
||||
|
||||
|
||||
def seq_scope():
|
||||
"""Create a scope that allows multiple consecutive statements.
|
||||
|
||||
The IRBuilder requires a parent frame when having multiple consecutive
|
||||
top-level statements (e.g., multiple loops). This function creates a
|
||||
dummy attr frame that serves as a parent scope.
|
||||
|
||||
Returns
|
||||
-------
|
||||
frame.AttrFrame
|
||||
A dummy attribute frame that wraps multiple statements.
|
||||
|
||||
Examples
|
||||
--------
|
||||
Without seq_scope, multiple consecutive loops fail:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
with IRBuilder() as ib:
|
||||
with T.serial(0, 10) as i:
|
||||
T.evaluate(i)
|
||||
with T.serial(0, 5) as j: # This would fail!
|
||||
T.evaluate(j)
|
||||
|
||||
With seq_scope, multiple consecutive statements work:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
with IRBuilder() as ib:
|
||||
with seq_scope():
|
||||
with T.serial(0, 10) as i:
|
||||
T.evaluate(i)
|
||||
with T.serial(0, 5) as j:
|
||||
T.evaluate(j)
|
||||
result = ib.get()
|
||||
"""
|
||||
return T.attr(tirx.const(0, "int32"), "pragma_scope", tirx.StringImm("seq"))
|
||||
|
||||
|
||||
def _unravel_index(index, shape):
|
||||
"""Convert a flat index to multi-dimensional indices.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
index : Expr
|
||||
The flat index.
|
||||
shape : Tuple
|
||||
The shape of the buffer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[Expr]
|
||||
The multi-dimensional indices.
|
||||
"""
|
||||
indices = []
|
||||
for i, dim in enumerate(reversed(shape)):
|
||||
if i == len(shape) - 1:
|
||||
# Outermost dimension: use remaining quotient directly (no modulo)
|
||||
indices.append(index)
|
||||
else:
|
||||
indices.append(index % dim)
|
||||
index = index // dim
|
||||
return list(reversed(indices))
|
||||
|
||||
|
||||
class _BufferProxy:
|
||||
"""Proxy for flat indexing on multi-dimensional buffers.
|
||||
|
||||
This class wraps a TIR Buffer and provides flat indexing that gets
|
||||
automatically converted to multi-dimensional indices. It also supports
|
||||
assignment syntax via __setitem__.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
buf : Buffer
|
||||
The TIR buffer to wrap.
|
||||
|
||||
Examples
|
||||
--------
|
||||
.. code-block:: python
|
||||
|
||||
buf = tvm.tirx.decl_buffer([2, 3], "float32")
|
||||
ptr = buffer_proxy(buf)
|
||||
|
||||
# Read with flat index (converted to [0, 1])
|
||||
val = ptr[1]
|
||||
|
||||
# Write with flat index
|
||||
ptr[1] = 42.0
|
||||
|
||||
# Multi-dimensional access still works
|
||||
val = ptr[0, 2]
|
||||
"""
|
||||
|
||||
def __init__(self, buf):
|
||||
self._buffer = buf
|
||||
self.dtype = buf.dtype
|
||||
self.shape = buf.shape
|
||||
self.name = buf.name
|
||||
self.data = buf.data
|
||||
|
||||
def _normalize_index(self, index):
|
||||
"""Convert flat index to multi-dimensional indices if needed."""
|
||||
try:
|
||||
index = [*index]
|
||||
except TypeError:
|
||||
index = [index]
|
||||
if len(index) == 1 and len(self._buffer.shape) != 1:
|
||||
index = _unravel_index(index[0], self._buffer.shape)
|
||||
return index
|
||||
|
||||
def __getitem__(self, index):
|
||||
index = self._normalize_index(index)
|
||||
return tirx.BufferLoad(self._buffer, index)
|
||||
|
||||
def __setitem__(self, index, value):
|
||||
index = self._normalize_index(index)
|
||||
T.buffer_store(self._buffer, value, index)
|
||||
|
||||
|
||||
def buffer_proxy(buf: Buffer) -> _BufferProxy:
|
||||
"""Create a buffer proxy for flat indexing on multi-dimensional buffers.
|
||||
|
||||
This provides flat indexing that gets converted to multi-dimensional indices.
|
||||
It also supports assignment syntax via __setitem__.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
buf : Buffer
|
||||
The TIR buffer to wrap.
|
||||
|
||||
Returns
|
||||
-------
|
||||
_BufferProxy
|
||||
A proxy object that supports flat indexing and assignment.
|
||||
|
||||
Examples
|
||||
--------
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.tirx.script.builder.utils import buffer_proxy
|
||||
|
||||
buf = tvm.tirx.decl_buffer([2, 3], "float32")
|
||||
ptr = buffer_proxy(buf)
|
||||
|
||||
# Flat indexing (index 1 -> indices [0, 1])
|
||||
val = ptr[1]
|
||||
|
||||
# Assignment syntax
|
||||
ptr[1] = 42.0
|
||||
"""
|
||||
return _BufferProxy(buf)
|
||||
Reference in New Issue
Block a user