tvm.runtime.relax_vm
Relax 虚拟机。
- class tvm.runtime.relax_vm.VMInstrumentReturnKind(value)
一个枚举。
- class tvm.runtime.relax_vm.VirtualMachine(rt_mod: Module | Executable, device: Device | List[Device], memory_cfg: str | Dict[Device, str] | None = None, profile: bool = False)
Relax VM 运行时。
- save_function(func_name: str, saved_name: str, *args: List[Any], include_return: bool = True, **kwargs: Dict[str, Any]) None
便利函数。从模块中获取一个函数，并保存一个 PackedFunc，当调用时，将使用给定的参数调用该函数。PackedFunc 可以使用 saved_name 从模块访问。包含此项是为了方便计时试验：调用返回的 PackedFunc 将比通常通过 VM 运行具有更少的字典查找开销。
如果保存的名称已被占用，则可以覆盖它，但它不能覆盖 Relax 源代码中定义的函数的名称。
这实际上是在创建一个闭包，但该函数具有不同的名称，以避免与 invoke_closure 混淆（它们并非旨在一起使用）。
- set_input(func_name: str, *args: Any, **kwargs: Any) None
设置函数的输入。当通过 RPC 使用 VM 时，此接口通过在内部将参数中的 NDArray 转换为 DLTensor 来工作，DLTensor 在 RPC 中受支持，远程可能只有最小的 C 运行时。
注意：如果使用 set_input，则函数必须使用 invoke_stateful 调用，并且结果必须使用 get_outputs 获取。
- invoke_stateful(func_name: str) None
使用 set_input 设置的参数从 VM 模块调用命名函数。在没有首先使用 set_input 的情况下调用 invoke_stateful 是错误的（即使它是为了设置 0 个输入）；相反，如果已调用 set_input，则在不使用 invoke_stateful 的情况下调用该函数是错误的。
可以通过调用 get_outputs 获取调用的结果。
- 参数:
func_name (str) – 要调用的函数的名称。
- get_outputs(func_name: str) Object | Tuple[Any]
在调用 invoke_stateful 之后，获取给定名称的函数输出的值。
在没有首先调用 invoke_stateful 的情况下调用此函数是错误的。
- set_instrument(instrument: PackedFunc) None
设置一个 instrumentation 函数。
如果存在 instrument，则将在每个 Call 指令之前/之后调用该函数。该函数具有以下签名
def instrument( func: Union[VMClosure, PackedFunc], func_symbol: str, before_run: bool, ret_value: any, *args) -> bool: pass
instrument 接受以下参数： - func：要调用的函数对象。 - func_symbol：函数的符号名称。 - before_run：是在调用之前还是之后。 - ret_value：调用的返回值，仅在运行后有效。 - args：传递给调用的参数。
instrument 函数可以选择一个整数，该整数对应于后续运行的操作方向。有关更多详细信息，请参见 VMInstrumentReturnKind。
- 参数:
instrument (tvm.runtime.PackedFunc) – 一个 instrumentation 函数，每次 VM 调用指令都会调用它。
另请参阅
VMInstrumentReturnKindVM 中可能的返回值。
- time_evaluator(func_name: str, dev: Device, number: int = 10, repeat: int = 1, min_repeat_ms: int = 0, cooldown_interval_ms: int = 0, repeats_to_cooldown: int = 1, f_preproc: str = '') Callable[[...], BenchmarkResult]
返回一个评估器，用于测量模块中函数的时间。这遵循 tvm.runtime.module 中 time_evaluator 的相同约定。这可以与 save_function() 结合使用，以便计时避免额外的字典查找。
- 参数:
func_name (str) – 模块中函数的名称。
dev (Device) – 我们应该在其上运行此功能的设备。
number (int) – 为了取平均值，运行此函数的次数。我们将这些运行称为一次测量 repeat。
repeat (int, 可选) – 重复测量的次数。总共，该函数将被调用 (1 + number x repeat) 次，其中第一个是预热并将被丢弃。返回的结果包含 repeat 成本，每个成本都是 number 成本的平均值。
min_repeat_ms (int, 可选) – 一次 repeat 的最小持续时间，以毫秒为单位。默认情况下，一次 repeat 包含 number 次运行。如果设置此参数，则将动态调整参数 number 以满足一次 repeat 的最小持续时间要求。即，当一次 repeat 的运行时间低于此时间时，number 参数将自动增加。
cooldown_interval_ms (int, 可选) – 由 repeats_to_cooldown 定义的重复次数之间的冷却间隔，以毫秒为单位。
repeats_to_cooldown (int, 可选) – 冷却激活之前的重复次数。
f_preproc (str, 可选) – 我们希望在执行时间评估器之前执行的预处理函数名称。
注意
该函数将被调用 (1 + number x repeat) 次，第一次调用将被丢弃，以防出现延迟初始化。
示例
与 VM 函数的正常使用（如果函数返回元组，则可能无法通过 RPC 工作）
target = tvm.target.Target("llvm", host="llvm") ex = tvm.compile(TestTimeEvaluator, target) vm = relax.VirtualMachine(mod, tvm.cpu()) timing_res = vm.time_evaluator("func_name", tvm.cpu())(arg0, arg1, ..., argn)
与状态 API 一起使用
target = tvm.target.Target("llvm", host="llvm") ex = tvm.compile(TestTimeEvaluator, target) vm = relax.VirtualMachine(mod, tvm.cpu()) vm.set_input("func_name", arg0, arg1, ..., argn) timing_res = vm.time_evaluator("invoke_stateful", tvm.cpu())("func_name")
通过 save_function 保存的闭包（这会在计时部分中减少字典查找）
target = tvm.target.Target("llvm", host="llvm") ex = tvm.compile(TestTimeEvaluator, target) vm = relax.VirtualMachine(mod, tvm.cpu()) vm.save_function("func_name", "func_name_saved", arg0, arg1, ..., argn) timing_res = vm.time_evaluator("func_name_saved", tvm.cpu())()
- 返回:
ftimer – 该函数接受与 func 相同的参数，并返回 BenchmarkResult。ProfileResult 报告 repeat 时间成本，以秒为单位。
- 返回类型:
函数