【架构分析】BladeDISC 源码分析

目录

概述

架构

源码分析

TensorRT 对比

---

概述

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 运行过程进行源码的详细分析如下

BladeDISC 核心工作时序分析

TensorRT 对比

TensorRT 对比 BladeDISC

TensorRT作为Tensorflow 支持的另外一个AI Compiler（参考我对TensorRT的分析文章 链接 ）从AI Compiler对接框架需要提供的功能层次上与BladeDISC做了对比如上图