tvm.relax.block_builder

用于构建 Relax AST 的开发者 API。

class tvm.relax.block_builder.FunctionScope(block_builder, name, params, attrs, is_pure): 函数的辅助作用域

class tvm.relax.block_builder.DataflowScope(block_builder): 数据流块的辅助作用域

class tvm.relax.block_builder.TestingScope(block_builder, def_vars): 用于测试目的的辅助作用域

class tvm.relax.block_builder.BlockBuilder(mod: IRModule | None = None)

用于测试和开发的构建 Relax IR 的构建器。

示例

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 也可用于使用 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()

static current() → BlockBuilder | None: 返回当前的 BlockBuilder。

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

注释 Relax 函数。

参数:

name (str, 可选) – 函数的名称
params (tvm.relax.Var | Tuple | List[tvm.relax.Var], 可选) – 函数的参数。如果 params 为 None，则表示将函数参数的初始化推迟到 emit_func_output。
attrs (Dict[str, Object], 可选) – 函数属性
pure (bool, 可选) – 函数是否被注释为 pure。
private (bool, 可选) – 函数是否被注释为 private。如果函数是 private 的，它将没有全局符号属性。如果它不是 private 的且不是内部函数，那么它将具有全局符号属性（映射到函数的名称）

返回:

ret – 用于构建 Relax 函数节点的 FunctionScope。

返回类型:

FunctionScope

testing_scope(def_vars: List[Var]) → TestingScope

启动用于单元测试目的的作用域。

参数:: def_vars (List[tir.Var]) – 标记为在作用域中定义的符号变量列表。
返回:: ret – 用于设置构建器以进行 emit 和其他目的的 TestingScope。
返回类型:: TestingScope

dataflow() → DataflowScope

注释 Relax 数据流块。

返回:: ret – 用于构建 Relax 数据流块的 DataflowScope。
返回类型:: DataflowScope

emit(expr: RelaxExpr, name_hint: str = '') → Var

Emit 一个 expr。这将推断 expr 的形状和类型，创建变量，并将 expr 绑定到变量。

参数:

expr (tvm.relax.Expr) – 要 emit 的 Expr。
name_hint (str) – 绑定变量的名称提示。

返回:

ret – 新创建的变量，它被绑定到输入 expr。

返回类型:

tvm.relax.Var

call_te(func: Callable, *args: Any, **kwargs: Any) → RelaxExpr

根据 te 函数生成 call 节点。此函数将参数从 relax 表达式转换为 te 张量，回调函数应返回一个 te 张量或 te 张量列表。请参阅 emit_te 中的详细示例

参数:

func (Callable) – 返回 te 张量或 te 张量列表的函数。
args (Any, 可选) – 传递给函数的参数。
kwargs (Any, 可选) –
传递给函数的关键字参数。请注意，以下关键字参数是保留的
- ’primfunc_name_hint’ 用于将名称提示传递给生成的 PrimFunc。
- ’primfunc_attrs’ 保留用于传递要添加到创建的 PrimFunc 的函数属性。

返回:

ret – 新创建的 call 节点

返回类型:

tvm.relax.Call

call_te_with_grad(func: Callable, *args: Any, te_grad_name: str, te_grad_kwargs: Dict[str, Object] | None = None, **kwargs: Any) → RelaxExpr

根据 te 函数生成 call 节点。此方法将生成 call_tir_with_grad 节点，即绑定到 te 梯度函数（由 te_grad_name 引用）的 call_tir 节点。

参数:

func (Callable) – 返回 te 张量或 te 张量列表的函数。
args (Any, 可选) – 传递给函数的参数。
te_grad_name (str) – 与 call_tir_with_grad 节点关联的 te 梯度函数的注册名称。必须作为关键字参数提供。
te_grad_kwargs (Dict[str, Object], 可选) – 传递给 te 梯度函数的关键字参数。可选地作为关键字参数提供。默认值：{}。
kwargs (Any, 可选) –
传递给函数的关键字参数。请注意，以下关键字参数是保留的
- ’primfunc_name_hint’ 用于将名称提示传递给生成的 PrimFunc。
- ’primfunc_attrs’ 保留用于传递要添加到创建的 PrimFunc 的函数属性。

返回:

ret – 新创建的 call 节点

返回类型:

tvm.relax.Call

emit_te(func: Callable, *args: Any, **kwargs: Any) → Var

根据 te 函数 Emit 一个 call 节点。此函数将参数从 relax 表达式转换为 te 张量，回调函数应返回一个 te 张量或 te 张量列表。

参数:

