tvm.relax.block_builder

tvm.relax.block_builder#

Developer API of constructing Relax AST.

class tvm.relax.block_builder.BlockBuilder(mod: IRModule = None)[源代码]#

A builder to build Relax IR for testing and dev.

示例

m = tir.Var("m", "int32")
n = tir.Var("n", "int32")
x = rx.Var("x", rx.TensorStructInfo([m, n], "float16"))
y = rx.Var("y", rx.TensorStructInfo([n], "float16")
bb = rx.BlockBuilder()
with bb.function([x, y], "func"):
    with bb.dataflow() as df:
        lv0 = bb.emit(rx.add(x, y))
        lv1 = bb.emit(rx.multiply(lv0, y))
        gv0 = bb.emit_output(lv1)
    bb.emit_func_output(gv0)
mod = bb.get()

BlockBuilder can also be used to construct neural networks with nn.Module API

from tvm.relax.testing import nn

n = tir.Var("n", "int64")
input_size = 784
hidden_sizes = [128, 32]
output_size = 10
bb = rx.BlockBuilder()

with bb.function("main"):
    model = nn.Sequential(
        nn.Linear(input_size, hidden_sizes[0]),
        nn.ReLU(),
        nn.Linear(hidden_sizes[0], hidden_sizes[1]),
        nn.ReLU(),
        nn.Linear(hidden_sizes[1], output_size),
        nn.LogSoftmax(),
    )
    data = nn.Placeholder((n, input_size), name="data")
    output = model(data)
    params = [data] + model.parameters()
    builder.emit_func_output(output, params=params)
mod = bb.get()
add_func(func: BaseFunc, func_name: str) → GlobalVar[源代码]#

Add a Relax function or a TIR PrimFunc to the IRModule being built.

参数:

  • func (BaseFunc) -- The function to be added.

  • func_name (str) -- The name of the function to be added.

返回:

gvar -- The global var bound to the added function.

返回类型:

GlobalVar

begin_scope(params: List[Var] | None = None) → None[源代码]#

Begin a new scope, with optional parameters that are visible within the scope.

参数:

params (Optional[List[Var]]) -- Parameters that are visible within the scope.

备注

This function should be called when new scope is introduced (function, seq) to properly track the variable availability and help the best effort deduction.

call_te(func: Callable, *args: Any, **kwargs: Any) → RelaxExpr[源代码]#

Generate a call node according to the te function. This function converts arguments from relax expression to te tensor, The callback func should return a te tensor or a list of te tensors. Please see detailed example in emit_te

参数:

  • func (Callable) -- A function that returns a te tensor or a list of te tensors.

  • args (Any, optional) -- arguments passed to the function.

  • kwargs (Any, optional) --

    The keyword arguments passed to the function. Note that the following keyword args are reserved:

    • 'primfunc_name_hint' for passing name hint to the PrimFunc that gets generated.

    • 'primfunc_attrs' is reserved for passing func attributes to be added to the PrimFunc that gets created.

返回:

ret -- A newly created call node

返回类型:

tvm.relax.Call

call_te_with_grad(func: Callable, *args: Any, te_grad_name: str, te_grad_kwargs: Dict[str, Object] = None, **kwargs: Any) → RelaxExpr[源代码]#

Generate a call node according to the te function. This method will generate a call_tir_with_grad node, i.e. a call_tir node bound with a te gradient function (refered by te_grad_name).

参数:

  • func (Callable) -- A function that returns a te tensor or a list of te tensors.

  • args (Any, optional) -- arguments passed to the function.

  • te_grad_name (str) -- The registered name of the te gradient function associated with the call_tir_with_grad node. Must be provided as a keyword argument.

  • te_grad_kwargs (Dict[str, Object], optional) -- The keyword arguments passed to the te gradient function. Optionally provided as a keyword argument. Default: {}.

  • kwargs (Any, optional) --

    The keyword arguments passed to the function. Note that the following keyword args are reserved:

    • 'primfunc_name_hint' for passing name hint to the PrimFunc that gets generated.

    • 'primfunc_attrs' is reserved for passing func attributes to be added to the PrimFunc that gets created.

返回:

ret -- A newly created call node

返回类型:

tvm.relax.Call

static current() → BlockBuilder | None[源代码]#

Returns the current BlockBuilder.

current_block_is_dataflow() → bool[源代码]#

Check if the block being built is DataflowBlock or not.

返回:

ret -- A boolean that indicates if the block being built is DataflowBlock or not.

返回类型:

bool

dataflow() → DataflowScope[源代码]#

Annotate a Relax dataflow block.

返回:

ret -- A DataflowScope for building a Relax dataflow block.

返回类型:

DataflowScope

emit(expr: RelaxExpr, name_hint: str = '') → Var[源代码]#

Emit an expr. This infers the shape and type of the expr, create a variable, and bind the expr to the variable.

参数:

  • expr (tvm.relax.Expr) -- The Expr to be emitted.

  • name_hint (str) -- Name hint for the bound variable.

返回:

ret -- A newly created variable that gets bound to the input expr.

返回类型:

tvm.relax.Var

emit_func_output(output: RelaxExpr | Tuple | List[RelaxExpr], params: Var | Tuple | List[Var] | None = None) → GlobalVar[源代码]#

Emit output for the function.

参数:

  • output (Expr | Tuple | List[Expr]) -- The output of the current block/function.

  • params (tvm.relax.Var | Tuple | List[tvm.relax.Var], optional) -- The parameters of the function to be built. If params is None, it means the params have been initialized in the function with scope.

返回:

gvar -- A GlobalVar representing the function

返回类型:

tvm.ir.GlobalVar

emit_normalized(binding: Binding) → None[源代码]#

Emit an already normalized binding.

参数:

binding (Binding) -- The binding to be emitted.

emit_output(output: RelaxExpr | Tuple | List[RelaxExpr], name_hint: str = '') → Var[源代码]#

Emit output for the current dataflow block or function.

参数:

  • output (Expr | Tuple | List[Expr]) -- The output of the current block/function.

  • name_hint (str) -- Name hint for the bound variable.

返回:

ret -- The return variable which gets bound to the output.

返回类型:

tvm.relax.Var

emit_te(func: Callable, *args: Any, **kwargs: Any) → Var[源代码]#

Emit a call node according to the te function. This function converts arguments from relax expression to te tensor, The callback func should return a te tensor or a list of te tensors.

参数:

  • func (Callable) -- A function that returns a te tensor or a list of te tensors.

  • args (Any, optional) -- arguments passed to the function.

  • kwargs (Any, optional) -- The keyword arguments passed to the function. Note that the key "primfunc_name_hint" is reserved for passing name hint to the PrimFunc that gets generated.

返回:

ret -- A newly created variable that gets bound to the call code.

返回类型:

tvm.relax.Var

示例

bb = rx.BlockBuilder()
n, m = tir.Var("n", "int64"), tir.Var("m", "int64")
x = rx.Var("x", rx.TensorStructInfo([n, m], "float32"))
y = rx.Var("y", rx.TensorStructInfo([n, m], "float32"))

def te_func(args, args_dict, msg):
    A = args[0]
    B = args_dict["B"]
    return te.compute((128, 128), lambda i, j: A[i, j] + B[i, j])

with bb.function([x, y], "rx_func"):
    out = bb.emit_te(te_func, [x], {"B": y}, msg="hello")
    bb.emit_func_output(out)

will result in TVMScript

@tvm.script.ir_module
class Module:
    @T.prim_func
    def te_func(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle,
                var_compute: T.handle) -> None:
        # function attr dict
        T.func_attr({"tir.noalias": True})
        m = T.int64()
        n = T.int64()
        rxplaceholder = T.match_buffer(var_rxplaceholder, [n, m], dtype="float32")
        rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [n, m], dtype="float32")
        compute = T.match_buffer(var_compute, [128, 128], dtype="float32")
        # body
        # with T.block("root")
        for i0, i1 in T.grid(128, 128):
            with T.block("compute"):
                i, j = T.axis.remap("SS", [i0, i1])
                T.reads([rxplaceholder[i, j], rxplaceholder_1[i, j]])
                T.writes([compute[i, j]])
                compute[i, j] = rxplaceholder[i, j] + rxplaceholder_1[i, j]

    @R.function
    def rx_func(x: Tensor((n, m), "float32"), y: Tensor((n, m), "float32")) -> Tensor:
        # block 0
        gv = relax.call_tir("te_func", (x, y), R.Tensor((128, 128), "float32"))
        return gv

