TFRT：全新的 TensorFlow 运行时

TFRT 是一个全新的 TensorFlow 运行时，旨在提供统一且可扩展的基础架构层，在多种领域专用硬件上实现一流的性能。它能够高效利用多线程主机 CPU，支持完全异步的编程模型，并专注于低级别的效率优化。

TFRT 将惠及广泛的用户群体，但如果您属于以下任何一类，可能会特别感兴趣：

• 正在探索复杂新模型并希望向 TensorFlow 添加自定义操作的研究人员
• 希望在生产环境中部署模型时提升性能的应用开发者
• 计划将硬件集成到 TensorFlow 中的硬件制造商，包括边缘设备和数据中心设备

…或者您只是对先进的机器学习基础架构和底层运行时技术充满好奇！

要了解 TFRT 的早期进展与成果，请观看我们的 TensorFlow 开发者峰会 2020 演示，其中我们针对 ResNet 50 的小批量 GPU 推理任务进行了性能基准测试；以及我们的 MLIR 开放设计深度解析演示，该演示详细介绍了 TFRT 的核心组件、底层抽象及总体设计原则。

注意： TFRT 目前仍处于早期开发阶段，尚未准备好供广泛使用。

快速入门

简而言之： 本节将介绍如何为 TFRT 搭建开发环境，以及构建和测试 TFRT 组件的具体步骤。

TFRT 当前支持 Ubuntu 16.04。未来还将支持 macOS、Windows 等平台。构建和测试 TFRT 需要 Bazel 和 clang 工具链。对于 GPU 后端，则需要 NVIDIA 的 CUDA Toolkit 和 cuDNN 库。

为了说明 TFRT 的构建与测试流程，我们将编译并运行用于图执行的两个二进制文件。

回顾我们在开发者峰会上的演示：在图执行过程中，TensorFlow 用户会将通过高级 TensorFlow API 创建的计算图传递给 TFRT，随后 TFRT 调用基于 MLIR 的图编译器对该图进行优化和降级，最终生成适用于 TFRT 图执行的 BEF 格式文件（MLIR 是我们用来表示 TFRT 主机程序的编译基础设施）。下图所示的简化 TensorFlow 训练栈中，蓝色箭头即展示了这一流程。

接下来介绍的两个二进制文件专注于图执行工作流的后端部分。当图编译器完成 TensorFlow 计算图的优化，并将其转换为 MLIR 表示的低级别 TFRT 主机程序后，tfrt_translate 会从该主机程序生成 BEF 文件，而 bef_executor 则负责运行这个 BEF 文件。从 TFRT 主机程序经由 tfrt_translate 到 bef_executor 的完整流程，如下面展开的 TensorFlow 训练栈图所示。请注意，TFRT 主机程序与 BEF 文件之间的蓝色箭头代表了 tfrt_translate 的作用。这两个程序均位于 tools 目录下。

tfrt_translate

tfrt_translate 程序实现了 MLIR 与 BEF 之间的双向转换，类似于汇编器和反汇编器的功能。

bef_executor

bef_executor 程序是 BEF 文件的执行驱动程序。它读取 BEF 文件，设置运行时环境，并异步执行文件中的一个或多个函数。

先决条件

安装 Bazel

要构建 TFRT，您需要先安装 Bazel。TFRT 使用 Bazel 4.0 版本进行构建和验证。请按照 Bazel 安装指南 安装 Bazel，并通过以下命令验证安装是否成功：

$ bazel --version
bazel 4.0.0

安装 clang

请按照 clang 安装指南 安装 clang。推荐使用自动安装脚本同时安装 clang、lldb 和 lld。TFRT 使用 clang 11.1 版本进行构建和验证。

如果您已安装多个版本的 clang，请确保正确版本被设为默认。在基于 Ubuntu 的系统上，可以使用 update-alternatives 来选择默认版本。以下示例假设您已安装 clang-11：

$ sudo update-alternatives --install /usr/bin/clang clang /usr/bin/clang-11 11
$ sudo update-alternatives --install /usr/bin/clang++ clang++ /usr/bin/clang++-11 11

通过以下命令验证安装是否成功：

$ clang --version
clang version 11.1.0

安装 libstdc++

TFRT 需要 libstdc++ 8 或更高版本。请检查 clang 所选的 GCC 版本：

$ clang++ -v |& grep "Selected GCC"
Selected GCC installation: /usr/bin/../lib/gcc/x86_64-linux-gnu/10

