天天看點

大規模深度學習預測場景下 codegen 的思考與應用背景簡介PULL vs. PUSH,通用 vs. 優化解決技術方案一些讨論總結與展望

更多關于 RTP 系統的介紹請見 深度預測平台RTP介紹

背景簡介

RTP 系統

RTP 系統(即 Rank Service),是一個面向搜尋和推薦的 ranking 需求,支援多種模型的線上 inference 服務。RTP 支援 LR、GBDT 以及 tensorflow 等多種模型及模型格式,并依托 suez 線上服務架構,将樣本組裝和模型預測一氣呵成,提升了業務疊代效率,并在性能和穩定性方面予業務以充分保障。

inference 計算

RTP 上的業務大多數是打分類的場景,它的計算流程由一個 Compute Graph 描述(TF Graph)。從計算邏輯上來看,它可以分為這樣三部分:

  • 存儲讀取。根據商品 ID,利用 suez 存儲能力,join 讀取 Item 字段内容;擷取 Query 和 User 資訊。
  • 特征生成。根據上一步的原始字段值,根據使用者的特征配置,加工生成特征(如 match,combo 等),進行離散化以及 embedding 等。以經典的 MLP 為例,我們會最終生成一個大的 feature 矩陣。
  • 模型計算。即大規模浮點運算部分,執行矩陣乘和激活等 Graph 節點,獲得預測結果。
大規模深度學習預測場景下 codegen 的思考與應用背景簡介PULL vs. PUSH,通用 vs. 優化解決技術方案一些讨論總結與展望

對于模型計算部分來說,我們利用圖優化和 GPU/FPGA 異構計算等手段,在 latency 和 utility 等方面都拿到了不錯的效果。本文将重點讨論存儲讀取和特征生成部分的問題和挑戰。

PULL vs. PUSH,通用 vs. 優化

背景

RTP 引擎源于 HA3 搜尋引擎,熟悉 HA3、寫過插件的人都會多多少少聽說過 matchdoc 和 expression。

搜尋引擎在進行倒排 match 時,命中的文檔一般 相距 都比較遠。matchdoc 結構将命中的文檔 ID 以及起字段值緩存到一塊連續記憶體中,使後面的 ranking 階段更加 cache 友好。這便是 matchdoc 的價值。

expression 則是常見的計算表達式 lib,提供了基本的一進制二進制表達式運算功能,也支援使用者自定義 function。

pull,通用

在 2016 年,RTP 是這樣計算 feature 的:

大規模深度學習預測場景下 codegen 的思考與應用背景簡介PULL vs. PUSH,通用 vs. 優化解決技術方案一些讨論總結與展望

這樣的實作具有很好的通用性和靈活性,然而在性能方面 expression 的封裝卻不盡人意。過度的虛函數開銷,以及訪存不連續帶來的 cache miss,通用性帶來的 overhead 難以令人接受。

push,優化

在 2017 年,RTP 是這樣計算 feature 的:

大規模深度學習預測場景下 codegen 的思考與應用背景簡介PULL vs. PUSH,通用 vs. 優化解決技術方案一些讨論總結與展望

我們去掉了 expression,而隻用 matchdoc 結構存儲每一步的中間結果,并且調整了 FOR 的結構,這些帶來了 20-30% 的性能收益。

然而,這些改動也耗費了巨大的人力,并且喪失了表達式的一般通用性,feature 處理成為了一個定制的 lib。這樣的計算方式,每一步都要緩存所有 feature 結果,這可能并不是必要的。同理而言,對于 matchdoc 中有些隻需要使用一次的字段,也并不需要提前緩存。而列結構的計算和存儲模型,是否就一定會比行結構好呢?

問題

push 或 pull 是衆多優化中沖突的一個縮影,而最具有問題的代表性。總結以上沖突,有以下幾點:

  • 優化喪失了通用性,添加新功能會比之前困難。
  • 對于 類型 和 數值 方面的枚舉優化,破壞了整潔的代碼結構。
  • 運算邏輯不變,但改變執行方式需要大量的代碼改動,令人望而卻步。

解決

codegen

expression 不好,在于性能,但其他方面都很優秀,通用性夠強,其本質是 按解釋執行。

PUSH model

不好,在于定制化沒有通用性,而性能上占優勢。更抽象的看,各自分别強在 描述 和 執行 上。

