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
发出当前数据流块或函数的输出。
- 参数:
- 返回: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
生成具有指定前缀的唯一名称。
add_func(func:BaseFunc, func_name:str) → GlobalVar
向正在构建的 IRModule 添加 Relax 函数或 TIR PrimFunc。
update_func(gv:GlobalVar, updated_func:BaseFunc) → None
向正在构建的 IRModule 添加 Relax 函数或 TIR PrimFunc。
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 了解详情