func (Callable) – 返回 te 张量或 te 张量列表的函数。
args (Any, 可选) – 传递给函数的参数。
kwargs (Any, 可选) – 传递给函数的关键字参数。请注意，键 “primfunc_name_hint” 保留用于将名称提示传递给生成的 PrimFunc。

返回:

ret – 新创建的变量，它被绑定到 call 代码。

返回类型:

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)

将导致 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)

将导致 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

match_cast(value: RelaxExpr, struct_info: StructInfo, name_hint: str = '') → Var

Emit 一个 MatchCast。

参数:

value (tvm.relax.Expr) – 要 emit 的 MatchCast 的值。
struct_info (StructInfo) – 要匹配的 struct info。
name_hint (str) – match cast 的名称

返回:

ret – 新创建的变量，它被绑定为 cast 的结果。

返回类型:

tvm.relax.Var

emit_output(output: RelaxExpr | Tuple | List[RelaxExpr], name_hint: str = '') → Var

为当前数据流块或函数 Emit 输出。

参数:

output (Expr | Tuple | List[Expr]) – 当前块/函数的输出。
name_hint (str) – 绑定变量的名称提示。

返回:

ret – 返回变量，它被绑定到输出。

返回类型:

tvm.relax.Var

为函数 Emit 输出。

参数:

output (Expr | Tuple | List[Expr]) – 当前块/函数的输出。
params (tvm.relax.Var | Tuple | List[tvm.relax.Var], 可选) – 要构建的函数的参数。如果 params 为 None，则表示参数已在具有作用域的函数中初始化。

返回:

gvar – 表示函数的 GlobalVar

返回类型:

tvm.ir.GlobalVar

normalize(expr: RelaxExpr) → RelaxExpr

Normalize 一个 Expr 以完成其形状和类型。

参数:: expr (Expr) – 输入 expr。
返回:: ret – 具有 normalized 形状和类型的 expr。
返回类型:: Expr

get() → IRModule

返回中间 IRModule。用于在构建过程中需要 IRModule 的情况。

返回:: ret – 正在构建的具有 Relax 和 TIR 函数的 IRModule。
返回类型:: tvm.IRModule

finalize() → IRModule

完成构建过程并返回结果 IRModule。

可能重命名 IRModule 中的 GlobalVar，以确保名称唯一性和不变量：每个公共函数都具有与其 “global_symbol” 属性相同的名称。

请注意，此方法应仅在构建过程结束时调用一次，因为它可能会使此构建器先前返回的全局变量无效。另请参阅 tvm.relax.transform.NormalizeGlobalVar。

返回:: ret – 正在构建的具有 Relax 和 TIR 函数的 IRModule。
返回类型:: tvm.IRModule

get_unique_name(name_prefix: str) → str

生成具有指定前缀的唯一名称。

参数:: name_hint (str) – 名称前缀。
返回:: ret – 生成的名称。
返回类型:: str

add_func(func: BaseFunc, func_name: str) → GlobalVar

将 Relax 函数或 TIR PrimFunc 添加到正在构建的 IRModule。

参数:

func (BaseFunc) – 要添加的函数。
func_name (str) – 要添加的函数的名称。

返回:

gvar – 绑定到添加函数的全局变量。

返回类型:

GlobalVar

update_func(gv: GlobalVar, updated_func: BaseFunc) → None

将 Relax 函数或 TIR PrimFunc 添加到正在构建的 IRModule。

参数:

gv (GlobalVar) – 引用要更新的函数的全局变量。
updated_func (BaseFunc) – 更新后的函数。

current_block_is_dataflow() → bool

检查正在构建的块是否为 DataflowBlock。

返回:: ret – 一个布尔值，指示正在构建的块是否为 DataflowBlock。
返回类型:: bool

emit_normalized(binding: Binding) → None

发出一个已经标准化的绑定。

参数:: binding (Binding) – 要发出的绑定。

lookup_binding(var: Var) → RelaxExpr | None

在绑定表中查找一个 var。

参数:: var (relax.Var) – 输入 var。
返回:: expr – 绑定到输入 var 的 Expr。
返回类型:: Expr

begin_scope(params: List[Var] | None = None) → None

开始一个新的作用域，带有可选的参数，这些参数在作用域内可见。

参数:: params (Optional[List[relax.Var]]) – 在作用域内可见的参数。

注意

当引入新的作用域（函数、seq）时，应该调用此函数，以正确跟踪变量的可用性并帮助尽力进行推导。

end_scope() → None: 结束当前作用域。请参阅 begin_scope 以了解详细信息