架構運算的代碼通用化,在于一套系統、binary 可以在适用于所有場景,做到利用架構定義的内容描述自己的算法,而 運作時 我們的系統可以根據描述進行運算。對于 RTP 的某一個算法模型(版本)來說,輸入描述、特征運算規則及網絡結構,都是在配置裡完全已知的,不會随着 Query 而變。對于這些已知的資訊,我們實際上可以對各個場景各自進行 hard code,将運作時解釋的東西提前到 編譯時 解決。将運算描述轉換成運算代碼,這個過程便是 codegen。

codegen(代碼生成) 并不是一項具體技術,它是一種方法和思想。C++、Java 等語言中的宏和模闆,protobuf 根據 proto 生成具體語言的執行代碼等等,這些都是 codegen 的例子。它們有些是語言内建的機制,有些是 library 的功能。簡單的說,codegen 是語言到語言、代碼到代碼的轉換,實作上則是一種編譯器。

對于從存儲到特征這個場景來說,我們需要将存儲的 DSL 和特征 DSL 聯合起來,編譯成高效的運作代碼。利用 codegen,我們可以在編譯時保留通用的 expression 描述,不失拓展性和靈活性;而運作時運作的實際代碼,也是高效的。

algo/sched 分離思想

解決了一個問題,但還沒有解決所有問題。有了編譯階段,但具體怎樣生成代碼呢?是按照 pull 方式還是 push 方式呢?何時需要 cache?一般來說,編譯期 expression(AST)描述了完整的運算過程,但一套規則并不能完美适配所有場景,取得最優性能。

運算描述(algorithm)和執行政策(schedule)分離的思想,在

Halide

一文中已經有了詳盡的闡釋。從這個角度看,matchdoc 隻是一種具象的執行政策,而利用 codegen,我們可以任意地在編譯器排程代碼,尋找訪存的最優解。而 expression 單純負責描述計算,具體的政策下放給編譯引擎,或是優化人員進行人工調優。

codegen 和 algo/sched 分離兩者是相輔相成的。沒有 codegen,便完全沒有了 algo/sched 分離的能力;而 algo/sched 分離是 codegen 所獲得最重要的能力,它的價值甚至比運作時提升到編譯時更為重要。它可以讓我們靈活應對任何一種存儲和計算形式,以預設政策為開始,自動調節和人工調優雙管齊下。

codegen 要求原本的計算原語需要被拆分開,擁有自由組合的能力。algo/sched 分離則有着更高的要求,它要求運算粒度要被拆分的足夠細,拆分的越充分,可以排程的空間和可能性便越高。這也正契合全圖化和算子化的方向,功能拆分越細,則靈活度越高,業務越自由,變更成本越小。而 codegen 和 algo/sched 分離的價值,不但會平衡抽象和拆分帶來的性能損失,而且會在更廣泛的空間上探索性能最優解。

相關工作

在深度學習方面,

Tensorflow XLA

算是 codegen 鼻祖,後來的

tvm

TensorComprehension

則在方法上更進一步。tvm 注重于多架構和多語言,提供了一些 highlevel 的 schedule api;TensorComprehension 則更注重于利用 polyhedral 方法搜尋和優化運算本身。兩者都某種程度上同源于 Halide。這些架構重點關注于網絡運算,在樣本生成和特征生成方面并沒有涉足。

在存儲和運算方面,

Impala

本身具有一定的 codegen 優化的能力,

PostgreSQL

的 codegen 則剛剛起步。它們一定程度上解決了按解釋執行的問題,也具有一定人工排程的能力。

PlinyCompute

是更一般的計算系統,平衡了抽象和性能,但使用的基本方法依然是基本的表達式計算與行列混合式存儲。類似的,在

dremio

這樣的 OLAP 場景下,

Gandiva

利用抽象表達式描述計算和訪存 (

Apache Arrow

),通過 codegen 進一步描述轉化成運作時高效代碼。但 Gandiva 描述力有限,而我們的存儲也比 Arrow 更複雜。

Julia

Ravi

等利用 LLVM-Jit 技術試圖解決解釋性語言的性能問題,但本身不具備 codegen 以及自由排程的能力。

GraalVM

打破了語言互操作的界限,将問題統一到 Java bitcode 和 JVM 中,而

