TensorFlow MLIR 团队发布

TensorFlow 生态系统包含许多编译器和优化器，它们在软件和硬件堆栈的多个级别上运行。作为 TensorFlow 的日常用户，这种多级堆栈可能会在使用不同类型的硬件（GPU、TPU、移动设备）时表现为难以理解的编译器和运行时错误。

这些组件，从图开始，可以概括如下

实际上比这更复杂

在这个图中，我们可以看到 TensorFlow 图 [1] 可以通过多种不同的方式运行。这包括

• 将它们发送到 TensorFlow 执行器，该执行器调用手写的 op 内核
  
  
• 将它们转换为 XLA 高级优化器表示（XLA HLO），进而可以调用 LLVM 编译器用于 CPU 或 GPU，或者继续使用 XLA 用于 TPU。（或者两者的某种组合！）
  
  
• 将它们转换为 TensorRT、nGraph 或其他针对特定硬件指令集的编译器格式
  
  
• 将图转换为 TensorFlow Lite 格式，然后在 TensorFlow Lite 运行时中执行，或者进一步转换为在 GPU 或 DSP 上运行，通过 Android 神经网络 API (NNAPI) 或相关技术。

此外，还有其他更复杂的路径，包括每一层内的多轮优化，例如 Grappler 框架，它在 TensorFlow 中优化张量布局和操作。

虽然这些众多的编译器和表示实现大大提高了性能，但这种异构环境可能会导致最终用户出现问题，例如在这些系统之间的边界处产生令人困惑的错误消息。此外，新的硬件和软件堆栈创建者必须为每条新路径重建优化和转换流程。

考虑到所有这些，我们想宣布 MLIR 或 Multi-Level Intermediate Representation。这是一种表示格式和编译器实用程序库，位于模型表示和生成特定于硬件的代码的低级编译器/执行器之间。使用 MLIR，我们希望能够在优化编译器设计和实现方面进行新的探索，并以生产级组件为后盾。

我们预计 MLIR 会让许多群体感兴趣，包括

• 希望优化机器学习模型性能和内存消耗的编译器研究人员和实施者
  
  
• 希望找到一种方法将他们的硬件连接到 TensorFlow 的硬件制造商，例如 TPU、手机中的便携式神经硬件和其他定制 ASIC
  
  
• 编写语言绑定的人员，他们希望利用优化编译器和硬件加速。

什么是 MLIR？

MLIR 本质上是一个灵活的现代优化编译器基础架构。这意味着它包含用于中间表示 (IR) 的规范以及一个用于对该表示进行转换的代码工具包。（在编译器术语中，当您从更高级别的表示移动到更低级别的表示时，这些转换可以称为“降低”，我们将在后面使用这个术语。）

MLIR 受到 LLVM 的很大影响，并且毫不犹豫地借鉴了其中的许多优秀理念。它有一个灵活的类型系统，并允许表示、分析和转换组合了同一编译单元中多个抽象级别的图。这些抽象包括 TensorFlow 操作、嵌套的多面体循环区域，甚至 LLVM 指令以及固定硬件操作和类型。

MLIR 方言

为了区分不同的硬件和软件目标，MLIR 有“方言”，包括

• TensorFlow IR，它代表 TensorFlow 图中所有可能的事物
  
  
• XLA HLO IR，它旨在利用 XLA 的编译能力（输出到 TPU 等）
  
  
• 一个实验性的仿射方言，它专注于 多面体表示 和优化
  
  
• LLVM IR，它与 LLVM 自己的表示之间存在一对一映射，允许 MLIR 通过 LLVM 生成 GPU 和 CPU 代码
  
  
• TensorFlow Lite，它将转换为在移动平台上运行的代码
  

每个方言都包含一组定义的操作，这些操作对它们施加了不变式，例如：“这是一个二元运算符，并且输入和输出具有相同的类型。”

添加到 MLIR

MLIR 没有固定的/内置的全局已知操作列表（没有“内联函数”）。方言可以定义完全自定义的类型，这就是 MLIR 如何模拟 LLVM IR 类型系统（它具有第一类聚合）、对 ML 优化加速器（如量化类型）重要的域抽象，甚至 Swift 或 Clang 类型系统（它们基于 Swift/Clang 声明节点）在未来。

如果您想连接一个新的低级编译器，您将创建一个新的方言以及 TensorFlow 图方言和您的方言之间的降低。这为硬件和编译器制造商铺平了道路。您甚至可以在同一模型中针对不同级别的方言；更高级别的优化器将尊重 IR 中不熟悉的部件，并等待更低级别处理它。

对于编译器研究人员和框架制造商，MLIR 允许您在各个级别组合转换，您甚至可以在 IR 中定义自己的操作和抽象 - 允许您最好地模拟您尝试解决的问题域。这样，MLIR 更像是一个纯粹的编译器基础架构而不是 LLVM。

虽然 MLIR 充当 ML 的编译器，但我们也看到它使编译器中使用机器学习技术成为可能！这尤其重要，因为开发数值库的工程师的扩展速度无法与 ML 模型或硬件的多样化速度相匹配。MLIR 的可扩展性促进了对代码降低策略的探索，并在抽象之间执行逐步降低。

接下来会发生什么

我们已经开放了 GitHub 存储库，欢迎您的兴趣（查看我们的 教程！）。我们将在未来几个月内发布更多此工具包，包括 TensorFlow 和 TF Lite 方言的规范。我们期待告诉你更多关于这个；有关详细信息，请参阅 Chris Lattner 在 c4ml 上的演讲 以及我们的 GitHub 上的 README。如果您想了解与 MLIR 相关的所有内容，请加入 我们的新邮件列表，该列表将在短期内专注于我们在发布更多此项目时的公告。敬请关注！

脚注

[1] 在 TensorFlow 2.0 中，图可以是隐式的；急切执行可以单独运行操作、按组运行或作为完整图运行（例如 Keras 顺序）。无论如何，这些图或图片段都必须进行优化和执行。