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

全图化引擎又称算子执行引擎,它的介绍能够参考从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编译框架TVMWeld等等。

借用weld的概念图表达一下异构计算编译器设计的愿景:让数据分析,深度学习,图像算法等能用统一易用的编程接口充分发挥异构计算平台的算力。

 

总结

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

参考

原文连接 

相关文章
相关标签/搜索