示例

bb = relax.BlockBuilder()
n = tir.Var("n", "int64")
x = relax.Var("x", relax.TensorStructInfo([n], "float32"))
y = relax.Var("y", relax.TensorStructInfo([n + 1], "float32"))

def te_func(A):
    C = te.compute((n + 1), lambda i: A[i])
    return C

with bb.function("rx_func", [x, y]):
    x1 = bb.emit_te(te_func, y)
    bb.emit_func_output(x1)

will result in TVMScript

@tvm.script.ir_module
class Module:
    @T.prim_func
    def te_func(var_rxplaceholder: T.handle, var_compute: T.handle, n: T.int64) -> None:
        rxplaceholder = T.match_buffer(var_rxplaceholder, [n + T.int64(1)],
                                       dtype="float32")
        compute = T.match_buffer(var_compute, [n + T.int64(1)], dtype="float32")
        # body
        # with T.block("root")
        for i0 in T.serial(0, n + T.int64(1)):
            with T.block("compute"):
                i = T.axis.spatial(n + T.int64(1), i0)
                T.reads([rxplaceholder[i]])
                T.writes([compute[i]])
                compute[i] = rxplaceholder[i]

    @R.function
    def rx_func(x: Tensor((n,), "float32"), y: Tensor(((n + 1),), "float32"))
        -> Tensor(None, "float32", ndim=-1):
        # block 0
        gv = relax.call_tir(te_func, (y,), R.Tensor((n + 1,), "float32"), (n,))
        return gv
