淺析 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