浅析 TensorFlow Runtime 技术

时间 2020-12-26

标签 python git github 编程后端 api promise 服务器 session 多线程栏目 Python 繁體版

原文原文链接

关于 TF Runtime 的疑问？

什么是TFRT ?

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

TFRT 的输入是什么？

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

在TF原生框架中，执行的流程是：Python Layers → GradDef (DAG) → 执行OpNode (ThreadPool并行)github
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的互转的缘由。好比：api

%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）promise

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

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

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

同步 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

Core Runtime

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

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

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框架的？

参考资料

【官方文档】—TFRT Host Runtime Design