github-featured/apache--tvm
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
apache--tvm/python/tvm/tirx/script/builder/ir.py
T
Tianqi Chen a6c7377bae [REFACTOR][TIRx] Keep AttrStmt node values unboxed (#20030) 
AttrStmt.node is an ffi::Any field, so converting POD arguments to
PrimExpr makes its representation depend on the caller rather than the
declared container type.

This change preserves values passed through AttrStmt and T.attr, uses
raw zero sentinel nodes consistently, and updates the printer
canonicalization.
2026-07-20 06:21:51 +08:00

3558 lines
102 KiB
Python
Raw Permalink Blame History

# 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.script.ir_builder.ir import meta_var
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,
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_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,
_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,
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.
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,
_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,
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.
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,
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,
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.
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 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,
_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:
# Canonical Vars are callable when they denote functions. Here a Var is
# already a resolved type annotation, rather than a deferred annotation factory.
if callable(type_annotation) and not isinstance(type_annotation, Expr):
type_annotation = type_annotation()
if isinstance(type_annotation, ir.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, ir.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 ``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=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(0, k, convert(v))) # type: ignore[attr-defined]
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 Return(value: Expr) -> None: # pylint: disable=invalid-name
"""Create a return node."""
return _ffi_api.Return(value) # 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,
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.
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,
_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,
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, ir.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:
C = TypeVar("C")
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)
# 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)
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",
"Return",
"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",
"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"]
Reference in New Issue View Git Blame Copy Permalink