wehub-resource-sync
3631 lines
104 KiB
Python
3631 lines
104 KiB
Python
# 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"""
|
|
|
|
import contextlib
|
|
import functools
|
|
import inspect
|
|
import threading
|
|
from collections.abc import Callable
|
|
from functools import partial
|
|
from numbers import Integral
|
|
from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, Union
|
|
|
|
# isort: off
|
|
from typing import Literal
|
|
|
|
# isort: on
|
|
|
|
from tvm_ffi.core import String
|
|
|
|
from tvm import DataType, ir
|
|
from tvm import tirx as tir
|
|
from tvm.ir import Call, Type, is_prim_expr
|
|
from tvm.ir import register_op_attr as _register_op_attr
|
|
from tvm.ir.base import deprecated
|
|
from tvm.runtime import convert
|
|
from tvm.script.ir_builder.base import IRBuilder
|
|
from tvm.target import Target
|
|
|
|
# pylint: disable=unused-import
|
|
from tvm.target.codegen import llvm_lookup_intrinsic_id
|
|
from tvm.tirx import Buffer, BufferRegion, Expr, IndexMap, type_annotation
|
|
from tvm.tirx import _ffi_api as _tirx_ffi_api
|
|
from tvm.tirx import op as _tir_op
|
|
from tvm.tirx.exec_scope import ExecScope, ScopeIdDef, Var
|
|
|
|
# import tirx.expr for direct ir construction to pass structural_equal comparison
|
|
from tvm.tirx.expr import (
|
|
EQ,
|
|
GE,
|
|
GT,
|
|
LE,
|
|
LT,
|
|
NE,
|
|
Add,
|
|
And,
|
|
Broadcast,
|
|
BufferLoad,
|
|
CallEffectKind,
|
|
Cast,
|
|
CommReducer,
|
|
Div,
|
|
FloatImm,
|
|
FloorDiv,
|
|
FloorMod,
|
|
IntImm,
|
|
IterVar,
|
|
Max,
|
|
Min,
|
|
Mod,
|
|
Mul,
|
|
Not,
|
|
Or,
|
|
ProducerLoad,
|
|
Ramp,
|
|
Reduce,
|
|
Select,
|
|
Shuffle,
|
|
StringImm,
|
|
Sub,
|
|
)
|
|
from tvm.tirx.layout import (
|
|
ComposeLayout,
|
|
Iter,
|
|
Layout,
|
|
R,
|
|
S,
|
|
SwizzleLayout,
|
|
TileLayout,
|
|
wg_local_layout,
|
|
)
|
|
|
|
from . import _ffi_api, frame, utils
|
|
from .external_kernel import call_kernel
|
|
|
|
# pylint: enable=unused-import
|
|
|
|
|
|
def cast(value, dtype, span=None):
|
|
"""Cast an expression to the requested data type."""
|
|
return _tirx_ffi_api._cast(dtype, value, span) # type: ignore[attr-defined]
|
|
|
|
|
|
def _current_s_tir() -> bool:
|
|
"""Return True if the innermost enclosing PrimFuncFrame has ``s_tir=True``.
|
|
|
|
Gates the parser's default layout fill: ``s_tir=True`` PrimFuncs leave
|
|
``layout=None`` (so s_tir-style passes that don't touch layout round-trip
|
|
cleanly); ``s_tir=False`` (default, tirx) get ``DefaultLayout(shape)``.
|
|
"""
|
|
from tvm.script.ir_builder.base import IRBuilder # local import to avoid cycle
|
|
|
|
if not IRBuilder.is_in_scope():
|
|
return False
|
|
builder = IRBuilder.current()
|
|
for f in reversed(list(builder.frames)):
|
|
if isinstance(f, frame.PrimFuncFrame):
|
|
return bool(f.s_tir)
|
|
return False
|
|
|
|
|
|
def _get_layout(layout: str | Layout | None, shape: list[Expr], scope: str) -> Layout | None:
|
|
if layout is None:
|
|
return None
|
|
if isinstance(layout, Layout):
|
|
return layout
|
|
assert isinstance(layout, str)
|
|
if layout == "default":
|
|
if _current_s_tir():
|
|
return None
|
|
if scope in ["trn.sbuf", "trn.psum"]:
|
|
return None
|
|
return TileLayout(S[tuple(shape)])
|
|
shape = tuple(shape)
|
|
if scope == "trn.sbuf":
|
|
layout = TileLayout.trainium(layout, shape)
|
|
elif scope == "trn.psum":
|
|
layout = TileLayout.trainium(layout, shape).to_psum()
|
|
return layout
|
|
|
|
|
|
def _normalize_prim_type(dtype) -> ir.PrimType:
|
|
if isinstance(dtype, ir.PrimType):
|
|
return dtype
|
|
dtype_str = getattr(dtype, "_dtype_str", None)
|
|
if dtype_str is not None:
|
|
return ir.PrimType(dtype_str)
|
|
if callable(dtype):
|
|
value = dtype()
|
|
ty = getattr(value, "ty", None)
|
|
if isinstance(ty, ir.PrimType):
|
|
return ty
|
|
return ir.PrimType(dtype)
|
|
|
|
|
|
def _get_elem_offset(elem_offset, byte_offset, dtype: str):
|
|
assert elem_offset is None or byte_offset is None, (
|
|
"elem_offset and byte_offset cannot be set at the same time"
|
|
)
|
|
if elem_offset is not None:
|
|
return elem_offset
|
|
if byte_offset is None:
|
|
return None
|
|
return byte_offset * 8 // (DataType(_normalize_prim_type(dtype).dtype).bits)
|
|
|
|
|
|
_block_name_suffix = threading.local()
|
|
_meta_construction_state = threading.local()
|
|
_THIS_FILE = __file__
|
|
|
|
|
|
class _MetaResourceRecord:
|
|
"""Resource created while constructing a meta_class instance."""
|
|
|
|
def __init__(
|
|
self, value: Any, filename: str, lineno: int, colno: int | None, code: str
|
|
) -> None:
|
|
self.value = value
|
|
self.filename = filename
|
|
self.lineno = lineno
|
|
self.colno = colno
|
|
self.code = code
|
|
|
|
|
|
class _MetaConstructionScope:
|
|
"""Thread-local construction scope for a single meta_class __init__ call."""
|
|
|
|
def __init__(self, instance: Any, cls: type) -> None:
|
|
self.instance = instance
|
|
self.cls = cls
|
|
self.created: list[_MetaResourceRecord] = []
|
|
|
|
def record(self, value: Any, frame_info: inspect.FrameInfo) -> None:
|
|
positions = getattr(frame_info, "positions", None)
|
|
colno = None
|
|
if positions is not None and positions.col_offset is not None:
|
|
colno = positions.col_offset + 1
|
|
code = frame_info.code_context[0].strip() if frame_info.code_context else ""
|
|
self.created.append(
|
|
_MetaResourceRecord(
|
|
value=value,
|
|
filename=frame_info.filename,
|
|
lineno=frame_info.lineno,
|
|
colno=colno,
|
|
code=code,
|
|
)
|
|
)
|
|
|
|
|
|
def _meta_construction_stack() -> list[_MetaConstructionScope]:
|
|
stack = getattr(_meta_construction_state, "stack", None)
|
|
if stack is None:
|
|
stack = []
|
|
_meta_construction_state.stack = stack
|
|
return stack
|
|
|
|
|
|
def _current_meta_construction_scope() -> _MetaConstructionScope | None:
|
|
stack = _meta_construction_stack()
|
|
return stack[-1] if stack else None
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _with_meta_construction_scope(instance: Any, cls: type):
|
|
scope = _MetaConstructionScope(instance, cls)
|
|
stack = _meta_construction_stack()
|
|
stack.append(scope)
|
|
try:
|
|
yield scope
|
|
finally:
|
|
stack.pop()
|
|
|
|
|
|
def _record_meta_resource(value: Any, skip_frames: int = 2) -> None:
|
|
scope = _current_meta_construction_scope()
|
|
if scope is not None:
|
|
stack = inspect.stack(context=1)
|
|
frame_info = None
|
|
for candidate in stack[2:]:
|
|
if candidate.filename != _THIS_FILE:
|
|
frame_info = candidate
|
|
break
|
|
if frame_info is None:
|
|
frame_info = stack[min(skip_frames + 1, len(stack) - 1)]
|
|
scope.record(value, frame_info)
|
|
|
|
|
|
def _get_sblock_name_suffix() -> str:
|
|
"""Get the current block name suffix for macro expansion."""
|
|
return getattr(_block_name_suffix, "value", "")
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def block_name_suffix_context(block_suffix: str):
|
|
"""Context manager to set block name suffix during macro expansion.
|
|
|
|
Parameters
|
|
----------
|
|
block_suffix : str
|
|
The suffix to append to block names (e.g., "_1", "_2").
|
|
|
|
Yields
|
|
------
|
|
None
|
|
"""
|
|
old_suffix = getattr(_block_name_suffix, "value", "")
|
|
_block_name_suffix.value = block_suffix
|
|
try:
|
|
yield
|
|
finally:
|
|
_block_name_suffix.value = old_suffix
|
|
|
|
|
|
def buffer(
|
|
shape: list[Expr] | tuple[Expr] | Expr | Integral,
|
|
dtype: str = "float32",
|
|
data: Var = None,
|
|
strides: list[Expr] | None = None,
|
|
elem_offset: Expr = None,
|
|
byte_offset: Expr = None,
|
|
scope: str = "global",
|
|
align: int = 0,
|
|
offset_factor: int = 0,
|
|
buffer_type: str = "",
|
|
axis_separators: list[int] | None = None,
|
|
layout: str | Layout | None = "default",
|
|
allocated_addr: int | tuple[int, ...] | None = None,
|
|
buffer_name: str = "",
|
|
) -> Buffer:
|
|
"""The buffer declaration function.
|
|
|
|
Parameters
|
|
----------
|
|
shape : Union[List[Expr], Tuple[Expr], Expr, Integral]
|
|
The type of the buffer prior to flattening.
|
|
|
|
dtype : str
|
|
The data type in the content of the buffer.
|
|
|
|
data : Var
|
|
The pointer to the head of the data.
|
|
|
|
strides : List[Expr]
|
|
The strides of each dimension.
|
|
|
|
elem_offset : Expr
|
|
The offset in terms of number of dtype elements (including lanes).
|
|
|
|
scope : str
|
|
The optional storage scope of buffer data pointer.
|
|
|
|
align : int
|
|
The alignment requirement of data pointer in bytes.
|
|
|
|
offset_factor : int
|
|
The factor of elem_offset field.
|
|
|
|
buffer_type : str
|
|
The buffer type.
|
|
|
|
axis_separators : List[int]
|
|
The separators between input axes when generating flattened output axes.
|
|
|
|
buffer_name : str
|
|
The name of the buffer.
|
|
|
|
Returns
|
|
-------
|
|
res : Buffer
|
|
The declared buffer.
|
|
"""
|
|
shape = (shape,) if is_prim_expr(shape) or isinstance(shape, Integral) else shape
|
|
if strides is not None:
|
|
strides = [Var(s, "int32") if isinstance(s, str) else s for s in strides]
|
|
else:
|
|
strides = []
|
|
if allocated_addr is None:
|
|
allocated_addr = []
|
|
if not isinstance(allocated_addr, list | tuple):
|
|
allocated_addr = [allocated_addr]
|
|
return _ffi_api.Buffer( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
shape,
|
|
dtype,
|
|
buffer_name,
|
|
data,
|
|
strides,
|
|
_get_elem_offset(elem_offset, byte_offset, dtype),
|
|
scope,
|
|
align,
|
|
offset_factor,
|
|
buffer_type,
|
|
axis_separators,
|
|
_get_layout(layout, shape, scope),
|
|
allocated_addr,
|
|
)
|
|
|
|
|
|
@deprecated("T.buffer_decl(...)", "T.Buffer(...)")
|
|
def buffer_decl(*args, **kwargs):
|
|
return buffer(*args, **kwargs)
|
|
|
|
|
|
def prim_func(
|
|
is_private: bool = False,
|
|
s_tir: bool = False,
|
|
persistent: bool = False,
|
|
*,
|
|
private: bool | None = None,
|
|
) -> frame.PrimFuncFrame:
|
|
"""The primitive function statement.
|
|
|
|
Parameters
|
|
----------
|
|
is_private : bool
|
|
Whether the PrimFunc is annotated as private.
|
|
s_tir : bool
|
|
Whether this PrimFunc uses s_tir (apache-derived TIR) semantics:
|
|
parser fills layout=None on buffers, ScriptComplete wraps body in a
|
|
root SBlock. Default (False) selects tirx semantics: parser fills
|
|
``DefaultLayout(shape)`` and no root-block wrapping.
|
|
persistent : bool
|
|
Whether this is a persistent kernel.
|
|
private : bool
|
|
Alias for ``is_private`` (used in decorator syntax).
|
|
|
|
Returns
|
|
-------
|
|
res : frame.PrimFuncFrame
|
|
The PrimFuncFrame.
|
|
"""
|
|
if private is not None:
|
|
is_private = private
|
|
return _ffi_api.PrimFunc(is_private, s_tir, persistent) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def arg(name: str, obj: Var | Buffer) -> Var | Buffer:
|
|
"""The PrimFunc arguments adding function.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the argument.
|
|
|
|
var : Union[Var, Buffer]
|
|
The argument of Var or Buffer.
|
|
|
|
Returns
|
|
-------
|
|
res : Union[Var, Buffer]
|
|
The argument.
|
|
"""
|
|
return _ffi_api.Arg(name, obj) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def func_name(name: str) -> None:
|
|
"""The PrimFunc naming statement.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the PrimFunc.
|
|
"""
|
|
_ffi_api.FuncName(name) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def func_attr(attrs: dict[str, Any]) -> None:
|
|
"""The PrimFunc annotation statement.
|
|
|
|
Parameters
|
|
----------
|
|
attrs : Dict[str, Any]
|
|
The annotations of the PrimFunc.
|
|
"""
|
|
_ffi_api.FuncAttrs(attrs) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def func_ret(ret_type: Type | None) -> Type:
|
|
"""The PrimFunc return type statement.
|
|
|
|
Parameters
|
|
----------
|
|
ret_type : Type
|
|
The return type of the PrimFunc.
|
|
|
|
Returns
|
|
-------
|
|
res : Type
|
|
The return type.
|
|
"""
|
|
if ret_type is None:
|
|
ret_type = Type.missing()
|
|
return _ffi_api.FuncRet(ret_type) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def match_buffer(
|
|
param: Var | BufferLoad | BufferRegion,
|
|
shape: list[Expr] | tuple[Expr] | Expr | Integral = None,
|
|
dtype: str = "float32",
|
|
data: Var = None,
|
|
strides: list[Expr] | None = None,
|
|
elem_offset: Expr = None,
|
|
scope: str = "global",
|
|
align: int = -1,
|
|
offset_factor: int = 0,
|
|
buffer_type: str = "default",
|
|
axis_separators: list[int] | None = None,
|
|
layout: str | Layout | None = "default",
|
|
) -> Buffer:
|
|
"""The buffer match function.
|
|
|
|
Note
|
|
----
|
|
This function will perform different behavior, depending on the type of param.
|
|
If the param is a var in function parameter, it will create a buffer from DLTensor.
|
|
Else if the param is a subregion of other buffers, then create a subregion match inside a block.
|
|
|
|
Example
|
|
-------
|
|
Match buffer from function parameter
|
|
.. code-block:: python
|
|
A = T.match_buffer(a, (128, 128), dtype="float32")
|
|
|
|
Match buffer from Buffer subregion
|
|
.. code-block:: python
|
|
A = T.match_buffer(B[0:128, i * 128 : i * 128 + 128], (128, 128), dtype="float32")
|
|
|
|
Parameters
|
|
----------
|
|
param : Union[Var, BufferLoad, BufferRegion]
|
|
The parameter of the PrimFunc to match.
|
|
|
|
shape : Union[List[Expr], Tuple[Expr], Expr, Integral]
|
|
The type of the buffer prior to flattening.
|
|
|
|
dtype : str
|
|
The data type in the content of the buffer.
|
|
|
|
data : Var
|
|
The pointer to the head of the data.
|
|
|
|
strides : List[Expr]
|
|
The strides of each dimension.
|
|
|
|
elem_offset : Expr
|
|
The offset in terms of number of dtype elements (including lanes).
|
|
|
|
scope : str
|
|
The optional storage scope of buffer data pointer.
|
|
|
|
align : int
|
|
The alignment requirement of data pointer in bytes.
|
|
|
|
offset_factor : int
|
|
The factor of elem_offset field.
|
|
|
|
buffer_type : str
|
|
The buffer type.
|
|
|
|
axis_separators : List[int]
|
|
The separators between input axes when generating flattened output axes.
|
|
|
|
layout: Optional[Union[str, Layout]]
|
|
The layout of the buffer.
|
|
|
|
Returns
|
|
-------
|
|
res : Buffer
|
|
The matched buffer.
|
|
"""
|
|
if shape is None:
|
|
if isinstance(param, BufferRegion):
|
|
dtype = param.buffer.dtype
|
|
shape = [region.extent for region in param.region]
|
|
else:
|
|
raise ValueError("Shape must be specified when binding input param")
|
|
shape = (shape,) if is_prim_expr(shape) or isinstance(shape, Integral) else shape
|
|
if strides is not None:
|
|
idx_dtype = shape[0].ty if is_prim_expr(shape[0]) else "int32"
|
|
strides = [Var(s, idx_dtype) if isinstance(s, str) else s for s in strides]
|
|
else:
|
|
strides = []
|
|
return _ffi_api.MatchBuffer( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
param,
|
|
shape,
|
|
dtype,
|
|
data,
|
|
strides,
|
|
elem_offset,
|
|
scope,
|
|
align,
|
|
offset_factor,
|
|
buffer_type,
|
|
axis_separators,
|
|
_get_layout(layout, shape, scope),
|
|
)
|
|
|
|
|
|
def sblock(name: str = "", no_realize: bool = False, exec_scope: str = "") -> frame.SBlockFrame:
|
|
"""The sblock declaration statement.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the sblock.
|
|
|
|
no_realize : bool
|
|
The flag whether to construct SBlockRealize or SBlock.
|
|
|
|
exec_scope : str
|
|
The execution scope of the block.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.SBlockFrame
|
|
The SBlockFrame.
|
|
"""
|
|
if isinstance(name, list):
|
|
# tir+
|
|
return _ffi_api.ScopeSlice(name, no_realize)
|
|
block_suffix = _get_sblock_name_suffix()
|
|
if block_suffix and name:
|
|
name = name + block_suffix
|
|
return _ffi_api.Block(name, no_realize, exec_scope) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def device_entry() -> None:
|
|
"""Mark the device-region entry within the enclosing PrimFunc body.
|
|
|
|
Flat marker (no ``with``). Subsequent statements in the function body
|
|
accumulate into an ``AttrStmt("tirx.device_entry", True, body=...)``;
|
|
the wrapping is closed by the PrimFunc frame at function end.
|
|
|
|
Anything written before this marker is host code (e.g. ``T.match_buffer``);
|
|
anything after is device code.
|
|
|
|
Example::
|
|
|
|
@T.prim_func
|
|
def kernel(...):
|
|
A = T.match_buffer(...)
|
|
T.device_entry() # device region starts here
|
|
bx = T.cta_id([SM_COUNT]) # standalone scope-id def
|
|
...
|
|
"""
|
|
attr_frame = _ffi_api.DeviceEntry() # type: ignore[attr-defined] # pylint: disable=no-member
|
|
attr_frame.__enter__()
|
|
# No return: the frame is registered on the IRBuilder stack; the
|
|
# PrimFunc frame's exit drains it.
|
|
|
|
|
|
def elected():
|
|
"""Stub that rejects the removed ``T.elected()`` sugar.
|
|
|
|
Write the explicit form instead::
|
|
|
|
if T.ptx.elect_sync():
|
|
... # thread is the default scope
|
|
"""
|
|
raise RuntimeError(
|
|
"T.elected() is no longer available. Write explicitly: "
|
|
"`if T.ptx.elect_sync(): ...` (thread is the default scope)"
|
|
)
|
|
|
|
|
|
def scope_id(extents: list[Expr | int] | None, parent: str, cur: str) -> Var | list[Var]:
|
|
ret = _ffi_api.ScopeId(extents, parent, "T.scope_id", cur) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
if len(ret) == 1:
|
|
return ret[0]
|
|
return ret
|
|
|
|
|
|
def cluster_id(extents: list[Expr | int] | None = None) -> Var | list[Var]:
|
|
"""Define a kernel→cluster scope id. Pass ``None`` (the default) to defer the
|
|
extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure."""
|
|
ret = _ffi_api.ClusterId(extents, "kernel") # type: ignore[attr-defined] # pylint: disable=no-member
|
|
if len(ret) == 1:
|
|
return ret[0]
|
|
return ret
|
|
|
|
|
|
def cta_id(extents: list[Expr | int] | None = None, preferred=None) -> Var | list[Var]:
|
|
"""Define a kernel→cta scope id. Pass ``None`` (the default) to defer the
|
|
extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure."""
|
|
ret = _ffi_api.CtaId(extents, "kernel", preferred) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
if len(ret) == 1:
|
|
return ret[0]
|
|
return ret
|
|
|
|
|
|
def cta_id_in_cluster(extents: list[Expr | int] | None = None, preferred=None) -> Var | list[Var]:
|
|
"""Define a cluster→cta scope id. Pass ``None`` (the default) to defer the
|
|
extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure."""
|
|
ret = _ffi_api.CtaId(extents, "cluster", preferred) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
if len(ret) == 1:
|
|
return ret[0]
|
|
return ret
|
|
|
|
|
|
def cta_id_in_pair() -> Var:
|
|
ret = _ffi_api.CtaIdInPair() # type: ignore[attr-defined] # pylint: disable=no-member
|
|
return ret[0]
|
|
|
|
|
|
def warpgroup_id(extents: list[Expr | int] | None = None) -> Var | list[Var]:
|
|
"""Define a cta→warpgroup scope id. Pass ``None`` (the default) to defer
|
|
the extent; it will be inferred at LowerTIRx from sibling closure."""
|
|
ret = _ffi_api.WarpgroupId(extents, "cta") # type: ignore[attr-defined] # pylint: disable=no-member
|
|
if len(ret) == 1:
|
|
return ret[0]
|
|
return ret
|
|
|
|
|
|
def warp_id(extents: list[Expr | int] | None = None) -> Var | list[Var]:
|
|
"""Define a cta→warp scope id. Pass ``None`` (the default) to defer the
|
|
extent; it will be inferred at LowerTIRx from sibling closure."""
|
|
ret = _ffi_api.WarpId(extents, "cta") # type: ignore[attr-defined] # pylint: disable=no-member
|
|
if len(ret) == 1:
|
|
return ret[0]
|
|
return ret
|
|
|
|
|
|
def warp_id_in_wg(extents: list[Expr | int] | None = None) -> Var | list[Var]:
|
|
"""Define a warpgroup→warp scope id. Pass ``None`` (the default) to defer
|
|
the extent; it will be inferred at LowerTIRx from sibling closure."""
|
|
ret = _ffi_api.WarpId(extents, "warpgroup") # type: ignore[attr-defined] # pylint: disable=no-member
|
|
if len(ret) == 1:
|
|
return ret[0]
|
|
return ret
|
|
|
|
|
|
def lane_id(extents: list[Expr | int] | None = None) -> Var | list[Var]:
|
|
"""Define a warp→thread scope id. Pass ``None`` (the default) to defer the
|
|
extent; it will be inferred at LowerTIRx from sibling closure."""
|
|
ret = _ffi_api.ThreadId(extents, "warp") # type: ignore[attr-defined] # pylint: disable=no-member
|
|
if len(ret) == 1:
|
|
return ret[0]
|
|
return ret
|
|
|
|
|
|
def thread_id(extents: list[Expr | int] | None = None) -> Var | list[Var]:
|
|
"""Define a cta→thread scope id. Pass ``None`` (the default) to defer the
|
|
extent; it will be inferred at LowerTIRx from sibling closure."""
|
|
ret = _ffi_api.ThreadId(extents, "cta") # type: ignore[attr-defined] # pylint: disable=no-member
|
|
if len(ret) == 1:
|
|
return ret[0]
|
|
return ret
|
|
|
|
|
|
def thread_id_in_wg(extents: list[Expr | int] | None = None) -> Var | list[Var]:
|
|
"""Define a warpgroup→thread scope id. Pass ``None`` (the default) to defer
|
|
the extent; it will be inferred at LowerTIRx from sibling closure."""
|
|
ret = _ffi_api.ThreadId(extents, "warpgroup") # type: ignore[attr-defined] # pylint: disable=no-member
|
|
if len(ret) == 1:
|
|
return ret[0]
|
|
return ret
|
|
|
|
|
|
def init() -> frame.BlockInitFrame:
|
|
"""The block initialization statement.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.BlockInitFrame
|
|
The BlockInitFrame.
|
|
"""
|
|
return _ffi_api.Init() # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def where(predicate: Expr | int) -> None:
|
|
"""The block predicate statement.
|
|
|
|
Parameters
|
|
----------
|
|
predicate : Union[Expr, Literal[0, 1]]
|
|
The predicate condition.
|
|
"""
|
|
if isinstance(predicate, bool):
|
|
predicate = IntImm("bool", predicate)
|
|
if isinstance(predicate, int):
|
|
if predicate in [0, 1]:
|
|
predicate = IntImm("bool", predicate)
|
|
else:
|
|
raise ValueError(f"Invalid value for predicate: {predicate}")
|
|
_ffi_api.Where(predicate) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def reads(*buffer_slices: list[BufferRegion | BufferLoad]) -> None:
|
|
"""The block buffer region reading statement.
|
|
|
|
Parameters
|
|
----------
|
|
buffer_slices : List[Union[BufferRegion, BufferLoad]]
|
|
The array of buffer regions to read.
|
|
"""
|
|
if len(buffer_slices) == 1:
|
|
if isinstance(buffer_slices[0], tuple):
|
|
buffer_slices = list(buffer_slices[0])
|
|
elif isinstance(buffer_slices[0], list):
|
|
buffer_slices = buffer_slices[0] # type: ignore[assignment]
|
|
else:
|
|
buffer_slices = [buffer_slices[0]]
|
|
else:
|
|
buffer_slices = list(buffer_slices) # type: ignore[assignment]
|
|
_ffi_api.Reads(buffer_slices) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def writes(*buffer_slices: list[BufferRegion | BufferLoad]) -> None:
|
|
"""The block buffer region writing statement.
|
|
|
|
Parameters
|
|
----------
|
|
buffer_slices : List[Union[BufferRegion, BufferLoad]]
|
|
The array of buffer regions to write.
|
|
"""
|
|
if len(buffer_slices) == 1:
|
|
if isinstance(buffer_slices[0], tuple):
|
|
buffer_slices = list(buffer_slices[0])
|
|
elif isinstance(buffer_slices[0], list):
|
|
buffer_slices = buffer_slices[0] # type: ignore[assignment]
|
|
else:
|
|
buffer_slices = [buffer_slices[0]]
|
|
else:
|
|
buffer_slices = list(buffer_slices) # type: ignore[assignment]
|
|
_ffi_api.Writes(buffer_slices) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def sblock_attr(attrs: dict[str, Any]) -> None:
|
|
"""The block annotation statement (for non-tirx SBlock usage).
|
|
|
|
Parameters
|
|
----------
|
|
attrs : Dict[str, Any]
|
|
The annotation of the block.
|
|
"""
|
|
return _ffi_api.BlockAttrs(attrs) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def alloc_buffer(
|
|
shape: list[Expr] | tuple[Expr] | Expr | Integral,
|
|
dtype: str = "float32",
|
|
data: Var | None = None,
|
|
strides: list[Expr] | None = None,
|
|
elem_offset: Expr | None = None,
|
|
byte_offset: Expr | None = None,
|
|
scope: str = "global",
|
|
align: int = -1,
|
|
offset_factor: int = 0,
|
|
buffer_type: str = "default",
|
|
axis_separators: list[int] | None = None,
|
|
layout: str | Layout | None = "default",
|
|
allocated_addr: int | tuple[int, ...] | None = None,
|
|
annotations: dict[str, Any] | None = None,
|
|
) -> Buffer:
|
|
"""Statement-level buffer allocation (creates an AllocBuffer IR node).
|
|
|
|
Emits an AllocBuffer statement and returns the Buffer directly::
|
|
|
|
buf = T.alloc_buffer((128, 128))
|
|
|
|
For SBlock-level buffer allocation (added to SBlock.alloc_buffers),
|
|
use T.sblock_alloc_buffer() instead.
|
|
|
|
Parameters
|
|
----------
|
|
shape : Union[List[Expr], Tuple[Expr], Expr, Integral]
|
|
The shape of the buffer to allocate.
|
|
dtype : str
|
|
The data type of the buffer elements.
|
|
scope : str
|
|
The storage scope of the buffer (e.g., "global", "shared").
|
|
data : Optional[Var]
|
|
Optional explicit data pointer.
|
|
strides : Optional[List[Expr]]
|
|
Optional strides.
|
|
elem_offset : Optional[Expr]
|
|
Optional element offset.
|
|
byte_offset : Optional[Expr]
|
|
Optional byte offset.
|
|
align : int
|
|
Alignment requirement in bytes.
|
|
offset_factor : int
|
|
Offset factor.
|
|
buffer_type : str
|
|
Buffer type.
|
|
axis_separators : Optional[List[int]]
|
|
Optional axis separators.
|
|
layout : Optional[Union[str, Layout]]
|
|
Optional layout.
|
|
allocated_addr : Optional[Union[int, Tuple[int, ...]]]
|
|
Optional pre-allocated address metadata.
|
|
annotations : Optional[Dict[str, Any]]
|
|
Optional annotations for the allocation.
|
|
|
|
Returns
|
|
-------
|
|
res : Buffer
|
|
The allocated buffer.
|
|
"""
|
|
shape = (shape,) if is_prim_expr(shape) or isinstance(shape, Integral) else shape
|
|
buf = buffer(
|
|
shape=shape,
|
|
dtype=dtype,
|
|
data=data,
|
|
strides=strides,
|
|
elem_offset=elem_offset,
|
|
byte_offset=byte_offset,
|
|
scope=scope,
|
|
align=align,
|
|
offset_factor=offset_factor,
|
|
buffer_type=buffer_type,
|
|
axis_separators=axis_separators,
|
|
layout=layout,
|
|
allocated_addr=allocated_addr,
|
|
buffer_name="",
|
|
)
|
|
_record_meta_resource(buf, skip_frames=2)
|
|
|
|
# AllocBuffer.annotations holds typed IR values. The C++ side stores
|
|
# alignment / shape-like ints as ``IntImm(int32, ...)``; if the user
|
|
# (or a parsed-source round-trip) passes a bare Python int, normalize
|
|
# it so structural equality is preserved against the LowerOpaqueBlock
|
|
# output. Booleans must stay as IntImm("bool", ...).
|
|
def _normalize_ann_value(v):
|
|
if isinstance(v, bool):
|
|
return tir.IntImm("bool", int(v))
|
|
if isinstance(v, int):
|
|
return tir.IntImm("int32", v)
|
|
if isinstance(v, float):
|
|
return tir.FloatImm("float32", v)
|
|
return v
|
|
|
|
norm_annotations = {k: _normalize_ann_value(v) for k, v in (annotations or {}).items()}
|
|
_ffi_api.AddToParent(tir.AllocBuffer(buf, norm_annotations)) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
return buf
|
|
|
|
|
|
def wg_reg_tile(elem_per_thread: int, dtype: str = "float32") -> Buffer:
|
|
"""Warpgroup-wide ``(128, elem_per_thread)`` register tile in local scope.
|
|
|
|
Sugar for the recurring pattern::
|
|
|
|
T.alloc_buffer(
|
|
(128, elem_per_thread), dtype,
|
|
layout=wg_local_layout(elem_per_thread),
|
|
scope="local",
|
|
)
|
|
|
|
Used to stage a tcgen05 load: each of the 128 threads in a warpgroup
|
|
owns one row of ``elem_per_thread`` contiguous elements.
|
|
"""
|
|
return alloc_buffer(
|
|
(128, elem_per_thread),
|
|
dtype,
|
|
layout=wg_local_layout(elem_per_thread),
|
|
scope="local",
|
|
)
|
|
|
|
|
|
def sblock_alloc_buffer(
|
|
shape: list[Expr] | tuple[Expr] | Expr | Integral,
|
|
dtype: str = "float32",
|
|
data: Var = None,
|
|
strides: list[Expr] | None = None,
|
|
elem_offset: Expr = None,
|
|
scope: str = "global",
|
|
align: int = -1,
|
|
offset_factor: int = 0,
|
|
buffer_type: str = "default",
|
|
axis_separators: list[int] | None = None,
|
|
layout: str | Layout | None = "default",
|
|
allocated_addr: int | tuple[int, ...] | None = None,
|
|
) -> Buffer:
|
|
"""SBlock-level buffer allocation function.
|
|
|
|
Parameters
|
|
----------
|
|
shape : Union[List[Expr], Tuple[Expr], Expr, Integral]
|
|
The type of the buffer prior to flattening.
|
|
dtype : str
|
|
The data type in the content of the buffer.
|
|
data : Var
|
|
The pointer to the head of the data.
|
|
strides : List[Expr]
|
|
The strides of each dimension.
|
|
elem_offset : Expr
|
|
The offset in terms of number of dtype elements (including lanes).
|
|
scope : str
|
|
The optional storage scope of buffer data pointer.
|
|
align : int
|
|
The alignment requirement of data pointer in bytes.
|
|
offset_factor : int
|
|
The factor of elem_offset field.
|
|
buffer_type : str
|
|
The buffer type.
|
|
axis_separators : List[int]
|
|
The separators between input axes when generating flattened output axes.
|
|
|
|
layout: Optional[Union[str, Layout]]
|
|
The layout of the buffer.
|
|
|
|
allocated_addr: Optional[Union[int, Tuple[int]]]
|
|
The address of the allocated buffer. Might be multi-dimensional.
|
|
There can be pooled storage scopes on some devices. For example,
|
|
the Trainium device has a pooled storage scope for the SRAN buffers. ("trn.sbuf")
|
|
CUDA has a pooled storage scope for the shared memory ("shared.dyn")
|
|
|
|
Returns
|
|
-------
|
|
res : Buffer
|
|
The allocated buffer.
|
|
"""
|
|
shape = (shape,) if is_prim_expr(shape) or isinstance(shape, Integral) else shape
|
|
if strides is not None:
|
|
strides = [Var(s, "int32") if isinstance(s, str) else s for s in strides]
|
|
else:
|
|
strides = []
|
|
if axis_separators is None:
|
|
axis_separators = []
|
|
if allocated_addr is None:
|
|
allocated_addr = []
|
|
if not isinstance(allocated_addr, list | tuple):
|
|
allocated_addr = [allocated_addr]
|
|
alloc_frame = _ffi_api.SBlockAllocBuffer( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
shape,
|
|
dtype,
|
|
data,
|
|
strides,
|
|
elem_offset,
|
|
scope,
|
|
align,
|
|
offset_factor,
|
|
buffer_type,
|
|
axis_separators,
|
|
_get_layout(layout, shape, scope),
|
|
allocated_addr,
|
|
)
|
|
if isinstance(alloc_frame, frame.AllocBufferFrame):
|
|
alloc_frame.add_callback(partial(alloc_frame.__exit__, None, None, None))
|
|
buf = alloc_frame.__enter__()
|
|
else:
|
|
buf = alloc_frame
|
|
_record_meta_resource(buf, skip_frames=2)
|
|
return buf
|
|
|
|
|
|
def _as_range(dom: ir.Range | list[Expr]) -> ir.Range:
|
|
"""The range constructor.
|
|
|
|
Parameters
|
|
----------
|
|
dom : Union[Range, List[Expr]]
|
|
The domain.
|
|
|
|
Returns
|
|
-------
|
|
res : Range
|
|
The Range.
|
|
"""
|
|
if isinstance(dom, ir.Range):
|
|
return dom
|
|
if isinstance(dom, list | tuple):
|
|
from tvm.arith import Analyzer # pylint: disable=import-outside-toplevel
|
|
|
|
extent = Analyzer().simplify(dom[1] - dom[0])
|
|
if isinstance(extent, tir.IntImm):
|
|
return ir.Range.from_min_extent(dom[0], extent)
|
|
return ir.Range(dom[0], dom[1])
|
|
if is_prim_expr(dom):
|
|
return ir.Range(IntImm(dom.ty, 0), dom)
|
|
return ir.Range(0, dom)
|
|
|
|
|
|
class axis: # pylint: disable=invalid-name
|
|
"""The axis class"""
|
|
|
|
@staticmethod
|
|
def spatial(
|
|
dom: ir.Range | list[Expr] | tuple[Expr],
|
|
binding: Expr,
|
|
dtype: str = "int32",
|
|
) -> Var:
|
|
"""The spatial block axis defining function.
|
|
|
|
Parameters
|
|
----------
|
|
dom : Union[Range, List[Expr], Tuple[Expr]]
|
|
The domain of the iteration variable.
|
|
|
|
binding : Expr
|
|
The binding value of the iteration variable.
|
|
|
|
dtype : str
|
|
The data type of the iteration variable.
|
|
|
|
Returns
|
|
-------
|
|
res : Var
|
|
The iteration variable.
|
|
"""
|
|
return _ffi_api.AxisSpatial( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
_as_range(dom), binding, dtype
|
|
)
|
|
|
|
@staticmethod
|
|
def reduce(
|
|
dom: ir.Range | list[Expr] | tuple[Expr],
|
|
binding: Expr,
|
|
dtype: str = "int32",
|
|
) -> Var:
|
|
"""The reduced block axis defining function.
|
|
|
|
Parameters
|
|
----------
|
|
dom : Union[Range, List[Expr], Tuple[Expr]]
|
|
The domain of the iteration variable.
|
|
|
|
binding : Expr
|
|
The binding value of the iteration variable.
|
|
|
|
dtype : str
|
|
The data type of the iteration variable.
|
|
|
|
Returns
|
|
-------
|
|
res : Var
|
|
The iteration variable.
|
|
"""
|
|
return _ffi_api.AxisReduce( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
_as_range(dom), binding, dtype
|
|
)
|
|
|
|
@staticmethod
|
|
def scan(
|
|
dom: ir.Range | list[Expr] | tuple[Expr],
|
|
binding: Expr,
|
|
dtype: str = "int32",
|
|
) -> Var:
|
|
"""The scanning block axis defining function.
|
|
|
|
Parameters
|
|
----------
|
|
dom : Union[Range, List[Expr], Tuple[Expr]]
|
|
The domain of the iteration variable.
|
|
|
|
binding : Expr
|
|
The binding value of the iteration variable.
|
|
|
|
dtype : str
|
|
The data type of the iteration variable.
|
|
|
|
Returns
|
|
-------
|
|
res : Var
|
|
The iteration variable.
|
|
"""
|
|
return _ffi_api.AxisScan( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
_as_range(dom), binding, dtype
|
|
)
|
|
|
|
@staticmethod
|
|
def opaque(
|
|
dom: ir.Range | list[Expr] | tuple[Expr],
|
|
binding: Expr,
|
|
dtype: str = "int32",
|
|
) -> Var:
|
|
"""The opaque block axis defining function.
|
|
|
|
Parameters
|
|
----------
|
|
dom : Union[Range, List[Expr], Tuple[Expr]]
|
|
The domain of the iteration variable.
|
|
|
|
binding : Expr
|
|
The binding value of the iteration variable.
|
|
|
|
dtype : str
|
|
The data type of the iteration variable.
|
|
|
|
Returns
|
|
-------
|
|
res : Var
|
|
The iteration variable.
|
|
"""
|
|
return _ffi_api.AxisOpaque( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
_as_range(dom), binding, dtype
|
|
)
|
|
|
|
@staticmethod
|
|
def remap(kinds: str, bindings: list[Expr], dtype: str = "int32") -> list[Var] | Var:
|
|
"""The block axis remapping function.
|
|
|
|
Parameters
|
|
----------
|
|
kinds : str
|
|
The types of the iteration variables.
|
|
|
|
bindings : List[Expr]
|
|
The binding values of the iteration variables.
|
|
|
|
dtype : str
|
|
The data types of the iteration variables.
|
|
|
|
Returns
|
|
-------
|
|
res : Var
|
|
The iteration variables.
|
|
"""
|
|
iter_vars = _ffi_api.AxisRemap( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
kinds, bindings, dtype
|
|
)
|
|
return iter_vars[0] if len(iter_vars) == 1 else iter_vars
|
|
|
|
S = spatial # pylint: disable=invalid-name
|
|
R = reduce # pylint: disable=invalid-name
|
|
|
|
|
|
def serial(
|
|
start: Expr,
|
|
stop: Expr = None,
|
|
*,
|
|
annotations: dict[str, Any] | None = None,
|
|
step: Expr | None = None,
|
|
unroll: bool | None = None,
|
|
) -> frame.ForFrame:
|
|
"""The serial For statement.
|
|
|
|
Parameters
|
|
----------
|
|
start : Expr
|
|
The minimum value of iteration.
|
|
|
|
stop : Expr
|
|
The maximum value of iteration.
|
|
|
|
annotations : Dict[str, Any]
|
|
The optional annotations of the For statement.
|
|
|
|
step : Expr
|
|
The optional step value of iteration.
|
|
|
|
unroll : bool, optional
|
|
If True, adds ``{"pragma_unroll": True}`` annotation, which asks CUDA codegen
|
|
to emit ``#pragma unroll`` while preserving the loop as a C++ ``for``.
|
|
If False, adds ``{"disable_unroll": True}`` annotation.
|
|
Shorthand for ``annotations={"disable_unroll": True}``.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.ForFrame
|
|
The ForFrame.
|
|
"""
|
|
if unroll is not None:
|
|
annotations = dict(annotations) if annotations else {}
|
|
if unroll:
|
|
annotations["pragma_unroll"] = True
|
|
else:
|
|
annotations["disable_unroll"] = True
|
|
if stop is None:
|
|
stop = start
|
|
if is_prim_expr(start):
|
|
start = IntImm(start.ty, 0)
|
|
else:
|
|
start = 0
|
|
return _ffi_api.Serial(start, stop, annotations, step) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def parallel(
|
|
start: Expr,
|
|
stop: Expr = None,
|
|
*,
|
|
annotations: dict[str, Any] | None = None,
|
|
step: Expr | None = None,
|
|
) -> frame.ForFrame:
|
|
"""The parallel For statement.
|
|
|
|
Parameters
|
|
----------
|
|
start : Expr
|
|
The minimum value of iteration.
|
|
|
|
stop : Expr
|
|
The maximum value of iteration.
|
|
|
|
annotations : Dict[str, Any]
|
|
The optional annotations of the For statement.
|
|
|
|
step : Expr
|
|
The optional step value of iteration.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.ForFrame
|
|
The ForFrame.
|
|
"""
|
|
if stop is None:
|
|
stop = start
|
|
if is_prim_expr(start):
|
|
start = IntImm(start.ty, 0)
|
|
else:
|
|
start = 0
|
|
return _ffi_api.Parallel(start, stop, annotations, step) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def vectorized(
|
|
start: Expr,
|
|
stop: Expr = None,
|
|
*,
|
|
annotations: dict[str, Any] | None = None,
|
|
step: Expr | None = None,
|
|
) -> frame.ForFrame:
|
|
"""The vectorized For statement.
|
|
|
|
Parameters
|
|
----------
|
|
start : Expr
|
|
The minimum value of iteration.
|
|
|
|
stop : Expr
|
|
The maximum value of iteration.
|
|
|
|
annotations : Dict[str, Any]
|
|
The optional annotations of the For statement.
|
|
|
|
step : Expr
|
|
The optional step value of iteration.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.ForFrame
|
|
The ForFrame.
|
|
"""
|
|
if stop is None:
|
|
stop = start
|
|
if is_prim_expr(start):
|
|
start = IntImm(start.ty, 0)
|
|
else:
|
|
start = 0
|
|
return _ffi_api.Vectorized(start, stop, annotations, step) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def unroll(
|
|
start: Expr,
|
|
stop: Expr = None,
|
|
*,
|
|
annotations: dict[str, Any] | None = None,
|
|
step: Expr | None = None,
|
|
) -> frame.ForFrame:
|
|
"""The unrolled For statement.
|
|
|
|
Parameters
|
|
----------
|
|
start : Expr
|
|
The minimum value of iteration.
|
|
|
|
stop : Expr
|
|
The maximum value of iteration.
|
|
|
|
annotations : Dict[str, Any]
|
|
The optional annotations of the For statement.
|
|
|
|
step : Expr
|
|
The optional step value of iteration.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.ForFrame
|
|
The ForFrame.
|
|
"""
|
|
if stop is None:
|
|
stop = start
|
|
if is_prim_expr(start):
|
|
start = IntImm(start.ty, 0)
|
|
else:
|
|
start = 0
|
|
return _ffi_api.Unroll(start, stop, annotations, step) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def thread_binding(
|
|
start: Expr,
|
|
stop: Expr = None,
|
|
thread: str | None = None,
|
|
*,
|
|
annotations: dict[str, Any] | None = None,
|
|
) -> frame.ForFrame:
|
|
"""The thread-binding For statement.
|
|
|
|
Parameters
|
|
----------
|
|
start : Expr
|
|
The minimum value of iteration.
|
|
|
|
stop : Expr
|
|
The maximum value of iteration.
|
|
|
|
thread : str
|
|
The thread for loop variable to bind.
|
|
|
|
annotations : Dict[str, Any]
|
|
The optional annotations of the For statement.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.ForFrame
|
|
The ForFrame.
|
|
"""
|
|
if thread is None:
|
|
if not isinstance(stop, str):
|
|
raise ValueError("Thread cannot be None for thread_binding")
|
|
thread = stop
|
|
stop = start
|
|
if is_prim_expr(start):
|
|
start = IntImm(start.ty, 0)
|
|
else:
|
|
start = 0
|
|
elif stop is None:
|
|
stop = start
|
|
if is_prim_expr(start):
|
|
start = IntImm(start.ty, 0)
|
|
else:
|
|
start = 0
|
|
return _ffi_api.ThreadBinding( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
start, stop, thread, annotations
|
|
)
|
|
|
|
|
|
def grid(*extents: tuple[Expr | tuple[Expr, Expr]]) -> frame.ForFrame:
|
|
"""The grid For statement.
|
|
|
|
Parameters
|
|
----------
|
|
extents : Tuple[Union[Expr, Tuple[Expr, Expr]]]
|
|
If a single Expr is provided, it is used as the extent of the iteration.
|
|
If a tuple of two Expr is provided, the first is the start of the iteration,
|
|
and the second is the extent of the iteration.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.ForFrame
|
|
The ForFrame.
|
|
"""
|
|
# Convert integer extents to IntImm
|
|
# TODO(@bohan): fix this after FFI refactor
|
|
processed_extents = []
|
|
for extent in extents:
|
|
if isinstance(extent, tuple):
|
|
start, extent = extent
|
|
start = IntImm("int32", start) if isinstance(start, int) else start
|
|
extent = IntImm("int32", extent) if isinstance(extent, int) else extent
|
|
processed_extents.append((start, extent))
|
|
else:
|
|
processed_extents.append(IntImm("int32", extent) if isinstance(extent, int) else extent)
|
|
extents = tuple(processed_extents)
|
|
return _ffi_api.Grid(extents) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def Assert(condition: Expr, message, error_kind: str = "RuntimeError") -> frame.AssertFrame: # pylint: disable=invalid-name
|
|
"""Create an assertion statement.
|
|
|
|
Parameters
|
|
----------
|
|
condition : Expr
|
|
The Expr to test.
|
|
|
|
message : str or list[str]
|
|
The error message when the assertion fails. Can be a single string
|
|
or a list of string parts (fragments stored separately in the IR
|
|
for binary size reduction through string reuse).
|
|
|
|
error_kind : str
|
|
The error kind (e.g. "RuntimeError", "TypeError", "ValueError").
|
|
|
|
Returns
|
|
-------
|
|
res : frame.AssertFrame
|
|
The result AssertFrame.
|
|
"""
|
|
if isinstance(condition, bool):
|
|
condition = IntImm("bool", condition)
|
|
if not isinstance(message, list | tuple):
|
|
message = [message]
|
|
return _ffi_api.Assert(condition, error_kind, message) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def Bind( # pylint: disable=invalid-name
|
|
value: Expr,
|
|
type_annotation: Type | None = None, # pylint: disable=redefined-outer-name
|
|
*,
|
|
var: Var | None = None, # pylint: disable=redefined-outer-name
|
|
) -> Var:
|
|
"""Create a Bind (variable binding).
|
|
|
|
Emits a flat Bind statement to the current frame and returns the bound variable.
|
|
|
|
Parameters
|
|
----------
|
|
value : Expr
|
|
The value to be bound.
|
|
type_annotation : Optional[Type] = None
|
|
The type annotation of the binding. Usually it is used for fine-grained var typing,
|
|
particularly, PointerType.
|
|
var : Optional[Var] = None
|
|
The variable to bind. If not specified, a new variable will be created.
|
|
|
|
Returns
|
|
-------
|
|
var : Var
|
|
The bound variable.
|
|
"""
|
|
if type_annotation is not None:
|
|
if callable(type_annotation):
|
|
type_annotation = type_annotation()
|
|
if isinstance(type_annotation, Var):
|
|
type_annotation = type_annotation.ty
|
|
return _ffi_api.Bind(value, type_annotation, var) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def Let( # pylint: disable=invalid-name
|
|
expr: Expr,
|
|
where: dict[Var, Expr], # pylint: disable=redefined-outer-name
|
|
) -> Expr:
|
|
"""Create a Let expression binding"""
|
|
assert len(where) == 1, "T.Let only allows `where` to have exactly one element"
|
|
var, value = next(iter(where.items())) # pylint: disable=redefined-outer-name
|
|
return tir.Let(var, value, expr)
|
|
|
|
|
|
bind = Bind
|
|
|
|
|
|
class LetAnnotation:
|
|
"""Marker for explicit LetStmt. Created by T.let or T.let[type].
|
|
Usage in TVMScript:
|
|
x: T.let[T.int32] = expr # LetStmt with explicit type
|
|
x: T.let = expr # LetStmt with auto-typed RHS
|
|
"""
|
|
|
|
def __init__(self, type_spec=None):
|
|
self.type_spec = type_spec
|
|
|
|
def __class_getitem__(cls, item):
|
|
return LetAnnotation(item)
|
|
|
|
def __getitem__(self, item):
|
|
return LetAnnotation(item)
|
|
|
|
def as_var(self, rhs_dtype=None):
|
|
"""Resolve to a tir.Var."""
|
|
if self.type_spec is not None:
|
|
if isinstance(self.type_spec, Var):
|
|
return self.type_spec # Already a Var (e.g. T.handle(...))
|
|
elif callable(self.type_spec):
|
|
return self.type_spec() # e.g. T.int32() -> Var
|
|
elif isinstance(self.type_spec, Type):
|
|
return Var("", self.type_spec)
|
|
else:
|
|
raise TypeError(f"Invalid type for T.let: {self.type_spec}")
|
|
elif rhs_dtype is not None:
|
|
rhs_ty = rhs_dtype if isinstance(rhs_dtype, Type) else ir.PrimType(rhs_dtype)
|
|
return Var("", rhs_ty)
|
|
else:
|
|
raise TypeError("T.let requires either a type or an RHS value")
|
|
|
|
|
|
let = LetAnnotation() # Singleton for T.let (no subscript)
|
|
|
|
|
|
class LocalVectorAnnotation:
|
|
"""Marker for local vector/tensor allocation via type annotation subscript.
|
|
|
|
Created when a DtypeConstructor is subscripted, e.g. ``T.float32[N]`` or
|
|
``T.float32[M, N]``. The parser's ``visit_ann_assign`` recognises this
|
|
object and lowers it to ``T.alloc_local(shape=..., dtype=...)``.
|
|
"""
|
|
|
|
__slots__ = ("dtype", "shape")
|
|
|
|
def __init__(self, dtype: str, shape: tuple):
|
|
self.dtype = dtype
|
|
self.shape = shape
|
|
|
|
|
|
class DtypeConstructor:
|
|
"""Callable + subscriptable dtype object.
|
|
|
|
Replaces the plain functions previously returned by ``func_gen``.
|
|
|
|
* ``T.float32()`` — same FFI call as before (returns ``Var``).
|
|
* ``T.float32[N]`` — returns ``LocalVectorAnnotation("float32", (N,))``.
|
|
* ``T.float32[M, N]`` — returns ``LocalVectorAnnotation("float32", (M, N))``.
|
|
* ``x: T.float32`` — parser calls this object, gets a ``Var``.
|
|
"""
|
|
|
|
def __init__(self, ffi_name: str, dtype_str: str):
|
|
self._ffi_name = ffi_name
|
|
self._dtype_str = dtype_str
|
|
|
|
def __call__(
|
|
self,
|
|
expr: "None | Expr | Literal['inf', '-inf', 'nan'] | int | float" = None,
|
|
) -> "Expr":
|
|
if isinstance(expr, str):
|
|
expr = float(expr)
|
|
return getattr(_ffi_api, self._ffi_name)(expr)
|
|
|
|
def __getitem__(self, shape):
|
|
if isinstance(shape, tuple):
|
|
return LocalVectorAnnotation(self._dtype_str, shape)
|
|
return LocalVectorAnnotation(self._dtype_str, (shape,))
|
|
|
|
def __repr__(self):
|
|
return f"DtypeConstructor({self._dtype_str!r})"
|
|
|
|
|
|
def allocate(
|
|
extents: list[Expr],
|
|
dtype: str,
|
|
scope: str = "global",
|
|
condition: Expr = None,
|
|
annotations=None,
|
|
) -> frame.AllocateFrame:
|
|
"""Allocate node.
|
|
|
|
Parameters
|
|
----------
|
|
extents : List[Expr]
|
|
The extents of the allocate.
|
|
|
|
dtype : str
|
|
The data type of the buffer.
|
|
|
|
scope : str
|
|
The storage scope.
|
|
|
|
condition : Expr
|
|
The condition.
|
|
|
|
annotations: Optional[Mapping[str, Object]]
|
|
Additional annotation hints.
|
|
"""
|
|
if isinstance(condition, bool):
|
|
condition = IntImm("bool", condition)
|
|
return _ffi_api.Allocate( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
extents, dtype, scope, condition, annotations
|
|
)
|
|
|
|
|
|
def attr(
|
|
node_or_dict: Any, attr_key: str | None = None, value: Expr | str | None = None
|
|
) -> Union[frame.AttrFrame, "utils._FrameScope"]:
|
|
"""Create an attribute node, or multiple attribute nodes from a dict.
|
|
|
|
Usage 1 — single attr::
|
|
|
|
with T.attr(node, key, value):
|
|
...
|
|
|
|
Usage 2 — dict sugar (node defaults to ``T.int32(0)``)::
|
|
|
|
with T.attr({"key1": value1, "key2": value2}):
|
|
...
|
|
|
|
Parameters
|
|
----------
|
|
node_or_dict : Any
|
|
If a dict, each key-value pair becomes an AttrStmt with
|
|
``node=T.int32(0)``. Otherwise the node to annotate.
|
|
|
|
attr_key : str, optional
|
|
Attribute type key (required when ``node_or_dict`` is not a dict).
|
|
|
|
value : Union[Expr, str], optional
|
|
The attribute value (required when ``node_or_dict`` is not a dict).
|
|
|
|
Returns
|
|
-------
|
|
res : Union[frame.AttrFrame, _FrameScope]
|
|
A single AttrFrame, or a _FrameScope wrapping multiple AttrFrames.
|
|
"""
|
|
if isinstance(node_or_dict, dict):
|
|
frames = []
|
|
for k, v in node_or_dict.items():
|
|
if isinstance(v, bool):
|
|
v = IntImm("bool", v)
|
|
frames.append(
|
|
_ffi_api.Attr( # type: ignore[attr-defined]
|
|
convert(IntImm("int32", 0)), k, convert(v)
|
|
)
|
|
)
|
|
if len(frames) == 1:
|
|
return frames[0]
|
|
return utils._FrameScope(frames)
|
|
else:
|
|
if attr_key is None or value is None:
|
|
raise ValueError("T.attr(node, attr_key, value) requires all three arguments")
|
|
node_or_dict = convert(node_or_dict)
|
|
value = convert(value)
|
|
return _ffi_api.Attr(node_or_dict, attr_key, value) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def hint(message: str = "", **attrs) -> frame.HintFrame:
|
|
"""Universal directive primitive for the sketch language.
|
|
|
|
Parameters
|
|
----------
|
|
message : str
|
|
Free-form directive string that the agent interprets.
|
|
**attrs
|
|
Optional structured key-value attributes for known patterns.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.HintFrame
|
|
Usable as context manager (with T.hint("msg"):) or bare statement (T.hint("msg")).
|
|
"""
|
|
return _ffi_api.Hint(message, attrs or {}) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def While(condition: Expr) -> frame.WhileFrame: # pylint: disable=invalid-name
|
|
"""Create a while node.
|
|
|
|
Parameters
|
|
----------
|
|
condition : Expr
|
|
The termination condition of the loop.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.WhileFrame
|
|
The result WhileFrame.
|
|
"""
|
|
if isinstance(condition, bool):
|
|
condition = IntImm("bool", condition)
|
|
return _ffi_api.While(condition) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def Break() -> None: # pylint: disable=invalid-name
|
|
"""Create a break node."""
|
|
return _ffi_api.Break() # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def Continue() -> None: # pylint: disable=invalid-name
|
|
"""Create a continue node."""
|
|
return _ffi_api.Continue() # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def If(condition: Expr) -> frame.IfFrame: # pylint: disable=invalid-name
|
|
"""Create an if node.
|
|
|
|
Parameters
|
|
----------
|
|
condition : Expr
|
|
The condition of if statement, executes the true branch if the condition is true,
|
|
otherwise jump into the false branch.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.IfFrame
|
|
The result IfFrame.
|
|
"""
|
|
if isinstance(condition, bool):
|
|
condition = IntImm("bool", condition)
|
|
return _ffi_api.If(condition) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def Then() -> frame.ThenFrame: # pylint: disable=invalid-name
|
|
"""Create a then.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.ThenFrame
|
|
The result ThenFrame.
|
|
"""
|
|
return _ffi_api.Then() # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def Else() -> frame.ElseFrame: # pylint: disable=invalid-name
|
|
"""Create an else.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.ElseFrame
|
|
The result ElseFrame.
|
|
"""
|
|
return _ffi_api.Else() # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def decl_buffer(
|
|
shape,
|
|
dtype="float32",
|
|
data=None,
|
|
strides=None,
|
|
elem_offset=None,
|
|
byte_offset=None,
|
|
scope="global",
|
|
align=0,
|
|
offset_factor=0,
|
|
buffer_type="",
|
|
axis_separators=None,
|
|
layout="default",
|
|
allocated_addr=None,
|
|
) -> Buffer:
|
|
"""Create a buffer declaration node.
|
|
|
|
When ``data`` is provided, creates a DeclBuffer (alias to existing data).
|
|
When ``data`` is None, creates an AllocBuffer (new allocation).
|
|
|
|
Parameters
|
|
----------
|
|
shape : Union[List[Expr], Tuple[Expr], Expr, Integral]
|
|
The type of the buffer prior to flattening.
|
|
|
|
dtype : str
|
|
The data type in the content of the buffer.
|
|
|
|
data : Var
|
|
The pointer to the head of the data.
|
|
|
|
strides : List[Expr]
|
|
The strides of each dimension.
|
|
|
|
elem_offset : Expr
|
|
The offset in terms of number of dtype elements (including lanes).
|
|
|
|
byte_offset : Expr
|
|
The offset in terms of number of bytes.
|
|
|
|
scope : str
|
|
The optional storage scope of buffer data pointer.
|
|
|
|
align : int
|
|
The alignment requirement of data pointer in bytes.
|
|
|
|
offset_factor : int
|
|
The factor of elem_offset field.
|
|
|
|
buffer_type : str
|
|
The buffer type.
|
|
|
|
axis_separators : List[int]
|
|
The separators between input axes when generating flattened output axes.
|
|
|
|
layout : Layout
|
|
The layout of the buffer.
|
|
|
|
Returns
|
|
-------
|
|
res : Buffer
|
|
The declared buffer.
|
|
"""
|
|
shape = (shape,) if is_prim_expr(shape) or isinstance(shape, Integral) else shape
|
|
if strides is not None:
|
|
strides = [Var(s, "int32") if isinstance(s, str) else s for s in strides]
|
|
else:
|
|
strides = []
|
|
dtype = _normalize_prim_type(dtype)
|
|
decl_frame = _ffi_api.DeclBuffer( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
shape,
|
|
dtype,
|
|
"",
|
|
data,
|
|
strides,
|
|
_get_elem_offset(elem_offset, byte_offset, dtype),
|
|
scope,
|
|
align,
|
|
offset_factor,
|
|
buffer_type,
|
|
axis_separators,
|
|
_get_layout(layout, shape, scope),
|
|
allocated_addr,
|
|
)
|
|
if isinstance(decl_frame, frame.DeclBufferFrame):
|
|
decl_frame.add_callback(partial(decl_frame.__exit__, None, None, None))
|
|
buf = decl_frame.__enter__()
|
|
else:
|
|
buf = decl_frame
|
|
_record_meta_resource(buf, skip_frames=2)
|
|
return buf
|
|
|
|
|
|
alloc_shared = functools.partial(alloc_buffer, scope="shared")
|
|
alloc_local = functools.partial(alloc_buffer, scope="local")
|
|
smem = alloc_shared
|
|
tmem = functools.partial(alloc_buffer, scope="tmem")
|
|
|
|
|
|
def alloc_tcgen05_ldst_frag(instr_shape, tensor_shape, dtype):
|
|
"""Allocate a register fragment for ``tcgen05.{ld,st}`` atoms.
|
|
|
|
Sizes the per-thread storage, allocates ``local`` scope memory, and returns
|
|
a 2-D view of shape ``tensor_shape`` with a matching ``tcgen05_atom_layout``.
|
|
Pass the result to ``Tx.copy_async`` (with a ``(128, W)``-shaped TMEM
|
|
buffer) to trigger the corresponding dispatch path.
|
|
|
|
Parameters
|
|
----------
|
|
instr_shape : str
|
|
``"32x32b"`` (M=128 fragment, 128 row warpgroup tile, layout
|
|
``(128, K):(1@tid_in_wg, 1)``); or ``"16x64b"`` / ``"16x128b"`` /
|
|
``"16x256b"`` (M=64 fragments, 64 row warpgroup tile with the
|
|
per-shape per-lane register decomposition).
|
|
tensor_shape : tuple[int, int]
|
|
Logical fragment shape ``(frag_rows, K)`` in element units. ``frag_rows``
|
|
is ``128`` for ``.32x32b`` and ``64`` for the ``.16x*b`` shapes.
|
|
dtype : str
|
|
``"float32"``, ``"float16"``, or ``"bfloat16"``.
|
|
|
|
Returns
|
|
-------
|
|
Buffer
|
|
2-D view of shape ``tensor_shape`` whose layout matches
|
|
``tcgen05_atom_layout(instr_shape, tensor_shape, dtype)``.
|
|
|
|
Examples
|
|
--------
|
|
M=128 readback (existing dispatch):
|
|
``frag = T.alloc_tcgen05_ldst_frag("32x32b", (128, 64), "float32")``
|
|
``Tx.copy_async(frag[:, :], tmem[:, 0:64])``
|
|
|
|
M=64 readback (.16x64b dispatch):
|
|
``frag = T.alloc_tcgen05_ldst_frag("16x64b", (64, 64), "float32")``
|
|
``Tx.copy_async(frag[:, :], tmem[0:64, 0:64])``
|
|
"""
|
|
from tvm.tirx.layout import tcgen05_atom_layout # local import to avoid cycle
|
|
|
|
rows, cols = tensor_shape
|
|
bits = DataType(dtype).bits
|
|
# Per-warpgroup total bits = 64 rows x K cols x bits. Divided across 128
|
|
# threads gives per-thread bits; convert to element count.
|
|
per_thread_bits = (rows * cols * bits) // 128
|
|
if per_thread_bits % bits != 0:
|
|
raise ValueError(
|
|
f"alloc_tcgen05_ldst_frag tensor_shape={tensor_shape} dtype={dtype!r} "
|
|
f"does not evenly divide across 128 threads"
|
|
)
|
|
per_thread_elems = per_thread_bits // bits
|
|
|
|
layout = tcgen05_atom_layout(instr_shape, tensor_shape, dtype)
|
|
flat = alloc_local((per_thread_elems,), dtype)
|
|
return flat.view(rows, cols, layout=layout)
|
|
|
|
|
|
def alloc_cast_frag(src, dtype):
|
|
"""Allocate a register frag holding ``src`` value-cast to ``dtype``.
|
|
|
|
Inherits ``src``'s logical shape and its ``(lane, register)`` layout — only
|
|
the element dtype changes — so ``Tx.cast(dst, src)`` is a per-thread
|
|
element-wise cast with no cross-lane movement. ``.permute(...)`` the result
|
|
to the axis order a downstream consumer (e.g. ``stmatrix`` via
|
|
``Tx.copy(dispatch="ldstmatrix")``) expects.
|
|
|
|
Parameters
|
|
----------
|
|
src : Buffer
|
|
Source register frag (e.g. from ``alloc_tcgen05_ldst_frag``).
|
|
dtype : str
|
|
Destination element dtype.
|
|
|
|
Returns
|
|
-------
|
|
Buffer
|
|
Fresh ``local`` frag, ``src.shape`` shaped, ``src.layout``, dtype-cast.
|
|
"""
|
|
rows, cols = src.shape
|
|
per_thread_elems = (rows * cols) // 128
|
|
flat = alloc_local((per_thread_elems,), dtype)
|
|
return flat.view(rows, cols, layout=src.layout)
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
ScalarT = TypeVar("ScalarT")
|
|
|
|
# Keep type checking/linting simple by treating wrapper as identity.
|
|
def scalar_wrapper(x: ScalarT) -> ScalarT:
|
|
return x
|
|
|
|
else:
|
|
|
|
class scalar_wrapper:
|
|
"""Internal wrapper to allow IRBuilder auto-naming on scalar assignment."""
|
|
|
|
def __init__(self, scalar: BufferLoad):
|
|
assert isinstance(scalar, BufferLoad)
|
|
self.scalar = scalar
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
return getattr(self.scalar, name)
|
|
|
|
def __add__(self, other):
|
|
return self.scalar + other
|
|
|
|
def __radd__(self, other):
|
|
return other + self.scalar
|
|
|
|
def __sub__(self, other):
|
|
return self.scalar - other
|
|
|
|
def __rsub__(self, other):
|
|
return other - self.scalar
|
|
|
|
def __mul__(self, other):
|
|
return self.scalar * other
|
|
|
|
def __rmul__(self, other):
|
|
return other * self.scalar
|
|
|
|
def __truediv__(self, other):
|
|
return self.scalar / other
|
|
|
|
def __rtruediv__(self, other):
|
|
return other / self.scalar
|
|
|
|
def __floordiv__(self, other):
|
|
return self.scalar // other
|
|
|
|
def __rfloordiv__(self, other):
|
|
return other // self.scalar
|
|
|
|
def __mod__(self, other):
|
|
return self.scalar % other
|
|
|
|
def __rmod__(self, other):
|
|
return other % self.scalar
|
|
|
|
def __lt__(self, other):
|
|
return self.scalar < other
|
|
|
|
def __le__(self, other):
|
|
return self.scalar <= other
|
|
|
|
def __gt__(self, other):
|
|
return self.scalar > other
|
|
|
|
def __ge__(self, other):
|
|
return self.scalar >= other
|
|
|
|
def __eq__(self, other):
|
|
return self.scalar == other
|
|
|
|
def __ne__(self, other):
|
|
return self.scalar != other
|
|
|
|
def __and__(self, other):
|
|
return self.scalar & other
|
|
|
|
def __rand__(self, other):
|
|
return other & self.scalar
|
|
|
|
def __or__(self, other):
|
|
return self.scalar | other
|
|
|
|
def __ror__(self, other):
|
|
return other | self.scalar
|
|
|
|
def __xor__(self, other):
|
|
return self.scalar ^ other
|
|
|
|
def __rxor__(self, other):
|
|
return other ^ self.scalar
|
|
|
|
def __neg__(self):
|
|
return -self.scalar
|
|
|
|
def __invert__(self):
|
|
return ~self.scalar
|
|
|
|
|
|
def alloc_scalar(dtype: str = "float32", scope: str = "global") -> BufferLoad:
|
|
"""Allocate a zero-dimensional buffer (scalar)."""
|
|
buf = alloc_buffer(shape=(1,), dtype=dtype, scope=scope, layout=TileLayout(S[1]))
|
|
assert isinstance(buf, Buffer)
|
|
scalar = buf[0]
|
|
if _current_meta_construction_scope() is not None:
|
|
return scalar
|
|
return scalar_wrapper(scalar)
|
|
|
|
|
|
def decl_scalar(dtype, data, scope, elem_offset=None, byte_offset=None) -> BufferLoad:
|
|
"""Declare a zero-dimensional buffer (scalar) from a pointer."""
|
|
buf = decl_buffer(
|
|
shape=(1,),
|
|
dtype=dtype,
|
|
data=data,
|
|
scope=scope,
|
|
elem_offset=_get_elem_offset(elem_offset, byte_offset, dtype),
|
|
strides=None,
|
|
align=-1,
|
|
offset_factor=0,
|
|
buffer_type="default",
|
|
axis_separators=None,
|
|
layout=TileLayout(S[1]),
|
|
)
|
|
assert isinstance(buf, Buffer)
|
|
scalar = buf[0]
|
|
if _current_meta_construction_scope() is not None:
|
|
return scalar
|
|
return scalar_wrapper(scalar)
|
|
|
|
|
|
def shared_scalar(dtype: str = "float32") -> BufferLoad:
|
|
"""Allocate a zero-dimensional buffer in shared memory."""
|
|
return alloc_scalar(dtype=dtype, scope="shared")
|
|
|
|
|
|
def local_scalar(dtype: str = "float32") -> BufferLoad:
|
|
"""Allocate a zero-dimensional buffer in local memory."""
|
|
return alloc_scalar(dtype=dtype, scope="local")
|
|
|
|
|
|
def _is_meta_class_instance(value: Any) -> bool:
|
|
return getattr(type(value), "_is_meta_class", False)
|
|
|
|
|
|
def _sanitize_meta_name_part(value: Any, fallback: str) -> str:
|
|
if isinstance(value, str) and value.isidentifier():
|
|
return value
|
|
if isinstance(value, str):
|
|
sanitized = "".join(c if c.isalnum() or c == "_" else "_" for c in value)
|
|
if sanitized and sanitized[0].isalpha():
|
|
return sanitized
|
|
return fallback
|
|
|
|
|
|
def _meta_resource_for_value(value: Any) -> Any | None:
|
|
if isinstance(value, scalar_wrapper):
|
|
return value.scalar.buffer
|
|
if isinstance(value, BufferLoad):
|
|
return value.buffer
|
|
if isinstance(value, Buffer):
|
|
return value
|
|
return None
|
|
|
|
|
|
def _resource_in(resource: Any, resources: list[Any]) -> bool:
|
|
return any(_same_meta_resource(resource, other) for other in resources)
|
|
|
|
|
|
def _name_meta_value(
|
|
prefix: str,
|
|
value: Any,
|
|
visited: set[int] | None = None,
|
|
owned_resources: list[Any] | None = None,
|
|
named_resources: list[Any] | None = None,
|
|
) -> None:
|
|
if visited is None:
|
|
visited = set()
|
|
if named_resources is None:
|
|
named_resources = []
|
|
obj_id = id(value)
|
|
if obj_id in visited:
|
|
return
|
|
visited.add(obj_id)
|
|
|
|
resource = _meta_resource_for_value(value)
|
|
if resource is not None:
|
|
if owned_resources is not None and not _resource_in(resource, owned_resources):
|
|
return
|
|
if _resource_in(resource, named_resources):
|
|
return
|
|
IRBuilder.name(prefix, resource)
|
|
named_resources.append(resource)
|
|
return
|
|
if isinstance(value, Var | IterVar):
|
|
if owned_resources is not None:
|
|
return
|
|
IRBuilder.name(prefix, value)
|
|
return
|
|
if _is_meta_class_instance(value):
|
|
existing_prefix = getattr(value, "_tirx_meta_name", None)
|
|
if existing_prefix is not None and existing_prefix != prefix:
|
|
return
|
|
object.__setattr__(value, "_tirx_meta_name", prefix)
|
|
instance_owned_resources = getattr(value, "_tirx_meta_owned_resources", [])
|
|
for field_name, field_value in vars(value).items():
|
|
if field_name.startswith("_tirx_"):
|
|
continue
|
|
_name_meta_value(
|
|
f"{prefix}_{field_name}",
|
|
field_value,
|
|
visited,
|
|
instance_owned_resources,
|
|
named_resources,
|
|
)
|
|
return
|
|
if isinstance(value, list | tuple):
|
|
for i, item in enumerate(value):
|
|
_name_meta_value(f"{prefix}_{i}", item, visited, owned_resources, named_resources)
|
|
return
|
|
if isinstance(value, dict):
|
|
for i, (key, item) in enumerate(value.items()):
|
|
part = _sanitize_meta_name_part(key, f"item{i}")
|
|
_name_meta_value(f"{prefix}_{part}", item, visited, owned_resources, named_resources)
|
|
|
|
|
|
def _same_meta_resource(lhs: Any, rhs: Any) -> bool:
|
|
same_as = getattr(lhs, "same_as", None)
|
|
if same_as is not None:
|
|
try:
|
|
return bool(same_as(rhs))
|
|
except TypeError:
|
|
pass
|
|
return lhs is rhs
|
|
|
|
|
|
def _collect_meta_resources(value: Any, visited: set[int] | None = None) -> list[Any]:
|
|
if visited is None:
|
|
visited = set()
|
|
obj_id = id(value)
|
|
if obj_id in visited:
|
|
return []
|
|
visited.add(obj_id)
|
|
|
|
resource = _meta_resource_for_value(value)
|
|
if resource is not None:
|
|
return [resource]
|
|
if _is_meta_class_instance(value):
|
|
owned = []
|
|
for field_name, field_value in vars(value).items():
|
|
if field_name.startswith("_tirx_"):
|
|
continue
|
|
owned.extend(_collect_meta_resources(field_value, visited))
|
|
return owned
|
|
if isinstance(value, list | tuple):
|
|
owned = []
|
|
for item in value:
|
|
owned.extend(_collect_meta_resources(item, visited))
|
|
return owned
|
|
if isinstance(value, dict):
|
|
owned = []
|
|
for item in value.values():
|
|
owned.extend(_collect_meta_resources(item, visited))
|
|
return owned
|
|
return []
|
|
|
|
|
|
def _format_unowned_meta_resource_error(cls: type, record: _MetaResourceRecord, total: int) -> str:
|
|
count = "" if total == 1 else f" ({total} total)"
|
|
location = f"{record.filename}:{record.lineno}"
|
|
if record.colno is not None:
|
|
location = f"{location}:{record.colno}"
|
|
message = [
|
|
f"TIRx meta_class constructor created an unowned resource{count}.",
|
|
f" class: {cls.__name__}",
|
|
f" location: {location}",
|
|
]
|
|
if record.code:
|
|
message.extend(["", f" {record.code}", " ^ resource must be assigned to self.<field>"])
|
|
message.extend(
|
|
[
|
|
"",
|
|
"Resources created in a meta_class constructor must be reachable from the",
|
|
"constructed instance.",
|
|
"unowned resource at "
|
|
f"{location}: assign it to self.<field>, or move the allocation into a "
|
|
"parser-owned assignment.",
|
|
]
|
|
)
|
|
return "\n".join(message)
|
|
|
|
|
|
def _validate_meta_construction_scope(scope: _MetaConstructionScope) -> None:
|
|
if not scope.created:
|
|
object.__setattr__(scope.instance, "_tirx_meta_owned_resources", [])
|
|
return
|
|
created_resources = [record.value for record in scope.created]
|
|
owned_resources = _collect_meta_resources(scope.instance)
|
|
missing = [
|
|
record
|
|
for record in scope.created
|
|
if not any(_same_meta_resource(record.value, owned) for owned in owned_resources)
|
|
]
|
|
if missing:
|
|
raise ValueError(_format_unowned_meta_resource_error(scope.cls, missing[0], len(missing)))
|
|
object.__setattr__(scope.instance, "_tirx_meta_owned_resources", created_resources)
|
|
|
|
|
|
def name_meta_class_value(prefix: str, value: Any) -> None:
|
|
"""Name all TIR resources owned by a meta_class instance."""
|
|
_name_meta_value(prefix, value)
|
|
|
|
|
|
def launch_thread(
|
|
thread: IterVar | str, # pylint: disable=redefined-outer-name
|
|
extent: Expr,
|
|
) -> frame.LaunchThreadFrame:
|
|
"""Launch a thread.
|
|
|
|
Parameters
|
|
----------
|
|
thread : Union[IterVar, str]
|
|
The iteration variable.
|
|
|
|
extent : Expr
|
|
The extent of environment thread.
|
|
|
|
Returns
|
|
-------
|
|
res : frame.LaunchThreadFrame
|
|
The result LaunchThreadFrame.
|
|
|
|
Examples
|
|
--------
|
|
|
|
.. code-block:: python
|
|
|
|
from tvm.script.ir_builder import tirx as T
|
|
brow = T.env_thread("blockIdx.y")
|
|
T.launch_thread(brow, 1)
|
|
|
|
"""
|
|
|
|
if isinstance(thread, str):
|
|
thread = String(thread)
|
|
return _ffi_api.LaunchThread(thread, extent) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def env_thread(thread_tag: str, dtype: str = "int32") -> IterVar:
|
|
"""Bind a var to thread env
|
|
|
|
Parameters
|
|
----------
|
|
thread_tag : str
|
|
The thread type tag.
|
|
|
|
dtype : str
|
|
The data type of the thread env.
|
|
|
|
Returns
|
|
-------
|
|
res : IterVar
|
|
The result iteration variable gets bound to the thread env.
|
|
|
|
"""
|
|
return _ffi_api.EnvThread(thread_tag, dtype) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def buffer_store(
|
|
buffer: Buffer, # pylint: disable=redefined-outer-name
|
|
value: Expr,
|
|
indices: list[Expr | slice],
|
|
predicate: Expr | None = None,
|
|
) -> None:
|
|
"""Buffer store node.
|
|
|
|
Parameters
|
|
----------
|
|
buffer : Buffer
|
|
The buffer.
|
|
|
|
value : Expr
|
|
The value to be stored.
|
|
|
|
indices : List[Union[Expr, slice]]
|
|
The indices location to be stored.
|
|
|
|
predicate : Optional[Expr]
|
|
A vector mask of boolean values indicating which lanes of a vector are to be
|
|
stored. The number lanes of the mask must be equal to the number of lanes in
|
|
value.
|
|
"""
|
|
from tvm.arith import Analyzer # pylint: disable=import-outside-toplevel
|
|
|
|
if not isinstance(indices, list | tuple | ir.Array):
|
|
indices = [indices]
|
|
|
|
expr_indices = []
|
|
for index in indices:
|
|
if isinstance(index, slice):
|
|
step = 1 if index.step is None else index.step
|
|
lanes = Analyzer().simplify( # pylint: disable=redefined-outer-name
|
|
(index.stop - index.start + step - 1) // step
|
|
)
|
|
if lanes == 1:
|
|
expr_indices.append(index.start)
|
|
else:
|
|
expr_indices.append(ramp(index.start, step, lanes))
|
|
else:
|
|
expr_indices.append(index)
|
|
if isinstance(value, bool) and buffer.dtype == "bool":
|
|
value = IntImm("bool", value)
|
|
return _ffi_api.BufferStore( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
buffer, value, expr_indices, predicate
|
|
)
|
|
|
|
|
|
def evaluate(value: Expr) -> None:
|
|
"""Evaluate the input expression.
|
|
|
|
Parameters
|
|
----------
|
|
value: Expr
|
|
The input expression to evaluate.
|
|
"""
|
|
if isinstance(value, str):
|
|
value = StringImm(value)
|
|
if isinstance(value, bool):
|
|
value = IntImm("bool", value)
|
|
return _ffi_api.Evaluate(value) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def _ffi_name_to_dtype(name: str) -> str:
|
|
"""Convert an FFI type name to its TVM dtype string.
|
|
|
|
Examples: "Float32" -> "float32", "Int8x4" -> "int8x4",
|
|
"Float8E4M3" -> "float8_e4m3", "Float8E4M3B11FNUZ" -> "float8_e4m3b11fnuz".
|
|
"""
|
|
import re
|
|
|
|
# Insert underscore before E-notation in float8 names (E3M4, E4M3, etc.)
|
|
s = re.sub(r"(?<=[a-z0-9])E(\d)", r"_e\1", name, flags=re.IGNORECASE)
|
|
return s.lower()
|
|
|
|
|
|
def func_gen(name: str):
|
|
"""Generate a DtypeConstructor for each Expr dtype.
|
|
|
|
Parameters
|
|
----------
|
|
name: str
|
|
The ffi function name to call, e.g. "Float32", "Int32".
|
|
"""
|
|
return DtypeConstructor(name, _ffi_name_to_dtype(name))
|
|
|
|
|
|
def static_assert(x: Any, message: str = ""):
|
|
assert x, message
|
|
|
|
|
|
def add_to_parent(stmt: tir.Stmt) -> None:
|
|
"""Add a statement to the parent frame."""
|
|
_ffi_api.AddToParent(stmt) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
# pylint: disable=invalid-name
|
|
int8 = func_gen("Int8")
|
|
int16 = func_gen("Int16")
|
|
int32 = func_gen("Int32")
|
|
int64 = func_gen("Int64")
|
|
int8x2 = func_gen("Int8x2")
|
|
int16x2 = func_gen("Int16x2")
|
|
int32x2 = func_gen("Int32x2")
|
|
int64x2 = func_gen("Int64x2")
|
|
int8x4 = func_gen("Int8x4")
|
|
int16x4 = func_gen("Int16x4")
|
|
int32x4 = func_gen("Int32x4")
|
|
int64x4 = func_gen("Int64x4")
|
|
int8x8 = func_gen("Int8x8")
|
|
int16x8 = func_gen("Int16x8")
|
|
int32x8 = func_gen("Int32x8")
|
|
int64x8 = func_gen("Int64x8")
|
|
int8x16 = func_gen("Int8x16")
|
|
int16x16 = func_gen("Int16x16")
|
|
int32x16 = func_gen("Int32x16")
|
|
int64x16 = func_gen("Int64x16")
|
|
int8x32 = func_gen("Int8x32")
|
|
int16x32 = func_gen("Int16x32")
|
|
int32x32 = func_gen("Int32x32")
|
|
int64x32 = func_gen("Int64x32")
|
|
int8x64 = func_gen("Int8x64")
|
|
int16x64 = func_gen("Int16x64")
|
|
int32x64 = func_gen("Int32x64")
|
|
int64x64 = func_gen("Int64x64")
|
|
|
|
uint8 = func_gen("UInt8")
|
|
uint16 = func_gen("UInt16")
|
|
uint32 = func_gen("UInt32")
|
|
uint64 = func_gen("UInt64")
|
|
uint8x2 = func_gen("UInt8x2")
|
|
uint16x2 = func_gen("UInt16x2")
|
|
uint32x2 = func_gen("UInt32x2")
|
|
uint64x2 = func_gen("UInt64x2")
|
|
uint8x4 = func_gen("UInt8x4")
|
|
uint16x4 = func_gen("UInt16x4")
|
|
uint32x4 = func_gen("UInt32x4")
|
|
uint64x4 = func_gen("UInt64x4")
|
|
uint8x8 = func_gen("UInt8x8")
|
|
uint16x8 = func_gen("UInt16x8")
|
|
uint32x8 = func_gen("UInt32x8")
|
|
uint64x8 = func_gen("UInt64x8")
|
|
uint8x16 = func_gen("UInt8x16")
|
|
uint16x16 = func_gen("UInt16x16")
|
|
uint32x16 = func_gen("UInt32x16")
|
|
uint64x16 = func_gen("UInt64x16")
|
|
uint8x32 = func_gen("UInt8x32")
|
|
uint16x32 = func_gen("UInt16x32")
|
|
uint32x32 = func_gen("UInt32x32")
|
|
uint64x32 = func_gen("UInt64x32")
|
|
uint8x64 = func_gen("UInt8x64")
|
|
uint16x64 = func_gen("UInt16x64")
|
|
uint32x64 = func_gen("UInt32x64")
|
|
uint64x64 = func_gen("UInt64x64")
|
|
|
|
float16 = func_gen("Float16")
|
|
float32 = func_gen("Float32")
|
|
float64 = func_gen("Float64")
|
|
float16x2 = func_gen("Float16x2")
|
|
float32x2 = func_gen("Float32x2")
|
|
float64x2 = func_gen("Float64x2")
|
|
float16x4 = func_gen("Float16x4")
|
|
float32x4 = func_gen("Float32x4")
|
|
float64x4 = func_gen("Float64x4")
|
|
float16x8 = func_gen("Float16x8")
|
|
float32x8 = func_gen("Float32x8")
|
|
float64x8 = func_gen("Float64x8")
|
|
float16x16 = func_gen("Float16x16")
|
|
float32x16 = func_gen("Float32x16")
|
|
float64x16 = func_gen("Float64x16")
|
|
float16x32 = func_gen("Float16x32")
|
|
float32x32 = func_gen("Float32x32")
|
|
float64x32 = func_gen("Float64x32")
|
|
float16x64 = func_gen("Float16x64")
|
|
float32x64 = func_gen("Float32x64")
|
|
float64x64 = func_gen("Float64x64")
|
|
|
|
# Float8 variants
|
|
float8_e3m4 = func_gen("Float8E3M4")
|
|
float8_e3m4x2 = func_gen("Float8E3M4x2")
|
|
float8_e3m4x4 = func_gen("Float8E3M4x4")
|
|
float8_e3m4x8 = func_gen("Float8E3M4x8")
|
|
float8_e3m4x16 = func_gen("Float8E3M4x16")
|
|
float8_e3m4x32 = func_gen("Float8E3M4x32")
|
|
float8_e3m4x64 = func_gen("Float8E3M4x64")
|
|
|
|
float8_e4m3 = func_gen("Float8E4M3")
|
|
float8_e4m3x2 = func_gen("Float8E4M3x2")
|
|
float8_e4m3x4 = func_gen("Float8E4M3x4")
|
|
float8_e4m3x8 = func_gen("Float8E4M3x8")
|
|
float8_e4m3x16 = func_gen("Float8E4M3x16")
|
|
float8_e4m3x32 = func_gen("Float8E4M3x32")
|
|
float8_e4m3x64 = func_gen("Float8E4M3x64")
|
|
|
|
float8_e4m3b11fnuz = func_gen("Float8E4M3B11FNUZ")
|
|
float8_e4m3b11fnuzx2 = func_gen("Float8E4M3B11FNUZx2")
|
|
float8_e4m3b11fnuzx4 = func_gen("Float8E4M3B11FNUZx4")
|
|
float8_e4m3b11fnuzx8 = func_gen("Float8E4M3B11FNUZx8")
|
|
float8_e4m3b11fnuzx16 = func_gen("Float8E4M3B11FNUZx16")
|
|
float8_e4m3b11fnuzx32 = func_gen("Float8E4M3B11FNUZx32")
|
|
float8_e4m3b11fnuzx64 = func_gen("Float8E4M3B11FNUZx64")
|
|
|
|
float8_e4m3fn = func_gen("Float8E4M3FN")
|
|
float8_e4m3fnx2 = func_gen("Float8E4M3FNx2")
|
|
float8_e4m3fnx4 = func_gen("Float8E4M3FNx4")
|
|
float8_e4m3fnx8 = func_gen("Float8E4M3FNx8")
|
|
float8_e4m3fnx16 = func_gen("Float8E4M3FNx16")
|
|
float8_e4m3fnx32 = func_gen("Float8E4M3FNx32")
|
|
float8_e4m3fnx64 = func_gen("Float8E4M3FNx64")
|
|
|
|
float8_e4m3fnuz = func_gen("Float8E4M3FNUZ")
|
|
float8_e4m3fnuzx2 = func_gen("Float8E4M3FNUZx2")
|
|
float8_e4m3fnuzx4 = func_gen("Float8E4M3FNUZx4")
|
|
float8_e4m3fnuzx8 = func_gen("Float8E4M3FNUZx8")
|
|
float8_e4m3fnuzx16 = func_gen("Float8E4M3FNUZx16")
|
|
float8_e4m3fnuzx32 = func_gen("Float8E4M3FNUZx32")
|
|
float8_e4m3fnuzx64 = func_gen("Float8E4M3FNUZx64")
|
|
|
|
float8_e5m2 = func_gen("Float8E5M2")
|
|
float8_e5m2x2 = func_gen("Float8E5M2x2")
|
|
float8_e5m2x4 = func_gen("Float8E5M2x4")
|
|
float8_e5m2x8 = func_gen("Float8E5M2x8")
|
|
float8_e5m2x16 = func_gen("Float8E5M2x16")
|
|
float8_e5m2x32 = func_gen("Float8E5M2x32")
|
|
float8_e5m2x64 = func_gen("Float8E5M2x64")
|
|
|
|
float8_e5m2fnuz = func_gen("Float8E5M2FNUZ")
|
|
float8_e5m2fnuzx2 = func_gen("Float8E5M2FNUZx2")
|
|
float8_e5m2fnuzx4 = func_gen("Float8E5M2FNUZx4")
|
|
float8_e5m2fnuzx8 = func_gen("Float8E5M2FNUZx8")
|
|
float8_e5m2fnuzx16 = func_gen("Float8E5M2FNUZx16")
|
|
float8_e5m2fnuzx32 = func_gen("Float8E5M2FNUZx32")
|
|
float8_e5m2fnuzx64 = func_gen("Float8E5M2FNUZx64")
|
|
|
|
float8_e8m0fnu = func_gen("Float8E8M0FNU")
|
|
float8_e8m0fnux2 = func_gen("Float8E8M0FNUx2")
|
|
float8_e8m0fnux4 = func_gen("Float8E8M0FNUx4")
|
|
float8_e8m0fnux8 = func_gen("Float8E8M0FNUx8")
|
|
float8_e8m0fnux16 = func_gen("Float8E8M0FNUx16")
|
|
float8_e8m0fnux32 = func_gen("Float8E8M0FNUx32")
|
|
float8_e8m0fnux64 = func_gen("Float8E8M0FNUx64")
|
|
|
|
# Float6 variants
|
|
float6_e2m3fn = func_gen("Float6E2M3FN")
|
|
float6_e2m3fnx2 = func_gen("Float6E2M3FNx2")
|
|
float6_e2m3fnx4 = func_gen("Float6E2M3FNx4")
|
|
float6_e2m3fnx8 = func_gen("Float6E2M3FNx8")
|
|
float6_e2m3fnx16 = func_gen("Float6E2M3FNx16")
|
|
float6_e2m3fnx32 = func_gen("Float6E2M3FNx32")
|
|
float6_e2m3fnx64 = func_gen("Float6E2M3FNx64")
|
|
|
|
float6_e3m2fn = func_gen("Float6E3M2FN")
|
|
float6_e3m2fnx2 = func_gen("Float6E3M2FNx2")
|
|
float6_e3m2fnx4 = func_gen("Float6E3M2FNx4")
|
|
float6_e3m2fnx8 = func_gen("Float6E3M2FNx8")
|
|
float6_e3m2fnx16 = func_gen("Float6E3M2FNx16")
|
|
float6_e3m2fnx32 = func_gen("Float6E3M2FNx32")
|
|
float6_e3m2fnx64 = func_gen("Float6E3M2FNx64")
|
|
|
|
# Float4 variants
|
|
float4_e2m1fn = func_gen("Float4E2M1FN")
|
|
float4_e2m1fnx2 = func_gen("Float4E2M1FNx2")
|
|
float4_e2m1fnx4 = func_gen("Float4E2M1FNx4")
|
|
float4_e2m1fnx8 = func_gen("Float4E2M1FNx8")
|
|
float4_e2m1fnx16 = func_gen("Float4E2M1FNx16")
|
|
float4_e2m1fnx32 = func_gen("Float4E2M1FNx32")
|
|
float4_e2m1fnx64 = func_gen("Float4E2M1FNx64")
|
|
|
|
bfloat16 = func_gen("BFloat16")
|
|
|
|
# Shorthand aliases
|
|
f16 = float16
|
|
f32 = float32
|
|
f64 = float64
|
|
bf16 = bfloat16
|
|
i8 = int8
|
|
i16 = int16
|
|
i32 = int32
|
|
i64 = int64
|
|
u8 = uint8
|
|
u16 = uint16
|
|
u32 = uint32
|
|
u64 = uint64
|
|
# pylint: enable=invalid-name
|
|
|
|
|
|
def boolean(expr: Expr | None = None) -> Expr:
|
|
"""Construct a new tirx.Var with type boolean or cast expression to type boolean.
|
|
|
|
Parameters
|
|
----------
|
|
expr: Expr
|
|
The expression to be cast.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The new tirx.Var with type boolean or casted expression with type boolean.
|
|
"""
|
|
return _ffi_api.Boolean(expr) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def handle(
|
|
dtype: str | None = None,
|
|
storage_scope: str = "global",
|
|
) -> Var:
|
|
"""Create a TIR var that represents a pointer.
|
|
|
|
Parameters
|
|
----------
|
|
dtype: str | None
|
|
The data type of the pointer. If omitted, construct an opaque handle.
|
|
|
|
storage_scope: str
|
|
The storage scope of the pointer.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The new tirx.Var with type handle or casted expression with type handle.
|
|
"""
|
|
if dtype in ("TensorMap", "tensormap", "CUtensorMap", "cuTensorMap"):
|
|
return _ffi_api.TensorMap() # type: ignore[attr-defined] # pylint: disable=no-member
|
|
return _ffi_api.Handle( # type: ignore[attr-defined] # pylint: disable=no-member
|
|
dtype,
|
|
storage_scope,
|
|
)
|
|
|
|
|
|
def TensorMap() -> Var: # pylint: disable=invalid-name
|
|
"""Create a TIRx var that represents a CUDA tensor-map descriptor.
|
|
|
|
The host/runtime ABI passes a handle to descriptor storage. CUDA kernel
|
|
codegen lowers this type to ``const __grid_constant__ CUtensorMap`` when it
|
|
appears as a kernel parameter.
|
|
"""
|
|
return _ffi_api.TensorMap() # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def void(expr: Expr | None = None) -> Expr:
|
|
"""Construct a new tirx.Var with type void or cast expression to type void.
|
|
|
|
Parameters
|
|
----------
|
|
expr: Expr
|
|
The expression to be cast.
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The new tirx.Var with type void or casted expression with type void.
|
|
"""
|
|
return _ffi_api.Void(expr) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
@deprecated("T.var", "T.{dtype}")
|
|
def var(dtype: str, name: str = "") -> Var:
|
|
"""Construct a new tirx.Var.
|
|
|
|
Parameters
|
|
----------
|
|
dtype: str
|
|
The dtype of the Var.
|
|
|
|
name: str
|
|
The name of the Var.
|
|
|
|
Returns
|
|
-------
|
|
res : Var
|
|
The result tirx.Var.
|
|
"""
|
|
return Var(name, dtype) # pylint: disable=no-member
|
|
|
|
|
|
def ptr(dtype: str, storage_scope: str = "global") -> Var:
|
|
"""The pointer declaration function.
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type of the pointer.
|
|
|
|
storage_scope : str
|
|
The storage scope of the pointer.
|
|
|
|
Returns
|
|
-------
|
|
res : Var
|
|
The pointer.
|
|
"""
|
|
return _ffi_api.Ptr(dtype, storage_scope) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
@deprecated("T.buffer_var", "T.handle")
|
|
def buffer_var(dtype: str, storage_scope: str = "global") -> Var:
|
|
"""The pointer declaration function.
|
|
|
|
Parameters
|
|
----------
|
|
dtype : str
|
|
The data type of the pointer.
|
|
|
|
storage_scope : str
|
|
The storage scope of the pointer.
|
|
|
|
Returns
|
|
-------
|
|
res : Var
|
|
The pointer.
|
|
"""
|
|
return _ffi_api.Ptr(dtype, storage_scope) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def min(a: Expr, b: Expr) -> Expr: # pylint: disable=redefined-builtin
|
|
"""Compute the minimum value of two expressions.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand
|
|
|
|
b : Expr
|
|
The right hand operand
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result expression.
|
|
"""
|
|
return _ffi_api.min(a, b) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def max(a: Expr, b: Expr) -> Expr: # pylint: disable=redefined-builtin
|
|
"""Compute the maximum value of two expressions.
|
|
|
|
Parameters
|
|
----------
|
|
a : Expr
|
|
The left hand operand
|
|
|
|
b : Expr
|
|
The right hand operand
|
|
|
|
Returns
|
|
-------
|
|
res : Expr
|
|
The result expression.
|
|
"""
|
|
return _ffi_api.max(a, b) # type: ignore[attr-defined] # pylint: disable=no-member
|
|
|
|
|
|
def iter_var(v: Var | str, dom: ir.Range, iter_type: str, thread_tag: str) -> IterVar:
|
|
"""The iteration variable.
|
|
|
|
Parameters
|
|
----------
|
|
var : Union[Var, str]
|
|
The internal variable that is used for iteration.
|
|
|
|
dom : Range
|
|
The domain of the iteration.
|
|
|
|
iter_type : str
|
|
The iteration type.
|
|
|
|
thread_tag : str
|
|
The thread type tag.
|
|
|
|
Returns
|
|
-------
|
|
res : IterVar
|
|
The iteration variable.
|
|
"""
|
|
iter_type = getattr(IterVar, iter_type)
|
|
return IterVar(dom, v, iter_type, thread_tag)
|
|
|
|
|
|
def comm_reducer(combiner: Callable, identity: list[Expr]) -> CommReducer:
|
|
"""
|
|
Create a CommReducer from lambda inputs/outputs and the identities
|
|
|
|
Parameters
|
|
----------
|
|
combiner : Callable
|
|
A binary function which takes two Expr as input to return a Expr.
|
|
|
|
identity : List[Expr]
|
|
A list of types of output Expr.
|
|
|
|
Returns
|
|
-------
|
|
res : CommReducer
|
|
The CommReducer.
|
|
"""
|
|
params = inspect.signature(combiner).parameters
|
|
num_args = len(params)
|
|
args = []
|
|
for name, i in zip(params.keys(), identity + identity):
|
|
if isinstance(i, int):
|
|
args.append(Var(name, "int32"))
|
|
else:
|
|
args.append(Var(name, i.ty))
|
|
res = combiner(*args)
|
|
if not isinstance(res, tuple):
|
|
res = (res,)
|
|
return CommReducer(args[: num_args // 2], args[num_args // 2 :], res, identity)
|
|
|
|
|
|
def index_map(
|
|
mapping: Callable,
|
|
*,
|
|
inverse_index_map: Callable | None = None,
|
|
index_dtype: str = "int64",
|
|
) -> IndexMap:
|
|
"""Create a TIR Index mapping"""
|
|
return IndexMap.from_func(mapping, inverse_index_map=inverse_index_map, index_dtype=index_dtype)
|
|
|
|
|
|
def target(
|
|
target_config: dict | str,
|
|
host: dict | str | Target | None = None,
|
|
) -> Target:
|
|
"""
|
|
Create a target
|
|
|
|
Parameters
|
|
----------
|
|
target_config : Union[Dict, str]
|
|
The target configuration.
|
|
|
|
host : Optional[Union[Dict, str, Target]]
|
|
The target configuration.
|
|
|
|
Returns
|
|
-------
|
|
res : Target
|
|
The target.
|
|
"""
|
|
if not isinstance(target_config, str | dict):
|
|
raise ValueError(
|
|
f"T.target expected a config dict or string, but got {type(target_config)}"
|
|
)
|
|
if host is not None and not isinstance(host, str | dict | Target):
|
|
raise ValueError(
|
|
"T.target expected the host to be "
|
|
"a config dict, string, or T.target, "
|
|
f"but got {type(host)}"
|
|
)
|
|
if isinstance(target_config, dict) and "host" in target_config and host is not None:
|
|
raise ValueError(
|
|
"T.target expects to either receive the host "
|
|
"as part of the target's config dictionary, "
|
|
"or as a separate argument, but not both."
|
|
)
|
|
return Target(target_config, host)
|
|
|
|
|
|
def Range(begin: Expr, end: Expr) -> ir.Range: # pylint: disable=invalid-name
|
|
"""
|
|
Create a Range object.
|
|
|
|
Parameters
|
|
----------
|
|
begin : Expr
|
|
The begin value of the range.
|
|
|
|
end : Optional[Expr]
|
|
The end value of the range.
|
|
"""
|
|
return ir.Range(begin, end)
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
T = TypeVar("T")
|
|
C = TypeVar("C")
|
|
|
|
# When type checking (and by extension, for linters like Pylint), treat
|
|
# meta_var as an identity function.
|
|
def meta_var(x: T) -> T:
|
|
return x
|
|
|
|
def meta_class(cls: C) -> C:
|
|
return cls
|
|
|
|
else:
|
|
|
|
def _install_meta_class(cls):
|
|
if cls.__dict__.get("_tirx_meta_class_installed", False):
|
|
cls._is_meta_class = True
|
|
return cls
|
|
|
|
original_init = getattr(cls, "__init__", object.__init__)
|
|
original_setattr = getattr(cls, "__setattr__", object.__setattr__)
|
|
original_init_subclass = getattr(cls, "__init_subclass__", None)
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
with _with_meta_construction_scope(self, type(self)) as scope:
|
|
original_init(self, *args, **kwargs)
|
|
_validate_meta_construction_scope(scope)
|
|
|
|
def __setattr__(self, name, value):
|
|
if isinstance(value, scalar_wrapper):
|
|
value = value.scalar
|
|
original_setattr(self, name, value)
|
|
|
|
@classmethod
|
|
def __init_subclass__(subcls, **kwargs):
|
|
if original_init_subclass is not None:
|
|
original_init_subclass(**kwargs)
|
|
_install_meta_class(subcls)
|
|
|
|
cls.__init__ = __init__
|
|
cls.__setattr__ = __setattr__
|
|
cls.__init_subclass__ = __init_subclass__
|
|
cls._is_meta_class = True
|
|
cls._tirx_meta_class_installed = True
|
|
return cls
|
|
|
|
def meta_class(cls):
|
|
"""Decorator for utility classes used inside @T.prim_func.
|
|
|
|
Instances of decorated classes are treated as parser meta values.
|
|
"""
|
|
return _install_meta_class(cls)
|
|
|
|
class meta_var:
|
|
"""A meta variable used in TVMScript metaprogramming.
|
|
|
|
The value does not appear in the final TIR and only exists in the parser.
|
|
|
|
Parameters
|
|
----------
|
|
value: Any
|
|
The meta variable.
|
|
"""
|
|
|
|
def __init__(self, value: Any) -> None:
|
|
self.value = value
|
|
|
|
def __iter__(self):
|
|
# Return a generator that yields wrapped items.
|
|
return (meta_var(i) for i in self.value)
|
|
|
|
|
|
# pylint: disable=invalid-name
|
|
|
|
|
|
T = TypeVar("T")
|
|
P = ParamSpec("P")
|
|
|
|
|
|
def _op_wrapper(func: Callable[P, T]) -> Callable[P, T]:
|
|
@functools.wraps(func)
|
|
def wrapped(*args, **kwargs) -> T:
|
|
if "dtype" in kwargs:
|
|
kwargs.pop("dtype")
|
|
return func(*args, **kwargs)
|
|
|
|
# Expose underlying tir op name for printer registration
|
|
try:
|
|
wrapped.__tir_op_name__ = getattr(func, "__name__", None)
|
|
except Exception: # pragma: no cover
|
|
pass
|
|
return wrapped
|
|
|
|
|
|
def _dtype_forward(func):
|
|
@functools.wraps(func)
|
|
def wrapped(*args, **kwargs):
|
|
if "dtype" in kwargs:
|
|
args = (kwargs.pop("dtype"), *args)
|
|
return func(*args, **kwargs)
|
|
|
|
# Expose underlying tir op name for printer registration
|
|
try:
|
|
wrapped.__tir_op_name__ = getattr(func, "__name__", None)
|
|
except Exception: # pragma: no cover
|
|
pass
|
|
return wrapped
|
|
|
|
|
|
class WebGPUNamespace:
|
|
"""The WebGPU intrinsics submodule."""
|
|
|
|
@staticmethod
|
|
def subgroup_shuffle(var, lane):
|
|
if isinstance(var, Buffer):
|
|
var = var[0]
|
|
return _tir_op.call_intrin(var.ty, "tirx.webgpu.subgroup_shuffle", var, lane)
|
|
|
|
@staticmethod
|
|
def subgroup_shuffle_up(var, delta):
|
|
if isinstance(var, Buffer):
|
|
var = var[0]
|
|
return _tir_op.call_intrin(var.ty, "tirx.webgpu.subgroup_shuffle_up", var, delta)
|
|
|
|
@staticmethod
|
|
def subgroup_shuffle_down(var, delta):
|
|
if isinstance(var, Buffer):
|
|
var = var[0]
|
|
return _tir_op.call_intrin(var.ty, "tirx.webgpu.subgroup_shuffle_down", var, delta)
|
|
|
|
|
|
webgpu = WebGPUNamespace()
|
|
|
|
|
|
#
|
|
# Register printer namespace mapping from the builder namespaces so the
|
|
# TVMScript printer emits dotted names that match parser namespaces.
|
|
#
|
|
def _register_script_namespace_printer_names(ns_obj, dotted_prefix):
|
|
def register_printer_name(op_name, script_name):
|
|
try:
|
|
ir.Op.get(op_name)
|
|
except Exception:
|
|
return
|
|
try:
|
|
_register_op_attr(op_name, "TScriptPrinterName", script_name, level=20)
|
|
except Exception:
|
|
pass
|
|
|
|
def visit(ns_obj, dotted_prefix):
|
|
# If the namespace object itself maps to an op via __call__
|
|
call_op = getattr(ns_obj, "__tir_call_op_name__", None)
|
|
if call_op:
|
|
flat_name = f"tirx.{call_op}"
|
|
for op_name in {flat_name, _tir_op._canonical_device_intrin_name(flat_name)}:
|
|
register_printer_name(op_name, dotted_prefix)
|
|
# Walk attributes to find wrapped ops and sub-namespaces
|
|
for name in dir(ns_obj):
|
|
if name.startswith("_"):
|
|
continue
|
|
try:
|
|
val = getattr(ns_obj, name)
|
|
except Exception:
|
|
continue
|
|
# Sub-namespace: recurse
|
|
if hasattr(val, "__dict__") and val.__class__.__name__.endswith("Namespace"):
|
|
visit(val, f"{dotted_prefix}.{name}")
|
|
continue
|
|
# Wrapped op (callable with attached __tir_op_name__)
|
|
op_name = getattr(val, "__tir_op_name__", None)
|
|
if callable(val) and op_name:
|
|
flat_name = f"tirx.{op_name}"
|
|
script_name = f"{dotted_prefix}.{name}"
|
|
for full_op_name in {flat_name, _tir_op._canonical_device_intrin_name(flat_name)}:
|
|
register_printer_name(full_op_name, script_name)
|
|
|
|
visit(ns_obj, dotted_prefix)
|
|
|
|
|
|
def register_script_namespace(name: str, namespace: object) -> object:
|
|
"""Register a TVMScript namespace on the TIRx builder facade."""
|
|
globals()[name] = namespace
|
|
if "__all__" in globals() and name not in __all__:
|
|
__all__.append(name)
|
|
|
|
import sys # pylint: disable=import-outside-toplevel
|
|
|
|
for module_name in [
|
|
"tvm.tirx.script.builder",
|
|
"tvm.tirx.script.parser",
|
|
"tvm.tirx.script",
|
|
"tvm.script.tirx",
|
|
]:
|
|
module = sys.modules.get(module_name)
|
|
if module is None:
|
|
continue
|
|
setattr(module, name, namespace)
|
|
module_all = getattr(module, "__all__", None)
|
|
if isinstance(module_all, list) and name not in module_all:
|
|
module_all.append(name)
|
|
|
|
_register_script_namespace_printer_names(namespace, name)
|
|
return namespace
|
|
|
|
|
|
def _register_tir_namespace_printer_names():
|
|
try:
|
|
_register_script_namespace_printer_names(webgpu, "webgpu")
|
|
except Exception:
|
|
# Best-effort registration; avoid import-time hard failure
|
|
pass
|
|
|
|
|
|
# Execute registration on import so printer picks up dotted names
|
|
_register_tir_namespace_printer_names()
|
|
|
|
|
|
abs = _op_wrapper(_tir_op.abs) # pylint: disable=redefined-builtin
|
|
acos = _op_wrapper(_tir_op.acos)
|
|
acosh = _op_wrapper(_tir_op.acosh)
|
|
address_of = _op_wrapper(_tir_op.address_of)
|
|
asin = _op_wrapper(_tir_op.asin)
|
|
asinh = _op_wrapper(_tir_op.asinh)
|
|
atan = _op_wrapper(_tir_op.atan)
|
|
atan2 = _op_wrapper(_tir_op.atan2)
|
|
atanh = _op_wrapper(_tir_op.atanh)
|
|
bitwise_and = _op_wrapper(_tir_op.bitwise_and)
|
|
bitwise_not = _op_wrapper(_tir_op.bitwise_not)
|
|
bitwise_or = _op_wrapper(_tir_op.bitwise_or)
|
|
bitwise_xor = _op_wrapper(_tir_op.bitwise_xor)
|
|
ceil = _op_wrapper(_tir_op.ceil)
|
|
clz = _op_wrapper(_tir_op.clz)
|
|
copysign = _op_wrapper(_tir_op.copysign)
|
|
cos = _op_wrapper(_tir_op.cos)
|
|
cosh = _op_wrapper(_tir_op.cosh)
|
|
erf = _op_wrapper(_tir_op.erf)
|
|
exp = _op_wrapper(_tir_op.exp)
|
|
exp2 = _op_wrapper(_tir_op.exp2)
|
|
exp10 = _op_wrapper(_tir_op.exp10)
|
|
filter = _op_wrapper(_tir_op.filter) # pylint: disable=redefined-builtin
|
|
selector = _op_wrapper(_tir_op.selector)
|
|
floor = _op_wrapper(_tir_op.floor)
|
|
ceildiv = _op_wrapper(_tir_op.ceildiv)
|
|
floordiv = _op_wrapper(_tir_op.floordiv)
|
|
floormod = _op_wrapper(_tir_op.floormod)
|
|
fmod = _op_wrapper(_tir_op.fmod)
|
|
fma = _op_wrapper(_tir_op.fma)
|
|
hypot = _op_wrapper(_tir_op.hypot)
|
|
if_then_else = _op_wrapper(_tir_op.if_then_else)
|
|
infinity = _op_wrapper(_tir_op.infinity)
|
|
isfinite = _op_wrapper(_tir_op.isfinite)
|
|
isinf = _op_wrapper(_tir_op.isinf)
|
|
isnan = _op_wrapper(_tir_op.isnan)
|
|
isnullptr = _op_wrapper(_tir_op.isnullptr)
|
|
ldexp = _op_wrapper(_tir_op.ldexp)
|
|
likely = _op_wrapper(_tir_op.likely)
|
|
log = _op_wrapper(_tir_op.log)
|
|
log1p = _op_wrapper(_tir_op.log1p)
|
|
log2 = _op_wrapper(_tir_op.log2)
|
|
log10 = _op_wrapper(_tir_op.log10)
|
|
lookup_param = _op_wrapper(_tir_op.lookup_param)
|
|
max_value = _op_wrapper(_tir_op.max_value)
|
|
min_value = _op_wrapper(_tir_op.min_value)
|
|
nearbyint = _op_wrapper(_tir_op.nearbyint)
|
|
nextafter = _op_wrapper(_tir_op.nextafter)
|
|
popcount = _op_wrapper(_tir_op.popcount)
|
|
pow = _op_wrapper(_tir_op.pow) # pylint: disable=redefined-builtin
|
|
q_multiply_shift = _op_wrapper(_tir_op.q_multiply_shift)
|
|
q_multiply_shift_per_axis = _op_wrapper(_tir_op.q_multiply_shift_per_axis)
|
|
ret = _op_wrapper(_tir_op.ret)
|
|
continue_loop = _op_wrapper(_tir_op.continue_loop)
|
|
break_loop = _op_wrapper(_tir_op.break_loop)
|
|
round = _op_wrapper(_tir_op.round) # pylint: disable=redefined-builtin
|
|
rsqrt = _op_wrapper(_tir_op.rsqrt)
|
|
shift_left = _op_wrapper(_tir_op.shift_left)
|
|
shift_right = _op_wrapper(_tir_op.shift_right)
|
|
sigmoid = _op_wrapper(_tir_op.sigmoid)
|
|
sin = _op_wrapper(_tir_op.sin)
|
|
sinh = _op_wrapper(_tir_op.sinh)
|
|
sqrt = _op_wrapper(_tir_op.sqrt)
|
|
tan = _op_wrapper(_tir_op.tan)
|
|
tanh = _op_wrapper(_tir_op.tanh)
|
|
thread_return = _op_wrapper(_tir_op.thread_return)
|
|
trunc = _op_wrapper(_tir_op.trunc)
|
|
truncdiv = _op_wrapper(_tir_op.truncdiv)
|
|
truncmod = _op_wrapper(_tir_op.truncmod)
|
|
tvm_access_ptr = _op_wrapper(_tir_op.tvm_access_ptr)
|
|
ptr_byte_offset = _op_wrapper(_tir_op.ptr_byte_offset)
|
|
tvm_throw_last_error = _op_wrapper(_tir_op.tvm_throw_last_error)
|
|
print_buffer = _op_wrapper(_tir_op.print_buffer)
|
|
tvm_stack_alloca = _op_wrapper(_tir_op.tvm_stack_alloca)
|
|
tvm_stack_make_shape = _op_wrapper(_tir_op.tvm_stack_make_shape)
|
|
tvm_stack_make_array = _op_wrapper(_tir_op.tvm_stack_make_array)
|
|
call_packed = _op_wrapper(_tir_op.call_packed)
|
|
call_cpacked = _op_wrapper(_tir_op.call_cpacked)
|
|
call_packed_lowered = _op_wrapper(_tir_op.call_packed_lowered)
|
|
call_cpacked_lowered = _op_wrapper(_tir_op.call_cpacked_lowered)
|
|
tvm_tuple = _op_wrapper(_tir_op.tvm_tuple)
|
|
handle_add_byte_offset = _op_wrapper(_tir_op.handle_add_byte_offset)
|
|
tvm_struct_set = _op_wrapper(_tir_op.tvm_struct_set)
|
|
tvm_struct_get = _tir_op.tvm_struct_get
|
|
tvm_thread_invariant = _op_wrapper(_tir_op.tvm_thread_invariant)
|
|
tvm_thread_allreduce = _op_wrapper(_tir_op.tvm_thread_allreduce)
|
|
tvm_load_matrix_sync = _op_wrapper(_tir_op.tvm_load_matrix_sync)
|
|
tvm_mma_sync = _op_wrapper(_tir_op.tvm_mma_sync)
|
|
tvm_bmma_sync = _op_wrapper(_tir_op.tvm_bmma_sync)
|
|
tvm_fill_fragment = _op_wrapper(_tir_op.tvm_fill_fragment)
|
|
tvm_store_matrix_sync = _op_wrapper(_tir_op.tvm_store_matrix_sync)
|
|
tvm_storage_sync = _tir_op.tvm_storage_sync
|
|
tvm_kernel_replace_point = _op_wrapper(_tir_op.tvm_kernel_replace_point)
|
|
tvm_global_barrier_kinit = _tir_op.tvm_global_barrier_kinit
|
|
tvm_warp_shuffle = _tir_op.tvm_warp_shuffle
|
|
tvm_warp_shuffle_up = _tir_op.tvm_warp_shuffle_up
|
|
tvm_warp_shuffle_down = _tir_op.tvm_warp_shuffle_down
|
|
tvm_warp_shuffle_xor = _tir_op.tvm_warp_shuffle_xor
|
|
tvm_warp_activemask = _tir_op.tvm_warp_activemask
|
|
cooperative_tensor_fill = _op_wrapper(_tir_op.cooperative_tensor_fill)
|
|
cooperative_tensor_load = _op_wrapper(_tir_op.cooperative_tensor_load)
|
|
cooperative_tensor_store = _op_wrapper(_tir_op.cooperative_tensor_store)
|
|
cooperative_tensor_multiply_accumulate = _op_wrapper(_tir_op.cooperative_tensor_multiply_accumulate)
|
|
assume = _op_wrapper(_tir_op.assume)
|
|
undef = _op_wrapper(_tir_op.undef)
|
|
TVMBackendAllocWorkspace = _op_wrapper(_tir_op.TVMBackendAllocWorkspace)
|
|
TVMBackendFreeWorkspace = _op_wrapper(_tir_op.TVMBackendFreeWorkspace)
|
|
start_profile_intrinsic = _op_wrapper(_tir_op.start_profile_intrinsic)
|
|
end_profile_intrinsic = _op_wrapper(_tir_op.end_profile_intrinsic)
|
|
anylist_getitem = _op_wrapper(_tir_op.anylist_getitem)
|
|
anylist_resetitem = _op_wrapper(_tir_op.anylist_resetitem)
|
|
anylist_setitem_call_packed = _op_wrapper(_tir_op.anylist_setitem_call_packed)
|
|
anylist_setitem_call_cpacked = _op_wrapper(_tir_op.anylist_setitem_call_cpacked)
|
|
vscale = _op_wrapper(_tir_op.vscale)
|
|
ignore_loop_partition = _op_wrapper(_tir_op.ignore_loop_partition)
|
|
|
|
reinterpret = _dtype_forward(_tir_op.reinterpret)
|
|
call_extern = _dtype_forward(_tir_op.call_extern)
|
|
call_intrin = _dtype_forward(_tir_op.call_intrin)
|
|
call_llvm_intrin = _dtype_forward(_tir_op.call_llvm_intrin)
|
|
call_llvm_pure_intrin = _dtype_forward(_tir_op.call_llvm_pure_intrin)
|
|
call_pure_extern = _dtype_forward(_tir_op.call_pure_extern)
|
|
vectorlow = _dtype_forward(_tir_op.vectorlow)
|
|
vectorhigh = _dtype_forward(_tir_op.vectorhigh)
|
|
vectorcombine = _dtype_forward(_tir_op.vectorcombine)
|
|
get_active_lane_mask = _dtype_forward(_tir_op.get_active_lane_mask)
|
|
dp4a = _dtype_forward(_tir_op.dp4a)
|
|
|
|
|
|
broadcast = Broadcast
|
|
ramp = Ramp
|
|
fabs = abs
|
|
tvm_call_packed = call_packed
|
|
tvm_call_cpacked = call_cpacked
|
|
tvm_call_packed_lowered = call_packed_lowered
|
|
tvm_call_cpacked_lowered = call_cpacked_lowered
|
|
|
|
# pylint: enable=invalid-name
|
|
|
|
bases = [
|
|
"float8_e3m4",
|
|
"float8_e4m3",
|
|
"float8_e4m3b11fnuz",
|
|
"float8_e4m3fn",
|
|
"float8_e4m3fnuz",
|
|
"float8_e5m2",
|
|
"float8_e5m2fnuz",
|
|
"float8_e8m0fnu",
|
|
"float6_e2m3fn",
|
|
"float6_e3m2fn",
|
|
"float4_e2m1fn",
|
|
"float16",
|
|
"float32",
|
|
"float64",
|
|
]
|
|
lanes = [1, 2, 4, 8, 16, 32, 64]
|
|
|
|
float_types = []
|
|
for base in bases:
|
|
for lane in lanes:
|
|
suffix = f"x{lane}" if lane != 1 else ""
|
|
float_types.append(f"{base}{suffix}")
|
|
|
|
__all__ = [
|
|
*float_types,
|
|
"int8",
|
|
"int16",
|
|
"int32",
|
|
"int64",
|
|
"int8x2",
|
|
"int16x2",
|
|
"int32x2",
|
|
"int64x2",
|
|
"int8x4",
|
|
"int16x4",
|
|
"int32x4",
|
|
"int64x4",
|
|
"int8x8",
|
|
"int16x8",
|
|
"int32x8",
|
|
"int64x8",
|
|
"int8x16",
|
|
"int16x16",
|
|
"int32x16",
|
|
"int64x16",
|
|
"int8x32",
|
|
"int16x32",
|
|
"int32x32",
|
|
"int64x32",
|
|
"int8x64",
|
|
"int16x64",
|
|
"int32x64",
|
|
"int64x64",
|
|
"uint8",
|
|
"uint16",
|
|
"uint32",
|
|
"uint64",
|
|
"uint8x2",
|
|
"uint16x2",
|
|
"uint32x2",
|
|
"uint64x2",
|
|
"uint8x4",
|
|
"uint16x4",
|
|
"uint32x4",
|
|
"uint64x4",
|
|
"uint8x8",
|
|
"uint16x8",
|
|
"uint32x8",
|
|
"uint64x8",
|
|
"uint8x16",
|
|
"uint16x16",
|
|
"uint32x16",
|
|
"uint64x16",
|
|
"uint8x32",
|
|
"uint16x32",
|
|
"uint32x32",
|
|
"uint64x32",
|
|
"uint8x64",
|
|
"uint16x64",
|
|
"uint32x64",
|
|
"uint64x64",
|
|
"float8_e4m3fn",
|
|
"float8_e5m2",
|
|
"float4_e2m1fn",
|
|
"float16",
|
|
"float32",
|
|
"float64",
|
|
"float4_e2m1fnx2",
|
|
"float8_e4m3fnx4",
|
|
"float8_e5m2x4",
|
|
"float4_e2m1fnx4",
|
|
"float16x2",
|
|
"float32x2",
|
|
"float64x2",
|
|
"float16x4",
|
|
"float32x4",
|
|
"float64x4",
|
|
"float8_e4m3fnx8",
|
|
"float8_e5m2x8",
|
|
"float4_e2m1fnx8",
|
|
"float16x8",
|
|
"float32x8",
|
|
"float64x8",
|
|
"float8_e4m3fnx16",
|
|
"float8_e5m2x16",
|
|
"float4_e2m1fnx16",
|
|
"float16x16",
|
|
"float32x16",
|
|
"float64x16",
|
|
"float8_e4m3fnx32",
|
|
"float8_e5m2x32",
|
|
"float4_e2m1fnx32",
|
|
"float16x32",
|
|
"float32x32",
|
|
"float64x32",
|
|
"float8_e4m3fnx64",
|
|
"float8_e5m2x64",
|
|
"float4_e2m1fnx64",
|
|
"float16x64",
|
|
"float32x64",
|
|
"float64x64",
|
|
"bfloat16",
|
|
"buffer",
|
|
"buffer_decl",
|
|
"prim_func",
|
|
"arg",
|
|
"func_name",
|
|
"func_attr",
|
|
"func_ret",
|
|
"match_buffer",
|
|
"sblock",
|
|
"block_name_suffix_context",
|
|
"init",
|
|
"where",
|
|
"reads",
|
|
"writes",
|
|
"sblock_attr",
|
|
"alloc_buffer",
|
|
"sblock_alloc_buffer",
|
|
"wg_reg_tile",
|
|
"axis",
|
|
"serial",
|
|
"parallel",
|
|
"vectorized",
|
|
"unroll",
|
|
"thread_binding",
|
|
"grid",
|
|
"Assert",
|
|
"attr",
|
|
"hint",
|
|
"While",
|
|
"Break",
|
|
"Continue",
|
|
"If",
|
|
"Then",
|
|
"Else",
|
|
"decl_buffer",
|
|
"launch_thread",
|
|
"env_thread",
|
|
"buffer_store",
|
|
"evaluate",
|
|
"boolean",
|
|
"handle",
|
|
"void",
|
|
"var",
|
|
"ptr",
|
|
"min",
|
|
"max",
|
|
"iter_var",
|
|
"comm_reducer",
|
|
"index_map",
|
|
"target",
|
|
"buffer_var",
|
|
"abs",
|
|
"fabs",
|
|
"acos",
|
|
"acosh",
|
|
"address_of",
|
|
"asin",
|
|
"asinh",
|
|
"atan",
|
|
"atan2",
|
|
"atanh",
|
|
"bitwise_and",
|
|
"bitwise_not",
|
|
"bitwise_or",
|
|
"bitwise_xor",
|
|
"ceil",
|
|
"clz",
|
|
"copysign",
|
|
"cos",
|
|
"cosh",
|
|
"erf",
|
|
"exp",
|
|
"exp2",
|
|
"exp10",
|
|
"floor",
|
|
"ceildiv",
|
|
"floordiv",
|
|
"floormod",
|
|
"fmod",
|
|
"fma",
|
|
"filter",
|
|
"selector",
|
|
"hypot",
|
|
"if_then_else",
|
|
"infinity",
|
|
"isfinite",
|
|
"isinf",
|
|
"isnan",
|
|
"isnullptr",
|
|
"ldexp",
|
|
"likely",
|
|
"log",
|
|
"log1p",
|
|
"log2",
|
|
"log10",
|
|
"lookup_param",
|
|
"max_value",
|
|
"min_value",
|
|
"nearbyint",
|
|
"nextafter",
|
|
"popcount",
|
|
"pow",
|
|
"q_multiply_shift",
|
|
"q_multiply_shift_per_axis",
|
|
"ret",
|
|
"continue_loop",
|
|
"break_loop",
|
|
"reinterpret",
|
|
"round",
|
|
"rsqrt",
|
|
"shift_left",
|
|
"shift_right",
|
|
"sigmoid",
|
|
"sin",
|
|
"sinh",
|
|
"sqrt",
|
|
"tan",
|
|
"tanh",
|
|
"thread_return",
|
|
"trunc",
|
|
"truncdiv",
|
|
"truncmod",
|
|
"tvm_access_ptr",
|
|
"ptr_byte_offset",
|
|
"tvm_throw_last_error",
|
|
"print_buffer",
|
|
"tvm_stack_alloca",
|
|
"tvm_stack_make_shape",
|
|
"tvm_stack_make_array",
|
|
"call_packed",
|
|
"call_cpacked",
|
|
"call_packed_lowered",
|
|
"call_cpacked_lowered",
|
|
"call_extern",
|
|
"call_intrin",
|
|
"call_llvm_intrin",
|
|
"call_llvm_pure_intrin",
|
|
"call_pure_extern",
|
|
"tvm_tuple",
|
|
"handle_add_byte_offset",
|
|
"tvm_struct_set",
|
|
"tvm_struct_get",
|
|
"tvm_thread_invariant",
|
|
"tvm_thread_allreduce",
|
|
"tvm_load_matrix_sync",
|
|
"tvm_mma_sync",
|
|
"tvm_bmma_sync",
|
|
"tvm_fill_fragment",
|
|
"tvm_store_matrix_sync",
|
|
"tvm_storage_sync",
|
|
"tvm_kernel_replace_point",
|
|
"tvm_global_barrier_kinit",
|
|
"tvm_warp_shuffle",
|
|
"tvm_warp_shuffle_up",
|
|
"tvm_warp_shuffle_down",
|
|
"tvm_warp_shuffle_xor",
|
|
"tvm_warp_activemask",
|
|
"cooperative_tensor_fill",
|
|
"cooperative_tensor_load",
|
|
"cooperative_tensor_store",
|
|
"cooperative_tensor_multiply_accumulate",
|
|
"vectorlow",
|
|
"vectorhigh",
|
|
"vectorcombine",
|
|
"dp4a",
|
|
"assume",
|
|
"undef",
|
|
"tvm_call_packed",
|
|
"tvm_call_cpacked",
|
|
"tvm_call_packed_lowered",
|
|
"tvm_call_cpacked_lowered",
|
|
"TVMBackendAllocWorkspace",
|
|
"TVMBackendFreeWorkspace",
|
|
"start_profile_intrinsic",
|
|
"end_profile_intrinsic",
|
|
"meta_var",
|
|
"anylist_getitem",
|
|
"anylist_resetitem",
|
|
"anylist_setitem_call_packed",
|
|
"anylist_setitem_call_cpacked",
|
|
"llvm_lookup_intrinsic_id",
|
|
"type_annotation",
|
|
"broadcast",
|
|
"ramp",
|
|
"cast",
|
|
# tvm.tirx.expr
|
|
"Var",
|
|
"Reduce",
|
|
"FloatImm",
|
|
"IntImm",
|
|
"StringImm",
|
|
"Cast",
|
|
"Add",
|
|
"Sub",
|
|
"Mul",
|
|
"Div",
|
|
"Mod",
|
|
"FloorDiv",
|
|
"FloorMod",
|
|
"Min",
|
|
"Max",
|
|
"EQ",
|
|
"NE",
|
|
"LT",
|
|
"LE",
|
|
"GT",
|
|
"GE",
|
|
"And",
|
|
"Or",
|
|
"Not",
|
|
"Select",
|
|
"BufferLoad",
|
|
"ProducerLoad",
|
|
"Ramp",
|
|
"Broadcast",
|
|
"Shuffle",
|
|
"Call",
|
|
"CallEffectKind",
|
|
"let",
|
|
"Bind",
|
|
"bind",
|
|
"LetAnnotation",
|
|
"LocalVectorAnnotation",
|
|
"DtypeConstructor",
|
|
"Let",
|
|
"IterVar",
|
|
"CommReducer",
|
|
"Range",
|
|
"vscale",
|
|
"get_active_lane_mask",
|
|
"call_kernel",
|
|
"ignore_loop_partition",
|
|
]
|
|
|
|
__all__ += [
|
|
"ComposeLayout",
|
|
"ExecScope",
|
|
"Iter",
|
|
"Layout",
|
|
"R",
|
|
"S",
|
|
"ScopeIdDef",
|
|
"SwizzleLayout",
|
|
"TensorMap",
|
|
"TileLayout",
|
|
"Var",
|
|
"add_to_parent",
|
|
"alloc_cast_frag",
|
|
"alloc_local",
|
|
"alloc_scalar",
|
|
"alloc_shared",
|
|
"alloc_tcgen05_ldst_frag",
|
|
"cluster_id",
|
|
"cta_id",
|
|
"cta_id_in_cluster",
|
|
"cta_id_in_pair",
|
|
"decl_scalar",
|
|
"device_entry",
|
|
"lane_id",
|
|
"local_scalar",
|
|
"meta_class",
|
|
"register_script_namespace",
|
|
"scalar_wrapper",
|
|
"scope_id",
|
|
"shared_scalar",
|
|
"smem",
|
|
"static_assert",
|
|
"thread_id",
|
|
"thread_id_in_wg",
|
|
"tmem",
|
|
"warp_id",
|
|
"warp_id_in_wg",
|
|
"warpgroup_id",
|
|
"webgpu",
|
|
]
|
|
|
|
# Shorthand dtype aliases
|
|
__all__ += ["bf16", "f16", "f32", "f64", "i8", "i16", "i32", "i64", "u8", "u16", "u32", "u64"]
|