模型编译深度学习推理引擎【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址https://gitcode.com/gh_mirrors/tv/tvm点击查看免费下载tvm.script.ir_builder是 TVM 中一套方言无关dialect-agnostic的编程式 IR 构造框架它以线程局部的帧栈frame stack为核心让开发者用 Pythonwith上下文以命令式风格直接构建 IRModule、Relax 函数与 TIR 块而无需先书写 TVMScript 文本再解析。本文将以该模块的 API 文档为骨架深入仓库源码与测试用例系统讲解IRBuilder/IRBuilderFrame的核心机制、IR 层ir、Relax 层relax、分布式扩展relax.distributed与 TIRx 层tirx的完整编程接口以及它在实际代码生成与测试场景中的典型用法。本文对应的官方 API 参考文档为 docs/reference/api/python/script/ir_builder.rst该文档通过 Sphinxautomodule指令聚合了五个公开子模块的成员签名与 docstring本文在此基础上结合源码python/tvm/script/ir_builder/、include/tvm/script/ir_builder/与测试用例逐层展开。一、理解核心抽象IRBuilder 与 IRBuilderFrame 的框架栈模型tvm.script.ir_builder的顶层命名空间导出了两个基础对象IRBuilder方言无关的构建器与IRBuilderFrame构建器栈帧。二者在 python/tvm/script/ir_builder/base.py 中定义为tvm.runtime.Object的注册子类底层对应 C 端 include/tvm/script/ir_builder/base.h 中的IRBuilderNode与IRBuilderFrameNode。1.1 IRBuilder线程局部的作用域管理器IRBuilder的惯用方式是放进with作用域在其中调用各方言的方法退出时取出构建结果。其关键成员如下方法作用底层 FFIIRBuilder()构造一个构建器实例_ffi_api.IRBuilder__enter__/__exit__进入/退出with作用域使构建器可被IRBuilder.current()获取IRBuilderEnter/IRBuilderExitIRBuilder.current()静态获取当前线程局部作用域中的构建器IRBuilderCurrentIRBuilder.is_in_scope()静态查询当前线程局部作用域是否存在构建器IRBuilderIsInScopeget()取出构建完成的 IR 对象IRBuilderGetwith_source_span(span)为嵌套作用域内构造的 IR 节点附加前端源码spanIRBuilderPushSourceSpan/IRBuilderPopSourceSpanIRBuilder.name(s, v)静态给对象命名返回同一对象IRBuilderNameIRBuilder.name_many(s, vs)静态批量命名要求名称列表与对象列表等长逐项调用name从 include/tvm/script/ir_builder/base.h 的节点定义可以看出每个IRBuilderNode内部维护了三块状态ffi::ArrayIRBuilderFrame frames上下文帧栈从栈顶向下查找最近的指定类型帧ffi::Optionalffi::ObjectRef resultIR 构造的最终产出std::vectorSpan source_spans活动的前端源码跨度从外到内。C 端提供模板方法FindFrameTFrame()自栈顶向下查找指定类型帧见 base.h、GetLastFrameTFrame()仅检查栈顶是否为指定类型见 base.h以及GetTObjectRef()取出结果并做类型校验见 base.h。这些机制正是上下文相关 IR 构造的基础例如同一个match_buffer调用在PrimFuncFrame下是函数签名中的 buffer 声明而在SBlockFrame下则生成MatchBufferRegion参见IRBuilderFrame的 docstring 示例与 base.h 中的 C 注释。with_source_span的嵌套语义值得注意当多个不同的源码区间发生嵌套如 TVMScript 内联展开时会被保留为SequentialSpan用于前端报错定位。1.2 IRBuilderFrame携带回调的栈帧IRBuilderFrame是构建器帧的基类。它同样支持with作用域__enter__调用IRBuilderFrameEnter将自身压入当前构建器的帧栈__exit__在没有异常发生时才调用IRBuilderFrameExit出栈——从 base.py 可以看到若with作用域因异常退出出栈逻辑会被跳过避免污染栈状态add_callback(callback)注册一个在退出with作用域时被调用的回调底层对应 base.h 中IRBuilderFrameNode的std::vectorffi::TypedFunctionvoid() callbacks成员该字段刻意未注册进反射因为它不应被遍历访问。各方言的帧如IRModuleFrame、Relax 的FunctionFrame/BindingBlockFrame、TIRx 的PrimFuncFrame/SBlockFrame都是IRBuilderFrame的子类python/tvm/script/ir_builder/ir/frame.py 中IRModuleFrame的注册即是例证。二、包结构与方言懒加载机制python/tvm/script/ir_builder/init.py 说明了该包的模块组织原则IR 层是基础层不注册为方言它的 builder 是真实子模块tvm.script.ir_builder.ir始终可以直接导入其余方言 builder 按需懒加载tvm.script.ir_builder.relax、tvm.script.ir_builder.relax.distributed、tvm.script.ir_builder.tirx等通过__getattr__在tvm.script._DIALECT_REGISTRY中查找方言名再动态导入dialect_module_path.builder例如tvm.tirx.script.builder结果缓存到模块全局变量后续访问不再走__getattr__。这意味着你既可以用from tvm.script.ir_builder import relax as relax_builder显式导入也可以直接import tvm.script.ir_builder后通过属性访问触发懒加载。该机制由 python/tvm/script/init.py 中的_DIALECT_REGISTRY与_DialectRedirectFinder配合实现负责处理深层语句式导入。三、IR 层tvm.script.ir_builder.ir面向 IRModule 的基础构建器IR 层暴露为I命名空间在 TVMScript 中常写作from tvm.script import ir as I。其全部 API 集中在 python/tvm/script/ir_builder/ir/ir.py入口均转发至_ffi_api3.1 IRModule 帧与模块级操作函数作用参数说明ir_module()开启一个IRModuleFrame返回帧对象无参数module_attrs(attrs, allow_overwriteFalse)设置模块级 attrsattrs: Dict[str, Object]allow_overwrite控制是否允许覆盖已有 attrmodule_get_attr(attr_key)读取指定 attr返回Optional[Object]不存在时返回Nonemodule_set_attr(attr_key, attr_value, allow_overwriteFalse)设置/覆盖指定 attrattr_value可为Nonemodule_global_infos(global_infos)设置模块级 global infosglobal_infos: Dict[str, List[GlobalInfo]]3.2 函数声明与定义decl_function(func_name, func_signature) - GlobalVar声明一个尚未给出实现体、仅指定签名参数与返回类型/形状的函数常用于跨函数调用场景。若func_signature不是BaseFunc实例会抛出ValueErrordef_function(func_name, func)为之前声明过的函数补全实现体。二者配合可以在一个 IRModule 中先声明后定义从而支持相互递归或前向引用。3.3 GlobalInfo 相关工具函数作用dummy_global_info()创建DummyGlobalInfo表达式常用于尚未确定设备信息的占位vdevice(targetNone, vdevice_id0, memory_scopeglobal)创建虚拟设备VDeviceglobal infovdevice_id默认 0memory_scope默认globallookup_vdevice(target_kindNone, device_index-1)从模块 globalinfo 的 vdevice 列表中按 target 类型如llvm、cuda与设备索引检索VDevicelookup_name(name) - bool检查是否存在指定名字的全局变量3.4meta_var解析期元编程标记meta_var(value)是 TVMScript 解析期parser-time专用的元编程值包装器对它的赋值会被解包而不产生 IR binding且支持迭代解包__iter__将列表元素逐个包装。在 python/tvm/script/ir_builder/ir/ir.py 中它被定义为运行时类meta_var类型检查环境下则是同名函数其 docstring 明确说明在 Relax 方言中它是默认原始绑定发射primitive binding emission的显式退出开关。各方言命名空间会提供指向同一实现的兼容别名。四、Relax 层tvm.script.ir_builder.relax构建 Relax 函数的完整 APIRelax 方言 builder 位于 python/tvm/relax/script/builder/ir.py约 990 行是实际使用最频繁的构建器。它复导出tvm.relax.op下大量算子add、matmul、nn.conv2d、nn.batch_norm、nn.relu等并提供了函数级与绑定块级 API。4.1 函数帧 API函数作用关键参数function(is_pureTrue, is_privateFalse)开启函数帧返回FunctionFrameis_pure标注函数纯性is_private标注私有性arg(name, ty)向最近函数帧添加参数返回Varty为 Relax 类型如R.Tensorfunc_name(name)指定最近函数帧的名字—func_attr(attrs)指定函数 attrsDict[str, Object]—func_ret_type(ret_ty)指定函数返回类型func_ret_ty为向后兼容别名—func_ret_value(value)指定函数返回值表达式—4.2 绑定块BindingBlockAPI函数作用dataflow()开启 dataflow 绑定块帧返回BindingBlockFrameoutput(*vars)将 dataflow 块内的变量暴露为块外全局可见的变量emit(value, annotate_tyNone)发射一条绑定生成Var可选显式标注类型emit_te(func, *args, **kwargs)通过 TE张量表达式算子生成Callemit_match_cast(value, ty)发射MatchCast绑定返回Varemit_var_binding(binding)直接发射一条VarBindingemit_with_type(...)/emit_with_ty(...)带类型信息发射tuple(*fields)/shape(value)构造 Relax 元组与形状表达式此外to_vdevice(data, dst_vdevice)支持将dst_vdevice以字符串形式如cuda:0或llvm形式解析为VDevice后调用tvm.relax.op.to_vdevicerewriter(rewriter_mod)则可以从一个定义了pattern/replacement两个同名签名函数的 IRModule 或 TVMScript 类构造PatternMatchingRewriter用于声明式改写规则。4.3 实战示例用构建器生成一个卷积网络 IRModulepython/tvm/relax/backend/adreno/mod_utils.py 中的get_relax_conv2d_mod展示了标准三步法IRBuilder→ 方言帧 →builder.get()from tvm.script.ir_builder import IRBuilder from tvm.script.ir_builder import relax as relax_builder from tvm.script import relax as R with IRBuilder() as builder: with relax_builder.function(): R.func_name(main) data R.arg(data, R.Tensor(data_shape, dtype)) weight R.arg(weight, R.Tensor(weight_shape, dtype)) with R.dataflow() as frame: output R.emit(R.nn.conv2d(data, weight, out_dtypedtype, stridesstride, dilationdilation, paddingpadding, groupsgroups)) if has_bias: output R.emit(output bias) R.output(output) R.func_ret_value(frame.output_vars[0]) func builder.get() return tvm.IRModule({main: func})关键点with IRBuilder() as builder使构建器进入线程局部作用域with relax_builder.function()压入函数帧随后R.arg注册参数、R.func_name命名with R.dataflow()压入 dataflow 绑定块帧块内R.emit逐个发射算子绑定R.output声明块输出退出所有作用域后builder.get()返回构造好的relax.Function再手工包成tvm.IRModule。在 tests/python/relax/test_codegen_cutlass.py 中还有IRBuilder与relax_builder.function()组合构造 CUTLASS codegen 测试输入的更复杂示例函数内嵌套R.dataflow()、T.prim_func(s_tirTrue)的 TIR 内核等说明该方法可与 TIRx 层混合使用覆盖端到端 codegen 场景。五、分布式扩展tvm.script.ir_builder.relax.distributedrelax.distributed是 Relax 构建器的分布式扩展对应 python/tvm/relax/script/builder/distributed/ir.py用于构造带设备网格device mesh、分布式张量与分布式算子绑定的 Relax IR。它与主relax构建器共享IRBuilder/帧栈机制以with D.DeviceMesh(...)、with D.Tensor(...)等分布式专用帧/表达式配合relax_builder.function()使用构造出的 IR 可直接交给 Relax 的分布式 pass 管线如DistributedNormalize继续处理。该模块的 API 文档通过automodule自动收集其分布式类型DTensorStructInfo等与全局信息如DeviceMesh在 python/tvm/relax/distributed/ 中有完整定义。六、TIRx 层tvm.script.ir_builder.tirx底层的 TIR 构建能力TIRx 方言 builder 位于 python/tvm/tirx/script/builder/ir.py提供面向底层 TIR 的帧与语句构建接口。核心入口包括prim_func(...)见 ir.py开启PrimFuncFrame是 TIR 函数级作用域sblock(name, no_realizeFalse, exec_scope)见 ir.py开启SBlockFrameno_realize控制是否自动 realize bufferexec_scope指定执行作用域func_name(name)/func_attr(attrs)/func_ret(ret_type)见 ir.py函数命名、属性与返回类型设置sblock_attr(attrs)ir.py与sblock_alloc_buffer(...)ir.py块级属性与 buffer 分配block_name_suffix_context(block_suffix)ir.py为块名追加后缀的上下文管理器func_gen(name)ir.py函数生成器辅助。结合 python/tvm/tirx/script/builder/utils.py 的frame_scope(frames)辅助函数可以批量压入多个帧。TIRx builder 与 parserpython/tvm/tirx/script/parser/entry.py 的prim_func共享同一套帧模型因此 TVMScript 的T.prim_func语法与编程式T.prim_func(...)构造在语义上等价。七、综合实战混合使用 IR 层与方言构建器IR 层的I.ir_module()帧可以与任何方言帧组合实现先开模块、再填函数的完整流程。结合 python/tvm/script/ir_builder/ir/ir.py 的 API一个通用模式如下from tvm.script.ir_builder import IRBuilder from tvm.script.ir_builder import ir as I from tvm.script.ir_builder import relax as relax_builder from tvm.script import relax as R with IRBuilder() as builder: with I.ir_module() as mod_frame: # 模块级信息可选 I.module_attrs({tir.noalias: True}) with relax_builder.function(is_pureTrue): R.func_name(main) x R.arg(x, R.Tensor((16, 16), float32)) y R.emit(x * 2.0) R.func_ret_value(y) mod builder.get() # 返回 IRModule要点回顾I.ir_module()与relax_builder.function()形成嵌套帧栈出栈顺序严格反向I.module_attrs在IRModuleFrame存活期内生效若要构造跨函数调用可先用I.decl_function(helper, signature)声明再在后续用I.def_function(helper, func)补全实现最终builder.get()的返回类型由栈顶帧决定RelaxFunctionFrame出栈后得到relax.FunctionIRModuleFrame出栈后得到IRModuleC 端GetTObjectRef()会校验结果类型见 base.h。八、源码阅读索引若希望深入理解tvm.script.ir_builder的实现建议按以下路径阅读Python 绑定层python/tvm/script/ir_builder/base.pyIRBuilder/IRBuilderFrame、python/tvm/script/ir_builder/ir/ir.pyIR 层 API、python/tvm/script/ir_builder/init.py方言懒加载C 实现层include/tvm/script/ir_builder/base.h帧栈与Get/FindFrame模板、src/script/ir_builder/base.cc线程局部作用域与 FFI 注册Relax 方言python/tvm/relax/script/builder/ir.py、python/tvm/relax/script/builder/distributed/ir.pyTIRx 方言python/tvm/tirx/script/builder/ir.py、python/tvm/tirx/script/builder/utils.py测试用例tests/python/relax/backend/adreno/mod_utils.py、tests/python/relax/test_codegen_cutlass.py。结语tvm.script.ir_builder是 TVM 中程序化构造 IR的统一入口IRBuilder提供线程局部作用域IRBuilderFrame提供带回调的上下文帧栈ir层管理 IRModule 与全局信息relax、relax.distributed与tirx层分别覆盖高层函数构建、分布式扩展与底层 TIR 语句构造。掌握了帧栈模型与各层的函数签名你便可以在不依赖 TVMScript 文本解析的前提下以纯命令式 Python 精确生成任意结构的 IR——这正是自动化 codegen、测试输入生成与编译器前端集成的核心能力所在。赞分享模型编译深度学习推理引擎【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址https://gitcode.com/gh_mirrors/tv/tvm点击查看免费下载相关推荐TVM Relax BlockBuilder 完全指南用 Python 构建 Relax IR 的开发者 APITVM Relax BlockBuilder 完全指南用 Python 构建 Relax IR 的开发者 API 本指南基于 Apache TVM 开源仓库中模型编译深度学习推理引擎TVM TIRx 核心脚本 API 指南从 tvm.script.tirx 解析器到 IR Builder 的完整解析TVM TIRx 核心脚本 API 指南从 tvm.script.tirx 解析器到 IR Builder 的完整解析 导读 TIRx 是 Apache TV模型编译深度学习推理引擎TVM TIRx Python API 参考指南从内核编写、IR 检查到编译器扩展的完整实践TVM TIRx Python API 参考指南从内核编写、IR 检查到编译器扩展的完整实践 TIRx 是 Apache TVM 中面向现代加速器CUDA、模型编译深度学习推理引擎创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考