前言
MLIR 社区充满活力。但由于它是一个新的且快速发展的项目,因此可用的教程和文档并不多。没有权威的MLIR书籍。大多数围绕事物的推理都来自民间传说和技术性很强的 RFC。而且由于 MLIR 构建在 LLVM(该缩写词以前的意思是“低级虚拟机”)之上,因此现有的许多文档都通过类比 LLVM 来解释概念,这对于像我这样不熟悉内部结构的人来说没有帮助LLVM 的工作原理。最后,在我看来,目前的教程水平太高,无法让人们真正了解如何在框架中编写程序。为此,我希望降低使用 MLIR 的准入门槛。因此,这一系列博客文章将详细介绍 MLIR 的总体情况。
MLIR 和 LLVM 简史
关于 MLIR,您首先会注意到的是,它位于LLVM 项目的 monorepo中名为mlir/
文件夹下。 LLVM 是一种抽象的汇编语言,编译器开发人员可以将其作为后端,然后 LLVM 本身打包了许多优化和可以编译到的“真正”后端目标。比如说,如果您使用Rust 编程语言,并且希望编译为 x86、ARM 和 WebAssembly,而无需执行所有这些工作,则只需输出 LLVM 代码,然后运行 LLVM 的编译套件即可。
我不想过多了解 LLVM 的历史(更多详细信息请参阅此采访),而且我对它没有任何第一手知识,但从我可以收集到的 LLVM(以前代表“Low Level”) Virtual Machine”)是Chris Lattner在 2000 年代初期的博士项目,旨在成为下一代 C 编译器。 Chris 转到 Apple,负责 LLVM 以及基于 LLVM 构建的 Swift 等语言的工作。 2017 年,他加入 Google Brain,担任 TensorFlow 基础设施团队总监,他和他的团队构建了 MLIR来统一其生态系统中的孤立工具。
我们将在以后的文章中详细讨论 MLIR 到底是什么以及它提供什么。有关高级概述,请参阅 MLIR 论文。简而言之,它是一个构建编译器的框架,其基本理念是一个大编译器应该被分解为许多子语言之间的小编译器(编译器人称之为“中间表示”或“IR”),其中每个子语言-语言旨在使特定类型的优化更自然地表达。因此,MLIR 是多级中间表示的缩写。
MLIR 与 TensorFlow 相关,因为训练和推理都可以被视为程序,其指令类似于“2d 卷积”和“softmax”。优化这些指令并将其转换为较低级别的硬件指令(尤其是在 TPU 加速器上)的过程在很大程度上是编译器的问题。 MLIR 将过程分解为不同抽象级别的 IR,例如张量运算、线性代数和较低级别的控制流。
但 LLVM 无法直接用作 TensorFlow 编译器。对于 CPU 来说,它过于传统且过于专业化,在较低的抽象层上运行,并且存在附带的技术债务。但 LLVM 确实有许多可重用的部分,例如数据结构、错误处理和测试基础设施。再加上 Lattner 对他工作了近 20 年的项目的熟悉程度,将 MLIR 放入 monorepo 中可能会更容易启动它。
Build systems
现在,LLVM 和 MLIR 的官方构建系统是CMake。但出于几个原因我将使用Bazel。首先,我想引导感兴趣的读者了解 HEIR,这就是 HEIR 所使用的,因为它是 Google 旗下的项目。其次,尽管人们可能担心 Bazel 配置复杂或不受支持,但由于 MLIR 和 LLVM 已成为 Google 生产基础设施的关键,Google 帮助维护Bazel 与 CMake 配置并行的“覆盖”,并且 Google 有待命工程师负责确保 Google 的 MLIR 内部副本与 LLVM monorepo 保持同步,并确保任何构建问题都能得到及时修复。剩下的粗糙边缘对于像我这样不耐烦的假人来说足够简单了。
(下面这部分可跳过,去看看CMake)
这是 Bazel 的概述(其中部分内容重复于我之前的文章)。Bazel是 Google 内部构建系统“Blaze”的开源类似物,而Starlark是其受 Python 启发的脚本语言。关于 Bazel 有很多观点,这里不再赘述。您可以使用bazelisk
程序安装它。
首先是一些术语。要使用 Bazel,您需要执行以下操作。
- 定义一个
WORKSPACE文件,该文件定义项目的所有外部依赖项、如何获取其源代码以及应使用哪些 bazel 命令来构建它们。这可以被认为是顶级 CMakeLists,只不过除了声明项目目录树的根和项目名称之外,它不包含任何构建项目的指令。
- 在每个子目录中定义一组
BUILD
文件,声明可以从该目录(但不是其子目录)中的源文件构建的构建目标。这类似于子目录中的 CMakeLists 文件。每个构建目标都可以声明对其他构建目标的依赖关系,bazel build
确保首先构建依赖关系,并在会话中缓存构建结果。许多项目在项目根目录中都有一个BUILD
文件,用于公开项目的公共库和 API。 - 使用
cc_library
、cc_binary
和cc_test
等内置 bazel规则将文件分组到可以使用bazel build
构建的库、也可以使用bazel run
运行的可执行二进制文件以及也可以使用bazel test
运行的测试。大多数 bazel 规则归结为使用特定参数调用一些可执行程序,例如gcc
或javac
,同时还跟踪文件系统上“封闭”位置中构建工件的累积依赖项集。 - 定义执行自定义程序的新 bazel 规则,并声明静态依赖关系图的依赖关系和输出。 MLIR 的自定义规则围绕
tblgen
程序,该程序是 MLIR 的自定义模板语言,用于生成 C++ 代码。 - 编写将内置 bazel 命令链接在一起的任何其他 bazel宏。宏看起来像调用各个 bazel 规则并可能在它们之间传递数据的 Python 函数。它们编写在
.bzl
文件(包含 Starlark 代码)中,由bazel
直接解释。当我们谈论 MLIR 的测试框架lit
时,我们会看到一个很好的 bazel 宏示例,但本文包含一个简单的示例,用于在WORKSPACE
文件(也是 Starlark)中设置 LLVM 依赖项。
一般来说,bazel
分两个阶段构建目标。首先是分析阶段,它加载所有BUILD
文件和导入的.bzl
文件,并扫描所有被调用的规则。特别是,它运行宏,因为它需要知道宏调用什么规则(并且规则可以由控制流保护,或者它们的参数可以动态生成等)。但它本身并不运行构建规则。在此过程中,它可以构建完整的依赖关系图,并报告有关拼写错误、缺少依赖关系、循环等的错误。分析阶段完成后,它会按依赖关系顺序运行底层规则,并缓存结果。仅当规则所依赖的文件或其底层依赖项发生变化时,Bazel 才会再次运行规则。
WORKSPACE 和 llvm-project 依赖项
添加.gitignore来过滤掉 Bazel 的构建目录后,此提交设置了一个初始WORKSPACE
文件和两个 bazel 文件,它们执行不寻常的两步舞蹈来配置 LLVM 代码库。工作区文件如下所示:
workspace(name = "mlir_tutorial")load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe")load("//bazel:import_llvm.bzl", "import_llvm")import_llvm("llvm-raw")load("//bazel:setup_llvm.bzl", "setup_llvm")setup_llvm("llvm-project")
这不是一种正常的依赖关系。正常的依赖关系可能如下所示:
http_archive(name = "abc",build_file = "//bazel:abc.BUILD",sha256 = "7fa5a448a4309fb4d6cf856c3fe4cc4be46b09dd552a05d5cfacd75f8d9504ad",urls = ["https://github.com/berkeley-abc/abc/archive/eb44a80bf2eb8723231e72bb095c97d1e4834d56.zip",],
)
上面告诉 bazel:从给定的 URL 中提取 zip 文件,仔细检查它的哈希和,然后(因为依赖的项目不是用 bazel 构建的)我会告诉你在我的存储库中的哪里可以找到你应该的 BUILD 文件用于构建它。如果项目有 BUILD 文件,我们可以省略build_file
并且它会正常工作。
现在,LLVM 具有 bazel 构建文件,但它们隐藏在项目的utils/bazel
子目录中。 Bazel 要求其特殊文件位于正确的位置,并且 bazel 配置旨在与 CMake 配置同步。因此utils/bazel
目录有一个llvm_configure
bazel 宏,它执行一个正确符号链接所有内容的 python 脚本。有关上游系统的更多信息可以在此处找到。
因此,要运行这个宏,我们必须下载 LLVM 代码作为存储库,将其放入import_llvm.bzl
文件中,并调用该宏,将其放入setup_llvm.bzl
中。为什么有两个文件? bazel 的一个明显的怪癖是,您无法从下载依赖项的同一个WORKSPACE
文件中的依赖项的 bazel 文件中load()
宏。
还值得一提的是,import_llvm.bzl
是我放置硬编码提交哈希的位置,该哈希将该项目固定到特定的 LLVM 版本。
相关文章:
【AI编译器】MLIR(7)— Dialect转换 - 知乎
【AI编译器】MLIR(6)— Dialect使用Traits - 知乎
【AI编译器】MLIR(5)— 定义一个Dialect - 知乎
【AI编译器】MLIR(4)— Tablegen写pass - 知乎
【AI编译器】MLIR(3) — 编写第一个pass - 知乎
【AI编译器】MLIR(2) — 运行和测试Lowering - 知乎
【AI编译器】MLIR(1) — 入门 - 知乎