end_scope() → None[源代码]#

End the current scope. Please see begin_scope for details

finalize() → IRModule[源代码]#

Finalize the building process and return the result IRModule.

Possibly rename GlobalVars in the IRModule to ensure name uniqueness and the invariant: every public function has the same name as its "global_symbol" attribute.

Note this method should be called only once at the end of the building process, since it may invalidate global vars previously returned by this builder. See also tvm.relax.transform.NormalizeGlobalVar.

返回:

ret -- An IRModule with Relax and TIR functions being built.

返回类型:

tvm.IRModule

function(name: str, params: Var | Tuple | List[Var] | None = None, attrs: Dict[str, Object] | None = None, pure: bool = True, private: bool = False) → FunctionScope[源代码]#

Annotate a Relax function.

参数:

  • name (str, optional) -- The name of the function

  • params (tvm.relax.Var | Tuple | List[tvm.relax.Var], optional) -- The parameters of the function. If params is None, it means deferring initialization of function parameters until emit_func_output.

  • attrs (Dict[str, Object], optional) -- The function attrs

  • pure (bool, optional) -- Whether the function is annotated as pure.

  • private (bool, optional) -- Whether the function is annotated as private. If the function is private, it will not have a global symbol attribute. If it is not private and not an inner function, then it will have a global symbol attribute (mapped to the function's name)

返回:

ret -- A FunctionScope for building a Relax function node.

返回类型:

FunctionScope

get() → IRModule[源代码]#

Return intermediate IRModule. For the situation where the IRModule is needed in the middle of a building process.

返回:

ret -- An IRModule with Relax and TIR functions being built.

返回类型:

tvm.IRModule

get_unique_name(name_prefix: str) → str[源代码]#

Generate a unique name with a specified prefix.

参数:

name_hint (str) -- The name prefix.

返回:

ret -- The generated name.

返回类型:

str

lookup_binding(var: Var) → RelaxExpr | None[源代码]#

Lookup a var in the binding table.

参数:

var (Var) -- The input var.

返回:

expr -- The Expr bound to the input var.

返回类型:

Expr

match_cast(value: RelaxExpr, struct_info: StructInfo, name_hint: str = '') → Var[源代码]#

Emit a MatchCast.

参数:

  • value (tvm.relax.Expr) -- The value of the MatchCast to be emitted.

  • struct_info (StructInfo) -- The struct info to be matched.

  • name_hint (str) -- The name of the match cast

返回:

ret -- A newly created variable that get bounds to be the casted result.

返回类型:

tvm.relax.Var

normalize(expr: RelaxExpr) → RelaxExpr[源代码]#

Normalize an Expr to complete its shape and type.

参数:

expr (Expr) -- The input expr.

返回:

ret -- The expr with normalized shape and type.

返回类型:

Expr

testing_scope(def_vars: List[Var]) → TestingScope[源代码]#

Start a scope for unit-testing purposes.

参数:

def_vars (List[tir.Var]) -- List of symbolic variables that are marked as defined in scope.

返回:

ret -- A TestingScope to setup builder for emit and other purposes.

返回类型:

TestingScope

update_func(gv: GlobalVar, updated_func: BaseFunc) → None[源代码]#

Add a Relax function or a TIR PrimFunc to the IRModule being built.

参数:

  • gv (GlobalVar) -- The global var referring the function to be updated.

  • updated_func (BaseFunc) -- The updated function.

class tvm.relax.block_builder.DataflowScope(block_builder)[源代码]#

Auxiliary scope for Dataflow block

class tvm.relax.block_builder.FunctionScope(block_builder, name, params, attrs, is_pure)[源代码]#

Auxiliary scope for function

class tvm.relax.block_builder.TestingScope(block_builder, def_vars)[源代码]#

Auxiliary scope for testing purposes

目录