编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载导读本文以 TVM 官方架构文档 docs/arch/index.rst 为核心骨架面向希望理解 TVM 编译栈整体设计、乃至希望参与 TVM 开发的技术人员。文章先沿着一条典型的端到端编译流程剖析 Import导入、Transformation变换、Target Translation目标翻译、Runtime Execution运行时执行四个阶段及其核心数据结构IRModule、relay.Function、tir.PrimFunc、runtime.Module、runtime.PackedFunc随后从静态视角梳理代码库的各个逻辑模块runtime、node、ir、target、tir、arith、te、topi、relay、autotvm 等及其相互关系。读完本文你将掌握 TVM 编译管线的完整脉络理解如何通过 Python 与 C API 直接操纵 IR 与编译变换并能依据源码路径进一步深入阅读实现细节。本文所有源码与文档引用均基于当前仓库原文档中的架构示意图托管在外部资源上故此处不引用外部图片仅以文字精确描述各阶段与模块关系。一、一条端到端的编译流程原文档开篇即以Example Compilation Flow为线索将 TVM 将一个高层模型描述转化为可部署模块的过程划分为四个高层次的步骤Import导入前端frontend组件将模型摄入为一个IRModule其中包含一组函数函数内部即模型的表达。Transformation变换编译器将IRModule变换为另一个功能等价或在量化等场景下近似等价的IRModule。大量变换与目标target无关同时我们也允许 target 影响变换管线的配置。Target Translation目标翻译编译器将IRModule翻译codegen为 target 指定的可执行格式翻译结果封装为一个runtime.Module可以被导出、加载并在目标运行时环境中执行。Runtime Execution运行时执行用户重新加载runtime.Module并在支持的运行时环境中运行编译得到的函数。这四条主线的本质是整个编译过程就是在少数关键数据结构之间反复施加变换。下面先介绍这些关键数据结构再逐一展开每个阶段。二、关键数据结构Key Data Structures理解一个复杂系统的最佳方式之一是识别其关键数据结构以及操作变换这些数据结构的 API。TVM 的核心数据结构主要有两个层次2.1 IRModule —— 贯穿整个编译栈的主数据结构IRModuleintermediate representation module中间表示模块是贯穿整个编译栈的首要数据结构它包含一组函数的集合。目前支持两种主要函数变体relay::Function高层函数式程序表示通常对应一个端到端的模型。可以把它理解为一个计算图并额外支持控制流、递归和复杂数据结构。tir::PrimFunc低层程序表示包含循环嵌套选择、多维 load/store、线程化、向量/张量指令等元素。通常用于表示执行模型中某个可能被融合后的算子程序的算子级函数。在编译过程中一个 relay 函数会被 lower 为多个tir::PrimFunc以及一个调用这些tir::PrimFunc的顶层函数。从源码上看IRModule及相关 IR 数据结构均派生自运行时的对象系统见后文tvm/runtime与tvm/node部分因此可以被 Python 前端直接访问和操作。2.2 runtime.Module 与 runtime.PackedFunc —— 运行时的两个核心抽象runtime.Module封装编译结果。一个 runtime.Module 提供GetFunction方法按名称获取PackedFunc。在 include/tvm/runtime/module.h 中可以看到Module是ObjectRef的容器类提供GetFunction(name, query_imports)、Import(other)、静态的LoadFromFile(file_name, format)等接口其底层容器ModuleNode通过type_key()标识模块类型如llvm、cuda、c等并通过SaveToFile/SaveToBinary支持序列化。runtime.PackedFunc是类型擦除的统一函数接口。一个 PackedFunc 可以接受并返回以下类型POD 类型int、float、string、runtime.PackedFunc、runtime.Module、runtime.NDArray以及其他 runtime.Object 子类。include/tvm/runtime/packed_func.h 中定义了PackedFunc的operator()直接按普通函数方式调用与CallPacked打包格式调用同头文件还提供带类型检查的TypedPackedFuncR(Args...)包装C 开发中推荐优先使用。这两个机制使得 TVM 的模块化非常强大例如要在 CUDA 上执行某个addone函数可以用 LLVM 生成宿主端代码来计算启动参数如线程组大小再调用另一个由 CUDA driver API 支撑的CUDAModule中的 PackedFunc同样的机制也可用于 OpenCL 内核。2.3 两个层面数据结构的对应关系小结编译阶段的数据结构包含内容IRModule包含 relay.Function 与 tir.PrimFuncruntime.Module包含 runtime.PackedFunc绝大多数编译工作都是在这些关键数据结构之间进行的变换relay/transform与tir/transform是确定性的、基于规则的变换auto_scheduler与autotvm是基于搜索的变换见下文。三、变换Transformations优化与 Lowering每一个变换都服务于以下两种目的之一优化optimization将程序变换为等价的、可能更优的版本lowering将程序变换为更接近 target 的低层表示。3.1 relay/transform高层模型优化relay/transform包含一组优化模型的 pass。优化包括通用程序优化如常量折叠 constant folding、死代码消除 dead-code elimination以及张量计算特有 pass如 layout 变换、缩放因子折叠 scaling factor folding。在 relay 优化管线末尾会运行一个名为FuseOps的 pass将端到端函数如 MobileNet切分为子函数段如 conv2d-relu。这些被切分出的段即functions segments。切分将原问题一分为二对每个子函数单独进行编译与优化处理整体执行结构按顺序调用生成的子函数以完成整个模型的执行。从源码看src/relay/transforms/fuse_ops.cc 注册了两个可配置项relay.FuseOps.max_depthInteger控制最大融合深度与relay.FuseOps.link_paramsBool控制是否将参数链接进图FuseOps(fuse_opt_level)通过PassContext读取这些配置默认最大融合深度为kMaxFusedOps并依赖InferType前置 pass。这些配置均可在 Python 侧通过tvm.transform.PassContext(config{...})覆盖。3.2 整体执行结构图执行器、虚拟机与 AOTrelay/backend中提供了处理整体执行问题的几种不同方式对于形状已知、无控制流的简单模型可以 lower 为图执行器graph executor把执行结构存储在图里支持虚拟机virtual machine后端用于动态执行计划支持提前编译AOT把高层执行结构与生成的 primitive 函数一起编译进可执行文件。所有这些执行模式都统一封装在runtime.Module接口之下。从源码看src/relay/backend/build_module.cc 通过GetPackedFunc(relay.build_module._GraphExecutorCodegen)获取图执行器代码生成器其实现位于 src/relay/backend/graph_executor_codegen.ccAOT 代码生成器注册于 src/relay/backend/aot_executor_codegen.cc。虚拟机编译实现位于 src/relay/backend/vm/compiler.ccVMCompiler通过GetFunction暴露Lower、OptimizeModule、GetExecutable等接口虚拟机运行时实现位于 src/runtime/vm/executable.cc 等文件。3.3 tir/transformTIR 层变换与对下游编译器的依赖tir/transform包含 TIR 层函数的变换 pass。许多 tir pass 服务于 lowering 目的例如将多维访问 flatten 为一维指针访问将 intrinsics 展开为 target 特化的指令装饰函数入口以符合运行时调用约定。当然也有优化 pass例如访问索引化简access index simplification和死代码消除。值得强调的设计原则许多低层优化可以在 target 阶段由 LLVM、CUDA C 等下游编译器处理因此 TVM 将寄存器分配等低层优化留给下游编译器只聚焦于下游编译器不覆盖的优化。这决定了 Target Translation 阶段应当尽可能轻量绝大部分变换与 lowering 都应在此之前完成。四、基于搜索与学习式的变换Search-space and Learning-based Transformations上述变换都是确定性、基于规则的。TVM 栈的设计目标之一是为不同硬件平台提供高性能代码优化需要探索尽可能多的优化选择——包括多维张量访问、循环分块行为、专用加速器内存层级、线程化等。为任意硬件定义一套启发式规则很难因此 TVM 采取搜索 学习的路线定义一组可作用于程序的动作actions例如循环变换、内联、向量化这些动作被称为调度原语scheduling primitives它们的集合构成一个程序的优化搜索空间。系统在搜索空间中搜索不同的调度序列挑选最佳组合搜索过程通常由机器学习算法引导。搜索完成后可以记录某个可能融合的算子的最佳调度序列编译器之后只需查表并应用该调度序列。值得注意的是调度应用阶段与规则式变换完全一致因此可以与传统 pass 共享同一套接口约定。基于搜索的优化用于解决初始 TIR 函数生成问题这部分模块名为AutoTVM / auto_scheduler详见tvm/autotvm一节。TVM 团队计划在后续演进中将学习式变换扩展到更多领域。从代码组织看搜索策略位于 src/auto_scheduler/search_policyAutoTVM 的成本模型与特征提取位于 src/autotvm。五、Target Translation目标翻译目标翻译阶段将 IRModule 变换为对应 target 的可执行格式对 x86、ARM 等后端使用LLVM IRBuilder在内存中构建 LLVM IR也可以生成源码级语言如CUDA C和OpenCL还支持通过外部代码生成器把 Relay 函数子图直接翻译到特定 target。tvm/target模块包含所有这些代码生成器并提供一个统一的Target类来描述编译目标。目标翻译之前的变换同样会受到 target 影响——例如 target 的向量长度会改变向量化行为。编译管线可以根据 target 的属性信息以及注册到每个 target idcuda、opencl的 builtin 信息进行定制。设计上最终代码生成阶段必须尽可能轻量绝大部分变换与 lowering 应在目标翻译之前完成使目标翻译只做最后一公里的代码生成工作。六、Runtime Execution运行时执行TVM 运行时的主要目标是提供一套最小化 API让用户以自己偏好的语言Python、C、Rust、Go、Java、JavaScript 等加载并执行编译产物。原文档给出了两个 Python 示例此处完整保留并结合源码注释展开。6.1 算子级执行示例addoneimport tvm # Example runtime execution program in python, with type annotated mod: tvm.runtime.Module tvm.runtime.load_module(compiled_artifact.so) arr: tvm.runtime.NDArray tvm.nd.array([1, 2, 3], devicetvm.cuda(0)) fun: tvm.runtime.PackedFunc mod[addone] fun(arr) print(arr.numpy())要点tvm.runtime.load_module对应Module::LoadFromFileinclude/tvm/runtime/module.h按文件名与可选 format 加载宿主模块mod[addone]实际调用Module::GetFunction(name, query_importsfalse)按名称从当前模块及其依赖模块若query_importstrue获取 PackedFunc函数不存在时返回PackedFunc(nullptr)PackedFunc 可以直接像普通函数一样调用tvm.runtime.NDArray作为参数自动打包传递。6.2 端到端模型执行示例resnet18import tvm # Example runtime execution program in python, with types annotated factory: tvm.runtime.Module tvm.runtime.load_module(resnet18.so) # Create a stateful graph execution module for resnet18 on cuda(0) gmod: tvm.runtime.Module factoryresnet18) data: tvm.runtime.NDArray get_input_data() # set input gmodset_input # execute the model gmod[run]() # get the output result gmodget_output.numpy()要点这里factory是一个工厂模块factoryresnet18)返回一个有状态的图执行模块graph executor module其set_input/run/get_output正是图执行器运行时暴露的 PackedFunc 接口图执行器运行时的实现位于 src/runtime/graph_executor其调试版本 src/runtime/graph_executor/debug/graph_executor_debug.cc 还额外提供debug_get_output、profile等接口。核心结论runtime.Module与runtime.PackedFunc足以同时封装算子级程序如 addone与端到端模型这一统一接口正是 TVM 运行时以少驭多的关键。七、逻辑架构组件Logical Architecture Components上面的编译流程是运行时视角各组件在运行编译器时的交互本节则是静态总览视角梳理代码库的逻辑模块及其关系。7.1 tvm/support —— 基础设施工具support 模块包含最通用的基础设施工具例如通用 arena 分配器、socket、日志等。对应源码位于 src/support头文件位于 include/tvm/support。7.2 tvm/runtime —— 运行时基石runtime 是 TVM 栈的基础提供加载与执行编译产物的机制并定义一组稳定的标准 C API 与 Python、Rust 等前端语言对接。runtime::Object是除runtime::PackedFunc外运行时的主要数据结构它是一个带类型索引type index的引用计数基类支持运行时类型检查与向下转型downcasting。从 include/tvm/runtime/object.h 可以看到Object持有type_index_与ref_counter_子类通过注册_type_index自动填充类型索引与析构器。对象系统允许开发者向运行时引入新数据结构如 Array、Map、新的 IR 数据结构。编译器自身也重度依赖运行时机制所有 IR 数据结构都是runtime::Object的子类因此可以直接从 Python 前端访问和操作各种 API 通过 PackedFunc 机制暴露给前端。不同硬件后端的运行时支持定义在 runtime 的子目录中例如 runtime/opencl这些硬件特化运行时模块定义了设备内存分配与设备函数序列化等 API。runtime/rpc实现了对 PackedFunc 的 RPC 支持。可以用 RPC 机制把交叉编译的库发送到远程设备并基准测试执行性能RPC 基础设施为学习式优化提供了来自广泛硬件后端的数据收集能力。实现位于 src/runtime/rpc如 src/runtime/rpc/rpc_channel.ccRPCChannel 是连接本地与远程 RPC 会话的通信端点抽象。相关深度指南链接已转换为仓库根相对路径runtime、debugger、virtual_machine、introduction_to_module_serialization、device_target_interactions。7.3 tvm/node —— IR 数据结构的反射、序列化与结构等价node 模块在runtime::Object之上为 IR 数据结构增加了反射reflection、序列化、结构等价structural equivalence与哈希hashing等能力。得益于 node 模块在 Python 中可以直接按字段名访问 TVM IRNode 的任何字段x tvm.tir.Var(x, int32) y tvm.tir.Add(x, x) # a and b are fields of a tir.Add node # we can directly use the field name to access the IR structures assert y.a x我们也可以把任意 IR 节点序列化为 JSON 格式再加载回来。这种可保存、可存储、可检视的能力为让编译器更易用奠定了基础。相关实现位于 src/nodestructural_equal.cc、structural_hash.cc、serialization.cc、reflection.cc 等。7.4 tvm/ir —— 统一的数据结构与接口tvm/ir文件夹包含所有 IR 函数变体共用的统一数据结构与接口被tvm/relay和tvm/tir共享主要包括IRModuleTypePassContext 与 PassOp不同函数变体relay.Function 与 tir.PrimFunc可以共存于同一个 IRModule。虽然内容表示不同但它们使用相同的数据结构表示类型因此可以用同一数据结构表示这些变体的函数类型签名。统一类型系统允许一个函数变体调用另一个函数变体只要明确调用约定为未来的跨函数变体优化打开大门。同时提供统一的PassContext用于配置 pass 行为以及通用的组合 pass 来执行 pass 管线。原文档给出的配置示例# configure the behavior of the tir.UnrollLoop pass with tvm.transform.PassContext(config{tir.UnrollLoop: { auto_max_step: 10 }}): # code affected by the pass contextOp是表示所有系统定义的原语算子/intrinsics 的公共类开发者可以向系统注册新的 Op 及其附加属性例如该 Op 是否为 elementwise。Pass 基础设施的深入介绍见 pass_infra。7.5 tvm/target —— 代码生成器与 Target 描述target 模块包含所有把 IRModule 翻译为目标 runtime.Module 的代码生成器并提供统一的Target类描述目标。编译管线可以通过查询 target 的属性信息以及注册到每个 target idcuda、opencl的 builtin 信息来定制。相关讨论见 device_target_interactions。7.6 tvm/tir —— 低层程序表示TIR 定义低层程序表示使用tir::PrimFunc表示可被 TIR pass 变换的函数。除 IR 数据结构外tir 模块还通过公共 Op 注册表定义一组 builtin intrinsics 及其属性以及tir/transform中的变换 pass。7.7 tvm/arith —— 索引算术分析arith 模块与 TIR 紧密关联。低层代码生成中的关键问题之一是指针/索引算术性质分析——正负性、变量边界、描述迭代空间的整数集合。arith 模块提供一组主要是整数分析工具TIR pass 可以利用这些分析来化简和优化代码。对应源码位于 src/arith如 rewrite_simplify.cc、int_set.cc、presburger_set.cc 等。7.8 tvm/te —— 张量表达式 DSLte 即 tensor expression张量表达式是一个领域特定语言模块让我们可以通过书写张量表达式快速构造tir::PrimFunc变体。重要张量表达式本身不是可以存入 IRModule 的自包含函数而是可以拼接起来构建 IRModule 的 IR 片段。te/schedule提供一组调度原语来控制生成函数未来部分调度组件可能会被引入到tir::PrimFunc自身。相关主题见 inferbound 与 hybrid_script。7.9 tvm/topi —— 张量算子库存虽然每个用例都可以直接通过 TIR 或 TE 构造算子但这样做很繁琐。topiTensor operator inventory提供一组预定义算子以 TE 或 TIR 形式实现语义对齐 numpy 并覆盖常见深度学习工作负载同时提供一批常用调度模板用于在不同 target 平台上获得高性能实现。对应实现位于 include/tvm/topi 与 python/tvm/topi。7.10 tvm/relay —— 高层函数式 IRRelay 是表示完整模型的高层函数式 IR各种优化定义在relay.transform中。Relay 编译器定义了多种方言dialect每种方言支持特定风格的优化值得一提的包括QNN用于导入预量化模型VM用于 lower 到动态虚拟机memory用于内存优化。深入主题见 relay_intro、relay_op_strategy、convert_layout。7.11 tvm/autotvm —— 自动化搜索式程序优化AutoTVM 与 AutoScheduler 都是自动化、基于搜索的程序优化组件目前主要包括成本模型与特征提取cost models and feature extraction存储程序基准结果、用于构建成本模型的记录格式record format一组作用于程序变换的搜索策略search policies。自动化程序优化仍是活跃的研究领域因此设计上刻意模块化研究人员可以通过 Python 绑定快速修改某个组件或应用自己的算法自定义搜索并插入自己的算法。实现位于 src/auto_scheduler 与 src/autotvm配套基准说明见 benchmark。7.12 Frontends —— 模型导入前端前端负责把不同框架的模型摄入 TVM 栈命名空间为tvm.relay.frontend是模型导入 API 的入口。TensorFlow 前端指南见 frontend/tensorflow。7.13 Security 与 microTVMSecurity安全相关文档见 security。microTVM面向嵌入式/MCU 场景的微型运行时设计相关主题见 microtvm_design、microtvm_project_api、model_library_format。八、总结与延伸讨论Summary and Discussions回顾全文编译流程中的关键数据结构是IRModule包含 relay.Function 与 tir.PrimFuncruntime.Module包含 runtime.PackedFunc。编译的大部分内容都是这些关键数据结构之间的变换relay/transform与tir/transform是确定性的、基于规则的变换auto_scheduler与autotvm包含基于搜索的变换。最后强调编译流程示例只是 TVM 栈的一种典型用法。TVM 将这些关键数据结构与变换同时暴露给 Python 与 C API因此你可以像使用 numpy 一样使用 TVM——只是感兴趣的数据结构从numpy.ndarray换成了tvm.IRModule。典型的高级用法包括直接用 Python API 构造 IRModule组合自定义变换集合例如自定义量化流程直接用 Python API 操纵 IR。若需进一步深入原文档还推荐阅读开发者指南dev-how-to以及 device_target_interactions各物理设备与代码生成目标的交互。对于希望参与开发的读者可结合本文提到的源码路径include/tvm/runtime/module.h、include/tvm/runtime/packed_func.h、src/relay/transforms/fuse_ops.cc、src/relay/backend/build_module.cc 等逐模块阅读形成文档脉络 源码证据的完整认知闭环。赞分享编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载相关推荐Apache TVM 编译架构全解析从模型 IR 到可部署模块的端到端编译流程Apache TVM 编译架构全解析从模型 IR 到可部署模块的端到端编译流程 Apache TVM 是面向机器学习的开源编译框架其核心使命是把高层模型描述模型编译深度学习推理引擎Helix 编辑器源码架构指南从 Rope 到 Term 的模块化分层设计Helix 编辑器源码架构指南从 Rope 到 Term 的模块化分层设计 本篇技术指南基于 Helix一种后现代模态文本编辑器仓库根目录下的 docs/代码编辑器开发工具CLITriton编译器架构前端、中端和后端的模块化设计Triton编译器架构前端、中端和后端的模块化设计 概述 Triton是一个专为深度学习计算优化的编程语言和编译器其核心设计理念是通过模块化的编译器架构实现编译器编程语言人工智能深度学习高性能计算上一篇Proxyee 开源项目教程下一篇Awesome Hacking Resources黑客与渗透测试资源宝库创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 TVM 架构设计全解：从端到端编译流程到逻辑模块分层
编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载导读本文以 TVM 官方架构文档 docs/arch/index.rst 为核心骨架面向希望理解 TVM 编译栈整体设计、乃至希望参与 TVM 开发的技术人员。文章先沿着一条典型的端到端编译流程剖析 Import导入、Transformation变换、Target Translation目标翻译、Runtime Execution运行时执行四个阶段及其核心数据结构IRModule、relay.Function、tir.PrimFunc、runtime.Module、runtime.PackedFunc随后从静态视角梳理代码库的各个逻辑模块runtime、node、ir、target、tir、arith、te、topi、relay、autotvm 等及其相互关系。读完本文你将掌握 TVM 编译管线的完整脉络理解如何通过 Python 与 C API 直接操纵 IR 与编译变换并能依据源码路径进一步深入阅读实现细节。本文所有源码与文档引用均基于当前仓库原文档中的架构示意图托管在外部资源上故此处不引用外部图片仅以文字精确描述各阶段与模块关系。一、一条端到端的编译流程原文档开篇即以Example Compilation Flow为线索将 TVM 将一个高层模型描述转化为可部署模块的过程划分为四个高层次的步骤Import导入前端frontend组件将模型摄入为一个IRModule其中包含一组函数函数内部即模型的表达。Transformation变换编译器将IRModule变换为另一个功能等价或在量化等场景下近似等价的IRModule。大量变换与目标target无关同时我们也允许 target 影响变换管线的配置。Target Translation目标翻译编译器将IRModule翻译codegen为 target 指定的可执行格式翻译结果封装为一个runtime.Module可以被导出、加载并在目标运行时环境中执行。Runtime Execution运行时执行用户重新加载runtime.Module并在支持的运行时环境中运行编译得到的函数。这四条主线的本质是整个编译过程就是在少数关键数据结构之间反复施加变换。下面先介绍这些关键数据结构再逐一展开每个阶段。二、关键数据结构Key Data Structures理解一个复杂系统的最佳方式之一是识别其关键数据结构以及操作变换这些数据结构的 API。TVM 的核心数据结构主要有两个层次2.1 IRModule —— 贯穿整个编译栈的主数据结构IRModuleintermediate representation module中间表示模块是贯穿整个编译栈的首要数据结构它包含一组函数的集合。目前支持两种主要函数变体relay::Function高层函数式程序表示通常对应一个端到端的模型。可以把它理解为一个计算图并额外支持控制流、递归和复杂数据结构。tir::PrimFunc低层程序表示包含循环嵌套选择、多维 load/store、线程化、向量/张量指令等元素。通常用于表示执行模型中某个可能被融合后的算子程序的算子级函数。在编译过程中一个 relay 函数会被 lower 为多个tir::PrimFunc以及一个调用这些tir::PrimFunc的顶层函数。从源码上看IRModule及相关 IR 数据结构均派生自运行时的对象系统见后文tvm/runtime与tvm/node部分因此可以被 Python 前端直接访问和操作。2.2 runtime.Module 与 runtime.PackedFunc —— 运行时的两个核心抽象runtime.Module封装编译结果。一个 runtime.Module 提供GetFunction方法按名称获取PackedFunc。在 include/tvm/runtime/module.h 中可以看到Module是ObjectRef的容器类提供GetFunction(name, query_imports)、Import(other)、静态的LoadFromFile(file_name, format)等接口其底层容器ModuleNode通过type_key()标识模块类型如llvm、cuda、c等并通过SaveToFile/SaveToBinary支持序列化。runtime.PackedFunc是类型擦除的统一函数接口。一个 PackedFunc 可以接受并返回以下类型POD 类型int、float、string、runtime.PackedFunc、runtime.Module、runtime.NDArray以及其他 runtime.Object 子类。include/tvm/runtime/packed_func.h 中定义了PackedFunc的operator()直接按普通函数方式调用与CallPacked打包格式调用同头文件还提供带类型检查的TypedPackedFuncR(Args...)包装C 开发中推荐优先使用。这两个机制使得 TVM 的模块化非常强大例如要在 CUDA 上执行某个addone函数可以用 LLVM 生成宿主端代码来计算启动参数如线程组大小再调用另一个由 CUDA driver API 支撑的CUDAModule中的 PackedFunc同样的机制也可用于 OpenCL 内核。2.3 两个层面数据结构的对应关系小结编译阶段的数据结构包含内容IRModule包含 relay.Function 与 tir.PrimFuncruntime.Module包含 runtime.PackedFunc绝大多数编译工作都是在这些关键数据结构之间进行的变换relay/transform与tir/transform是确定性的、基于规则的变换auto_scheduler与autotvm是基于搜索的变换见下文。三、变换Transformations优化与 Lowering每一个变换都服务于以下两种目的之一优化optimization将程序变换为等价的、可能更优的版本lowering将程序变换为更接近 target 的低层表示。3.1 relay/transform高层模型优化relay/transform包含一组优化模型的 pass。优化包括通用程序优化如常量折叠 constant folding、死代码消除 dead-code elimination以及张量计算特有 pass如 layout 变换、缩放因子折叠 scaling factor folding。在 relay 优化管线末尾会运行一个名为FuseOps的 pass将端到端函数如 MobileNet切分为子函数段如 conv2d-relu。这些被切分出的段即functions segments。切分将原问题一分为二对每个子函数单独进行编译与优化处理整体执行结构按顺序调用生成的子函数以完成整个模型的执行。从源码看src/relay/transforms/fuse_ops.cc 注册了两个可配置项relay.FuseOps.max_depthInteger控制最大融合深度与relay.FuseOps.link_paramsBool控制是否将参数链接进图FuseOps(fuse_opt_level)通过PassContext读取这些配置默认最大融合深度为kMaxFusedOps并依赖InferType前置 pass。这些配置均可在 Python 侧通过tvm.transform.PassContext(config{...})覆盖。3.2 整体执行结构图执行器、虚拟机与 AOTrelay/backend中提供了处理整体执行问题的几种不同方式对于形状已知、无控制流的简单模型可以 lower 为图执行器graph executor把执行结构存储在图里支持虚拟机virtual machine后端用于动态执行计划支持提前编译AOT把高层执行结构与生成的 primitive 函数一起编译进可执行文件。所有这些执行模式都统一封装在runtime.Module接口之下。从源码看src/relay/backend/build_module.cc 通过GetPackedFunc(relay.build_module._GraphExecutorCodegen)获取图执行器代码生成器其实现位于 src/relay/backend/graph_executor_codegen.ccAOT 代码生成器注册于 src/relay/backend/aot_executor_codegen.cc。虚拟机编译实现位于 src/relay/backend/vm/compiler.ccVMCompiler通过GetFunction暴露Lower、OptimizeModule、GetExecutable等接口虚拟机运行时实现位于 src/runtime/vm/executable.cc 等文件。3.3 tir/transformTIR 层变换与对下游编译器的依赖tir/transform包含 TIR 层函数的变换 pass。许多 tir pass 服务于 lowering 目的例如将多维访问 flatten 为一维指针访问将 intrinsics 展开为 target 特化的指令装饰函数入口以符合运行时调用约定。当然也有优化 pass例如访问索引化简access index simplification和死代码消除。值得强调的设计原则许多低层优化可以在 target 阶段由 LLVM、CUDA C 等下游编译器处理因此 TVM 将寄存器分配等低层优化留给下游编译器只聚焦于下游编译器不覆盖的优化。这决定了 Target Translation 阶段应当尽可能轻量绝大部分变换与 lowering 都应在此之前完成。四、基于搜索与学习式的变换Search-space and Learning-based Transformations上述变换都是确定性、基于规则的。TVM 栈的设计目标之一是为不同硬件平台提供高性能代码优化需要探索尽可能多的优化选择——包括多维张量访问、循环分块行为、专用加速器内存层级、线程化等。为任意硬件定义一套启发式规则很难因此 TVM 采取搜索 学习的路线定义一组可作用于程序的动作actions例如循环变换、内联、向量化这些动作被称为调度原语scheduling primitives它们的集合构成一个程序的优化搜索空间。系统在搜索空间中搜索不同的调度序列挑选最佳组合搜索过程通常由机器学习算法引导。搜索完成后可以记录某个可能融合的算子的最佳调度序列编译器之后只需查表并应用该调度序列。值得注意的是调度应用阶段与规则式变换完全一致因此可以与传统 pass 共享同一套接口约定。基于搜索的优化用于解决初始 TIR 函数生成问题这部分模块名为AutoTVM / auto_scheduler详见tvm/autotvm一节。TVM 团队计划在后续演进中将学习式变换扩展到更多领域。从代码组织看搜索策略位于 src/auto_scheduler/search_policyAutoTVM 的成本模型与特征提取位于 src/autotvm。五、Target Translation目标翻译目标翻译阶段将 IRModule 变换为对应 target 的可执行格式对 x86、ARM 等后端使用LLVM IRBuilder在内存中构建 LLVM IR也可以生成源码级语言如CUDA C和OpenCL还支持通过外部代码生成器把 Relay 函数子图直接翻译到特定 target。tvm/target模块包含所有这些代码生成器并提供一个统一的Target类来描述编译目标。目标翻译之前的变换同样会受到 target 影响——例如 target 的向量长度会改变向量化行为。编译管线可以根据 target 的属性信息以及注册到每个 target idcuda、opencl的 builtin 信息进行定制。设计上最终代码生成阶段必须尽可能轻量绝大部分变换与 lowering 应在目标翻译之前完成使目标翻译只做最后一公里的代码生成工作。六、Runtime Execution运行时执行TVM 运行时的主要目标是提供一套最小化 API让用户以自己偏好的语言Python、C、Rust、Go、Java、JavaScript 等加载并执行编译产物。原文档给出了两个 Python 示例此处完整保留并结合源码注释展开。6.1 算子级执行示例addoneimport tvm # Example runtime execution program in python, with type annotated mod: tvm.runtime.Module tvm.runtime.load_module(compiled_artifact.so) arr: tvm.runtime.NDArray tvm.nd.array([1, 2, 3], devicetvm.cuda(0)) fun: tvm.runtime.PackedFunc mod[addone] fun(arr) print(arr.numpy())要点tvm.runtime.load_module对应Module::LoadFromFileinclude/tvm/runtime/module.h按文件名与可选 format 加载宿主模块mod[addone]实际调用Module::GetFunction(name, query_importsfalse)按名称从当前模块及其依赖模块若query_importstrue获取 PackedFunc函数不存在时返回PackedFunc(nullptr)PackedFunc 可以直接像普通函数一样调用tvm.runtime.NDArray作为参数自动打包传递。6.2 端到端模型执行示例resnet18import tvm # Example runtime execution program in python, with types annotated factory: tvm.runtime.Module tvm.runtime.load_module(resnet18.so) # Create a stateful graph execution module for resnet18 on cuda(0) gmod: tvm.runtime.Module factoryresnet18) data: tvm.runtime.NDArray get_input_data() # set input gmodset_input # execute the model gmod[run]() # get the output result gmodget_output.numpy()要点这里factory是一个工厂模块factoryresnet18)返回一个有状态的图执行模块graph executor module其set_input/run/get_output正是图执行器运行时暴露的 PackedFunc 接口图执行器运行时的实现位于 src/runtime/graph_executor其调试版本 src/runtime/graph_executor/debug/graph_executor_debug.cc 还额外提供debug_get_output、profile等接口。核心结论runtime.Module与runtime.PackedFunc足以同时封装算子级程序如 addone与端到端模型这一统一接口正是 TVM 运行时以少驭多的关键。七、逻辑架构组件Logical Architecture Components上面的编译流程是运行时视角各组件在运行编译器时的交互本节则是静态总览视角梳理代码库的逻辑模块及其关系。7.1 tvm/support —— 基础设施工具support 模块包含最通用的基础设施工具例如通用 arena 分配器、socket、日志等。对应源码位于 src/support头文件位于 include/tvm/support。7.2 tvm/runtime —— 运行时基石runtime 是 TVM 栈的基础提供加载与执行编译产物的机制并定义一组稳定的标准 C API 与 Python、Rust 等前端语言对接。runtime::Object是除runtime::PackedFunc外运行时的主要数据结构它是一个带类型索引type index的引用计数基类支持运行时类型检查与向下转型downcasting。从 include/tvm/runtime/object.h 可以看到Object持有type_index_与ref_counter_子类通过注册_type_index自动填充类型索引与析构器。对象系统允许开发者向运行时引入新数据结构如 Array、Map、新的 IR 数据结构。编译器自身也重度依赖运行时机制所有 IR 数据结构都是runtime::Object的子类因此可以直接从 Python 前端访问和操作各种 API 通过 PackedFunc 机制暴露给前端。不同硬件后端的运行时支持定义在 runtime 的子目录中例如 runtime/opencl这些硬件特化运行时模块定义了设备内存分配与设备函数序列化等 API。runtime/rpc实现了对 PackedFunc 的 RPC 支持。可以用 RPC 机制把交叉编译的库发送到远程设备并基准测试执行性能RPC 基础设施为学习式优化提供了来自广泛硬件后端的数据收集能力。实现位于 src/runtime/rpc如 src/runtime/rpc/rpc_channel.ccRPCChannel 是连接本地与远程 RPC 会话的通信端点抽象。相关深度指南链接已转换为仓库根相对路径runtime、debugger、virtual_machine、introduction_to_module_serialization、device_target_interactions。7.3 tvm/node —— IR 数据结构的反射、序列化与结构等价node 模块在runtime::Object之上为 IR 数据结构增加了反射reflection、序列化、结构等价structural equivalence与哈希hashing等能力。得益于 node 模块在 Python 中可以直接按字段名访问 TVM IRNode 的任何字段x tvm.tir.Var(x, int32) y tvm.tir.Add(x, x) # a and b are fields of a tir.Add node # we can directly use the field name to access the IR structures assert y.a x我们也可以把任意 IR 节点序列化为 JSON 格式再加载回来。这种可保存、可存储、可检视的能力为让编译器更易用奠定了基础。相关实现位于 src/nodestructural_equal.cc、structural_hash.cc、serialization.cc、reflection.cc 等。7.4 tvm/ir —— 统一的数据结构与接口tvm/ir文件夹包含所有 IR 函数变体共用的统一数据结构与接口被tvm/relay和tvm/tir共享主要包括IRModuleTypePassContext 与 PassOp不同函数变体relay.Function 与 tir.PrimFunc可以共存于同一个 IRModule。虽然内容表示不同但它们使用相同的数据结构表示类型因此可以用同一数据结构表示这些变体的函数类型签名。统一类型系统允许一个函数变体调用另一个函数变体只要明确调用约定为未来的跨函数变体优化打开大门。同时提供统一的PassContext用于配置 pass 行为以及通用的组合 pass 来执行 pass 管线。原文档给出的配置示例# configure the behavior of the tir.UnrollLoop pass with tvm.transform.PassContext(config{tir.UnrollLoop: { auto_max_step: 10 }}): # code affected by the pass contextOp是表示所有系统定义的原语算子/intrinsics 的公共类开发者可以向系统注册新的 Op 及其附加属性例如该 Op 是否为 elementwise。Pass 基础设施的深入介绍见 pass_infra。7.5 tvm/target —— 代码生成器与 Target 描述target 模块包含所有把 IRModule 翻译为目标 runtime.Module 的代码生成器并提供统一的Target类描述目标。编译管线可以通过查询 target 的属性信息以及注册到每个 target idcuda、opencl的 builtin 信息来定制。相关讨论见 device_target_interactions。7.6 tvm/tir —— 低层程序表示TIR 定义低层程序表示使用tir::PrimFunc表示可被 TIR pass 变换的函数。除 IR 数据结构外tir 模块还通过公共 Op 注册表定义一组 builtin intrinsics 及其属性以及tir/transform中的变换 pass。7.7 tvm/arith —— 索引算术分析arith 模块与 TIR 紧密关联。低层代码生成中的关键问题之一是指针/索引算术性质分析——正负性、变量边界、描述迭代空间的整数集合。arith 模块提供一组主要是整数分析工具TIR pass 可以利用这些分析来化简和优化代码。对应源码位于 src/arith如 rewrite_simplify.cc、int_set.cc、presburger_set.cc 等。7.8 tvm/te —— 张量表达式 DSLte 即 tensor expression张量表达式是一个领域特定语言模块让我们可以通过书写张量表达式快速构造tir::PrimFunc变体。重要张量表达式本身不是可以存入 IRModule 的自包含函数而是可以拼接起来构建 IRModule 的 IR 片段。te/schedule提供一组调度原语来控制生成函数未来部分调度组件可能会被引入到tir::PrimFunc自身。相关主题见 inferbound 与 hybrid_script。7.9 tvm/topi —— 张量算子库存虽然每个用例都可以直接通过 TIR 或 TE 构造算子但这样做很繁琐。topiTensor operator inventory提供一组预定义算子以 TE 或 TIR 形式实现语义对齐 numpy 并覆盖常见深度学习工作负载同时提供一批常用调度模板用于在不同 target 平台上获得高性能实现。对应实现位于 include/tvm/topi 与 python/tvm/topi。7.10 tvm/relay —— 高层函数式 IRRelay 是表示完整模型的高层函数式 IR各种优化定义在relay.transform中。Relay 编译器定义了多种方言dialect每种方言支持特定风格的优化值得一提的包括QNN用于导入预量化模型VM用于 lower 到动态虚拟机memory用于内存优化。深入主题见 relay_intro、relay_op_strategy、convert_layout。7.11 tvm/autotvm —— 自动化搜索式程序优化AutoTVM 与 AutoScheduler 都是自动化、基于搜索的程序优化组件目前主要包括成本模型与特征提取cost models and feature extraction存储程序基准结果、用于构建成本模型的记录格式record format一组作用于程序变换的搜索策略search policies。自动化程序优化仍是活跃的研究领域因此设计上刻意模块化研究人员可以通过 Python 绑定快速修改某个组件或应用自己的算法自定义搜索并插入自己的算法。实现位于 src/auto_scheduler 与 src/autotvm配套基准说明见 benchmark。7.12 Frontends —— 模型导入前端前端负责把不同框架的模型摄入 TVM 栈命名空间为tvm.relay.frontend是模型导入 API 的入口。TensorFlow 前端指南见 frontend/tensorflow。7.13 Security 与 microTVMSecurity安全相关文档见 security。microTVM面向嵌入式/MCU 场景的微型运行时设计相关主题见 microtvm_design、microtvm_project_api、model_library_format。八、总结与延伸讨论Summary and Discussions回顾全文编译流程中的关键数据结构是IRModule包含 relay.Function 与 tir.PrimFuncruntime.Module包含 runtime.PackedFunc。编译的大部分内容都是这些关键数据结构之间的变换relay/transform与tir/transform是确定性的、基于规则的变换auto_scheduler与autotvm包含基于搜索的变换。最后强调编译流程示例只是 TVM 栈的一种典型用法。TVM 将这些关键数据结构与变换同时暴露给 Python 与 C API因此你可以像使用 numpy 一样使用 TVM——只是感兴趣的数据结构从numpy.ndarray换成了tvm.IRModule。典型的高级用法包括直接用 Python API 构造 IRModule组合自定义变换集合例如自定义量化流程直接用 Python API 操纵 IR。若需进一步深入原文档还推荐阅读开发者指南dev-how-to以及 device_target_interactions各物理设备与代码生成目标的交互。对于希望参与开发的读者可结合本文提到的源码路径include/tvm/runtime/module.h、include/tvm/runtime/packed_func.h、src/relay/transforms/fuse_ops.cc、src/relay/backend/build_module.cc 等逐模块阅读形成文档脉络 源码证据的完整认知闭环。赞分享编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载相关推荐Apache TVM 编译架构全解析从模型 IR 到可部署模块的端到端编译流程Apache TVM 编译架构全解析从模型 IR 到可部署模块的端到端编译流程 Apache TVM 是面向机器学习的开源编译框架其核心使命是把高层模型描述模型编译深度学习推理引擎Helix 编辑器源码架构指南从 Rope 到 Term 的模块化分层设计Helix 编辑器源码架构指南从 Rope 到 Term 的模块化分层设计 本篇技术指南基于 Helix一种后现代模态文本编辑器仓库根目录下的 docs/代码编辑器开发工具CLITriton编译器架构前端、中端和后端的模块化设计Triton编译器架构前端、中端和后端的模块化设计 概述 Triton是一个专为深度学习计算优化的编程语言和编译器其核心设计理念是通过模块化的编译器架构实现编译器编程语言人工智能深度学习高性能计算上一篇Proxyee 开源项目教程下一篇Awesome Hacking Resources黑客与渗透测试资源宝库创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考