核心内容摘要
2025：捕捉爱的瞬间，见证不一样的“家”
TVM 现已更新到
0.
2
0 版本[TVM 中文文档]已经和新版本对齐。
Apache TVM 是一个深度的深度学习编译框架适用于 CPU、GPU 和各种机器学习加速芯片。
更多 TVM 中文文档可访问 →[Apache TVM]在线运行 TVM 学习教程链接是https://hyper.ai/notebooks/48919?utm_sourceDistributeutm_mediumDistribute-TVMutm_campaignDistribute-TVM-260126Relax 与 TVM IR 都包含一系列优化传递optimization passes用于改进模型在特定设备上的性能指标例如推理平均时间、内存占用或功耗。
这些优化包括标准优化与机器学习特定优化如常量折叠constant folding、死代码消除、算子布局变换、算子融合、缓冲区处理和循环变换等。
每个传递都是基于收集的分析结果进行的 IR-to-IR 转换。
然而随着 TVM 的快速发展越来越需要一种系统化且高效的方式来管理这些传递。
此外一个通用的框架能够在 TVM 栈的不同层次例如 Relax 和 tir之间管理传递这为开发者快速原型化和集成新传递铺平了道路。
本文档介绍了这种基础设施的设计它结合了生产级编译器中用于管理优化传递的方式以及现代深度学习框架用于构建层次化结构的风格。
例如许多现有的生产级编译器如 GCC 与 LLVM 采用「传递管理器pass manager」来高效管理传递执行。
最初传递数量较少时管理很简单但成熟编译器可能包含数百个独立传递。
外部用户往往希望添加自定义传递并能正确调度而无需手动修改固定顺序。
类似地现代深度学习框架如 Pytorch 与 MXNet Gluon也倾向于通过Sequential和Block实现类似「传递式」层构建机制。
借助这些构造框架能够轻松将模块或层添加到容器中从而快速搭建神经网络。
TVM 的传递基础设施设计灵感主要来自 LLVM 的层次化传递管理器 以及流行深度学习框架的模块化容器。
该系统的主要目标包括支持更灵活的优化编排让用户能自由构建自定义优化流水线。
提供便捷的调试机制。
让开发者无需手动解决传递之间的依赖。
简化新传递的实现方式例如允许用户直接用 Python 实现一个传递由系统自动管理其执行。
设计概述系统重点关注可扩展性使用户能快速添加新传递而不破坏兼容性。
其结构包括后端与前端后端实现核心逻辑前端则提供简单的 API 供用户创建与控制优化流程。
C 后端我们提供PassInfo对象来存储单个传递所需的基本信息name为传递名opt_level指示该传递在哪个优化级别启用required表示执行该传递前所需的其他传递详见include/tvm/ir/transform.h。
在注册传递时开发者可以指定传递名称、优化级别与依赖。
opt_level可帮助系统在给定优化级别下判断某个传递是否需要执行required字段用于自动解析传递依赖。
class PassInfoNode : public Object { ffi::String name; int opt_level; ffi::Arrayffi::String required; };PassContextPassContext携带优化传递所需的关键信息。
例如它包含错误报告系统方便优化作者诊断失败原因。
PassContext也取代了旧的BuildConfig用于配置编译选项如优化级别、必需/禁用传递等。
例如我们可以配置在opt_level3下执行所有传递并通过disabled_passxx禁用某些传递系统会聚合该级别的所有传递并排除被禁用的项。
PassContext还提供对所有传递进行检测instrumentation的能力见pass_instrument_cpp_backend。
该类支持 Pythonwith语法便于在给定配置下执行优化。
同时用户可以通过PassContext::Current()在线程安全的方式获取当前上下文 因为系统使用线程本地存储PassContextThreadLocalStore来保存上下文对象。
class PassContextNode : public Object { public: int opt_level{2}; tvm::ffi::Arraytvm::Expr required_pass; tvm::ffi::Arraytvm::Expr disabled_pass; mutable ffi::OptionalDiagnosticContext diag_ctx; ffi::Mapffi::String, Any config; ffi::Arrayinstrument::PassInstrument instruments; }; class PassContext : public NodeRef { public: TVM_DLL static PassContext Create(); TVM_DLL static PassContext Current(); TVM_DLL void InstrumentEnterPassContext(); TVM_DLL void InstrumentExitPassContext(); TVM_DLL bool InstrumentBeforePass(const IRModule mod, const PassInfo info) const; TVM_DLL void InstrumentAfterPass(const IRModule mod, const PassInfo info) const; /* 其他字段省略 */ private: // 进入 pass 上下文作用域 TVM_DLL void EnterWithScope(); // 离开 pass 上下文作用域 TVM_DLL void ExitWithScope(); // 用于支持 Python with 语法 friend class tvm::WithPassContext; }; struct PassContextThreadLocalEntry { /*! rief 默认 pass 上下文 */ PassContext default_context; /*! rief 当前 pass 上下文 */ std::stackPassContext context_stack; PassContextThreadLocalEntry() { default_context PassContext(make_nodePassContextNode()); } }; /*! rief 线程本地存储用于保存 pass 上下文 */ typedef dmlc::ThreadLocalStorePassContextThreadLocalEntry PassContextThreadLocalStore;Pass 构造传递Pass基础设施以分层结构设计可在 Relax/tir 程序的不同粒度上工作。
系统定义了一个纯虚类PassNode作为各种优化传递的基类。
此类包含多个必须在子类中实现的虚函数适用于模块级、函数级或顺序传递级别。
class PassNode : Object { virtual PassInfo Info() const 0; virtual Module operator()(const IRModule mod, const PassContext pass_ctx) const 0; };该函数对象定义了传递的执行方式 每个传递都在特定上下文PassContext下作用于一个IRModule 并以Module到Module的方式实现。
因此所有传递都以模块为单位更新整个 IR。
系统实现了多个PassNode子类来支持不同类型的优化 包括函数级传递、模块级传递与顺序传递sequential pass。
每个子类本身都可充当一个传递管理器例如它们可以收集所需传递并执行或基于元信息建立依赖图。
完整定义见src/ir/transform.cc。
模块级传递模块级传递主要用于全局或过程间优化IPO类似于 LLVM 中的模块传递。
Relax 中一些典型需要全局视图的优化如 A-normal form 转换、lambda 提升就属于此类。
在该级别用户可以在模块中添加或删除函数。
class ModulePassNode : PassNode { PassInfo pass_info; std::functionModule(Module, PassContext) pass_func; Module operator()(const Module mod, const PassContext pass_ctx) const final; // 其他成员/方法省略 };pass_info存储模块传递的相关信息pass_func定义实际优化逻辑。
例如在模块上执行死代码消除可在pass_func中实现它将删除模块中未使用的函数。
此字段被设计为「打包函数packed function」 因此优化逻辑既可用 C 实现也可用 Python 实现。
函数级传递函数级传递用于实现 Relax/tir 模块中函数内的优化。
它一次提取模块中的一个函数进行优化输出优化后的 RelaxFunction或 tirPrimFunc。
多数优化都属于此类如 Relax 的公共子表达式消除、推理简化或 tir 的向量化与内存扁平化。
函数级传递仅作用于单个函数Relax 或 tir因此无法通过此类传递添加或删除函数因为其不具备全局信息。
class FunctionPassNode : PassNode { PassInfo pass_info; std::functionFunction(Function, Module, PassContext) pass_func; Module operator()(const Module mod, const PassContext pass_ctx) const final; bool SkipFunction(const Function func) const; // 其他成员/方法省略 };pass_info与模块级传递相同。
pass_func接受函数与模块作为输入可在函数上执行优化 函数若被注解为SkipOptimization将被跳过。
顺序传递Sequential PassSequentialPass类似于 PyTorch 的nn.Sequential可包含多个顺序执行的传递。
class SequentialPassNode : PassNode { PassInfo pass_info; // 需要执行的传递列表 ffi::ArrayPass passes; bool PassEnabled(const PassInfo info) const; Module operator()(const Module mod, const PassContext pass_ctx) const final; };以下展示顺序传递的执行逻辑系统会按照传递添加的顺序依次执行。
Module SequentialNode::operator()(const Module module, const PassContext pass_ctx) const { Module mod module; for (const Pass pass : passes) { ICHECK(pass.defined()) Found undefined pass for optimization.; const PassInfo pass_info pass-Info(); if (!PassEnabled(pass_info)) continue; for (const auto it : pass_info-required) { const auto* name it.astvm::ir::StringImm(); ICHECK(name); mod GetPass(name-value)(mod, pass_ctx); } mod pass(mod, pass_ctx); } return mod; }在执行传递前系统会判断该传递是否启用首先检查是否被用户禁用其次查看是否被显式声明为必需。
若仍未确定则根据opt_level判断是否执行。
执行时系统会根据传递名从注册表中获取对应实现Pass GetPass(const std::string pass_name) { using tvm::runtime::Registry; std::string fpass_name relax.transform. pass_name; const std::optionaltvm::ffi::Function f tvm::ffi::Function::GetGlobal(fpass_name); ICHECK(f.has_value()) Cannot find fpass_name to create the pass pass_name; return (*f)(); }系统还提供辅助函数用于创建各类传递并暴露给 Python 前端Pass CreateFunctionPass( std::functionFunction(Function, IRModule, PassContext) pass_func, int opt_level, ffi::String name, ffi::Arrayffi::String required); Pass CreatePrimFuncPass( std::functionPrimFunc(PrimFunc, IRModule, PassContext) pass_func, int opt_level, ffi::String name, ffi::Arrayffi::String required); Pass CreateModulePass( std::functionIRModule(IRModule, PassContext) pass_func, int opt_level, ffi::String name, ffi::Arrayffi::String required); Pass Sequential(tvm::ffi::ArrayPass passes, PassInfo pass_info);传递注册前文介绍了不同粒度的传递和编译上下文。
下面展示如何注册一个传递。
以常量折叠constant folding为例 它用于在 Relax 函数中折叠常量实现位于 src/relax/transforms/fold_constant.cc。
该传递提供了Expr到Expr的转换 APIExpr FoldConstant(const Expr expr);要将其注册到传递基础设施中首先需要确定传递的粒度。
常量折叠作用于函数级因此通过CreateFunctionPass创建pass_func以打包函数形式返回用于对 [IRModule]{.title-ref} 中的每个函数调用该转换 API。
{}表示该传递没有前置依赖若有依赖开发者需明确列出。
同时注册名为relax.transform.FoldConstant的 API 入口使该传递可被 C 例如以上的GetPass与 Python 访问namespace transform { Pass FoldConstant() { auto pass_func [](Function f, IRModule m, PassContext pc) { return ConstantFolder::Fold(f, m); }; return CreateFunctionPass(pass_func, 0, FoldConstant, {}); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl tvm::ffi::reflection; refl::GlobalDef().def(relax.transform.FoldConstant, FoldConstant); } } // namespace transform为方便其他 C 模块调用在include/tvm/relax/transform.h中声明TVM_DLL Pass FoldConstant();传递检测Pass Instrument传递检测机制用于分析传递本身例如统计执行时间与内存占用或观察 IR 如何被改变。
我们在PassContext生命周期中引入四个检测点TVM_DLL void InstrumentEnterPassContext(); TVM_DLL void InstrumentExitPassContext(); TVM_DLL bool InstrumentBeforePass(const IRModule mod, const PassInfo info) const; TVM_DLL void InstrumentAfterPass(const IRModule mod, const PassInfo info) const;InstrumentEnterPassContext在进入PassContext作用域时调用。
InstrumentExitPassContext在离开PassContext或执行发生异常时调用。
当通过 :pytvm.transform.PassContext的override_instruments覆盖检测器时也会触发见pass_instrument_overriden。
InstrumentBeforePass在传递执行前调用 若该传递应执行则在执行后调用InstrumentAfterPass。
其伪代码如下if (pass_ctx.InstrumentBeforePass(ir_module, pass_info)) { new_ir_module run_pass(ir_module, pass_ctx); pass_ctx.InstrumentAfterPass(new_ir_module, pass_info); return new_ir_module; }PassInstrument接口允许你在上述四个阶段插入自定义逻辑。
可向单个PassContext注册多个检测器实例它们将按instruments指定的顺序依次调用。
接口定义如下namespace instrument { class PassInstrumentNode : public Object { public: ffi::String name; virtual void EnterPassContext() const 0; virtual void ExitPassContext() const 0; virtual bool ShouldRun(const IRModule mod, const transform::PassInfo info) const 0; virtual void RunBeforePass(const IRModule mod, const transform::PassInfo info) const 0; virtual void RunAfterPass(const IRModule mod, const transform::PassInfo info) const 0; /* 其他字段省略 */ }; class PassInstrument : public ObjectRef { public: TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PassInstrument, ObjectRef, PassInstrumentNode); }; } // namespace instrumentPython 前端提供了便捷方式来实现PassInstrument见pass_instrument_py_frontend。
在一个PassContext中某个PassInstrument实例的调用顺序如下with PassContext(instruments[pi]) # pi 为某个 PassInstrument 实现 pi.EnterPassContext() if pi.ShouldRun(Pass
: pi.RunBeforePass() Pass1() pi.RunAfterPass() if pi.ShouldRun(Pass
: pi.RunBeforePass() Pass2() pi.RunAfterPass() pi.ExitPassContext()以下简述PassInstrument与PassContext方法之间的关系详见 src/ir/transform.ccInstrumentEnterPassContextEnterPassContext()按传入instruments的顺序执行。
若执行中抛出异常PassContext会清空所有已注册的检测器。
然后对已成功执行EnterPassContext()的检测器依次调用ExitPassContext()。
例如注册了 A、B、C 三个检测器A 成功B 抛异常则 C 不会执行随后调用 A 的ExitPassContext()。
InstrumentExitPassContext各检测器的ExitPassContext()按instruments顺序执行。
若发生异常instruments会被清空。
抛出异常后注册的检测器不会执行ExitPassContext。
InstrumentBeforePass若该传递未被显式列为必需则会调用ShouldRun。
若未被ShouldRun阻塞则按顺序调用RunBeforePass。
该函数返回布尔值指示该传递是否应执行。
若发生异常将立即抛出Python 依靠上下文管理器安全退出确保各检测器的ExitPassContext被调用C 见 include/tvm/support/with.h。
InstrumentAfterPass按顺序调用RunAfterPass。
若发生异常将立即抛出依靠上下文管理器或With类include/tvm/support/with.h安全退出。
内置检测器系统内置若干检测器标注TODO的尚未实现PassTimingInstrument见 src/ir/instrument.cc用于分析各传递的执行时间。
PrintIRBeforeTODO在传递执行前打印 IR。
也可通过 :pytvm.transform.PrintIR{.interpreted-text role“func”} 在传递周围插入打印实现但使用检测器无需修改传递序列。
PrintAfterTODO在传递执行后打印 IR。
Python 前端前端仅需少量 API 即可创建并执行传递完整实现见python/tvm/relax/transform/transform.py与python/tvm/ir/transform.py。
后端将根据提供的信息决定如何创建 Pass 对象。
PassContextPython 前端为PassContext提供了包装以支持with语法并提供current静态方法tvm_ffi.register_object(transform.PassContext) class PassContext(tvm.runtime.Object): def __enter__(self): _transform.EnterPassContext(self) return self def __exit__(self, ptype, value, trace, config): _transform.ExitPassContext(self) staticmethod def current(): Return the current pass context. return _transform.GetCurrentPassContext()PassContext用于配置编译选项优化级别、必需/禁用传递等并可传入配置字典以便不同传递读取需要的数据如回退设备信息、循环展开步数/深度等。
若要从config中获取某项配置其键名需通过TVM_REGISTER_PASS_CONFIG_OPTION注册例如循环展开传递TVM_REGISTER_PASS_CONFIG_OPTION(tir.UnrollLoop, UnrollLoopConfig);详见src/tir/transforms/unroll_loop.cc。
Python 中的传递检测使用装饰器python/tvm/ir/instrument.py可以快速实现PassInstrument。
推荐使用装饰器方式而非继承enter_pass_ctx进入PassContext时执行exit_pass_ctx退出PassContext时执行should_run在传递执行前调用返回该传递是否应执行run_before_pass传递执行前调用run_after_pass传递执行后调用。
可通过 :pytvm.transform.PassContext的instruments参数注册实例。
更多示例见use pass instrument教程。
覆盖当前 PassContext 中的检测器override_instruments方法可覆盖当前PassContext中的instruments。
例如当未显式创建新PassContext而直接运行传递时仍可将检测器注册到全局上下文cur_pass_ctx tvm.transform.PassContext.current() # 覆盖 PassInstrument 实例 cur_pass_ctx.override_instruments([pass_inst]) mod pass_seq(mod) result pass_inst.get_result()注意调用override_instruments时旧检测器的exit_pass_ctx会被调用随后新检测器的enter_pass_ctx会被调用。