浅析 TensorFlow Runtime 技术

关于 TF Runtime 的疑问？

什么是TFRT ?

TensorFlow Runtime，简称 TFRT，它提供了统一的、可扩展的基础架构层，可以极致地发挥CPU多线程性能，支持全异步编程（无锁队列+异步化语义）。TFRT 可以减少开发、验证和部署企业级模型所需的时间。

TFRT 的输入是什么？

输入为Tensorflow GraphDef，TFRT 会调用基于MLIR的图编译器，执行图优化，并将其lower成 BEF —— 用于执行TFRT graph的二进制可执行格式。

• 在TF原生框架中，执行的流程是：Python Layers → GradDef (DAG) → 执行OpNode (ThreadPool并行)

• Runtime 的思路：Python Layers → GradDef (DAG) → Compile IR → Binary (BEF) → execute (BEFExecutor)

基础概念：

• Host Program in MLIR是graph的低阶中间表示
• BEF是一个BEFExecutor的可执行文件，读取BEF文件，然后异步执行里面的函数
• 两者通过tfrt_translate来转换，类似汇编器 Assembler

这里的 IR 是什么？

其实可以理解为是一套表示拓扑关系的代码，甚至是一个graph。通过拓扑递推，可以很容易转为一段IR代码。这也是为什么BEF支持IR与Graph的互转的原因。比如：

%1 = hex.constant.i32 1
%2 = hex.constant.i32 2
%3 = hex.add.i32 %1, %2
hex.print.i32 %3
# 实际可以表示为一个DAG图

和 XLA 的区别？

XLA 本质上并没有脱离图执行的框架，它只是通过 graph cluster 把部分子图通过 HLO 的转换走 JIT 执行，将子图包裹在一个XlaRunOp里，再与图的其他节点一起执行。所以只是把几个节点换成了一个更快的大节点。（看起来有点类似fuse）

官方文档里称BEF为 Kernel graph的实际载体，实际还是一个graph，即表示bef executor最终执行的实体依然是一个 graph（但不是TF原生意义的GraphDef）。

TFRT 基本执行单元是什么？执行的流程？

TFRT里的 kernel 概念，分为如下两种：

• 同步 Kernel

  • 完全在调用它的线程中执行，不会涉及到其他线程里的计算。它产生的AsyncValue状态都是available的

    int32_t TFRTAddI32(Argument<int32_t> arg0, Argument<int32_t> arg1) {
      // The thread that calls TFRTAddI32 performs this addition, and produces
      // an available AsyncValue.
      return *arg0 + *arg1;
    }
    

• 异步 Kernel

  • 包含两个部分的计算：①调用它所在线程的同步计算 ② 其他线程中的异步计算。它产生的AsyncValue状态是unavailable的（并不全是）

    void TFRTAddI32Async(Argument<int32_t> arg0, Argument<int32_t> arg1,
                        Result<int32_t> output, HostContext* host) {
      // Synchronously allocate an unavailable AsyncValue for ‘output’.
      auto result = output.Allocate();
    
      // Asynchronously make ‘output’ available.
      host->EnqueueWork([arg0 = *arg0, arg1 = *arg1,
                         result_ref = FormRef(result)] {
        // A ConcurrentWorkQueue thread performs this addition.
        result_ref->emplace(arg0 + arg1);
      });
    
      // Synchronously returns unavailable ‘output’.
    }
    

执行流程：

• 创建一个AsyncKernelFrame，包含输入参数和输入result
• 将Frame传递给kernel执行
• 所有的AsyncValue通过registers来跟踪

也提供了eager API （op-by-op）：CoreRuntime 和 CoreRuntimeOp

• CoreRuntime：

  • 执行OpHandler，借助内部类Impl来实现
  • 它可以调用MakeOp(op_name, op_handler)来创建一个CoreRuntimeOp直接运行

• CoreRuntimeOp

  • 持有一个llvm::unique_function<void<const OpInvocation&>>类型的函数指针fn_
  • 仿函数用于执行函数fn_

