tvm.script.printer 详解:TVM 统一打印器的 Doc 中间表示与 Python 输出管线

发布时间:2026/9/23 2:35:36
tvm.script.printer 详解:TVM 统一打印器的 Doc 中间表示与 Python 输出管线 模型编译深度学习推理引擎【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址https://gitcode.com/gh_mirrors/tv/tvm点击查看免费下载tvm.script.printer是 Apache TVM 的「统一打印器」Unified PrinterPython 接口负责将各类 TVM IR包括 TIR、Relax 等方言以**可往返roundtrippable**的方式打印为 TVMScript 文本。本文以其 API 参考文档docs/reference/api/python/script/printer.rst为骨架结合仓库源码完整介绍该模块的 Doc 中间表示体系、to_python_script输出控制参数以及 C 侧Script()入口与分发表vtable的底层实现。读完本文你将掌握如何通过 Doc 对象手工构造 TVMScript 输出以及如何配置打印器的缩进、行号、语法糖、下划线标注等行为。模块定位从 IR 对象到 TVMScript 文本的两阶段管线tvm.script.printer的模块说明非常精炼见 python/tvm/script/printer/init.pyTVMScript Unified Printer本包提供一组 API将受支持的 TVM IR 以可往返roundtrippable的方式打印为 TVMScript。所谓「可往返」指的是打印出的 TVMScript 文本可以再次被 TVMScript 解析器读回还原出等价的 IR 对象形成IR → 文本 → IR的闭环。这决定了打印器在设计上必须保留类型信息、结构信息与足够的语法细节而不能像普通调试输出那样随意丢弃语义。从源码结构看整个打印过程是一条清晰的两阶段管线第一阶段的核心数据结构见 include/tvm/script/printer/ir_docsifier.hIR → Doc由IRDocsifier将 IR 节点Var、Expr、Stmt、Module 等递归转换为一种与具体方言无关的中间表示——Doc对象图。转换过程通过「分发表」dispatch vtable按对象类型分发并借助Frame栈来维护作用域、变量绑定关系与元数据。Doc → 文本由DocPrinter将Doc对象图渲染为 Python 语法的 TVMScript 字符串。Python 侧暴露的入口正是 doc_printer.py 中的to_python_script()。tvm.script.printer的 Python 包装分为三个文件_ffi_api.py通过tvm_ffi.init_ffi_api(script.printer, __name__)加载script.printer命名空间下的所有 FFI 函数是 Python 与 C 实现之间的桥梁doc.py定义完整的 Doc 中间表示类型体系doc_printer.py提供将 Doc 转换为 Python 文本字符串的to_python_script()。Doc 中间表示打印器的「AST」Doc 体系是整个打印器的核心抽象它本质上是 Python 语法树AST的 TVM 版本用对象图来描述将要打印的代码结构。所有 Doc 都继承自Doc基类其中ExprDoc与StmtDoc分别对应 Python 的「表达式」与「语句」两大类。完整定义见 doc.py下表汇总了全部 Doc 类型及其用途分类Doc 类型对应 Python 语法关键字段基类Doc—所有 Doc 的基类表达式ExprDoc表达式提供attr()、call()、__getitem__便捷构造表达式LiteralDoc字面量valuestr / float / bool / int / None表达式IdDoc标识符name表达式AttrAccessDoc属性访问a.bvalue、name表达式IndexDoc下标访问a[i]value、indices表达式CallDoc函数调用f(a, b1)callee、args、kwargs_keys、kwargs_values表达式OperationDoc一元/二元/三元运算kindOperationKind、operands表达式LambdaDoclambda表达式args、body表达式TupleDoc/ListDoc/DictDoc元组 / 列表 / 字典字面量elements/keys、values表达式SliceDoc切片a[i:j:k]start、stop、step语句StmtBlockDoc语句块仅容器不进入 IRstmts语句AssignDoc赋值x y/x: Tlhs、rhs、annotation语句IfDoc/WhileDoc/ForDoc条件 / 循环predicate、then_branch、else_branch等语句ScopeDocwith语句lhs、rhs、body语句ExprStmtDoc表达式语句expr语句AssertDoc/ReturnDocassert/returntest、msg/value语句FunctionDoc/ClassDoc函数 / 类定义name、args、decorators、body等语句CommentDoc/DocStringDoc注释 / 文档字符串comment/docsExprDoc表达式 Doc 与链式构造ExprDoc是表达式 Doc 的基类也是打印器日常使用最频繁的构造入口。它提供了三个便捷方法让开发者可以像写 Python 一样组装表达式 Doc见 doc.pyattr(name)生成AttrAccessDoc表示属性访问即obj.attrcall(*args, **kwargs)生成CallDoc以自身为被调用者支持位置参数与关键字参数__getitem__(indices)生成IndexDoc表示下标访问支持元组形式的多维下标。此外ExprDoc还实现了__iter__并抛出RuntimeError这是有意为之根据 PEP-234若对象只实现__getitem__而未实现__iter__解释器会尝试通过__getitem__(0)、__getitem__(1)…来迭代它直到IndexError。显式实现__iter__可以避免这种令人困惑的错误信息让用户立即得知ExprDoc不能被当作可迭代对象使用。LiteralDoc是字面量 Doc其构造函数根据 Python 值类型自动分派到对应的 FFI 构造器None、str、float、bool、int分别对应LiteralDocNone、LiteralDocStr、LiteralDocFloat、LiteralDocBoolean、LiteralDocInt其余类型直接抛出TypeError。需要注意的是LiteralDoc的value字段类型为str | IntImm | FloatImm | None——在 IR 打印场景中整数与浮点字面量通常携带 TVM 的IntImm/FloatImm信息来自tvm.tirx以保证往返时的精度与类型。CallDoc的实现细节值得一提构造函数将关键字参数拆分为kwargs_keys与kwargs_values两个并列列表见 doc.py这一设计是为了与 FFI 侧的数组容器无缝对接避免在 C 侧维护 Python dict 的键值顺序。OperationKind运算符枚举OperationDoc用于表示一元、二元以及特殊运算符其kind字段的类型是OperationKind枚举见 doc.py。该枚举的命名规范直接对照 Python 标准库ast模块并镜像自 C 侧OperationDocNode::Kind位于include/tvm/script/printer/doc.h。枚举区间划分为_UnaryStart~_UnaryEnd一元运算符包括USub负号、Invert按位取反、Not逻辑非_BinaryStart~_BinaryEnd二元运算符覆盖Add、Sub、Mult、Div、FloorDiv、Mod、Pow、LShift、RShift、BitAnd、BitOr、BitXor、Lt、LtE、Eq、NotEq、Gt、GtE、And、Or、MatMul矩阵乘法_SpecialStart~_SpecialEnd特殊运算符目前仅包含IfThenElsePython 三元表达式a if cond else b。StmtDoc语句 Doc语句类 Doc 覆盖了 TVMScript 所需的全部 Python 语句形态doc.pyAssignDoclhs、rhs、annotation三个字段使其既能表达普通赋值也能表达带类型注解的声明如x: T.int32IfDoc/WhileDoc/ForDoc条件分支与循环ForDoc的lhs/rhs分别对应循环变量与迭代对象ScopeDoc对应 Python 的with语句语义为with rhs as lhs: body...这是 TVMScript 中T.prim_func、R.function等作用域语法的关键载体ExprStmtDoc把表达式当语句输出如函数调用语句AssertDoc/ReturnDoc断言与返回FunctionDoc/ClassDoc函数与类定义其中FunctionDoc支持decorators装饰器列表、return_type返回类型注解与type_params类型参数供泛型/多态打印使用CommentDoc/DocStringDoc注释与文档字符串用于输出# comment与docstringStmtBlockDoc一个特殊的容器 Doc注释明确指出它「永远不会出现在 IR 中」只作为临时容器在打印过程中容纳一组StmtDoc。输出控制to_python_script 与 PrinterConfig第二阶段的核心入口是to_python_script()见 doc_printer.py签名如下def to_python_script( doc: Doc, indent_spaces: int 4, print_line_numbers: bool False, num_context_lines: int | None None, path_to_underline: list[AccessPath] | None None, ) - str它接收一个Doc内部构造PrinterConfig并调用 FFI 函数DocToPythonScript完成渲染各参数含义参数默认值说明indent_spaces4输出文本的缩进空格数print_line_numbersFalse是否在输出中附带行号num_context_linesNone在需要下划线标注的文本周围打印的上下文行数path_to_underlineNone需要加下划线标注的对象路径AccessPath列表这些参数会直接映射到 C 侧PrinterConfigNode的对应字段见 include/tvm/script/printer/config.h。PrinterConfig是一个完整的配置对象除了上述四个参数外还包含更多打印行为开关其中部分字段如下binding_names绑定层级名称栈用于跨函数调用时保持变量名一致性show_meta是否显示元数据段默认为falseir_prefixIR 节点前缀默认ITIR 与 Relax 方言会分别覆盖为T/Rmodule_alias跨函数调用时当前模块的别名默认cls为空时直接使用模块名buffer_dtype默认 buffer 数据类型默认float32int_dtype整数字面量的默认数据类型默认int32float_dtype浮点字面量的默认数据类型默认为空即始终显式打印T.float32/T.float64包装verbose_expr是否冗长打印表达式syntax_sugar是否输出语法糖形式设为false可进行完整无糖打印默认trueshow_object_address变量名是否包含对象地址render_invisible_path_info对不可见的下划线路径是否渲染访问路径上下文默认truepath_to_underline/path_to_annotate以路径方式指定下划线 / 注解对象obj_to_underline/obj_to_annotate直接以对象方式指定下划线 / 注解对象extra_config方言级扩展配置的通用映射键按dialect.knob约定命名例如tirx.prefix、relax.prefix、relax.show_all_ty通过GetExtraConfigT(key, fallback)读取。从 src/script/printer/script_printer.cc 的PrinterConfig构造函数可以看到所有这些配置均支持通过一个config_dict字典批量传入例如name、show_meta、ir_prefix、module_alias、buffer_dtype、int_dtype、float_dtype、verbose_expr、indent_spaces、print_line_numbers、num_context_lines、path_to_underline等键都会被逐一解析并写入PrinterConfigNode。因此to_python_script的显式参数只是这个配置体系在 Python 层的便捷子集开发者若有更精细的控制需求可以直接构造PrinterConfig再调用底层 FFI。底层实现Script() 入口、分发表与 ReprPrint 回退整个打印器的 C 入口是自由函数tvm::Script(node, config)声明见 include/tvm/script/printer/printer.h实现在 src/script/printer/script_printer.ccstd::string Script(const ffi::ObjectRef node, const ffi::OptionalPrinterConfig cfg) { PrinterConfig config cfg.value_or(PrinterConfig()); if (!TVMScriptPrinter::vtable().CanDispatch(node)) { // Fall back to ffi::ReprPrint for types not registered with TVMScriptPrinter. return RenderFallbackWithInvisiblePathInfo(ffi::ReprPrint(ffi::Any(node)), config); } return TVMScriptPrinter::vtable()(node, config); }其核心逻辑是若节点类型已在TVMScriptPrinter::vtable()一个全局ObjectFunctor分发表中注册了打印函数则按类型分发打印否则回退到通用的ffi::ReprPrint并在需要时附加不可见路径信息Access path: ...提示见RenderFallbackWithInvisiblePathInfo。各方言TIR、Relax 等通过宏TVM_REGISTER_SCRIPT_AS_REPR(ObjectType, Method)完成注册printer.h。该宏在静态初始化块中做两件事将kRepr类型属性指向RedirectedReprPrinterMethod该函数尝试用tvm::Script格式化对象失败时回退为纯地址字符串实现于src/script/printer/config.cc从而让str()/ 打印某个 IR 对象时自动走统一打印器将具体的打印函数Method注册进TVMScriptPrinter::vtable()的分发表。在 IR → Doc 阶段IRDocsifier承担核心转换职责include/tvm/script/printer/ir_docsifier.hframes帧栈FrameNode记录当前作用域已生成的语句stmts、对应的 docsifier 指针以及退出帧时需执行的回调callbacksdispatch_tokens分发表令牌栈栈顶令牌决定当前使用哪个方言的转换函数obj2info变量表将 IR 变量映射到名字与 Doc 构造器DocCreator。Define方法在定义变量时会自动改名以避免与既有变量冲突metadata/global_infos元数据段与 GlobalInfo 打印所需的累积容器defined_names已占用的变量名集合配合SetCommonPrefix实现共享前缀信息的变量名优化。AsDoc()是 IR → Doc 的通用转换模板对None、bool、int、float、字符串、DataType、Device等基础 FFI 类型直接生成对应的LiteralDoc多行字符串会走AddMetadata放入元数据段其余ObjectRef则通过分发表递归转换并追加source_paths与下划线/注解装饰AddDocDecoration。此外FrameNode::AddDispatchToken在进入作用域时压入方言令牌并在帧退出时自动弹出保证了with作用域内外的打印语义正确隔离。实战手工构造 Doc 并输出 TVMScript虽然日常开发中打印 IR 只需直接对 IR 对象调用str()或tvm.script.from_ir但理解 Doc 体系后我们也可以手工构造 Doc 并验证输出。例如构造一条带类型注解的赋值语句x: T.int32 1并打印from tvm.script.printer.doc import IdDoc, LiteralDoc, AssignDoc from tvm.script.printer.doc_printer import to_python_script # 左值标识符 x注解与右值字面量 lhs IdDoc(x) annotation IdDoc(T.int32) rhs LiteralDoc(1) doc AssignDoc(lhslhs, rhsrhs, annotationannotation) print(to_python_script(doc, indent_spaces4))输出将是x: T.int32 1再如利用ExprDoc.call与attr的链式构造生成函数调用T.max(3, 5)from tvm.script.printer.doc import IdDoc, LiteralDoc callee IdDoc(T.max) call_doc callee.call(LiteralDoc(3), LiteralDoc(5)) print(to_python_script(call_doc)) # 输出: T.max(3, 5)更复杂的场景——组合语句块、IfDoc、ForDoc、ScopeDoc等——可以参照StmtBlockDoc的容器语义逐层嵌套最终由to_python_script一次性渲染为带缩进、可读的 TVMScript 文本。行号与下划线参数则用于调试定位设置print_line_numbersTrue可在输出中标注行号结合path_to_underline可让打印器在特定对象对应的文本上添加下划线提示。源码地图继续深入阅读Python 入口与 Doc 定义python/tvm/script/printer/init.py、python/tvm/script/printer/doc.py、python/tvm/script/printer/doc_printer.pyFFI 桥接python/tvm/script/printer/_ffi_api.pyC 配置与入口include/tvm/script/printer/config.h、include/tvm/script/printer/printer.h、src/script/printer/script_printer.ccIR → Doc 转换核心include/tvm/script/printer/ir_docsifier.h、include/tvm/script/printer/ir_docsifier_functor.hDoc → 文本渲染src/script/printer/doc_printer/base_doc_printer.cc、src/script/printer/doc_printer/python_doc_printer.cc方言级 IR 转换实现src/script/printer/ir/端到端往返测试tests/python/tvmscript/结合上述源码可以看出tvm.script.printer并非一个简单的字符串格式化工具而是一套以 Doc 为中间表示、以分发表驱动的多方言可扩展打印框架。无论你是想理解 IR 对象str()输出的来龙去脉还是计划为自定义 IR 方言接入 TVMScript 打印本模块都是统一的接入点。赞分享模型编译深度学习推理引擎【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址https://gitcode.com/gh_mirrors/tv/tvm点击查看免费下载相关推荐Prettier 技术内幕从 AST 到 Doc 中间表示的两阶段打印架构详解Prettier 技术内幕从 AST 到 Doc 中间表示的两阶段打印架构详解 Prettier 之所以能保证“解析结果与 printWidth 完美契合”开发工具格式化CLI中文大模型怎么选100开源LLM底座与微调资源一站导航中文大模型怎么选100开源LLM底座与微调资源一站导航 当你想给业务接入中文对话能力却在 ChatGLM、Qwen、百川之间反复横跳不确定该从哪个底座起大模型文档Element UI表格数据打印Table打印输出方案Element UI表格数据打印Table打印输出方案 在后台管理系统开发中经常需要将Element UI的Table组件数据导出为纸质文档或PDF。本文将前端UI组件设计系统创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考