跳到主要内容
版本：0.21.0

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,optional)：函数的名称。
    • params (tvm.relax.Var|Tuple|List [tvm.relax.Var]**,optional)：函数的参数。如果 params 为 None，则表示将函数参数的初始化推迟到 emit_func_output。
    • attrs (Dict[str, Object],optional)：函数 attrs。
    • pure (bool,optional)：函数是否被注释为纯函数。
    • private (bool,optional)：函数是否被注释为私有函数。如果函数是私有的，则它没有全局符号属性。如果它不是私有的且不是内部函数，则它将具有全局符号属性（映射到函数名称）。
  • 返回：ret：用于构建 Relax 函数节点的 FunctionScope。
  • 返回类型：FunctionScope

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

启动用于单元测试的范围。

  • 参数：def_vars (List[tir.Var])：标记为范围内定义的符号变量列表。
  • 返回：ret：用于设置用于发射和其他目的的构建器的 TestingScope。
  • 返回类型：TestingScope

dataflow() → DataflowScope​

注释 Relax 数据流块。

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

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

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

  • 参数：
    • expr (tvm.relax.Expr)：要发出的 Expr。
    • name_hint (str)：绑定变量的名称提示。
  • 返回：ret：与输入 expr 绑定的新创建的变量。
  • 返回类型：tvm.relax.Var

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

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

  • 参数：
    • func (Callable)*：*返回 te 张量或 te 张量列表的函数。
    • args (Any,optional)：传递给函数的参数。
    • kwargs (Any,optional)：
      • 传递给函数的关键字参数。请注意，以下关键字参数是保留的：
      • ’primfunc_name_hint’ 用语将名称提示传递给生成的PrimFunc。
      • ’primfunc_attrs’保留用于传递要添加到创造的 PrimFunc 的函数属性。
  • 返回：ret：新创建的调用节点。
  • 返回类型：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,optional)：传递给函数的参数。
    • te_grad_name (str)：与 call_tir_with_grad 节点关联的 te 梯度函数的注册名称。必须作为关键字参数提供。
    • te_grad_kwargs (Dict[str,Object],optional)：传递给 te 梯度函数的关键字参数。可选地，以关键字参数的形式提供。默认值：{}。
    • kwargs (Any,optional)：传递给函数的关键字参数。请注意，以下关键字参数是保留的：
      • ’primfunc_name_hint’ 用于将名称提示传递给生成的 PrimFunc。
      • ’primfunc_attrs’ 保留用于传递要添加到创建的 PrimFunc 的函数属性。
  • 返回：ret：新创建的调用节点。
  • 返回类型：tvm.relax.Call

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

根据 te 函数发出一个调用节点。该函数将参数从松弛表达式转换为 te 张量。回调函数应返回一个 te 张量或一个 te 张量列表。

  • 参数：
    • func(Callable)： 返回 te 张量或 te 张量列表的函数。
    • args (Any,optional)：传递给函数的参数。
    • kwargs (Any,optional)：传递给函数的关键字参数。请注意，“primfunc_name_hint”键保留用于将名称提示传递给生成的 PrimFunc。
  • 返回：ret：与调用代码绑定的新创建的变量。
  • 返回类型：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​

发出 MatchCast。

  • 参数：
    • value (tvm.relax.Expr)：要发出的 MatchCast 的值。
    • struct_info (StructInfo)：要匹配的结构信息。
    • name_hint (str)：匹配转换的名称。
  • 返回：ret：一个新创建的变量，其界限是转换结果。
  • 返回类型：tvm.relax.Var

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

发出当前数据流块或函数的输出。

  • 参数：
    • output (Expr|Tuple|List[Expr])：当前块/函数的输出。
    • name_hint (str)：绑定变量的名称提示。
  • 返回：ret：与输出绑定的返回变量。
  • 返回类型：tvm.relax.Var

emit_func_output(output:RelaxExpr|Tuple|List[RelaxExpr], params:Var|Tuple|List[Var] |None= None) → GlobalVar​

为函数发出输出。

  • 参数：
    • output (Expr|Tuple|List[Expr])：当前块/函数的输出。
    • params (tvm.relax.Var|Tuple|List[tvm.relax.Var],optional)：要构建的函数的参数。如果 params 为 None，则表示 params 已在函数中初始化并具有作用域。
  • 返回：gvar： 代表函数的 GlobalVar。
  • 返回类型：tvm.ir.GlobalVar

normalize(expr:RelaxExpr) → RelaxExpr​

规范化 Expr 以完成其形状和类型。

  • 参数：expr (Expr)：输入表达式。
  • 返回：ret：具有规范化形状和类型的 expr。
  • 返回类型： Expr

get() → IRModule​

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

  • 返回：ret：正在构建具有 Relax 和 TIR 功能的 IRModule。
  • 返回类型： tvm.IRModule。

finalize() → IRModule​

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

可能在 IRModule 中重命名 GlobalVars 以确保名称的唯一性和不变性：每个公共函数都与其「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​

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

  • 参数：
    • func (BaseFunc)：要添加的函数。
    • func_name (str)：要添加的函数的名称。
  • 返回：gvar：与添加的函数绑定的全局变量。
  • 返回类型：GlobalVar

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

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

  • 参数：
    • 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 (relax.Var)：输入变量。
  • 返回：expr：与输入变量绑定的 Expr。
  • 返回类型： Expr。

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

开始一个新的范围，带有范围内可见的可选参数。

  • 参数：params (Optional[List[relax.Var]**]*)：范围内可见的参数。

当引入新范围（函数、序列）时应调用此函数，以正确跟踪变量可用性并帮助尽力推断。

end_scope() → None​

结束当前作用域。请参阅 begin_scope 了解详情

编辑此页
最后由 wenjing-user 于 更新