全图化引擎(AI·OS)中的编译技术

66次阅读

共计 8292 个字符,预计需要花费 21 分钟才能阅读完成。

全图化引擎又称算子执行引擎,它的介绍可以参考从 HA3 到 AI OS – 全图化引擎破茧之路。本文从算子化的视角介绍了编译技术在全图化引擎中的运用主要内容有:
1. 通过脚本语言扩展通用算子上的用户订制能力,目前这些通用算子包括得分算子,过滤算子等。这一方面侧重于编译前端,我们开发了一种嵌入引擎的脚本语言 cava,解决了用户扩展引擎功能的一些痛点,包括插件的开发测试效率,兼容性,引擎版本升级效率等。
2. 通过代码技术优化全图化引擎性能,由于全图化引擎是基于 tensorflow 开发,它天生具备 tensorflow xla 编译能力,利用内核的熔丝提升性能,这部分内容可以参考 XLA 概述 .xla 主要面向 tensorflow 内置的内核,能发挥的场景是在线预测模型算分。但是对于用户自己开发的算子,XLA 很难发挥作用。本文第二部分主要介绍对于自定义算子我们是如何做的代码生成优化的。
通用算子上的脚本语言静脉
由于算子开发和组图逻辑对普通用户来说成本较高,全图化引擎内置了一些通用算子,比如说射手算子,过滤器算子。这些通用算子能加载 C ++ 插件,也支持用静脉脚本写的插件。关于静脉参考可以这篇文章了解一下。
和 C ++ 插件相比,静脉插件有如下特点:

1. 类 java 的语法。扩大了插件开发的受众,让熟悉 java 的同学能快速上手使用引擎。
2. 性能高.cava 是强类型,编译型语言,它能和 c ++ 无损交互。这保证了 cava 插件的执行性能,在单值场景使用 cava 写的插件和 c ++ 的插件性能相当。
3. 使用池管理内存.cava 的内存管理可定制,服务端应用每个请求一个池是最高效的内存使用策略。
4. 安全。对数组越界,对象访问,除零异常做了保护。
5. 支持 jit,编译快。支持 upc 时编译代码,插件的上线就和上线普通配置一样,极大的提升迭代效率
6. 兼容性:由于 cava 的编译过程和引擎版本是强绑定的,只要引擎提供的 cava 类库接口不变,cava 的插件的兼容性很容易得到保证。而 c ++ 插件兼容性很难保证,任何引擎内部对象内存布局的变动就可能带来兼容性问题。

射手算子中的静脉插件
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;
}
}
其中 cava scorer 的算分逻辑(process 函数)调用次数是 doc 级别的,它的执行性能和 c ++ 相比唯一的差距是多了安全保护(数组越界,对象访问,除零异常)。可以说 cava 是目前能嵌入 C ++ 系统执行的性能最好的脚本语言。
过滤算子中静脉插件
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)
另外由于静脉是编译执行的,和原生的解释执行的表达式相比有天然的性能优势。
关于静脉前端的展望
静脉是全图化引擎上面向用户需求的语言,有用户定制扩展逻辑的需求都可以考虑用通用算子 + 静脉插件配合的模式来支持,例如全图化 SQL 上的 UDF,规则引擎的匹配需求等等。
后续静脉会进一步完善语言前端功能,完善类库,尽可能兼容的 Java。依托苏伊士和全图化引擎支持更多的业务需求。
自定义算子的代码生成优化
过去几年,在 OLAP 领域 codegen 一直是一个比较热门的话题。原因在于大多数数据库系统采用的是 Volcano Model 模式。

其中的下一个()通常为虚函数调用,开销较大。全图化引擎中也有类似的代码生成场景,例如统计算子,过滤算子等。此外,和 XLA 类似,全图化引擎中也有一些场景可以通过算子融合优化性能。目前我们的代码生成工作主要集中在 CPU 上对局部算子做优化,未来期望能在 SQL 场景做全图编译,并且在异构计算的编译器领域有所发展。
单算子的代码生成优化
统计算子
例如统计语句:group_key:键,agg_fun:总和(VAL)#COUNT(),按键分组统计键出现的次数和缬氨酸的和在统计算子的实现中,键的取值有一次虚函数调用,sum 和 count 的计算是两次虚函数调用,sum count 计算出来的值和需要通过 matchdoc 存取,而 matchdoc 的访问有额外的开销:一次是定位到 matchdoc storage,一次是通过偏移定位到存取位置。
那么统计代码生成是怎么去除虚函数调用和 matchdoc 访问的呢?在运行时,我们可以根据用户的查询获取字段的类型,需要统计的功能等信息,根据这些信息我们可以把通用的统计实现特化成专用的统计实现。例如统计 sum 和 count 只需定义包含 sum count 字段的 AggItem 结构体,而不需要 matchdoc; 统计函数 sum 和 count 变成了结构体成员的 + = 操作。
假设键和 VAL 字段的类型都是整型,那么上面的统计语句最终的代码生成成的静脉代码如下:
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;
}
}
}

这里总计数的虚函数被替换成和 + + 和计数 + =,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
}