在上述示例中，路径末尾的 10 表明 clang 将使用 libstdc++10,这与 TFRT 兼容。

如果需要升级，最简单的方法是安装 gcc-8。运行以下命令进行安装：

$ sudo add-apt-repository -y ppa:ubuntu-toolchain-r/test
$ sudo apt-get update
$ sudo apt-get install -y gcc-8 g++-8

再次运行 clang++ -v 检查以确认安装是否成功。

GPU 先决条件

注意： 如果您不打算构建 GPU 后端，可以跳过本节。不过，请务必在 Bazel 的目标模式中排除 //backends/gpu/...。

构建和运行 GPU 后端需要额外安装一些组件。

使用 pip 安装 clang 的 Python 绑定：

$ pip install libclang

从 NVIDIA 的 .run 安装包中，将 NVIDIA 的 CUDA Toolkit v11.2（详情参见 安装指南）安装到指定目录：

$ wget http://developer.download.nvidia.com/compute/cuda/11.2.2/local_installers/cuda_11.2.2_460.32.03_linux.run
$ sudo sh cuda_11.2.2_460.32.03_linux.run --toolkit --installpath=<path>

将 CUDA 共享库的路径添加到系统配置中：

$ sudo echo '<path>/lib64' > '/etc/ld.so.conf.d/cuda.conf'
$ sudo ldconfig

接着安装 NVIDIA 的 cuDNN 库（详情参见 安装指南）：

$ wget http://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1604/x86_64/libcudnn8_8.0.4.30-1+cuda11.1_amd64.deb
$ sudo apt install ./libcudnn8_8.0.4.30-1+cuda11.1_amd64.deb

注意： 上述软件包专为 CUDA 11.1 设计，但也可兼容 CUDA 11.2。TFRT 使用 cuDNN 8.1（适用于 CUDA 11.2）进行构建和验证。获取该软件包需注册一个（免费的）NVIDIA 开发者账号。

构建和运行 TFRT

要构建 TFRT，请使用 cd 命令进入 TFRT 工作区的根目录（即包含 WORKSPACE 文件的目录）。.bazelrc 文件中包含一组构建配置。您可以在仓库根目录下创建一个 user.bazelrc 文件，添加可能有用的额外 Bazel 配置。使用以下命令构建 tfrt_translate 和 bef_executor：

$ bazel build //tools:bef_executor
$ bazel build //tools:tfrt_translate

上述命令会以 opt 编译模式构建二进制文件。有关更多构建选项，请参阅 Bazel 的文档。构建成功后，Bazel 会在输出末尾提示二进制文件的存放位置（默认为 bazel-bin）。

在 tfrt_translate 和 bef_executor 构建完成后，可以使用以下命令运行一个 .mlir 程序：

$ bazel-bin/tools/tfrt_translate -mlir-to-bef path/to/program.mlir | bazel-bin/tools/bef_executor

TFRT 提供了一系列 .mlir 测试程序。例如：

$ bazel-bin/tools/tfrt_translate -mlir-to-bef mlir_tests/bef_executor/async.mlir | bazel-bin/tools/bef_executor

任何输出都会打印到终端。

添加 GPU 支持

在 Bazel 命令中添加 --config=cuda 参数，即可将 GPU 后端链接到上述目标。

可以通过 --repo_env=CUDA_PATH=<path> 指定自定义的 CUDA Toolkit 路径。默认路径为 /usr/local/cuda。

测试

TFRT 利用 LLVM 的 LIT 测试框架和 FileCheck 工具来构建基于 MLIR 的检查测试。这些测试用于验证测试输出中是否出现特定的字符串标记。有关测试的更多介绍和指南，请参阅 这里。下面是一个示例测试：

// RUN: tfrt_translate -mlir-to-bef %s | bef_executor | FileCheck %s
// RUN: tfrt_opt %s | tfrt_opt

// CHECK-LABEL: --- Running 'basic_tensor'
func @basic_tensor() {
  %c0 = tfrt.new.chain

  %a = dht.create_uninitialized_tensor.i32.2 [3 : i64, 2 : i64]
  %c1 = dht.fill_tensor_with_constant.i32 %a, %c0 0 : i32

  // CHECK: shape = [3, 2], values = [0, 0, 0, 0, 0, 0]
  %c2 = dht.print_tensor %a, %c1

  tfrt.return
}

要运行测试，只需调用 bazel test 命令：

$ bazel test //mlir_tests/bef_executor:basics.mlir.test

位于 //backends/gpu/... 下的大多数测试需要使用 --config=cuda 进行构建，以便将 GPU 后端链接到 bef_executor：

