Functions and Modules

A function describes one piece of computation: its parameters and types say how it is called, and its body says what it does. In many cases, a program needs multiple functions that interact with each other, such as a host function that launches device kernels. A module collects these functions under their names, so calls can refer to one another and passes can transform the whole program.

In TVMScript, a module is written as a class marked with @T.ir_module, with one @T.function per member:

from tvm_tirx.script import tirx as T

n = T.dynamic("n")


@T.ir_module
class Module:
    @T.function
    def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
        for i in range(n):
            B[i] = A[i] + 1

    @T.function
    def add_two(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
        for i in range(n):
            B[i] = A[i] + 2

add_one and add_two add one and two to each element of A. The rest of this chapter looks at the functions and then at how the module ties them together.

Note

The runnable code in this chapter is under examples/ir_guides in the repository.

Functions

A tirx.Function has three fields of its own:

Field	What it holds
params	The parameter variables, in calling order. add_one has A and B.
ret_type	The type of the returned value. add_one returns nothing, an empty tuple type, and writes its result into B instead. A function that returns a value declares it with ->, as in def f(x: T.float32) -> T.float32.
body	The statements that run when the function is called.

It also has the fields every ir.BaseFunc shares: attrs, described below; ty, the function’s signature; and loc, its source location. ty is built from the parameter and return types:

add_one = Module["add_one"]
print(add_one.ty)

The output is:

T.Func([T.Tensor((n,), T.float32), T.Tensor((n,), T.float32)], T.Tuple())

The two tensor parameters come first, and T.Tuple() is the empty return type.

Attributes

attrs stores additional information about a function. The most common attribute is global_symbol, the name a public function is exported under in the compiled library. @T.function sets it from the Python name, so add_one is exported as "add_one"; a function marked @T.function(private=True) has no global_symbol and stays internal to the module.

In TVMScript, T.func_attr adds attributes in the function body:

@T.function
def add_one(A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
    T.func_attr({"note": "example"})
    for i in range(n):
        B[i] = A[i] + 1

The printer writes them back the same way. add_one.script() gives:

from __future__ import annotations

# from tvm_tirx.script import tirx as T

@T.function
def add_one[n](A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32)):
    T.func_attr({"note": "example"})
    for i in range(n):
        B[i] = A[i] + T.float32(1.0)

From Python, with_attr returns a new function with an attribute added and leaves the original unchanged:

add_one = Module["add_one"]
marked = add_one.with_attr("note", "example")
assert "note" not in add_one.attrs  # the original is unchanged

Modules

A module has three main fields:

Field	What it holds
functions	The module’s functions, keyed by GlobalVar.
attrs	Additional attribute information about the whole module, in the same form as a function’s attrs.
global_infos	Reserved for special use cases that need named objects shared across functions. Most modules leave it empty.

functions is keyed by GlobalVar, which the next section explains. The module also keeps an index from names to GlobalVar keys, which Module["add_one"] uses to look up a function by name.

Global Variables

ir.GlobalVar is an expression, listed in the expression table, that refers to a function in a module. The module’s functions map is keyed by GlobalVar: we can get the reference by name and look up its definition:

add_one_ref = Module.get_global_var("add_one")
add_one = Module[add_one_ref]

Keeping the reference separate from the definition lets IR refer to a function without containing its body. For example, a Call whose op is a GlobalVar calls that function, and a pass can rewrite or replace the definition without touching the places that refer to it. In practice, references between functions mostly connect host and device functions, such as a host function and the kernels it launches.

A function has two names, which coincide for add_one:

Name	Meaning	add_one	Private function
GlobalVar name	The function’s name within the module.	"add_one"	Present
global_symbol	The name a public function is exported under.	"add_one"	Absent

A private function has a module name but no exported symbol.