系列文章目录章节1什么是MLIR章节2MLIR的引入章节3MLIR完成一个端到端的流程目录系列文章目录一、tensorflow执行方式1. 前端构建计算图与 Op 注册2. 后端准备Kernel 注册与设备匹配3. 核心调度Executor 执行器4. 底层执行Kernel Compute二、使用XLA加速tensorflow1.核心原理从“解释执行”到“编译执行”2.XLA 加速带来的主要优势3.XLA和LLVM的关系三、XLA存在的问题四、MLIR 针对性解决的核心痛点提示以下是本篇文章正文内容下面案例可供参考一、tensorflow执行方式执行方式流程如下其中Graph 定义了“做什么”Executor 决定了“何时做”和“在哪做”而 Kernel 负责“具体怎么做”。三者协同工作实现了 TensorFlow 的张量计算能力。TensorFlow 从计算图Graph的操作节点到执行器Executor再到内核Kernel的执行流程是一个典型的“定义-优化-调度-执行”的流水线过程。以下是这一过程的详细拆解1. 前端构建计算图与 Op 注册当用户在 Python 中调用如tf.matmul或tf.add时并没有立即进行数值计算而是完成了以下工作Op 定义与注册TensorFlow 内部维护了一个全局的OpDefLibrary。每个操作如 MatMul都有一个对应的 OpDefinition定义了它的输入输出类型、属性如 transpose_a等。这些定义通过宏如REGISTER_OP在 C 层注册。节点插入 GraphPython API如gen_math_ops.py会将用户传入的参数转换为底层 C 调用最终在当前的默认Graph对象中创建一个 NodeDef。NodeDef 是图的序列化表示仅包含操作的名称、类型、输入依赖和属性不包含具体的计算逻辑。此时图中的边Edge建立了节点之间的数据依赖关系。2. 后端准备Kernel 注册与设备匹配计算逻辑是由 Kernel 实现的而不是 Op 本身。一个 Op 可以有多个 Kernel 实现分别对应不同的硬件设备CPU/GPU/TPU和数据类型float32/int32。Kernel 注册开发者使用REGISTER_KERNEL_BUILDER宏将具体的 C 计算类继承自OpKernel注册到系统中。例如MatMul在 CPU 上有一个 kernel在 GPU 上有另一个基于 cuBLAS 的 kernel。设备放置Device Placement在会话Session初始化或图执行前TensorFlow 的运行时会分析图中的每个 Node根据设备的可用性、内存限制和用户指定的with tf.device()上下文决定每个节点由哪个设备Device执行。实例化 Kernel一旦确定了设备运行时就会为该节点查找并实例化对应的OpKernel对象。这个对象包含了执行计算所需的所有状态和方法主要是Compute()函数。3. 核心调度Executor 执行器Executor 是 TensorFlow 运行时的核心组件负责驱动整个计算图的执行。它的主要职责是将静态的图结构转化为动态的执行流。图优化与分区Executor 接收优化后的 Graph。如果涉及多设备如 CPU 和 GPU它会将图分割成多个子图Partition每个子图对应一个特定的设备。构建执行计划Executor 会对节点进行拓扑排序确定执行的先后顺序。它会识别哪些节点可以并行执行即没有数据依赖关系的节点。异步调度模型Executor 维护一个任务队列。它不会等待一个节点完全结束才启动下一个而是采用异步机制当一个节点的所有输入张量Tensor就绪时该节点变为“可执行”状态。Executor 将该节点的计算任务封装为一个回调函数提交给目标设备的线程池Thread Pool。对于 GPU这通常意味着将 CUDA Kernel 启动指令放入 GPU 流Stream中。4. 底层执行Kernel Compute当调度器将任务分发给具体设备后真正的计算发生在 Kernel 层面调用 Compute()设备线程池中的工作线程会调用对应OpKernel的Compute(OpKernelContext* context)方法。获取输入与分配输出OpKernelContext提供了访问输入张量数据的接口并负责为输出张量分配内存通常在设备的显存或内存中。执行具体算法CPU Kernel直接执行 C 代码可能调用 Eigen 库或 MKL 进行矩阵运算。GPU Kernel通常是一个轻量级的包装器负责准备参数并启动实际的 CUDA/HIP 核函数如cublasSgemm。结果回传与同步计算完成后Kernel 将结果写入输出张量。Executor 会监控这些完成信号一旦某个节点的输出就绪它会通知所有依赖该输出的后续节点触发它们的执行。二、使用XLA加速tensorflow使用 XLA 加速是指利用 XLAAccelerated Linear Algebra加速线性代数 编译器对深度学习模型的数学运算进行即时编译和优化从而显著提升模型在特定硬件如 TPU、GPU上的训练和推理速度。简单来说它就像是一个“超级翻译官”兼“优化大师”将你写的 Python/PyTorch/JAX 代码转换成针对硬件高度优化的机器码而不是让硬件去逐条解释执行原始的 Python 指令。执行方式如下1.核心原理从“解释执行”到“编译执行”传统框架如原生 PyTorch 的 eager 模式或早期 TensorFlow通常采用命令式执行每行代码发出一个操作指令硬件立即执行。这种方式灵活但效率低因为每个小操作都有启动开销且无法看到全局优化空间。XLA 加速的核心机制是“惰性求值”与“图编译”构建计算图Tracing/Lazy Evaluation当你运行模型时XLA 不会立即执行每个算子而是先“记录”下所有的操作序列形成一个中间表示IR的计算图。这个过程被称为“追踪”Tracing。整体编译与优化XLA 编译器拿到整个计算图后会进行全局优化。最关键的技术是算子融合Operator Fusion它将多个小的、连续的操作如Add-Relu-BatchNorm合并成一个单一的 GPU/TPU 内核Kernel。生成高效机器码优化后的图被编译成针对特定硬件如 NVIDIA GPU 的 CUDA 代码或 Google TPU 的指令集的高效二进制代码。执行最后硬件直接执行这段优化后的代码。由于减少了内核启动次数和数据在显存中的读写次数性能大幅提升。2.XLA 加速带来的主要优势算子融合Operator Fusion这是 XLA 最核心的优化手段。通过减少 GPU/TPU 内核启动开销和中间张量的内存读写显著降低延迟并提高吞吐量。内存优化XLA 能更智能地管理显存复用缓冲区从而允许你训练更大的模型或使用更大的 Batch Size。硬件特异性优化XLA 针对不同硬件后端CPU、GPU、TPU生成专属的最优代码。例如在 TPU 上XLA 能充分利用其矩阵乘法单元MXU的高并行能力在 NVIDIA GPU 上它能更好地利用 Tensor Cores。3.XLA和LLVM的关系XLA 和 LLVM 是典型的「高层领域专用编译器」与「通用底层编译器基础设施」的深度协作关系二者定位完全互补XLA 不重复造轮子而是基于 LLVM 成熟的底层能力专注于AI张量计算的专属优化。XLA 的完整编译流水线天然以 LLVM 作为核心后端基座XLA 前端接收来自 TensorFlow/JAX 的高层张量计算图生成 XLA 专属的中间表示 HLOHigh Level Optimizer完成算子融合、自动并行、内存复用等AI领域专属的高层优化。HLO 转译层将高层的 HLO 张量操作逐步展开、映射为符合 LLVM 规范的中间表示LLVM IR比如一个HLO矩阵乘法会被拆解为一系列LLVM的加载、乘加、存储指令序列 。LLVM 后端接管生成的LLVM IR直接送入LLVM完整的工具链经过指令调度、寄存器分配、目标平台专属优化后最终输出x86、ARM、CUDA等不同硬件架构的原生机器码 。三、XLA存在的问题XLA 原生架构的核心不足IR层级断层优化空间受限XLA 只有HLO这一个核心中间表示高层张量图优化完成后直接硬着陆到LLVM IR或TPU硬件指令中间缺失了多层级过渡IR。这导致循环变换、多维数组分块、仿射优化这类关键的中间层变换无法优雅实现只能靠手写大量特殊Case的Emitter代码优化能力被严重锁死 。硬件适配成本极高扩展性差XLA 为每一种新硬件后端如某款新ASIC、边缘端NPU都需要从零开始手写整套代码生成逻辑复用率极低。传统XLA的GPU代码生成甚至重度依赖LLVM的元数据hack来传递优化信息维护成本极高新硬件接入周期长达数月甚至数年 。代码库碎片化维护难度指数级上升XLA的所有变换逻辑都是用C硬编码实现的没有统一的IR基础设施。随着算子数量、硬件种类增加手写的Emitter代码量爆炸式增长公共子表达式消除、常量折叠这类基础编译器优化都要重复造轮子无法复用成熟的通用编译器生态能力 。跨框架兼容性差生态封闭早期XLA深度绑定TensorFlow没有标准化的中间交换格式PyTorch、PaddlePaddle等第三方框架很难无缝接入XLA的加速能力无法形成跨框架的统一编译生态严重限制了XLA的普及范围。简单来说组合爆炸很多组件无法共用各个层次优化难以迁移就是从XLA HLO直接到LLVM IR跨度太大四、MLIR 针对性解决的核心痛点MLIR通过“多层级方言Dialect”架构完美补全了XLA的短板新增Linalg、Affine等中间层方言打通了从HLO张量图到底层循环优化的完整链路让XLA可以直接复用MLIR生态中成熟的循环分块、向量化优化Pass无需从零手写代码 。提供统一的IR基础设施新硬件厂商只需实现对应MLIR方言的Lowering逻辑就能自动复用整条编译流水线的所有优化能力硬件接入成本降低一个数量级。基于MLIR推出StableHLO标准化方言成为跨框架的统一中间表示让PyTorch、JAX等所有主流框架都能无缝接入XLA的加速能力彻底打破生态封闭的局面。简单来说XLA解决了“在TPU上把AI模型跑快”的问题而MLIR解决了“让XLA能适配所有硬件、所有框架同时大幅降低编译器维护成本”的问题二者结合形成了现在工业界最成熟的AI编译栈。简单来说语义统一组件可以复用共享生态各个层次可以协调优化层级之间跨度小方便实现 章节2：MLIR的引入
系列文章目录章节1什么是MLIR章节2MLIR的引入章节3MLIR完成一个端到端的流程目录系列文章目录一、tensorflow执行方式1. 前端构建计算图与 Op 注册2. 后端准备Kernel 注册与设备匹配3. 核心调度Executor 执行器4. 底层执行Kernel Compute二、使用XLA加速tensorflow1.核心原理从“解释执行”到“编译执行”2.XLA 加速带来的主要优势3.XLA和LLVM的关系三、XLA存在的问题四、MLIR 针对性解决的核心痛点提示以下是本篇文章正文内容下面案例可供参考一、tensorflow执行方式执行方式流程如下其中Graph 定义了“做什么”Executor 决定了“何时做”和“在哪做”而 Kernel 负责“具体怎么做”。三者协同工作实现了 TensorFlow 的张量计算能力。TensorFlow 从计算图Graph的操作节点到执行器Executor再到内核Kernel的执行流程是一个典型的“定义-优化-调度-执行”的流水线过程。以下是这一过程的详细拆解1. 前端构建计算图与 Op 注册当用户在 Python 中调用如tf.matmul或tf.add时并没有立即进行数值计算而是完成了以下工作Op 定义与注册TensorFlow 内部维护了一个全局的OpDefLibrary。每个操作如 MatMul都有一个对应的 OpDefinition定义了它的输入输出类型、属性如 transpose_a等。这些定义通过宏如REGISTER_OP在 C 层注册。节点插入 GraphPython API如gen_math_ops.py会将用户传入的参数转换为底层 C 调用最终在当前的默认Graph对象中创建一个 NodeDef。NodeDef 是图的序列化表示仅包含操作的名称、类型、输入依赖和属性不包含具体的计算逻辑。此时图中的边Edge建立了节点之间的数据依赖关系。2. 后端准备Kernel 注册与设备匹配计算逻辑是由 Kernel 实现的而不是 Op 本身。一个 Op 可以有多个 Kernel 实现分别对应不同的硬件设备CPU/GPU/TPU和数据类型float32/int32。Kernel 注册开发者使用REGISTER_KERNEL_BUILDER宏将具体的 C 计算类继承自OpKernel注册到系统中。例如MatMul在 CPU 上有一个 kernel在 GPU 上有另一个基于 cuBLAS 的 kernel。设备放置Device Placement在会话Session初始化或图执行前TensorFlow 的运行时会分析图中的每个 Node根据设备的可用性、内存限制和用户指定的with tf.device()上下文决定每个节点由哪个设备Device执行。实例化 Kernel一旦确定了设备运行时就会为该节点查找并实例化对应的OpKernel对象。这个对象包含了执行计算所需的所有状态和方法主要是Compute()函数。3. 核心调度Executor 执行器Executor 是 TensorFlow 运行时的核心组件负责驱动整个计算图的执行。它的主要职责是将静态的图结构转化为动态的执行流。图优化与分区Executor 接收优化后的 Graph。如果涉及多设备如 CPU 和 GPU它会将图分割成多个子图Partition每个子图对应一个特定的设备。构建执行计划Executor 会对节点进行拓扑排序确定执行的先后顺序。它会识别哪些节点可以并行执行即没有数据依赖关系的节点。异步调度模型Executor 维护一个任务队列。它不会等待一个节点完全结束才启动下一个而是采用异步机制当一个节点的所有输入张量Tensor就绪时该节点变为“可执行”状态。Executor 将该节点的计算任务封装为一个回调函数提交给目标设备的线程池Thread Pool。对于 GPU这通常意味着将 CUDA Kernel 启动指令放入 GPU 流Stream中。4. 底层执行Kernel Compute当调度器将任务分发给具体设备后真正的计算发生在 Kernel 层面调用 Compute()设备线程池中的工作线程会调用对应OpKernel的Compute(OpKernelContext* context)方法。获取输入与分配输出OpKernelContext提供了访问输入张量数据的接口并负责为输出张量分配内存通常在设备的显存或内存中。执行具体算法CPU Kernel直接执行 C 代码可能调用 Eigen 库或 MKL 进行矩阵运算。GPU Kernel通常是一个轻量级的包装器负责准备参数并启动实际的 CUDA/HIP 核函数如cublasSgemm。结果回传与同步计算完成后Kernel 将结果写入输出张量。Executor 会监控这些完成信号一旦某个节点的输出就绪它会通知所有依赖该输出的后续节点触发它们的执行。二、使用XLA加速tensorflow使用 XLA 加速是指利用 XLAAccelerated Linear Algebra加速线性代数 编译器对深度学习模型的数学运算进行即时编译和优化从而显著提升模型在特定硬件如 TPU、GPU上的训练和推理速度。简单来说它就像是一个“超级翻译官”兼“优化大师”将你写的 Python/PyTorch/JAX 代码转换成针对硬件高度优化的机器码而不是让硬件去逐条解释执行原始的 Python 指令。执行方式如下1.核心原理从“解释执行”到“编译执行”传统框架如原生 PyTorch 的 eager 模式或早期 TensorFlow通常采用命令式执行每行代码发出一个操作指令硬件立即执行。这种方式灵活但效率低因为每个小操作都有启动开销且无法看到全局优化空间。XLA 加速的核心机制是“惰性求值”与“图编译”构建计算图Tracing/Lazy Evaluation当你运行模型时XLA 不会立即执行每个算子而是先“记录”下所有的操作序列形成一个中间表示IR的计算图。这个过程被称为“追踪”Tracing。整体编译与优化XLA 编译器拿到整个计算图后会进行全局优化。最关键的技术是算子融合Operator Fusion它将多个小的、连续的操作如Add-Relu-BatchNorm合并成一个单一的 GPU/TPU 内核Kernel。生成高效机器码优化后的图被编译成针对特定硬件如 NVIDIA GPU 的 CUDA 代码或 Google TPU 的指令集的高效二进制代码。执行最后硬件直接执行这段优化后的代码。由于减少了内核启动次数和数据在显存中的读写次数性能大幅提升。2.XLA 加速带来的主要优势算子融合Operator Fusion这是 XLA 最核心的优化手段。通过减少 GPU/TPU 内核启动开销和中间张量的内存读写显著降低延迟并提高吞吐量。内存优化XLA 能更智能地管理显存复用缓冲区从而允许你训练更大的模型或使用更大的 Batch Size。硬件特异性优化XLA 针对不同硬件后端CPU、GPU、TPU生成专属的最优代码。例如在 TPU 上XLA 能充分利用其矩阵乘法单元MXU的高并行能力在 NVIDIA GPU 上它能更好地利用 Tensor Cores。3.XLA和LLVM的关系XLA 和 LLVM 是典型的「高层领域专用编译器」与「通用底层编译器基础设施」的深度协作关系二者定位完全互补XLA 不重复造轮子而是基于 LLVM 成熟的底层能力专注于AI张量计算的专属优化。XLA 的完整编译流水线天然以 LLVM 作为核心后端基座XLA 前端接收来自 TensorFlow/JAX 的高层张量计算图生成 XLA 专属的中间表示 HLOHigh Level Optimizer完成算子融合、自动并行、内存复用等AI领域专属的高层优化。HLO 转译层将高层的 HLO 张量操作逐步展开、映射为符合 LLVM 规范的中间表示LLVM IR比如一个HLO矩阵乘法会被拆解为一系列LLVM的加载、乘加、存储指令序列 。LLVM 后端接管生成的LLVM IR直接送入LLVM完整的工具链经过指令调度、寄存器分配、目标平台专属优化后最终输出x86、ARM、CUDA等不同硬件架构的原生机器码 。三、XLA存在的问题XLA 原生架构的核心不足IR层级断层优化空间受限XLA 只有HLO这一个核心中间表示高层张量图优化完成后直接硬着陆到LLVM IR或TPU硬件指令中间缺失了多层级过渡IR。这导致循环变换、多维数组分块、仿射优化这类关键的中间层变换无法优雅实现只能靠手写大量特殊Case的Emitter代码优化能力被严重锁死 。硬件适配成本极高扩展性差XLA 为每一种新硬件后端如某款新ASIC、边缘端NPU都需要从零开始手写整套代码生成逻辑复用率极低。传统XLA的GPU代码生成甚至重度依赖LLVM的元数据hack来传递优化信息维护成本极高新硬件接入周期长达数月甚至数年 。代码库碎片化维护难度指数级上升XLA的所有变换逻辑都是用C硬编码实现的没有统一的IR基础设施。随着算子数量、硬件种类增加手写的Emitter代码量爆炸式增长公共子表达式消除、常量折叠这类基础编译器优化都要重复造轮子无法复用成熟的通用编译器生态能力 。跨框架兼容性差生态封闭早期XLA深度绑定TensorFlow没有标准化的中间交换格式PyTorch、PaddlePaddle等第三方框架很难无缝接入XLA的加速能力无法形成跨框架的统一编译生态严重限制了XLA的普及范围。简单来说组合爆炸很多组件无法共用各个层次优化难以迁移就是从XLA HLO直接到LLVM IR跨度太大四、MLIR 针对性解决的核心痛点MLIR通过“多层级方言Dialect”架构完美补全了XLA的短板新增Linalg、Affine等中间层方言打通了从HLO张量图到底层循环优化的完整链路让XLA可以直接复用MLIR生态中成熟的循环分块、向量化优化Pass无需从零手写代码 。提供统一的IR基础设施新硬件厂商只需实现对应MLIR方言的Lowering逻辑就能自动复用整条编译流水线的所有优化能力硬件接入成本降低一个数量级。基于MLIR推出StableHLO标准化方言成为跨框架的统一中间表示让PyTorch、JAX等所有主流框架都能无缝接入XLA的加速能力彻底打破生态封闭的局面。简单来说XLA解决了“在TPU上把AI模型跑快”的问题而MLIR解决了“让XLA能适配所有硬件、所有框架同时大幅降低编译器维护成本”的问题二者结合形成了现在工业界最成熟的AI编译栈。简单来说语义统一组件可以复用共享生态各个层次可以协调优化层级之间跨度小方便实现