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

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

1. 通過指令碼語言擴充套件通用運算元上的使用者訂製能力，目前這些通用運算元包括scorer運算元，filter運算元等。這一方面側重於編譯前端，我們開發了一種嵌入引擎的指令碼語言cava，解決了使用者擴充套件引擎功能的一些痛點，包括外掛的開發測試效率，相容性，引擎版本升級效率等。

2. 通過codegen技術優化全圖化引擎效能，由於全圖化引擎是基於tensorflow開發，它天生具備tensorflow xla編譯能力，利用kernel的fuse提升效能，這部分內容可以參考XLA Overview。xla主要面向tensorflow內建的kernel，能發揮的場景是線上預測模型算分。但是對於使用者自己開發的運算元，xla很難發揮作用。本文第二部分主要介紹對於自定義運算元我們是如何做codegen優化的。

通用運算元上的指令碼語言cava

由於運算元開發和組圖邏輯對普通使用者來說成本較高，全圖化引擎內建了一些通用運算元，比如說scorer運算元，filter運算元。這些通用運算元能載入c++外掛，也支援用cava指令碼寫的外掛。關於cava可以參考這篇文章瞭解一下。

和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目前有如下場景在使用

• 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;
    }
}

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是編譯執行的，和原生的解釋執行的表示式相比有天然的效能優勢。

關於cava前端的展望

cava是全圖化引擎上面向使用者需求的語言，有使用者定製擴充套件邏輯的需求都可以考慮用通用運算元+cava外掛配合的模式來支援，例如全圖化sql上的udf，規則引擎的匹配需求等等。

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

自定義運算元的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的概念圖表達一下異構計算編譯器設計的願景：讓資料分析，深度學習，影像演算法等能用統一易用的程式設計介面充分發揮異構計算平臺的算力。

總結

編譯技術已經開始在引擎的使用者體驗，迭代效率，效能優化中發揮作用，後續會跟著全圖化引擎的演進不斷髮展。能做的事情很多，挑戰很大，有感興趣的同學可以聯絡我們探討交流。

參考

• Balancing vectorized query execution with bandwidth-optimized storage, Chapter 3
• Efficiently Compiling Efficient Query Plans for Modern Hardware
• TensorFlow編譯優化策略 – XLA
• Weld: Rethinking the Interface Between Data-Intensive Libraries
• TVM: An Automated End-to-End Optimizing Compiler for Deep Learning