ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

TVM Relay 后端深度解析:te_compiler、解释器、图执行器与虚拟机(tvm.relay.backend 全模块指南)

TVM Relay 后端深度解析:te_compiler、解释器、图执行器与虚拟机(tvm.relay.backend 全模块指南) 编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载导读tvm.relay.backend是 Apache TVM 中承接「Relay 前端 IR」与「后端运行时」的关键桥梁模块它负责把经过算子融合与优化后的 Relay 函数降级lower为可直接在 CPU/GPU/专用加速器上执行的底层实现。本篇指南以 TVM 仓库中的官方 API 文档 backend.rst 为骨架深入其引用的四个核心子模块——interpreter、te_compiler、graph_executor_codegen与vm并结合 python/tvm/relay/backend/ 与 src/relay/backend/ 下的真实源码逐一讲解 TE 编译器缓存机制、算子实现选择策略、参考解释器、图执行器代码生成和 Relay 虚拟机的完整调用链。读完本文你将能理解relay.build从 Relay 模块到可运行运行时模块的完整降级路径并掌握relay.create_executor、TECompiler、VMCompiler等关键 API 的实际用法。一、模块总览backend 包导出什么在 python/tvm/relay/backend/init.py 中tvm.relay.backend包只做了三件事Backend codegen modules for relay. from . import te_compiler from .executor import Executor from .runtime import Runtime导入te_compiler子模块TE 编译器引擎导出Executor执行器配置类见 executor.py导出Runtime运行时配置类见 runtime.py。这两个配置类是后续relay.build阶段控制用哪种执行器executor、跑在哪个运行时runtime的核心句柄类构造方式注册表名关键方法Executor(name, options)如tvm.relay.backend.Executor(graph, {link-params: 1})executorlist_registered()、list_registered_options()Runtime(name, options)如tvm.relay.backend.Runtime(cpp)runtimelist_registered()、list_registered_options()Executor.__getitem__/Runtime.__getitem__可用于读取对应配置项底层通过_backend.CreateExecutor/_backend.CreateRuntimeC FFI注册对象真正的可选项由 C 侧src/relay/backend/executor.cc、src/relay/backend/runtime.cc维护。二、tvm.relay.backend.interpreterRelay 参考解释器文档的automodule:: tvm.relay.backend.interpreter对应 interpreter.py其模块 docstring 明确写它是 The Python interface to the Relay reference interpreter。2.1 抽象基类ExecutorExecutor是执行 Relay 程序的抽象接口定义了三个方法_convert_args(expr, args, kwargs)把「位置参数 关键字参数」合并为与 Relay 函数形参顺序一致的参数序列。规则是先依次放置位置参数再用关键字参数补齐剩余形参若某个形参既被位置参数占用又出现在 kwargs 中会抛出 duplicate argument supplied 异常参数不足则抛 insufficient arguments。_make_executor(exprNone)构造一个实现了expr行为的 Python 可调用对象子类必须实现。evaluate(exprNone, bindsNone)评估 Relay 表达式。binds允许为自由变量提供额外绑定内部通过ScopeBuilder转成let绑定当expr是Function或GlobalVar时直接构造执行器普通表达式则会包装成无参函数执行。若expr为 None则默认取模块中的main函数。executor relay.create_executor(kinddebug, modmodule) func executor.evaluate(expr) # 复用同一个可调用对象 a func(args1) b func(args2)2.2 解释器实现InterpreterInterpreter(Executor)的构造参数为modtvm.IRModule、device运行时设备、target构建目标仅支持同构执行。其_make_executor通过 FFI 调用_backend.EvalFunction(self.mod, expr, self.device, self.target)把表达式编译为一个可重复应用的 packed function随后把 Python 侧的 numpy 数组 / NDArray / ADT / 元组等参数通过_arg_to_ast反射回 Relay ASTConstant、Tuple、RefCreate、Call、const再调用 packed function 求值。需要特别注意的是 docstring 中的 CAUTION解释器在每次evaluate调用时才准备模块而非在create_executor时一次性准备因此若用同一 executor 多次求值不同参数模块绑定会被重复处理。最佳实践是像上文示例那样把evaluate的结果保存为函数复用而不是每轮都调用evaluate。2.3 工厂函数relay.create_executor解释器通常不是直接实例化的而是通过 build_module.py 中的relay.create_executor工厂创建def create_executor(kinddebug, modNone, deviceNone, targetllvm, paramsNone):kind执行器类型可选debug解释器、graph图执行器、aotAOT 执行器、vm虚拟机modtvm.IRModule缺省为空模块device执行设备缺省时从第一个 target 推导默认设备target支持多目标heterogeneous但文档明确提示该 API 不允许多设备因此异构编译尚未支持params推理期不变的常量参数会通过bind_params_by_name绑定进main。官方 docstring 自带一个最小可运行示例构造x 1的 Relay 表达式后用kindvm求值输入[2]返回[3.]。三、tvm.relay.backend.te_compilerTE 编译引擎文档的automodule:: tvm.relay.backend.te_compiler对应 te_compiler.py模块 docstring 即 TE compiler engine (replacing legacy compile_engine)——它是取代旧compile_engine的新一代编译引擎。3.1 核心数据结构LoweredOutput承载outputste.Tensor 列表与implementOpImplementation两个字段是一次算子降级的输出结果。CCacheKey编译缓存的键由source_funcrelay.Function与targettvm.Target组成。CCacheValue编译缓存的值including usage statistics含使用统计用于后续缓存淘汰策略。3.2 算子实现选择get_valid_implementations与select_implementation这是 TE 编译引擎最核心的决策逻辑get_valid_implementations(op, attrs, inputs, out_type, target)读取算子注册的FTVMStrategy可用tvm.relay.op.register_strategy注册在 target 上下文中调用 strategy 得到OpStrategy遍历其specializations用tvm.arith.Analyzer的canonical_simplify化简各 specialization 的条件子句筛选出全部有效实现。select_implementation(op, attrs, inputs, out_type, target, use_autotvmTrue)在上面的基础上挑选最优实现若use_autotvmFalse或 AutoTVM / AutoScheduler / MetaSchedule 已启用直接选择plevel最高的实现若use_autotvmTrue先对每个可调优实现计算 AutoTVM workload查询DispatchContext得到配置与代价选代价最低的实现找不到调优记录时退回最高plevel实现并打印提示 One or more operators have not been tuned. Please tune your model for better performance.可通过 DEBUG 日志查看细节找不到任何有效实现时抛出RuntimeError(fNo valid {op} implementations for {target})。从源码结构看plevel优先级机制是 TVM 为同一算子注册多种后端实现如 CUDA 的 cublas、cutlass 与 generic 实现时做默认排序的依据。3.3 动态形状处理lower_call注册为relay.backend.lower_call在把 Call 表达式降级前会检查输入输出类型是否动态_ty.is_dynamic对动态形状的调用会选择use_autotvmFalse直接取最高优先级实现因为当前 TVM 尚未允许为动态形状自动生成多个 kernel源码中保留 TODO 注释。同时get_shape负责把IntImm维度按INDEX_DEFAULT_I64编译开关转换为 int64 或 int32把tvm.tir.Any动态维度转为te.size_var(any_dim, int32)。3.4TECompiler类与全局单例class TECompiler(Object): def lower(self, source_func, targetNone, mod_namedefault): ... def jit(self, source_func, targetNone): ... def clear(self): ... def items(self): ...lower把source_funcrelay.Function 或 CCacheKey降级为CachedFuncmod_name会先经过mangle_module_name修正出错时会把函数 ASTastext(show_meta_dataFalse)拼进异常信息方便定位。jit即时编译为tvm.runtime.PackedFunc。clear清空缓存。items列出缓存中的所有(CCacheKey, CCacheValue)对。get()返回全局 TE 编译器单例_backend._TECompilerGlobal()。此外模块还提供lower_to_primfunc(relay_func, target)调用全局函数relay.backend.LowerToPrimFunc把融合后的 Relay 原始函数直接转为tir.PrimFunc——这一步正是Relay → TIR的关键跳转。C 侧实现位于 src/relay/backend/te_compiler.cc、src/relay/backend/te_compiler_cache.cc缓存与 src/relay/backend/task_extraction.ccAutoTVM 任务抽取。四、tvm.relay.backend.graph_executor_codegen图执行器代码生成文档的automodule:: tvm.relay.backend.graph_executor_codegen对应 graph_executor_codegen.py。模块 docstring 把编译流程概括为三步把单个 Relay 表达式要求是函数编译为图语言程序函数形参对应计算图中的 placeholder/输入与模型参数函数体即计算图本身图语言由Node、NodeRef、InputNode、OpNode构成表示 TVM 图格式的程序通过一个 printer 把图格式转成JSON 字符串该字符串可被contrib.graph_executor或其他兼容 TVM runtime 的系统加载。GraphExecutorCodegen(mod, target)的codegen(ir_module, func)方法返回三元组graph_json : str—— 供运行时消费的图 JSONmod : IRModule—— 降级后的函数集合params : Dict[str, tvm.nd.NDArray]—— 额外的常量参数通过_list_params_name与_get_param_by_name从 C 侧取回并 copy 到新数组。底层 C 实现在 src/relay/backend/graph_executor_codegen.cc内存规划依赖 src/relay/backend/graph_plan_memory.cc静态内存池分配_GraphExecutorCodegenFFI 对象则由 python/tvm/relay/_build_module.py 暴露。这一路径是传统relay.buildtvm.contrib.graph_executor部署方案含 Android/iOS 等移动端的基石。五、tvm.relay.backend.vmRelay 虚拟机文档的automodule:: tvm.relay.backend.vm对应 vm.py为 Relay VM 提供 Python 接口包含编译与执行两部分。5.1 便捷函数compiledef compile(mod, targetNone, target_hostNone, paramsNone): compiler VMCompiler() if params: compiler.set_params(params) compiler.lower(mod, target, target_host) compiler.codegen() return compiler.get_exec()一次编译分四步设置参数 →lower降级为 VM 字节码 →codegen生成 kernel 库 → 取回tvm.runtime.vm.Executable同时包含库代码与字节码。5.2VMCompiler类set_params(params)/get_params()设置/取回常量参数。set_params会把 numpy 数组转为NDArray再包成relay.constget_params返回{name: ndarray}字典。lower(mod, target, target_host)通过Target.canon_multi_target_and_host规范化目标后调用 C_VMCompiler[lower]。注意它会先根据当前DispatchContext判断是否处于 fallback根上下文——若是则进入autotvm.tophub.context(raw_targets)即自动从TopHub拉取预调优参数否则使用EmptyContext。这就是 VM 编译默认能吃到社区调优结果的原因。codegen()生成 kernel 库。optimize(mod, target, target_host, params)仅优化返回优化后的IRModule与参数不产出可执行文件适合调试或查看优化结果。get_exec()返回tvm.runtime.vm.Executable。target_host参数的语义编译 CUDA 等设备程序时还需要 CPU 侧的宿主代码与驱动交互设置维度与参数target_host即宿主代码生成目标默认优先用llvm未启用时退回 stackvm 解释器。5.3VMExecutorVMExecutor(Executor)把 VM 适配进统一的执行器接口_make_executor中用self.mod[main] expr替换 main 后调用compile得到可执行文件实例化tvm.runtime.vm.VirtualMachine最终返回的包装函数把 Python 参数经过_convert_args转换后直接self.vm.run(*args)。它特别适合做实验与调试docstring 也说明 VM 可直接通过tvm.runtime.vm使用。C 侧编译实现在 src/relay/backend/vm/compiler.cc配套的 VM 专用 pass 包括lambda_lift.cc闭包转换、manifest_lifetimes.cc生命周期分析与removed_unused_funcs.cc无用函数移除。六、后端降级全景与验证路径综合以上四个子模块一次典型的relay.build后端流程可以归纳为Relay IRModule │ pass 优化算子融合 FuseOps 等 ▼ 融合后的原始函数primitive function │ te_compiler.lower_call / select_implementationstrategy plevel AutoTVM ▼ te.Tensor 计算 schedule → LowerToPrimFunc │ TECompilerCacheCCacheKey → CachedFunc ▼ TIR PrimFunc → 各目标代码生成LLVM/CUDA/OpenCL/... │ graph_executor_codegen 或 vm/compiler 或 aot_executor_codegen ▼ graph JSON lowered func params → 运行时模块在仓库中验证这条链路最直接的方式是运行测试tests/python/relay/下存在大量覆盖 interpreter、VM、graph executor 与 TECompiler 的用例例如通过python -m pytest tests/python/relay/test_vm.py -k compile可验证 VM 编译路径。对希望深入 C 实现的读者推荐按以下顺序阅读src/relay/backend/te_compiler.cc src/relay/backend/te_compiler_cache.cc降级主流程与缓存src/relay/backend/graph_executor_codegen.cc src/relay/backend/graph_plan_memory.cc图 JSON 生成与静态内存规划src/relay/backend/interpreter.cc参考解释器求值src/relay/backend/vm/compiler.ccVM 字节码编译。七、小结如何选择后端执行路径执行路径适用场景核心 APIdebugInterpreter原型验证、逐表达式调试不追求性能relay.create_executor(kinddebug)graphGraphExecutor静态形状、传统部署含移动端relay.buildtvm.contrib.graph_executorvmVirtualMachine动态形状、函数调用/控制流丰富的模型relay.vm.compile/VMCompileraotAOT Executor无操作系统嵌入式/MCU 场景编译期规划内存relay.build(..., executorExecutor(aot))tvm.relay.backend的设计核心在于解耦前端策略FTVMStrategy/ plevel决定有哪些实现、AutoTVM 决定选哪个实现、TECompiler决定如何缓存与复用降级结果最终由 interpreter / graph / vm / aot 四种执行器把降级产物映射到对应运行时。理解这条链路就等于掌握了 TVM 从模型 IR 到实际推理的全部关键环节。赞分享编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载相关推荐TVM Relay 虚拟机Virtual Machine架构深度解析设计、指令集、编译器与序列化TVM Relay 虚拟机Virtual Machine架构深度解析设计、指令集、编译器与序列化 导读 本文以 docs/arch/virtual_mac编译器深度学习模型优化Electrobun 与 Svelte 实战轻量级前端框架的桌面应用开发Electrobun 与 Svelte 实战轻量级前端框架的桌面应用开发 Electrobun 是一个基于 TypeScript 的超快速、轻量级跨平台桌面应桌面应用跨平台KDBush2D点的终极静态空间索引库让搜索速度提升10倍KDBush2D点的终极静态空间索引库让搜索速度提升10倍 KDBush 是一款基于扁平 KD 树的超快速 2D 点静态空间索引库专为高效处理海量二维点创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表