$ bazel test --config=cuda //backends/gpu/mlir_tests/core_runtime:get_device.mlir.test

可以使用 Bazel 的 目标模式 来运行多个测试：

$ bazel test -- //... -//third_party/... -//backends/gpu/...  # 所有 CPU 测试。
$ bazel test --config=cuda //backends/gpu/...                 # 所有 GPU 测试。

后续步骤

请尝试我们的 教程，以获得一些关于 TFRT 的实践体验。

有关 TFRT 设计的更多详细信息，请参阅 主机运行时设计。

仓库概览

TFRT 根目录下的三个关键目录是：

• lib/：包含 TFRT 核心基础设施代码
• backends/：包含设备特定的基础设施以及算子和内核实现
• include/：包含 TFRT 核心基础设施的公共头文件

顶级目录	子目录	描述
include/		TFRT 基础设施的公共头文件
lib/		TFRT 基础设施，适用于主机运行时和所有设备运行时
	basic_kernels/	通用基础设施内核，例如控制流内核
	bef_executor/	BEFFile 和 BEFExecutor 的实现
	bef_executor_driver/	用于运行输入 MLIR 文件的 BEFExecutor 的驱动代码
	bef_converter/	MLIR 和 BEF 之间的转换器（bef_to_mlir 和 mlir_to_bef）
	core_runtime/	TFRT 核心运行时基础设施
	distributed_runtime/	TFRT 分布式运行时基础设施
	data/	TFRT 用于 TensorFlow 输入管道的基础设施
	host_context/	主机端 TFRT 数据结构，例如 HostContext、AsyncValue、ConcurrentWorkQueue
	metrics/	ML 指标集成
	support/	基础工具，例如 hash_util、string_util
	tensor/	基础 Tensor 类及主机端 Tensor 实现
	test_kernels/	测试用内核实现
	tracing/	跟踪/性能分析支持
cpp_tests/		C++ 单元测试
mlir_tests/		基于 MLIR 的单元测试
utils/		各种实用工具，例如用于生成测试 ML 模型的脚本。
tools/		包括 bef_executor、tfrt_translate 等在内的二进制文件。
backends/common/		不同后端共享的库，例如 eigen、dnn_op_utils.h
	ops/	跨设备共享的算子实现库，例如元数据函数
	compat/eigen/	eigen 的适配库，被多个后端使用
	utils/	各种实用工具，例如用于生成 MLIR 测试代码的脚本。
backends/cpu/		CPU 设备基础设施以及 CPU 算子和内核
	include/	CPU 相关的公共头文件
	lib/core_runtime/	CPU 核心运行时基础设施，例如 cpu_device
	lib/ops	CPU 算子
	lib/kernels	CPU 内核
	cpp_tests/	CPU 基础设施的单元测试
	mlir_tests/	CPU 基于 MLIR 的测试
backends/gpu/		GPU 基础设施以及算子和内核的实现。在与 TFRT 其他基础设施的接口稳定之后，我们可能会将该目录拆分为一个单独的仓库。
	include/	GPU 相关的公共头文件
	lib/core_runtime/	GPU 核心运行时基础设施
	lib/memory	GPU 内存抽象
	lib/stream	GPU 流抽象及其封装
	lib/tensor	GPU 张量
	lib/ops	GPU 算子
	lib/kernels	GPU 内核
	lib/data	GPU 内核用于输入管道基础设施
	cpp_tests/	GPU 基础设施的单元测试
	mlir_tests/	GPU 基于 MLIR 的测试
	tools/	各种实用工具

贡献指南

如果您希望为 TFRT 做出贡献，请务必查看 贡献指南。本项目遵循 TensorFlow 的 行为准则。 参与其中即表示您同意遵守此行为准则。

注意： 目前 TFRT 尚未开放外部贡献。TFRT 开发人员正在构建工作流程和持续集成系统，以支持外部贡献的接收。一旦准备就绪，我们将更新此页面。

持续集成状态

联系方式

订阅 TFRT 邮件列表，以参与运行时相关的通用讨论。

我们使用 GitHub 问题来跟踪 bug 和功能请求。

许可证

Apache 2.0 许可证

TFRT (TensorFlow Runtime) 快速上手指南

TFRT 是 TensorFlow 的新一代运行时，旨在提供统一、可扩展的基础设施层，专注于底层效率和多线程性能。本指南将帮助您快速搭建开发环境并运行基础示例。