Truffle

架構在此基礎上提供了良好的語言建構方案。但在我們場景下的代碼都是非常 native 的,記憶體也是由開發者主動管理。已經優化的非常精細的代碼再交給 JVM 來 handle,帶來的任何 overhead 是無法接受的。

技術方案

大規模深度學習預測場景下 codegen 的思考與應用背景簡介PULL vs. PUSH,通用 vs. 優化解決技術方案一些讨論總結與展望

IR 表示及 C++ 混合程式設計

将原本的讀取存儲和生成特征的過程拆碎和抽象并不難,重要的是将他們用可排程的 expression 表示出來。LLVM IR 是方法之一,還有很多 DSL 可以應用。對于 RTP 應用,我們時常需要操縱複雜的 C++ 結構,不能完全使用某種現有 IR,是以我們需要預先定義一些基本的 C++ 函數和資料結構,利用 IR 的函數機制來進行運算,IR 也需要具備這種能力。

相當多的 codegen 機制是通過自定義資料結構配合 LLVM IR 實作的。不過,LLVM IR 并不擅長于處理 C++ 的類型,也缺乏 C++ 語言的機制(模闆,重載等)。對于 C++ 這樣的機制複雜的語言,隻有原生編譯器前端自己做的最好。

HalideIR 源于 Halide 項目,被用于 tvm 做中間表示。使用 HalideIR 是期望能夠借助統一的 IR 表示,日後使用 tvm 内置的 schedule 能力。HalideIR 提供了基本的類型能力,但還不夠。我們依托 C++ 類型在 HalideIR 上了一套簡單的函數注冊和重載機制,照顧到了編譯器的類型檢查,也提供了基本的類型運算能力。

對于 IR,或者語言的選擇上來說,求同大于存異,複用大于創造。可以充分利用現有生态和庫是第一要務,而語言本身的特性則是越簡單越好,越貼近底層越好,因為我們本身在做 codegen,隻要機制夠用即可,甚至不必特别關心目标編譯器生成代碼品質,畢竟做 codegen 就把一切握在了手中。

對于無法用 HalideIR 的邏輯,例如構造 C++ STL map,以及 map 查找等,我們定義了一些内置函數。而未來的 UDF 功能,也将依托于函數功能。進一步結合類型運算,我們定義了多種

Cast

方法,從簡單的 int 到 string,到複雜的格式 string 到 map,富類型運算也是 codegen 的一個副産品,UDF 接口可以是以變得更自然舒适,而優化也會是以更優雅。

schedule

曾經的 matchdoc 是一個 1000 loc 的 lib,而現在,它隻需要幾行代碼

auto tmpidx = Variable::make(type_of<int32_t>(), "tmpidx");
    // item 的下标
    auto tmpdocid = Load::make(type_of<int32_t>(), docids, tmpidx, const_true());
    // 存儲中的實際 id
    auto cacheValue = OverloadResolver::call("readField",
            {readerVar, tmpdocid});
    // 讀取字段值
    auto attrCache = Variable::make(cacheValue.type(), "buffer");
    auto alloc = Allocate::make(attrCache, cacheValue.type(), {},
                                const_true(), Evaluate::make(0),
                                OverloadResolver::call("alloca",
                                        {Sub::make(endIdx, beginIdx) * cacheValue.type().bytes()}), "nop");
    // 配置設定 cache 部分記憶體空間 buffer
    auto storeCache = Store::make(attrCache, cacheValue, Sub::make(tmpidx, beginIdx), const_true());
    // 将字段值 cache 寫入 buffer
    auto storeAll = For::make(tmpidx, beginIdx, endIdx,
                              ForType::Serial, DeviceAPI::None, storeCache);
    // 對 beginIdx -> endIdx 之間的 item 都執行 storeCache 操作,既緩存字段值           

借助 IR 表示,我們可以自由的選擇需要把哪些運算中間結果 cache,或者隻是一遍過。

而對于以前 FOR 形式的 copy,因為有了 IR 層,運作時的數字變成編譯時的常量,vectorize/SIMD 等手段便可任意使用。

在搜尋和推薦、廣告場景中,Query/User 部分是 1-batch 而 Item 次元計算是 n-batch 的,這個規律在 codegen 下也可以充分利用。

