全圖化引擎(AI·OS)中的編譯技術

時間 2019-12-04

標籤引擎編譯技術简体版

原文原文鏈接

全圖化引擎又稱算子執行引擎，它的介紹能夠參考從HA3到AI OS -- 全圖化引擎破繭之路。本文從算子化的視角介紹了編譯技術在全圖化引擎中的運用。主要內容有：html

1. 經過腳本語言擴展通用算子上的用戶訂製能力，目前這些通用算子包括scorer算子，filter算子等。這一方面側重於編譯前端，咱們開發了一種嵌入引擎的腳本語言cava，解決了用戶擴展引擎功能的一些痛點，包括插件的開發測試效率，兼容性，引擎版本升級效率等。前端

2. 經過codegen技術優化全圖化引擎性能，因爲全圖化引擎是基於tensorflow開發，它天生具有tensorflow xla編譯能力，利用kernel的fuse提高性能，這部份內容能夠參考XLA Overview。xla主要面向tensorflow內置的kernel，能發揮的場景是在線預測模型算分。可是對於用戶本身開發的算子，xla很難發揮做用。本文第二部分主要介紹對於自定義算子咱們是如何作codegen優化的。java

通用算子上的腳本語言cava

因爲算子開發和組圖邏輯對普通用戶來講成本較高，全圖化引擎內置了一些通用算子，好比說scorer算子，filter算子。這些通用算子能加載c++插件，也支持用cava腳本寫的插件。關於cava能夠參考這篇文章瞭解一下。c++

和c++插件相比，cava插件有以下特色：算法

1. 類java的語法。擴大了插件開發的受衆，讓熟悉java的同窗能快速上手使用引擎。
2. 性能高。cava是強類型，編譯型語言，它能和c++無損交互。這保證了cava插件的執行性能，在單值場景使用cava寫的插件和c++的插件性能至關。
3. 使用pool管理內存。cava的內存管理可定製，服務端應用每一個請求一個pool是最高效的內存使用策略。
4. 安全。對數組越界，對象訪問，除零異常作了保護。
5. 支持jit，編譯快。支持upc時編譯代碼，插件的上線就和上線普通配置同樣，極大的提高迭代效率
6. 兼容性：因爲cava的編譯過程和引擎版本是強綁定的，只要引擎提供的cava類庫接口不變，cava的插件的兼容性很容易獲得保證。而c++插件兼容性很難保證，任何引擎內部對象內存佈局的變更就可能帶來兼容性問題。

scorer算子中的cava插件

cava scorer目前有以下場景在使用sql

1. 主搜海選場景，算法邏輯能夠快速上線驗證
2. 賽馬引擎2.0的算分邏輯，賽馬引擎重構後引入cava算分替代原先的戰馬算分

樣例以下：