注意：TFRT 目前处于早期开发阶段，尚未准备好用于生产环境。

1. 环境准备

系统要求

• 操作系统: Ubuntu 16.04 (官方主要支持版本)。MacOS 和 Windows 支持尚在规划中。
• 构建工具: Bazel 4.0, Clang 11.1。
• GPU 支持 (可选): 若需构建 GPU 后端，需安装 NVIDIA CUDA Toolkit v11.2 和 cuDNN 8.1+。

前置依赖安装

安装 Bazel

请参照 Bazel 官方安装指南 安装 Bazel 4.0。 验证安装：

$ bazel --version
bazel 4.0.0

安装 Clang

推荐安装包含 clang, lldb, lld 的自动脚本包。TFRT 基于 Clang 11.1 构建。 参照 LLVM APT 仓库 进行安装。若系统中存在多个 Clang 版本，请使用 update-alternatives 设置默认版本：

$ sudo update-alternatives --install /usr/bin/clang clang /usr/bin/clang-11 11
$ sudo update-alternatives --install /usr/bin/clang++ clang++ /usr/bin/clang++-11 11

验证安装：

$ clang --version
clang version 11.1.0

配置 libstdc++

TFRT 需要 libstdc++8 或更高版本。检查当前 Clang 选择的 GCC 版本：

$ clang++ -v |& grep "Selected GCC"

如果版本过低，请安装 gcc-8：

$ sudo add-apt-repository -y ppa:ubuntu-toolchain-r/test
$ sudo apt-get update
$ sudo apt-get install -y gcc-8 g++-8

GPU 环境 (可选)

如果不构建 GPU 后端，可跳过此步骤（并在后续构建命令中排除 //backends/gpu/...）。

1. 安装 Clang Python 绑定：

   $ pip install libclang
   

2. 安装 CUDA Toolkit v11.2 (以 .run 文件为例)：

   $ wget http://developer.download.nvidia.com/compute/cuda/11.2.2/local_installers/cuda_11.2.2_460.32.03_linux.run
   $ sudo sh cuda_11.2.2_460.32.03_linux.run --toolkit --installpath=<path>
   

   注册共享库路径：

   $ sudo echo '<path>/lib64' > '/etc/ld.so.conf.d/cuda.conf'
   $ sudo ldconfig
   

3. 安装 cuDNN (需 NVIDIA 开发者账号下载对应 deb 包)：

   $ wget http://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1604/x86_64/libcudnn8_8.0.4.30-1+cuda11.1_amd64.deb
   $ sudo apt install ./libcudnn8_8.0.4.30-1+cuda11.1_amd64.deb
   

2. 安装与构建

克隆仓库后，进入根目录（包含 WORKSPACE 文件的目录）。

构建核心工具

构建图执行所需的两个核心二进制文件：tfrt_translate (MLIR 转 BEF) 和 bef_executor (执行器)。

$ bazel build //tools:bef_executor
$ bazel build //tools:tfrt_translate

注：默认使用 opt 优化模式编译。构建成功后，二进制文件位于 bazel-bin/tools/ 目录下。

启用 GPU 支持 (可选)

若已安装 CUDA 环境，需在构建命令中添加 --config=cuda：

$ bazel build --config=cuda //tools:bef_executor
$ bazel build --config=cuda //tools:tfrt_translate

若 CUDA 安装在非默认路径，可通过 --repo_env=CUDA_PATH=<path> 指定。

3. 基本使用

TFRT 的工作流是将 MLIR 程序编译为二进制可执行格式 (BEF)，然后由执行器运行。

运行示例

使用以下命令管道将 .mlir 文件转换为 BEF 并立即执行：

$ bazel-bin/tools/tfrt_translate -mlir-to-bef path/to/program.mlir | bazel-bin/tools/bef_executor

测试运行

项目自带了一系列 .mlir 测试程序，例如异步执行测试：

$ bazel-bin/tools/tfrt_translate -mlir-to-bef mlir_tests/bef_executor/async.mlir | bazel-bin/tools/bef_executor

任何输出结果将直接打印到终端。

运行测试套件

TFRT 使用 LLVM 的 LIT 和 FileCheck 进行测试。

运行单个测试：

$ bazel test //mlir_tests/bef_executor:basics.mlir.test

运行所有 CPU 测试（排除第三方和 GPU 测试）：

$ bazel test -- //... -//third_party/... -//backends/gpu/...

运行所有 GPU 测试（需配置 cuda）：

$ bazel test --config=cuda //backends/gpu/...