tvm_tirx.tirx.script.macro

Contents

tvm_tirx.tirx.script.macro#

tvm_tirx.tirx.script.macro(function=None, **options)#

Decorate a helper that constructs IR in its caller’s active frames.

Parameters:

  • function (Callable, optional) – The helper function. May be omitted to supply keyword options.

  • hygienic (bool, optional) – Whether the helper resolves symbols in its definition environment instead of its calling environment. Defaults to True. T.macro and R.macro capture values at definition time; T.inline refreshes captured closure cells when called.

  • **options – Keyword configuration for this helper decorator. hygienic controls name lookup as described above; other options are accepted but are not consumed or forwarded to builder hooks.

Returns:

result (Callable) – The construction helper, or a decorator when function is omitted.

Return type:

Callable[[…], Any]

Notes

T.inline follows Python lexical scoping with late binding of captured closure cells. Its return statements produce Python values, as do those of R.macro. T.macro emits returns in the active primitive function.

Examples

An inline helper can read values from its enclosing scope:

import tvm_tirx
from tvm_tirx.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.function
def use(A: T.Tensor((1024,), "int32"), B: T.Tensor((), "int32")) -> None:
    capture(A, B)       # Produces B[()] = A[128]
Contents