package test; import ha3.*; /* * 將多值字段值累加,並乘以query裏面傳遞的ratio,做爲最後的分數 * / class DefaultScorer { MInt32Ref mref; double ratio; boolean init(IApiProvider provider) { IRefManager refManger = provider.getRefManager(); mref = refManger.requireMInt32("ids"); KVMapApi kv = provider.getKVMapApi(); ratio = kv.getDoubleValue("ratio");//獲取kvpair內參數 return true; } double process(MatchDoc doc) { int score = 0; MInt32 mint = mref.get(doc); for (int i = 0; i < mint.size(); i++) { score = score + mint.get(i); } return score * ratio; } }

其中cava scorer的算分邏輯（process函數）調用次數是doc級別的，它的執行性能和c++相比惟一的差距是多了安全保護（數組越界，對象訪問，除零異常）。能夠說cava是目前能嵌入c++系統執行的性能最好的腳本語言。

filter算子中cava插件

filter算子中主要是表達式邏輯，例如filter = (0.5 * a + b) > 10。之前表達式的能力較弱，只能使用算術，邏輯和關係運算符。使用cava插件可進一步擴展表達式的能力，它支持類java語法，能夠定義變量，使用分支循環等。數據庫

計算 filter = (0.5 * a + b) > 10，用cava可定義以下： class MyFunc { public boolean init(FunctionProvider provider) { return true; } public boolean process(MatchDoc doc, double a, double b) { return (0.5 * a + b) > 10; } } filter = MyFunc(a, b)

另外因爲cava是編譯執行的，和原生的解釋執行的表達式相比有自然的性能優點。apache

關於cava前端的展望

cava是全圖化引擎上面向用戶需求的語言，有用戶定製擴展邏輯的需求均可以考慮用通用算子+cava插件配合的模式來支持，例如全圖化sql上的udf，規則引擎的匹配需求等等。編程

後續cava會進一步完善語言前端功能，完善類庫，儘量兼容java。依託suez和全圖化引擎支持更多的業務需求。api

自定義算子的codegen優化

過去幾年，在OLAP領域codegen一直是一個比較熱門的話題。緣由在於大多數數據庫系統採用的是Volcano Model模式。

其中的next()一般爲虛函數調用，開銷較大。全圖化引擎中也有相似的codegen場景，例如統計算子，過濾算子等。此外，和xla相似，全圖化引擎中也有一些場景能夠經過算子融合優化性能。目前咱們的codegen工做主要集中在cpu上對局部算子作優化，將來指望能在SQL場景作全圖編譯，而且在異構計算的編譯器領域有所發展。

單算子的codegen優化

1. 統計算子

例如統計語句：group_key:key,agg_fun:sum(val)#count()，按key分組統計key出現的次數和val的和。在統計算子的實現中，key的取值有一次虛函數調用，sum和count的計算是兩次虛函數調用，sum count計算出來的值又須要經過matchdoc存取，而matchdoc的訪問有額外的開銷：一次是定位到matchdoc storage，一次是經過偏移定位到存取位置。

那麼統計codegen是怎麼去除虛函數調用和matchdoc訪問的呢？在運行時，咱們能夠根據用戶的查詢獲取字段的類型，須要統計的function等信息，根據這些信息咱們能夠把通用的統計實現特化成專用的統計實現。例如統計sum和count只需定義包含sum count字段的AggItem結構體，而不須要matchdoc；統計function sum和count變成告終構體成員的+=操做。

假設key和val字段的類型都是int，那麼上面的統計語句最終codegen成的cava代碼以下：

class AggItem { long sum0; long count1; int groupKey; } class JitAggregator { public AttributeExpression groupKeyExpr; public IntAggItemMap itemMap; public AggItemAllocator allocator; public AttributeExpression sumExpr0; ... static public JitAggregator create(Aggregator aggregator) { .... } public void batch(MatchDocs docs, uint size) { for (uint i = 0; i < size; ++i) { MatchDoc doc = docs.get(i); //由c++實現，可被inline int key = groupKeyExpr.getInt32(doc); AggItem item = (AggItem)itemMap.get(key); if (item == null) { item = (AggItem)allocator.alloc(); item.sum0 = 0; item.count1 = 0; item.groupKey = key; itemMap.add(key, (Any)item); } int sum0 = sumExpr0.getInt32(doc); item.sum0 += sum0; item.count1 += 1; } } }

這裏sum count的虛函數被替換成sum += 和count += ，matchdoc的存取變成結構體成員的讀寫item.sum0和item.count0。通過llvm jit編譯優化以後生成的ir以下：

define void @_ZN3ha313JitAggregator5batchEP7CavaCtxPN6unsafe9MatchDocsEj(%"class.ha3::JitAggregator"* %this, %class.CavaCtx* %"@cavaCtx@", %"class.unsafe::MatchDocs"* %docs, i32 %size) { entry: %lt39 = icmp eq i32 %size, 0 br i1 %lt39, label %for.end, label %for.body.lr.ph for.body.lr.ph: ; preds = %entry %wide.trip.count = zext i32 %size to i64 br label %for.body for.body: ; preds = %for.inc, %for.body.lr.ph %lsr.iv42 = phi i64 [ %lsr.iv.next, %for.inc ], [ %wide.trip.count, %for.body.lr.ph ] %lsr.iv = phi %"class.unsafe::MatchDocs"* [ %scevgep, %for.inc ], [ %docs, %for.body.lr.ph ] %lsr.iv41 = bitcast %"class.unsafe::MatchDocs"* %lsr.iv to i64* // ... prepare call for groupKeyExpr.getInt32 %7 = tail call i32 %5(%"class.suez::turing::AttributeExpressionTyped.64"* %1, i64 %6) // ... prepare call for itemMap.get %9 = tail call i8* @_ZN6unsafe13IntAggItemMap3getEP7CavaCtxi(%"class.unsafe::IntAggItemMap"* %8, %class.CavaCtx* %"@cavaCtx@", i32 %7) %eq = icmp eq i8* %9, null br i1 %eq, label %if.then, label %if.end10 // if (item == null) { if.then: ; preds = %for.body // ... prepare call for allocator.alloc %15 = tail call i8* @_ZN6unsafe16AggItemAllocator5allocEP7CavaCtx(%"class.unsafe::AggItemAllocator"* %14, %class.CavaCtx* %"@cavaCtx@") // item.groupKey = key; %groupKey = getelementptr inbounds i8, i8* %15, i64 16 %16 = bitcast i8* %groupKey to i32* store i32 %7, i32* %16, align 4 // item.sum0 = 0; item.count1 = 0; call void @llvm.memset.p0i8.i64(i8* %15, i8 0, i64 16, i32 8, i1 false) // ... prepare call for itemMap.add tail call void @_ZN6unsafe13IntAggItemMap3addEP7CavaCtxiPNS_3AnyE(%"class.unsafe::IntAggItemMap"* %17, %class.CavaCtx* %"@cavaCtx@", i32 %7, i8* %15) br label %if.end10 if.end10: ; preds = %if.end, %for.body %item.0.in = phi i8* [ %15, %if.end ], [ %9, %for.body ] %18 = bitcast %"class.unsafe::MatchDocs"* %lsr.iv to i64* // ... prepare call for sumExpr0.getInt32 %26 = tail call i32 %24(%"class.suez::turing::AttributeExpressionTyped.64"* %20, i64 %25) // item.sum0 += sum0; item.count1 += 1; %27 = sext i32 %26 to i64 %28 = bitcast i8* %item.0.in to <2 x i64>* %29 = load <2 x i64>, <2 x i64>* %28, align 8 %30 = insertelement <2 x i64> undef, i64 %27, i32 0 %31 = insertelement <2 x i64> %30, i64 1, i32 1 %32 = add <2 x i64> %29, %31 %33 = bitcast i8* %item.0.in to <2 x i64>* store <2 x i64> %32, <2 x i64>* %33, align 8 br label %for.inc for.inc: ; preds = %if.then, %if.end10 %scevgep = getelementptr %"class.unsafe::MatchDocs", %"class.unsafe::MatchDocs"* %lsr.iv, i64 8 %lsr.iv.next = add nsw i64 %lsr.iv42, -1 %exitcond = icmp eq i64 %lsr.iv.next, 0 br i1 %exitcond, label %for.end, label %for.body for.end: ; preds = %for.inc, %entry ret void }

codegen的代碼中有很多函數是經過c++實現的，如docs.get(i)，itemMap.get(key)等。可是優化後的ir中並無docs.get(i)的函數調用，這是因爲常常調用的c++中實現的函數會被提早編譯成bc，由cava編譯器加載，通過llvm inline優化pass後被消除。

能夠認爲cava代碼和llvm ir基本能作到無損映射（cava中不容易實現邏輯可由c++實現，預編譯成bc加載後被inline），有了cava這一層咱們能夠用常規面向對象的編碼習慣來作codegen，不用關心llvm api細節，讓codegen門檻進一步下降。

這個例子中，統計規模是100w文檔1w個key時，線下測試初步結論是latency大約能降1倍左右（54ms->27ms），有待表達式計算進一步優化。

2. 過濾算子

在通用過濾算子中，表達式計算是典型的可被codegen優化的場景。例如ha3的filter語句：filter=(a + 2* b - c) > 0：

表達式計算是經過AttributeExpression實現的，AttributeExpression的evaluate是虛函數。對於單doc接口咱們能夠用和統計相似的方式，使用cava對錶達式計算作codegen。

對於批量接口，和統計不一樣的是，表達式的批量計算更容易運用向量化優化，利用cpu的simd指令，使計算效率有成倍的提高。可是並非全部的表達式都能使用一致的向量化優化方法，好比filter= a > 0 AND b < 0這類表達式，有短路邏輯，向量化會帶來沒必要要的計算。

所以表達式的編譯優化須要有更好的codegen抽象，咱們發現Halide能比較好的知足咱們的需求。Halide的核心思想：算法描述（作什麼 ir）和性能優化（怎麼作 schedule）解耦。這種解耦能讓咱們更靈活的定製優化策略，好比某些場景走向量化，某些場景走普通的codegen；更進一步，不一樣計算平臺上使用不一樣的優化策略也成爲可能。

3. 倒排召回算子

在seek算子中，倒排召回是經過QueryExecutor實現的，QueryExecutor的seek是虛函數。例如query= a AND b OR c。

QueryExecutor的And Or AndNot有比較複雜的邏輯，虛函數的開銷相對佔比沒有表達式計算那麼大，以前用vtune作過預估，seek虛函數調用開銷佔比約10%（數字不必定準確，inline效果無法評估）。和精確統計，表達式計算相比，query的組合空間巨大，seek的codegen得更多的考慮對高性價比query作編譯優化。

4. 海選與排序算子

在ha3引擎中海選和精排邏輯中有大量比較操做。例如sort=+RANK;id字句，對應的compare函數是Rank Compartor和Id Compartor的聯合比較。compare的函數調用可被codegen掉，而且還可和stl算法聯合inline。std::sort使用非inline的comp函數帶來的開銷能夠參考以下例子：

bool myfunction (int i,int j) { return (i<j); } int docCount = 200000; std::random_device rd; std::mt19937_64 mt(rd()); std::uniform_int_distribution<int> keyDist(0, 200000); std::vector<int> myvector1; for (int i = 0 ; i < docCount; i++) { myvector1.push_back(keyDist(mt)); } std::vector<int> myvector2 = myvector1; std::sort (myvector1.begin(), myvector1.end()); // cost 15.475ms std::sort (myvector2.begin(), myvector2.end(), myfunction); // cost 19.757ms

對20w隨機數排序，簡單的比較inline帶來30%的提高。固然在引擎場景，因爲比較邏輯複雜，這部分收益可能不會太多。

算子的fuse和codegen

算子的fuse是tensorflow xla編譯的核心思想，在全圖化場景咱們有一些自定義算子也能夠運用這個思想，例如feature generator。

fg特徵生成是模型訓練中很重要的一個環節。在線fg是以子圖+配置形式描述計算，這部分的codegen能使數據從索引直接計算到tensor上，省去了不少環節中間數據的拷貝。目前這部分codegen工做能夠參考這篇文章

關於編譯優化的展望

SQL場景全圖的編譯執行

數據庫領域Whole-stage Code Generation早被提出並應用，例如Apache Spark as a Compiler；還有如今比較火的GPU數據庫Mapd，把整個執行計劃編譯成架構無關的中間表示（llvm ir），藉助llvm編譯到不一樣的target執行。

從實現上看，SQL場景的全圖編譯執行對全圖化引擎還有更多意義，好比能夠省去tensorflow算子執行帶來的線程切換的開銷，能夠去除算子間matchdoc傳遞（matchdoc做爲通用的數據佈局性能較差）帶來的性能損耗。

面向異構計算的編譯器

隨着摩爾定律觸及天花板，將來異構計算必定是一個熱門的領域。SQL大規模數據分析和在線預測就是異構計算能夠發揮做用的典型場景，好比分析場景大數據量統計，在線預測場景深度模型大規模並行計算。cpu驅動其餘計算平臺如gpu fpga，相互配合各自作本身擅長的事情，在將來有多是常態。這須要爲開發人員提供更好的編程接口。

全圖化引擎已經領先了一步，集成了tensorflow計算框架，天生具有了異構計算的能力。但在編譯領域，通用的異構計算編程接口還遠未到成熟的地步。工業界和學術界有很多嘗試，好比tensorflow的xla編譯框架，TVM，Weld等等。

借用weld的概念圖表達一下異構計算編譯器設計的願景：讓數據分析，深度學習，圖像算法等能用統一易用的編程接口充分發揮異構計算平臺的算力。