tvm.relax.block_builder#
Developer API of constructing Relax AST.
- class tvm.relax.block_builder.BlockBuilder(mod=None)[源代码]#
A builder to build Relax IR for testing and dev.
Examples#
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()
- _normalize_python_tuple(expr)[源代码]#
Internal utility function to convert to relax.Tuple
The emit, emit_output, and emit_func_output can be called with python list or tuple objects. These objects should be converted to relax.Tuple prior to calling an FFI function, as they would otherwise be converted to tvm.runtime.Array. In addition, any nested tuple objects should be converted.
- 参数:
expr (RelayExpr | Sequence[RelayExpr])
- add_func(func, func_name)[源代码]#
Add a Relax function or a TIR PrimFunc to the IRModule being built.
Parameters#
- funcBaseFunc
The function to be added.
- func_namestr
The name of the function to be added.
Returns#
- gvarGlobalVar
The global var bound to the added function.
- 参数:
func (BaseFunc)
func_name (str)
- 返回类型:
GlobalVar
- begin_scope(params=None)[源代码]#
Begin a new scope, with optional parameters that are visible within the scope.
Parameters#
- params: Optional[List[Var]]
Parameters that are visible within the scope.
Note#
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, *args, **kwargs)[源代码]#
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
Parameters#
- funcCallable
A function that returns a te tensor or a list of te tensors.
- argsAny, optional
arguments passed to the function.
- kwargsAny, 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.
Returns#
- rettvm.relax.Call
A newly created call node
- call_te_with_grad(func, *args, te_grad_name, te_grad_kwargs=None, **kwargs)[源代码]#
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).
Parameters#
- funcCallable
A function that returns a te tensor or a list of te tensors.
- argsAny, optional
arguments passed to the function.
- te_grad_namestr
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_kwargsDict[str, Object], optional
The keyword arguments passed to the te gradient function. Optionally provided as a keyword argument. Default: {}.
- kwargsAny, 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.
Returns#
- rettvm.relax.Call
A newly created call node
- static current()[源代码]#
Returns the current BlockBuilder.
- 返回类型:
BlockBuilder | None
- current_block_is_dataflow()[源代码]#
Check if the block being built is DataflowBlock or not.
Returns#
- retbool
A boolean that indicates if the block being built is DataflowBlock or not.
- 返回类型:
- dataflow()[源代码]#
Annotate a Relax dataflow block.
Returns#
- ret: DataflowScope
A DataflowScope for building a Relax dataflow block.
- 返回类型:
- emit(expr, name_hint='')[源代码]#
Emit an expr. This infers the shape and type of the expr, create a variable, and bind the expr to the variable.
Parameters#
- exprtvm.relax.Expr
The Expr to be emitted.
- name_hintstr
Name hint for the bound variable.
Returns#
- rettvm.relax.Var
A newly created variable that gets bound to the input expr.
- emit_func_output(output, params=None)[源代码]#
Emit output for the function.
Parameters#
- outputExpr | Tuple | List[Expr]
The output of the current block/function.
- paramstvm.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.
Returns#
gvar: tvm.ir.GlobalVar
A GlobalVar representing the function
- emit_normalized(binding)[源代码]#
Emit an already normalized binding.
Parameters#
- binding: Binding
The binding to be emitted.
- 参数:
binding (Binding)
- 返回类型:
None
- emit_output(output, name_hint='')[源代码]#
Emit output for the current dataflow block or function.
Parameters#
- outputExpr | Tuple | List[Expr]
The output of the current block/function.
- name_hintstr
Name hint for the bound variable.
Returns#
- rettvm.relax.Var
The return variable which gets bound to the output.
- emit_te(func, *args, **kwargs)[源代码]#
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.
Parameters#
- funcCallable
A function that returns a te tensor or a list of te tensors.
- argsAny, optional
arguments passed to the function.
- kwargsAny, 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.
Returns#
- rettvm.relax.Var
A newly created variable that gets bound to the call code.
Example#
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
Example#
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
- finalize()[源代码]#
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.
Returns#
- rettvm.IRModule
An IRModule with Relax and TIR functions being built.
- 返回类型:
IRModule
- function(name, params=None, attrs=None, pure=True, private=False)[源代码]#
Annotate a Relax function.
Parameters#
- namestr, optional
The name of the function
- paramstvm.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.
- attrsDict[str, Object], optional
The function attrs
- purebool, optional
Whether the function is annotated as pure.
- privatebool, 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)
Returns#
- ret: FunctionScope
A FunctionScope for building a Relax function node.
- get()[源代码]#
Return intermediate IRModule. For the situation where the IRModule is needed in the middle of a building process.
Returns#
- rettvm.IRModule
An IRModule with Relax and TIR functions being built.
- 返回类型:
IRModule
- get_unique_name(name_prefix)[源代码]#
Generate a unique name with a specified prefix.
Parameters#
- name_hintstr
The name prefix.
Returns#
- retstr
The generated name.
- lookup_binding(var)[源代码]#
Lookup a var in the binding table.
Parameters#
- var: Var
The input var.
Returns#
- expr: Expr
The Expr bound to the input var.
- 参数:
var (Var)
- 返回类型:
RelayExpr | None
- match_cast(value, struct_info, name_hint='')[源代码]#
Emit a MatchCast.
Parameters#
- valuetvm.relax.Expr
The value of the MatchCast to be emitted.
- struct_infoStructInfo
The struct info to be matched.
- name_hintstr
The name of the match cast
Returns#
- rettvm.relax.Var
A newly created variable that get bounds to be the casted result.
- 参数:
value (RelayExpr)
struct_info (StructInfo)
name_hint (str)
- 返回类型:
- normalize(expr)[源代码]#
Normalize an Expr to complete its shape and type.
Parameters#
- exprExpr
The input expr.
Returns#
- retExpr
The expr with normalized shape and type.
- 参数:
expr (RelayExpr)
- 返回类型:
RelayExpr
- 参数:
mod (IRModule)