Function inline

python/tvm/tirx/script/parser/entry.py:164–208

Decorator for inline function definitions with Python LEGB scoping. @T.inline follows Python's lexical scoping with late binding: - At definition time, record which scopes are visible. - At call time, read current values from those scopes. Example:: import tvm from

(*args, definition_depth: int | None = None, defining_var_table=None)

def inline(*args, definition_depth: int | None = None, defining_var_table=None) -> Callable:
    """Decorator for inline function definitions with Python LEGB scoping.

    @T.inline follows Python's lexical scoping with late binding:
    - At definition time, record which scopes are visible.
    - At call time, read current values from those scopes.

    Example::

        import tvm
        from tvm.script import tirx as T

        x_value = 128

        @T.inline
        def capture(A, B):
            B[()] = A[x_value]          # x_value resolved from enclosing scope

        @T.prim_func(s_tir=True)
        def use(A: T.Buffer((1024,), "int32"), B: T.Buffer((), "int32")) -> None:
            capture(A, B)               # Produces B[()] = A[128]
    """

    def _decorator(func: Callable) -> Callable:
        source, closure_vars = scan_macro(func, utils.inspect_function_capture(func))
        obj = TIRInline(
            source,
            closure_vars,
            func,
            definition_depth=definition_depth,
            defining_var_table=defining_var_table,
        )

        def wrapper(*args, **kwargs):
            return obj(*args, **kwargs)

        return wrapper

    if len(args) == 0:
        setattr(_decorator, "dispatch_token", "tir.inline")
        return _decorator
    if len(args) == 1 and inspect.isfunction(args[0]):
        return _decorator(args[0])

    raise ValueError("Invalid use of T.inline. Usage: @T.inline or @T.inline()")


setattr(inline, "dispatch_token", "tir.inline")