计算框架相关的编译器知识(XLA, nGraph)

Tensorflow XLA的加速原理主要是融合Kernel

def  model_fn(x,y,z):    
    return tf.reduce_sum(x + y * z)  

如果运行模型时不使用 XLA,图表会启动三个内核,分别用于乘法、加法和减法。

官网原话:XLA 可以优化计算图,将加法、乘法和减法 “融合” 到单个 GPU Kernel中。此外,这种融合运算不会将 y*z 和 x+y*z 生成的中间值写入内存,而是将这些中间计算的结果直接 “流式传输” 给用户,并完整保存在 GPU 寄存器中。融合是 XLA 最重要的一种优化方式。内存带宽通常是硬件加速器上最稀缺的资源,因此删除内存运算是提升性能的最佳方法之一。

我的理解:计算融合到一个Kernel,减少Kernel启动次数,减少中间结果来回读写显存次数

Tensorflow使用XLA比不使用,在大部分模型下,能有10%~50%左右的加速提升;可以用户自己指定哪些操作做XLA编译优化,也可以让系统自动去找可优化的地方自动进行XLA优化;

AOT方式和JIT方式

两种方式下都会将整个计算图或则计算图的一部分直接编译成可执行代码。两则的区别也是比较明显的,除了编译时机不一样(AOT计算图在不会在运行阶段前被编译成可执行代码,JIT是在进入运行阶段后的适当的时机才会被编译成可执行代码),还有就是runtime(运行时)的参与程度。AOT中彻底不需要运行时的参与了,而JIT中还是需要运行时参与的,但是JIT会优化融合原计算图中的节点,加入XlaLaunch节点,来加速计算图的执行。

经典LLVM编译器架构:(扩展性支持好,即插即用的前端和后端)

计算框架相关的编译器知识(XLA, nGraph)_第1张图片

Tensorflow的XLA架构(类似LLVM):

计算框架相关的编译器知识(XLA, nGraph)_第2张图片

 

TVM

 

nGraph

原理上没有新东西提出来,只是提出了nGraph IR(Intermediate Representation),又一种DL IR。这些工作不外乎又是把问题变为了编译器的问题的又一种展现。

Tensorflow Model / ONNX / Caffe Model / ... ---> DL IR (nGraph IR / *.IR) ---> LLVM IR ---> CPU JIT / GPU / ...

如果把前面的Model看成一种语言或者DSL,就是DSL ---> DL IR ---> LLVM IR ---> Target ,然后你就在中间层疯狂的做优化,编译器优化开发也是这样做的。

在LLVM IR出现以前,很多编译器都有几层的IR表示,比如 C++ ----> 1st IR ----> OPT ----> 2nd IR ----> .... -> Target,只是LLVM出来以后,LLVM IR做了统一,编译器变为了 C++ ----> LLVM IR ---> OPT ----> LLVM IR ----> Target

在DL IR没有统一之前,我相信业界还会重复造轮子,然后来证明我的DL IR是最好的,你来用。就问题而言,这个问题已经变为了一个编译器问题。

 

你可能感兴趣的:(算法工程)