浅析 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是面向指定设备的(与平台绑定)

  1. 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++标准库中的 futurepromis 的组合。 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的引用计数

管理 AyncValueRegister 具体做哪些工作?

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 counter d的概念。

性能收益来源:

  • 它是 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
博客园精华区
我还没有学会写个人说明!
下一篇

《心灵奇旅》拿下豆瓣9.2分成年度开局最高片

你也可能喜欢

评论已经被关闭。

插入图片