目录
概述
BladeDISC 是Alibaba 推出的端到端AI Compiler,由于它基于 MLIR 的实现,所以它能够支持模型的动态shape,这也是为什么它的缩写DISC代表DynamIc Shape Compiler。本文基于BladeDISC 源码链接 ,对它的架构和工作原理进行分析,并从一个示例程序对它的运行过程做了源码分析。
架构
软件核心模块架构如上
- BladeDISC API模块 :通过TF API REGISTER_OPTIMIZATION 、REGISTER_OP、REGISTER_KERNEL_BUILDER 将BladeDisc优化器TaoOptimizationPass以及BladeDisc DiscLaunchOp 注册到了Tensorflow 框架,这样BladeDISC就和Tensorflow的图优化Pass建立起了关联(具体可以参考我的另外一篇关于Tensorflow计算图执行过程的分析文章 链接),在图优化Pass的执行阶段对计算图进行优化执行了(类似Tensorflow内置的XLA OptimizationPass)
- tao_bridge 模块:通过将计算图中可以被BladeDisc优化的子图分割、优化、替换为DiscLaunchOp节点。这里核心是3个图优化Pass (TaoMarkForCompilationPass、TaoEncapsulateSubgraphsPass、TaoBuildTaoOpPass)和DiscLaunchOp节点
- tao_compiler模块:负责用tao_compiler编译DiscLaunchOp节点表示的Tensorflow FunctionDef子图->MLIR->so,在计算图执行过程中再执行so中的main函数实现推理引擎加速
源码分析
示例代码
#BladeDISC/tao/tao_bridge/test/gpu/mlir/test_mlir.py
g = tf.Graph()
with g.as_default():
shape = [3, 3]
px = tf.placeholder(shape=shape, dtype=dtype, name="px")
py = tf.placeholder(shape=shape, dtype=dtype, name="py")
c = tf.constant([[1, 2, 3], [1, 2, 3], [1, 2, 3]],
dtype=dtype, shape=shape, name="c")
t1 = px + py
t2 = tf.matmul(t1, c)
t3 = tf.reduce_sum(t2)
t4 = t3 * t3
with super(TestTaoMlir, self).new_sess() as sess:
for i in range(0, 120):
np_x = np.ones(shape).astype(np_dtype)
np_y = np.ones(shape).astype(np_dtype)
r = sess.run(t4, {px: np_x, py: np_y})
通过对BladeDISC提供的test 程序 test_mlir.py 运行过程进行源码的详细分析如下
![](https://i-blog.csdnimg.cn/blog_migrate/967d55190bf6d1af7fb674e28bef95d0.jpeg)
TensorRT 对比
![](https://i-blog.csdnimg.cn/blog_migrate/60f5a9e506158fb15660c203e901008a.jpeg)
TensorRT作为Tensorflow 支持的另外一个AI Compiler(参考我对TensorRT的分析文章 链接 )从AI Compiler对接框架需要提供的功能层次上与BladeDISC做了对比如上图