代码生成的代码中有不少函数是通过 C ++ 实现的,如 docs.get(i)中,itemMap.get(键)等。但是优化后的 IR 中并没有 docs.get(I)的函数调用,这是由于经常调用的 c ++ 中实现的函数会被提前编译成 bc,由 cava 编译器加载,经过 llvm inline 优化 pass 后被消除。
可以认为 cava 代码和 llvm ir 基本能做到无损映射(cava 中不容易实现逻辑可由 c ++ 实现,预编译成 bc 加载后被内联),有了 cava 这一层我们可以用常规面向对象的编码习惯来做 codegen,不用关心 llvm api 细节,让 codegen 门槛进一步降低。
这个例子中,统计规模是 100 瓦特文档 1 瓦特个键时,线下测试初步结论是延迟大约能降 1 倍左右(54ms-> 27ms),有待表达式计算进一步优化。
2. 过滤算子
在通用过滤算子中,表达式计算是典型的可被 codegen 优化的场景。例如 ha3 的过滤语句:filter =(a + 2 * b – c)> 0:

表达式计算是通过 AttributeExpression 实现的,AttributeExpression 的评价是虚函数。对于单文档接口我们可以用和统计类似的方式,使用静脉对表达式计算做代码生成。
对于批量接口,和统计不同的是,表达式的批量计算更容易运用向量化优化,利用 CPU 的 SIMD 指令,使计算效率有成倍的提升。但是并不是所有的表达式都能使用一致的向量化优化方法,比如 filter = a> 0 AND b <0 这类表达式,有短路逻辑,向量化会带来不必要的计算。
因此表达式的编译优化需要有更好的 codegen 抽象,我们发现 Halide 能比较好的满足我们的需求.Halide 的核心思想:算法描述(做什么 ir)和性能优化(怎么做 schedule)解耦。种解耦能让我们更灵活的定制优化策略,比如某些场景走向量化,某些场景走普通的代码生成; 更进一步,不同计算平台上使用不同的优化策略也成为可能。
3. 倒排召回算子
在寻求算子中,倒排召回是通过 QueryExecutor 实现的,QueryExecutor 的 seek 是虚函数。例如 query = a AND b OR c。

QueryExecutor 的和或 ANDNOT 有比较复杂的逻辑,虚函数的开销相对占比没有表达式计算那么大,之前用 VTUNE 做过预估,求虚函数调用开销占比约 10%(数字不一定准确,内联效果没法评估)和精确统计,表达式计算相比,查询的组合空间巨大,寻求的代码生成得更多的考虑对高性价比的查询做编译优化。
海选与排序算子
在 HA3 引擎中海选和精排逻辑中有大量比较操作例如排序 = + RANK; ID 字句,对应的比较函数是秩 Compartor 和标识 Compartor 的联合比较.compare 的函数调用可被代码生成掉,并且还可和 STL 算法联合 inline.std :: 排序使用非在线的补偿函数带来的开销可以参考如下例子:
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

对 20 瓦特随机数排序,简单的比较直列带来 30%的提升。当然在引擎场景,由于比较逻辑复杂,这部分收益可能不会太多。
算子的保险丝和代码生成
算子的 fuse 是 tensorflow xla 编译的核心思想,在全图化场景我们有一些自定义算子也可以运用这个思想,例如特征生成器。
FG 生成特征的英文模型训练中很重要的一个环节。在线 FG 是以子图 + 配置形式描述计算,这部分的代码生成能使数据从索引直接计算到张量上,省去了很多环节中间数据的拷贝。目前这部分的代码生成可以工作参考这篇文章

关于编译优化的展望
SQL 场景全图的编译执行
数据库领域全阶段代码生成早被提出并应用,例如 Apache 的火花作为编译器 ; 还有现在比较火的 GPU 数据库 MAPD,把整个执行计划编译成架构无关的中间表示(LLVM IR),借助 LLVM 编译到不同的目标执行。
从实现上看,SQL 场景的全图编译执行对全图化引擎还有更多意义,比如可以省去 tensorflow 算子执行带来的线程切换的开销,可以去除算子间 matchdoc 传递(matchdoc 作为通用的数据布局性能较差)带来的性能损耗。
面向异构计算的编译器
随着摩尔定律触及天花板,未来异构计算一定是一个热门的领域.SQL 大规模数据分析和在线预测就是异构计算可以发挥作用的典型场景,比如分析场景大数据量统计,在线预测场景深度模型大规模并行计算.cpu 驱动其他计算平台如 gpu fpga,相互配合各自做自己擅长的事情,在未来有可能是常态。这需要为开发人员提供更好的编程接口。
全图化引擎已经领先了一步,集成了 tensorflow 计算框架,天生具备了异构计算的能力。但在编译领域,通用的异构计算编程接口还远未到成熟的地步。工业界和学术界有不少尝试,比如 tensorflow 的 xla 编译框架,TVM,Weld 等等。
借用焊接的概念图表达一下异构计算编译器设计的愿景:让数据分析,深度学习,图像算法等能用统一易用的编程接口充分发挥异构计算平台的算力。

总结
编译技术已经开始在引擎的用户体验,迭代效率,性能优化中发挥作用,后续会跟着全图化引擎的演进不断发展。能做的事情很多,挑战很大,感兴趣有同学的可以联系我们探讨交流。
参考

使用带宽优化存储平衡向量化查询执行,第 3 章
有效地编译现代硬件的高效查询计划
TensorFlow 编译优化策略 – XLA
Weld:重新思考数据密集型库之间的接口
TVM:用于深度学习的自动化端到端优化编译器

本文作者:sance 阅读原文
本文为云栖社区原创内容,未经允许不得转载。

正文完
 0