目前我們還沒有利用 tvm 的内置的 schedule 邏輯,規整矩陣上的 bound-inference 和我們 C++ 的資料結構上單值到多值的擴充還有沖突需要解決。schedule 還有很多可以探索的地方,它如編譯器的 transform pass,一方面有一些預定義的規則,而更廣闊的是根據場景去搜尋自己的最佳解決方案。思想上類似 Clang LTO,但使用上更自由和靈活。

codegen & execute

有了 IR,進一步生成可執行代碼已經不難了。當然我們不會直接生成彙編,還是要生成一種 中間語言,交給編譯器,讓它來最終生成可執行代碼。方案上有兩種選擇,一是生成 C/C++ 代碼,然後交出去完整編譯,二是生成 LLVM IR。LLVM IR 對于類型較少或比較固定,或者有統一抽象類型的語言來說,無疑是首選。但我們的系統需要混合 C++ 程式設計,直接和 LLVM IR 互動,需要處理複雜的 mangling 規則以及 type inference,帶來過多複雜性。而直接生成 C++ 代碼,則可以完全利用 C++ 本身的類型機制,适配重載和模闆,唯一缺點就是編譯時間會相對慢一些,不過相對舒服的多。

而生成 executable 也有兩種方案,一是送出給 LLVM-Jit,二是直接生成 so。送出給 LLVM-Jit 的弱勢主要在調試上,正常調試可以通過實作 GDB listener,但對 coredump 則無能為力;而使用 perf 也需要實作對應 listener,比較麻煩。但統一到 LLVM-Jit 上的好處就是未來會有更多的可能的優化手段,甚至做到運作時自動 inline 的境界,此時更多的是需要 VM 層協助。兩種方案我們都有實作,而選擇了編譯成 so 方案為預設方式。這樣既有代碼,也有 so,實作上類似于插件機制,友善調試和調查問題。但會多餘在磁盤上生成檔案,需要統一管理回收。

執行個體

來看看最終的結果吧

for (int32_t doc_idx = idxBegin; doc_idx < idxEnd; ++doc_idx) {
    int32_t docid = (( int32_t*)docids)[doc_idx];
    int32_t item_int32_attr1_value = buffer[(doc_idx - idxBegin)];
    if (!isInvalidValue(item_int32_attr1_value)) {
      SimpleBuffer feature_buffer = SimpleBuffer();
      (void)fillFeatureToBuffer(item_int32_attr1_value, feature_buffer);
      LookupResult<float> lookup_result = getValue(reader, (Fingerprint64(toString(feature_buffer)) % (uint64_t)4\
));
      if ((int)1) {
        (( float*)output)[((doc_idx * (int32_t)1) + (int32_t)0)] = getValue(lookup_result);
      } else {
        (( float*)output)[((doc_idx * (int32_t)1) + (int32_t)0)] = 0.000000e+00f;
      }
    } else {
      for (int32_t zero_idx = (int32_t)0; zero_idx < 1; ++zero_idx) {
        (( float*)output)[(((doc_idx * (int32_t)1) + (int32_t)0) + zero_idx)] = 0.000000e+00f;
      }
    }
  }           

if-else 分支條件被顯示優化成了 1。for 循環 bound 也變成了常數,下标等都是常數。不需要 cache 的結果也絕不用 buffer 儲存,需要 cache 的字段值則是從 buffer 中讀取。

效果

codegen 已經成功在雙十一淘寶、天貓搜尋等深度模型上的字段讀取和特征生成部分應用。在 CPU utility 方面,存儲讀取和特征生成方面有大幅節約。在 latency 方面,因為 latency 将整個特征部分包裹到一起,排程上有了按 item 次元 FOR 的能力,避免了原先純 Graph 方案中并行的 latency 傾斜問題,有效降低了整體 latency,為使用者體驗提供保障。

一些讨論

存儲和計算結構

RTP 字段目前是列存儲格式,雖然最外層按照 item 次元并行,但每個并行塊内還是在每個特征下 for item 的。而如果轉換為行存格式,則對于每個 item 按照行存順序生成所有特征更劃算一些。如何排布計算,如何排布存儲,codegen 能解決一半的問題,更多的需要整套系統的協調關聯。