如何整合硬件设备的？

借助 DeviceRuntime，让BEF只支持最底层的driver API的Op，从而尽量避免让每一种后端都单独实现一遍tf的各个Op。

如下图中使用的op直接对应到了cuda api：

Host Runtime的设计思路

Host Runtime 的位置?

host 指执行计算的机器设备，可能有，也可能没有硬件加速的资源。host 可以只是一个具有多GPU的服务器，或带有DSP和IPU的移动设备。

在TF原生的框架中，TF Core是按照 data-flow 进行op-by-op的执行，设计上有很多顺序同步执行的影子在里面。而 Host Runtime 通过重新编排计算逻辑，然后驱动 Device Runtime（如GPU、TPU）去加速计算，使得kernel的执行可以单独放在一个线程中，去异步执行，充分利用的多线程并行的优势。

为什么要做这件事？

• 期望能高效的eagerly执行op
  • TF对graph执行已经优化的很好了，毕竟都在C++端执行。但在earge模式下，python和runtime端之间的不必要的开销还是在存的。
• 统一图和op两个不同层次下多线程并行机制
• runtime 中异步是一等公民
  • a non-strict kernel/function may execute before all its inputs are ready.
• 更轻便地进行cross-kernel优化
  • TF 的op Kernel实现中封装了 Tensor 的内存申请之类的逻辑，这限制了cross-kernel中reuse buffe的优化。在 TFRT的kernel中，解耦了 shape计算和 tensor 内存申请的逻辑
• 实现模块化、可插拔式的新硬件支持机制
  • 期望解决之前为了接入新硬件而不得不hack整个代码库的痛点；能够建立一种模块化机制，直接提供完善的接入文档给硬件团队即可，变被动为主动。

如何去设计来实现上述目标么？

先回顾下背景： Core Runtime, Graph Lowering 和 Eager Execution

1. Core Runtime

   用来 eagerly 执行单个 op 或者整个graph function——包含GradDef 和 HLO。一个op graph通常是设备独立的。

2. Graph Lowering

Compiler passes 将一个op graph 转化为一个Kernel Graph，它是一个数据流计算的更低阶表示，为更快执行而设计，因此不适合做编译分析，但可以通过低阶方言（如MLIR）来表示。Kernel graph是面向指定设备的（与平台绑定）

3. Eager Execution

   Host Runtime支持eagerly 执行。但并不一定会涉及Graph/BEF的构造和BEFExecutor的使用。TF设计了两个方案：

   • Generic path：把 op 当做graph function来处理，可以很好处理组合 op 的情况，也可以复用graph function的那一整套代码。
   • Fast path：使用手写的C++或者预编译的 graph snippets 去完成op kernel的选取和调用（定制化优化？成本不高么？）

Kernel Graph 中的 Kernel 指什么？

TFRT里面也有 kernel 的概念，输入输出均为：AsyncValue——异步是一等公民的践行者。类似C++标准库中的 future 和 promis的组合。 graph中的所有data全部都会替换为AsyncValue。

执行流程：

• 创建一个AsyncKernelFrame，包含输入参数和输入result
• 将Frame传递给kernel执行
• 所有的AsyncValue通过registers来跟踪

// Kernel that adds two integers.
// AsyncKernelFrame holds the kernel’s arguments and results.
static void TFRTAdd(AsyncKernelFrame* frame) {
  // Fetch the kernel’s 0th argument.
  AsyncValue* arg1 = frame->GetArgAt(0);
  // Fetch the kernel’s 1st argument.
  AsyncValue* arg2 = frame->GetArgAt(1);

  int v1 = arg1->get<int>();
  int v2 = arg2->get<int>();

  // Set the kernel’s 0th result.
  frame->EmplaceResultAt<int>(0, v1 + v2);
}

TODO: Kernel中的内存申请接入机制

Kernel 类型分为如下两种：

