tvm_tirx.tirx.script.function.builder

tvm_tirx.tirx.script.function.builder#

function.builder(**options)#

Build a TIRx function with Python body semantics.

Use at the definition site with inspectable source. Assignments, loops, conditionals and returns execute as Python; explicit T.serial-family frames, T.if_/T.then_/T.else_, T.assign, T.bind_ and T.return_ construct IR. Signature annotations are re-evaluated from source, including postponed annotations; write expressions without quotes.

Shares T.function options, expression locations and module declarations. Statement locations are not injected: explicit frames, void emissions and IR emitted inside ordinary Python helpers may have unknown locations. Use T.call_global_var_(callee, args) for explicit IR function calls. The resulting IR prints as ordinary @T.function TVMScript.

Examples

Python loops execute while building; T.serial creates an IR loop:

from tvm_tirx.script import tirx as T

@T.function.builder
def add_rows(A: T.Tensor((2, 4), T.float32), B: T.Tensor((2, 4), T.float32)):
    for row in range(2):
        with T.serial(4) as col:
            T.assign(B[row, col], A[row, col] + 1)

Mixed modules share declarations, so explicit calls can refer to a later function:

@T.ir_module
class Module:
    @T.function.builder
    def main(A: T.Tensor((4,), T.int32)):
        with T.serial(4) as i:
            T.assign(A[i], T.call_global_var_(Module.add_one, [A[i]]))

    @T.function
    def add_one(value: T.int32) -> T.int32:
        return value + 1
Parameters:
  • function (LambdaType | None)

  • options (Any)

Return type:

Any