對于 embedding 操作,特征計算實際是在計算 embedding key,對于數組實作便是 offset。從這個角度看,特征計算類似于系統中的控制流,而 embedding vector 的拷貝則是實際的資料流。利用 codegen 方法,我們可以充分對利用各自的特點進行針對性優化,例如混合 item/feature 次元,在計算中做小塊的 SIMD。而針對于 embedding vector,則可以繞過 cache,甚至在異構計算中繞過記憶體。codegen 作為适配層,可以為這種優化提供充分的自由度。

UDF 能力及擴充

傳統的插件能力是在程式中定義一些固定的插入點,允許使用者通過實作 SDK 抽象接口來實作一些自定義拓展功能。HA3 引擎中的插件便是此路,它要求使用者分辨和定義哪些操作是 Query 次元的,哪些操作是 Item 次元的。

UDF(User Defined Function)則是更一般的插件方式,它允許使用者在幾乎任意位置插入自己的邏輯。

Flink UDF

相對引擎的插件形式就優雅許多,每個函數隻關心它真正計算需要的參數和計算邏輯,而不是和一些 context 打交道,算是個 pure function。

如果結合 codegen 和代碼上的語義分析,其實我們能夠分析出哪些運算是 Query 次元的,哪些運是 Item 次元的,可以對其進行合理的分拆和優化。使用者隻需要去定義運算本身,而計算的執行方式交給架構去分析和級聯,這不僅需要 codegen 的膠水能力,還需要計算架構上強大的分析能力。

通過使用者聲明的 UDF 接口,架構可以自動将資料轉換成需求的格式。而真正隻關心計算邏輯的使用者,也不用在乎這個運算是

elementwise

還是

broadcast

。不僅如此,我們甚至可以跨越語言的障礙,用 codegen 減少 UDF 中語言間互操作的 overhead,而使用者可以自由選擇喜歡擅長的程式設計語言而不必擔心性能問題。

混合的邊界

打碎運算和操作是獲得排程能力的基礎,但拆分的粒度應該如何把握?對于老代碼,我們可以利用先驗知識,定出 IR 和 C++ 自定義函數的邊界;對于新代碼,我們如何設計和把握 IR 和 C++ 之間的邊界,原則上是越細越好,而實際應用中并沒有明确的指南。對于引擎的使用者來說,找到 UDF 邏輯實作的邊界在哪裡,用圖表示,還是通過 C++ 等語言編寫,則是更困難的問題。

最自然的使用方式自然是用一種語言、一種表達方式解決問題。文法和語義分析是否能做到這一點,能否擁有靈活性的同時降低概念的複雜性,是我們需要繼續思考的問題。

總結與展望

本文主要讨論了在 RTP 的存儲讀取和特征生成場景中 codegen 的應用。利用 IR 和 C++ 混合程式設計,解決系統的抽象和性能問題,并提供 schedule 的優化能力。從這個角度看,expression 和 matchdoc 還有 fg lib 都隻是一個小的功能子集。

目前 codegen 隻覆寫了 RTP 的部分特征,還需要繼續改寫。還有更多的優化可以探索,例如行存的存儲格式、利用 prefetch、stream 等方法優化 cache 使用,這些曾經因為代碼結構不好做的事情,因為 codegen 而變得容易。有了 codegen,可以做更多的通用的和針對場景的優化。不過也有問題,而對于 UDF 的支援,IR 和 C++ 混合程式設計的邊界在哪裡?如何進一步和模型運算部分結合等,這些需要我們繼續去發展和解決。目前我們面對的主要是靜态的模型,而未來面對動态的 Query 執行的情況,codegen 也需要根據通路情況生成最優代碼,在 code(meta) 和 data 兩者之間平衡和統一,怎樣做樣的 cache 和執行政策則是一個更大的話題了。

脫離 RTP,放眼通用的存儲和計算架構,每個系統中都需要 codegen 能力。codegen 不但可以排程運算,還能夠适配存儲,甚至生成資料結構。分布式的計算引擎已經提供了很多 push/pull 的經驗,優化單機内的運算以及還有分布式的 io/latency 。但在每個具體場景中,優化方向和方法又不盡相同,如何做到大統一始終是一個難題。而如何自動地排程運算,甚至排程存儲以及微觀的存儲結構,更是很難。codegen 沒有給出答案,但給了可能。