# 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"]