• 同步 Kernel

  • 完全在调用它的线程中执行，不会涉及任何其他线程的计算。它产生的AsyncValue状态都是available的

    int32_t TFRTAddI32(Argument<int32_t> arg0, Argument<int32_t> arg1) {
      // The thread that calls TFRTAddI32 performs this addition, and produces
      // an available AsyncValue.
      return *arg0 + *arg1;
    }
    

• 异步 Kernel

  • 包含两个部分：①调用它所在线程的同步操作 ② 其他线程中的异步操作。它产生的``AsyncValue`状态是unavailable的（并不全是）

    void TFRTAddI32Async(Argument<int32_t> arg0, Argument<int32_t> arg1,
                        Result<int32_t> output, HostContext* host) {
      // Synchronously allocate an unavailable AsyncValue for ‘output’.
      auto result = output.Allocate();
    
      // Asynchronously make ‘output’ available.
      host->EnqueueWork([arg0 = *arg0, arg1 = *arg1,
                         result_ref = FormRef(result)] {
        // A ConcurrentWorkQueue thread performs this addition.
        result_ref->emplace(arg0 + arg1);
      });
    
      // Synchronously returns unavailable ‘output’.
    }
    

Kernel 的两种执行模式：

• Strict mode:

  • 此类Kernel被调用时，所有的AsyncValue均已是available。

• non Strict mode:

  • 只要有一个输入参数是available，就执行。比如三元操作，它其实只负责转发

  result = ternary(condition, true_result, false_result) //只要condition可用即可
  

  • 这类kernel实现难度较高

AsyncValue有什么用途？

前面提到：Kernel 的输入输出均为：AsyncValue，graph中的所有data也全部替换为了AsyncValue。

// A subset of interface functions in AsyncValue.
class AsyncValue {
 public:
  // Is the data available?
  bool IsAvailable() const;

  // Get the payload data as type T.
  // Assumes the data is already available, so get() never blocks.
  template <typename T> const T& get() const;

  // Store the payload data in-place.
  template <typename T, typename... Args>
  void emplace(Args&&... args);

  // Add a waiter callback that will run when the value becomes available.
  void AndThen(std::function<void()>&& waiter);
  // ...
};

AyncValuea有三个派生类：

• ConcreteAsyncValue<T>：用于表示和存放具体data
• ErrorAysncValue：用于处理异常传播和取消执行。BEFExecutor会监控每个Kernel执行返回的值，若果某个result值为此类型，则跳过所有依赖此值的下游op
• IndirectAsyncValue：有些情况下，某个result的dataType还不知道呢，但为了实现非阻塞机制，先创建一个IndirectSyncValue，保证non-strick Kernel的执行。它其实并不持有数据，而是持有了一个指向另一个AsyncValue的指针。

生命周期：通过引用计数实现：

• kernel会首先对results创建AyncValue（当dataType确定时）
• 一个AsyncValue的所有权会从kernel移交给BEFExecutor
• BEFExecutor将AsyncValue传递给所有使用它的下游 Op，并递增引用计数
• 每个下游Op Kernel完成计算后，递减此AsyncValue的引用计数

管理AyncValue的Register具体做哪些工作？

Register其实是一个指向AyncValue的指针，它也只操作指针，因此不涉及数据的移动和copy。

举个栗子：

available_value = upstream()
downstream(available_value, unavailable_value)

downstream需要等到两个参数都ready才会执行。当unavailable_value也available时，执行器从register加载数据，然后传递给downstream去执行

register有三种状态：

• Empty：初始状态，不指向任何AsyncValue
• Unavailable: 只用于异步kernel。同步kernel不会产生此状态。
• Available: 最终状态，且状态不可逆。

RunTime 如何实现异步加速的？

在 TFRT 中，执行Kernel的线程，与调度其他已ready的kernel的线程，可能属于同一个。TFRT 把后台调度kernel任务放到了一个ConcurrentWorkQueue中来异步执行。

但反向需要梯度才能执行，如何处理反向op以及IO阻塞问题呢？

TF采用了两个独立的线程池：

①专用线程池：存放长时非阻塞任务

• 固定线程数，每个硬件一个线程，避免线程资源抢占带来的开销。

②单独线程池：存放阻塞任务（如IO）

• 申请多一些线程数来处理IO任务
• 为了避免死锁，阻塞任务只能放在阻塞线程池里执行
• 要求Kernel的实现不能直接包含阻塞操作（例如？），更不能将部分阻塞操作放到非阻塞队列里。

图执行——Graph Executation

图执行时，host program 会把 graph 转换为MLIR表示的 Kernel graph。此处会应用一些compiler passes 将设备无关的 graph 转化为面向特定硬件平台的 kernel graph。

func @sample_function() -> i32 {
  %one = tfrt.constant.i32 1       // Make AsyncValue with value 1
  %two = tfrt.constant.i32 2       // Make AsyncValue with value 2
  %three = tfrt.add.i32 %one, %two // Make AsyncValue with value 3 (1+2)

  tfrt.print.i32 %three            // Print AsyncValue %three
  tfrt.return %three : i32         // Return AsyncValue %three
}

runtime 并不直接执行IR，而是通过mlir_to_bef将其转换为 BEF后再执行。通过 registers 跟踪和记录所有 AsyncValue 的状态。

如何解决control dependency问题？

在原生的TF中是通过tf.control_dependencies来对两个有顺序要求的Kernel添加依赖。在TFRT中，是通过Chain来实现。一个chain也是一个AsyncValue——可以是kernel的参数，也可以是result，这样的话，Chain要求consumer必须在producer之后，以此实现有序性。

func @control_dep1() {
  %a = dht.create_uninit_tensor.i32.2 [2 : i32, 2 : i32]
  %chain1 = dht.fill_tensor.i32 %a, 41
  %chain2 = dht.print_tensor.i32 %a, %chain1
 }

如何处理控制流的情况，如if ?

TFRT支持在Kernel中调用BEFExecutor（这一点跟Paddle目前的控制流处理思路有点类似）

void TFRTIf(AsyncKernelFrame* frame) {
  const auto* true_fn = &frame->GetConstantAt<Function>(0);
  const auto* false_fn = &frame->GetConstantAt<Function>(1);

  // First arg is the condition.
  ArrayRef<AsyncValue*> args = frame->GetArguments();
  AsyncValue* condition = args[0];

  // Execute true_fn or false_fn depending on ‘condition’.
  auto* fn = condition->get<bool>() ? true_fn : false_fn;
  fn->Execute(args.drop_front(),
              frame->GetResults(),
              frame->GetHostContext());
}

与底层的session的区别和联系？

貌似没啥关系。（待深入了解）

BEF文件里都包含了什么信息？

BEF 是runtime和compiler的桥梁，同时将compiler从runtime中解耦，从而可以独立应用编译优化策略。它支持保存到磁盘，重新加载执行（mmap bytes）。感觉和二进制文件很类似，因为它也包括很多section的概念。

BEF 包含了一些与硬件设备相关的信息：每个Kernel在哪种设备（CPU/GPU/TPU）上执行，以及哪些特殊的Kernel会被调用。

MLIR和BEF之间可以互相转换：

BEFExecutor的作用是什么？有特殊性能收益吗？

它是一个执行器，而非一个解释器，因为它没有program counterd的概念。

性能收益来源：

• 它是 lock-free 的
• 非阻塞执行：
  • 无论一个Value是否available，它都会执行下去。对于unvailable的value，执行器会将其推迟到AsyncValue::AndThen
  • 由于AyncValue都会由Register来跟踪，它一旦ready，会通知和唤起所有相关kernel

遗留问题

TFRT中公布的文档中很少涉及训练和反向op的内容，是否支持？

在官网给出的 mnist_training.md介绍中，提到了TFRT对训练的支持，但只是原型展示，并非最终版本。

• 单独重写了MNIST模型中所有的op，如matmul、relu、elem_add、argmax、reduce_mean
• 这里只重写relu_grad的kernel，其他op的反向kernel默认使用的是Tensorflow框架的？

参考资料

1. 【官方文档】—TFRT Host Runtime Design