章节2：MLIR的引入

发布时间：2026/10/27 18:41:07
章节2：MLIR的引入

系列文章目录

章节1：什么是MLIR

章节2：MLIR的引入

章节3：MLIR完成一个端到端的流程

目录

系列文章目录

一、tensorflow执行方式

1. 前端：构建计算图与 Op 注册

2. 后端准备：Kernel 注册与设备匹配

3. 核心调度：Executor 执行器

4. 底层执行：Kernel Compute

二、使用XLA加速tensorflow

1.核心原理：从“解释执行”到“编译执行”

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++ 层注册。

  • ‌节点插入 Graph‌：
    Python 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 维护一个任务队列。它不会等待一个节点完全结束才启动下一个，而是采用异步机制：

    1. 当一个节点的所有输入张量（Tensor）就绪时，该节点变为“可执行”状态。

    2. Executor 将该节点的计算任务封装为一个回调函数，提交给目标设备的线程池（Thread Pool）。

    3. 对于 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 加速‌是指利用 ‌XLA（Accelerated Linear Algebra，加速线性代数）‌ 编译器对深度学习模型的数学运算进行即时编译和优化，从而显著提升模型在特定硬件（如 TPU、GPU）上的训练和推理速度。

简单来说，它就像是一个“超级翻译官”兼“优化大师”，将你写的 Python/PyTorch/JAX 代码转换成针对硬件高度优化的机器码，而不是让硬件去逐条解释执行原始的 Python 指令。

执行方式如下：

1.核心原理：从“解释执行”到“编译执行”

传统框架（如原生 PyTorch 的 eager 模式或早期 TensorFlow）通常采用‌命令式执行‌：每行代码发出一个操作指令，硬件立即执行。这种方式灵活但效率低，因为每个小操作都有启动开销，且无法看到全局优化空间。

XLA 加速的核心机制是“惰性求值”与“图编译”：‌

  1. ‌构建计算图（Tracing/Lazy Evaluation）‌：
    当你运行模型时，XLA 不会立即执行每个算子，而是先“记录”下所有的操作序列，形成一个中间表示（IR）的计算图。这个过程被称为“追踪”（Tracing）。

  2. ‌整体编译与优化‌：
    XLA 编译器拿到整个计算图后，会进行全局优化。最关键的技术是‌算子融合（Operator Fusion）‌：它将多个小的、连续的操作（如Add->Relu->BatchNorm）合并成一个单一的 GPU/TPU 内核（Kernel）。

  3. ‌生成高效机器码‌：
    优化后的图被编译成针对特定硬件（如 NVIDIA GPU 的 CUDA 代码或 Google TPU 的指令集）的高效二进制代码。

  4. ‌执行‌：
    最后，硬件直接执行这段优化后的代码。由于减少了内核启动次数和数据在显存中的读写次数，性能大幅提升。

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 作为核心后端基座：

  1. ‌XLA 前端‌：接收来自 TensorFlow/JAX 的高层张量计算图，生成 XLA 专属的中间表示 ‌HLO（High Level Optimizer）‌，完成算子融合、自动并行、内存复用等AI领域专属的高层优化。

  2. ‌HLO 转译层‌：将高层的 HLO 张量操作，逐步展开、映射为符合 LLVM 规范的中间表示（LLVM IR），比如一个HLO矩阵乘法会被拆解为一系列LLVM的加载、乘加、存储指令序列 。

  3. ‌LLVM 后端接管‌：生成的LLVM IR直接送入LLVM完整的工具链，经过指令调度、寄存器分配、目标平台专属优化后，最终输出x86、ARM、CUDA等不同硬件架构的原生机器码 。

三、XLA存在的问题

XLA 原生架构的核心不足：

  1. ‌IR层级断层，优化空间受限‌
    XLA 只有HLO这一个核心中间表示，高层张量图优化完成后，直接硬着陆到LLVM IR或TPU硬件指令，中间缺失了多层级过渡IR。这导致循环变换、多维数组分块、仿射优化这类关键的中间层变换无法优雅实现，只能靠手写大量特殊Case的Emitter代码，优化能力被严重锁死 。

  2. ‌硬件适配成本极高，扩展性差‌
    XLA 为每一种新硬件后端（如某款新ASIC、边缘端NPU）都需要从零开始手写整套代码生成逻辑，复用率极低。传统XLA的GPU代码生成甚至重度依赖LLVM的元数据hack来传递优化信息，维护成本极高，新硬件接入周期长达数月甚至数年 。

  3. ‌代码库碎片化，维护难度指数级上升‌
    XLA的所有变换逻辑都是用C++硬编码实现的，没有统一的IR基础设施。随着算子数量、硬件种类增加，手写的Emitter代码量爆炸式增长，公共子表达式消除、常量折叠这类基础编译器优化都要重复造轮子，无法复用成熟的通用编译器生态能力 。

  4. ‌跨框架兼容性差，生态封闭‌
    早期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编译栈。

简单来说：语义统一，组件可以复用；共享生态，各个层次可以协调优化；层级之间跨度小，方便实现；