【云驻共创】MindSpore图算融合,打破算子边界

举报
SpiderMan 发表于 2022/08/29 10:47:02 2022/08/29
【摘要】 如今,通过人工或定制化的方式优化和融合算子可使网络的性能优化数倍,但人工定制化方式效率低且不能满足网络越来越多样化的需求。 所以MindSpore新特性图算融合就诞生了,它提供一种极简的算子达方式和泛化自动算子融合能力,将AI算力更极致地解放给用户。让我们进一步升入了解MindSpore是如何打破算子边界,让算子融合更高效。


目录

1.图算融合的背景

        1.1计算图和算子的认识

        1.2图算融合的优势

2.设计初心

        2.1图算融合概述

        2.2图算融合的设计初心

3.关键特性简介

        3.1统一图算前端表达

        3.2统一图算中间表达:算子白盒化

        3.3复用图层AD算法

                3.3.1自动生成反向算子

                3.3.2正反向算子的自动融合

        3.4 Bert上的图算融合

4.启方法汲编译流程演示

        4.1使用方法

                4.1.1样例脚本

        4.2自定义组合算子

                4.2.1样例脚本

5.实现过程小结

        5.1图算融合流程图

        5.2融合算子自动生成流程

        5.3小结

6.总结


    1.图算融合背景

    1.1计算图和算子认识

    包括MindSpore在内的AI框架,不约而同都采用计算图实现网络结构的表达。所以也间接证明了计算图在AI框架中强大的生命力。

    可以引用[tutorialspoint]中的一段话来介绍计算图:

    A computational graph is defined as a directed graph where the nodes correspond to mathematical operations. Computational graphs are a way of expressing and evaluating a mathematical expression.

    For example, here is a simple mathematical equation –

    P = x + y

    We can draw a computational graph of the above equation as follows.

    总结来说,计算图是一个以算子(Operators)为节点的有向图所表达的计算函数。在AI框架里面,这个计算函数对输入tensor依次调用执行有向图中的算子节点,并得到最终的输出tensor。

    引用Wikipedia来介绍一下算子:

    In mathematics, an operation is a function which takes zero or more input values (called operands) to a well-defined output value. The number of operands is the arity of the operation.

    所以,算子本质还是一个计算函数。在AI框架里面,这个计算函数还是对输入tensor依次进行若干内部处理,然后得到最终的输出tensor。

    算子的内部处理:通过对大量的算子进行分析,会发现,绝大部分算子都是可以拆分为若干更为基本的算子组成的的子图,比如Sigmoid,其内部计算可以用如下计算图表示:

    通过以上分析,不难发现,计算图和算子在计算本质上是一致的。算子是打包后的计算图,计算图是拆包后的算子。所以理论上,可以定义一个小规模的“基本算子”集合,然后通过一个或多个基本算子组合去等价表达任何的现有算子。从而进一步表达任意的现有计算图。

    用小规模的“基本算子”集合就可以表达任意现有计算图,但是现有主流AI框架还要定义大量的、本来可以用基本算子等价实现的复合算子,其主要原因是相比多个基本算子组成的计算图,使用复合算子可以具有更好的计算和内存局部性,从而获得更好的执行性能。

    因为对于基本算子计算图来说,相邻算子之间只能通过全局内存(或显存)进行数据传递。而对于复合算子来说,相邻的基本计算之间则可以通过局部内存或者寄存器进行数据传递。除了性能之外,在一些场景下,通过算子融合也能有效减少对全局内存的的实际占用。

    1.2图算融合的优势

    基于以上分析可见,相比复合算子,基本算子具有更好的通用组合表达能力,同时使得AI框架在算子支持方面更为精简。但相比基本算子,复合算子在性能方面存在优势。可以简单图示如下:

    对于以上问题,通过某些技术同时兼顾计算性能和组合表达能力,或者更具体地,使得上图的“计算性能”线变为水平。

    另一方面,在通过AI框架提供的算子表达计算网络时,通常是混合使用不同复合程度的算子。比如LayerNorm由12个基本算子组成,而Relu则只有1个基本算子。所以,对于这种不同复合程度的算子组成的计算图,从整体执行性能角度,可以将他们组合成更为合理的复合粒度算子,来获得最佳的整体计算图性能。

    在TVM、XLA等自动算子编译技术出现之前,AI框架主流采用手工融合的方式解决如上问题。主要思路是:

    • 识别常见的热点算子组合子图,比如:Add(Mul(x,y))。然后针对该算子子图手工实现对应融合算子;
    • 将融合算子注册到AI框架,并在AI框架中增加对应的优化pass,将匹配的算子子图替换为融合算子节点。

    这种方式的缺点是显而易见的。因为它只能针对若干热点场景进行融合,所以是无法做到通用和泛化的。基于以上问题背景,并结合MindSpore自身需求,我们提出图算融合解决方案。其主要思路是:

    • 以通用的pattern识别计算图中的融合场景,并生成融合算子子图;
    • 将生成的融合算子子图通过自动算子编译技术(AKG)生成对应的融合算子。

    通过图算融合之后,用户在进行网络定义时,不管使用组合表达能力更强的基本算子,还是复合算子。最终都将通过图算融合重新组织和融合,并替换为最优复合粒度的融合算子。


    2.设计初心

    2.1图算融合概述

    图算融合是MindSpore特有的网络性能优化技术。它可以通过自动分析和优化现有网络计算图逻辑,并结合目标硬件能力,对计算图进行计算化简和替代、算子拆分和融合、算子特例化编译等优化,以提升设备计算资源利用率,实现对网络性能的整体优化。相比传统优化技术,图算融合具有多算子跨边界联合优化、与算子编译跨层协同、基于Polyhedral的算子即时编译等独特优势。另外,图算融合只需要用户打开对应配置后,整个优化过程即可自动完成,不需要网络开发人员进行其它额外感知,使得用户可以聚焦网络算法实现。

    图算融合的适用场景包括:

    • 对网络执行时间具有较高性能要求的场景;
    • 通过拼接基本算子实现自定义组合算子,并希望对这些基本算子进行自动融合,以提升自定义组合算子性能的场景。

    MindSpore的三大特性:

    开发态友好:

    • 自动微分,网络+算子统一编程,函数式/算法原生表达,反向网络算子自动生成
    • 自动并行,模型自动切分实现最优效率的模型并行
    • 自动调优,动态图+静态图同套代码

    运行态高效:

    • On-Device执行,充分发挥异腾大算力
    • Pipeline优化,最大化并行线性度
    • 深度图优化,自适应硬件算力和精度

    部署态灵活:

    • 端-边-云按需协同计算,更好的保护隐私
    • 端-边-云统一架构,实现按需部署

    2.2图算融合设计初心

    图算融合的设计初心就是为了极简的算子表达和高效的自动融合,以下是设计初心的具体表述:

    • 统一前端表达:当前算子与图层没有统的前端表达用户新增算子需要额外学习TBE DSL,门槛较高。希望提供种通过ME原语在Python像写网络一样定义算子的能力。
    • 统一中间表达:当前算与图层没有统的中间表达,框架采取的均是提前写好硬件实现的算子并向框架注册算子。对于MSIR而言,每个算子的内部细节是不可知的,这样会丧失部分在图层进行细粒度跨算子编译优化的机会。同时也使得算子无法复用图层的通用优化。
    • 泛化融合能力:当前框架采用的均是采用手写融合pattern,手写融合算子人工注册的方式实现算子融合,这样的方式效率低,需要消耗大量人力(~天/pattern)不具备面尚用户多样化网络需求的泛化能力。因此需要提供可自动形成融合pattern自动生成高性能融合算子的能力(~分/pattern)。

    3.关键特性简介

    3.1统一图算前端表达

    前端:以定义原语的方式注册算子

    后端需要:算子信息、算子DSL

    使用图算方式定义算子

    只需描述算子的计算意图,所有性能优化的问题(format、schedule、tiling、fusion)都交给编译框架。让定义算子这种让人秃头的工作也变得简单生动起来。

    3.2统一图算中间表达:算子白盒化

    在IR中原黑盒算子原语被函数子图替

    打破算子和图层的信息边界

    • 图层知道算子间的依赖关系
    • 图层知道算子内的计算逻辑

    产生更多优化机会

    • 算子可复用图层的通用pass
    • 跨算子的细粒度的融合优化

    3.3复用图层AD算法

    3.3.1自动生成反向算子

    因为算子和图层统一使用MSIR表达,因此可以利用框架已有的自动求导的能力,对白盒算子进行自动求导。开发人员不用手动编写反向算子。

    3.3.2正反向算子的自动融合

    深度学习神经网络中算子的正反向经常有些相同的子计算,算子使用白盒表达后,就可以在图层自动分析消除正反算子中的相同的子计算。

    自动发现细粒度的算子融合机会,更高效的利用片上内存,减少外存数据的搬移。3.4 Bert上的图算融合

    在bert网络中虽然已经有很多定制的网络的pattern产生了定制融合算子,但是依然还有许多散落的基础黑盒算子,根据我们的融合算法产生了如下的融合pattern:

    ➢白加黑融合pattern:

    LambNextMV + Assign

    398 times

    LambUpdateWithLR + Assign

    398 times

    BiasAddGrad + Cast + Mul

    146 times




    ➢黑加黑融合pattern:

    RealDiv + Reshape + Assign

    402 times

    Mul + TensorAdd

    150 times

    Cast + TransData

    152 times

    Cast + Mul

    48 tims

    ReduceSum + Cast + Mul

    48 times

    Mul + ReduceSum

    3 times

    ReduceSum + Mul

    1 times

    TensorAdd + TensorAdd

    1 times

    Mul + ReduceSum + Cast + Mul 1 times

    1 times

    ReduceSum + TensorAdd

    1 times

    Cast + ExpandDim+Cast+Sub+Mul

    1 times

    Cast + RealDiv + Minimum + RealDiv + Sub + Pow + Mul + Sub + Mul +TensorAdd + Mul + Mul + TensorAdd

    1 times


    ➢融合后基础算子个数变化:

    Mul:413->160

    RealDiv:403->2

    Assign:1609->13

    Cast:649->250

    ReduceSum:56->2



    4.启用方法及编译流程演示

    4.1使用方法

    • MindSpore中的图算融合优化分布于网络图层编译和执行的多个步骤中,默认关闭状态,我们可以在训练脚本中为context指定参数enable_graph_kernel=True从而启用图算融合
    • from mindspore import context
      
      context.set context(enable graph kernel=True)
    • 通过定义继承与GraphKernel的子类定义图算融合的算子


    4.1.1样例脚本

    为了说明图算融合优化场景,构造了一个简单网络MyNet,包含一个乘法和加法计算。

    在打开图算融合进行优化之后,这两个计算便会自动合成一个融合算子:

    import numpy as np
    import mindspore.context as context
    from mindspore import Tensor
    from mindspore.nn import Cell
    import mindspore.ops as ops
    
    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
    # save graph ir to view fusion detail.
    context.set_context(save_graphs=True)
    # enable graph kernel optimization.
    context.set_context(enable_graph_kernel=True)
    
    class MyNet(Cell):
        def __init__(self):
            super(MyNet, self).__init__()
            self.add = ops.Add()
            self.mul = ops.Mul()
    
        def construct(self, x):
            a = self.mul(x, 2.0)
            res = self.add(a, 1.0)
            return res
    
    x = np.ones((4, 4)).astype(np.float32) * 0.5
    net = MyNet()
    result = net(Tensor(x))
    print("result: {}".format(result))
    

    输出结果:

    result: [[2. 2. 2. 2.]
     [2. 2. 2. 2.]
     [2. 2. 2. 2.]
     [2. 2. 2. 2.]]
    

    该计算图的融合结果如图1所示,其中左图为未使能图算融合时的对应计算图,右图为使能图算融合后的对应计算图。可以看到该网络中的加法和乘法被融合成一个算子。该融合过程可以通过查看中间IR,或者通过Profiling等工具跟踪算子执行过程进行验证。

    4.2自定义组合算子

    基于图算融合技术,用户可以很方便地实现高性能的自定义组合算子。相比其它自定义算子方式,这种方式具有对框架无侵入、简单易用等优点。

    其主要流程为:

    (1) 在脚本中用基本算子组合的方式实现自定义算子定义和使用;

    (2) 打开图算融合配置;

    (3) 图算融合对自定义组合算子中的基本算子自动进行算子融合,并生成高性能融合算子。

    4.2.1样例脚本

    构造一个简单网络MyNet,并在其中使用了自定义算子MyOp。

    代码样例如下:

    import numpy as np
    import mindspore.context as context
    from mindspore import Tensor
    from mindspore.nn import Cell
    import mindspore.ops.operations as P
    
    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
    # enable graph kernel optimization.
    context.set_context(enable_graph_kernel=True)
    
    class MyOp(Cell):
        """ my first custom OP composited by basic OPs """
        def __init__(self):
            super(MyOp, self).__init__()
            self.sub = P.Sub()
            self.mul = P.Mul()
    
        def construct(self, x, y):
            a = self.sub(x, y)
            return self.mul(a, x)
    
    class MyNet(Cell):
        def __init__(self):
            super(MyNet, self).__init__()
            self.mul = P.Mul()
            self.pow = P.Pow()
            self.my_op = MyOp()
    
        def construct(self, x, y):
            a = self.mul(x, 2.0)
            b = self.pow(a, 3.0)
            res = self.my_op(b, y)
            return res
    
    x = np.ones((4, 4)).astype(np.float32) * 0.2
    y = np.ones((4, 4)).astype(np.float32) * 0.3
    net = MyNet()
    result = net(Tensor(x), Tensor(y))
    print("result: {}".format(result))
    

    输出结果:

    result: [[-0.015104 -0.015104 -0.015104 -0.015104]
     [-0.015104 -0.015104 -0.015104 -0.015104]
     [-0.015104 -0.015104 -0.015104 -0.015104]
     [-0.015104 -0.015104 -0.015104 -0.015104]]
    

    该计算图的融合结果如图2所示,其中左图为未使能图算融合时的对应计算图,右图为使能图算融合后的对应计算图。可以看到不仅自定义算子MyOp中的基本算子进行了融合,并且与主图中的其他算子也进行了更大范围融合。该融合过程可以通过查看中间IR,或者通过Profiling等工具跟踪算子执行过程进行验


    5.实现过程

    5.1图算融合流程图

    5.2融合算子自动生成流程

    5.3小

    图算融合提供了一种简洁的算子前端表达方式。使用户可以在Python像写网络一样定义算子,仅需描述计算诉求,无需考虑性能优化问题。

    • 定义算子效率提升10X,代码量:百行—>十行

    图算融合提供了图层和算子的统的中间表达方式。使得算子可以复用图层的通用优化,同时打破了图层与算子之间的信息边界,创造出跨算子的优化机会。

    • 自动生成反向算子
    • 自动实现跨算子优化

    图算融合提供了一种泛化的算子融合能力。对多样化的用户网络可提供种通用的融合能力,更充分地释放AI算力。

    • 融合算子效率提升100X,人工定义—>自动发现、自动融合

    6.总

    在MindSpore,通过白盒化算子,初步打开了图算这一扇门。既提供了图算一体的DSL表达能力,又打开了图算混合优化的大门借此打破了算子边界可能多年以后,当AI遍布生活生产的每个角落,MindSpore这样的深度学习框架会成为一种必不可少,且每个人都能用的AI操作系统,它也就成为我们生活的一部分。

    MindSpore致力于AI开发生态的繁荣,开源开放可扩展架构,助力开发者灵活扩展第三方框架、第三方芯片支持能力,让开发者实现各种定制化需求。MindSpore将在门户网站、开源社区提供更多学习资源、支持与服务。

    我相信未来,也相信MindSpore


    本文参与华为云社区【内容共创】活动第19期。

    https://bbs.huaweicloud.com/blogs/370132

    任务26:MindSpore图算融合,打破算子边界


    【版权声明】本文为华为云社区用户原创内容,转载时必须标注文章的来源(华为云社区)、文章链接、文章作者等基本信息, 否则作者和本社区有权追究责任。如果您发现本社区中有涉嫌抄袭的内容,欢迎发送邮件进行举报,并提供相关证据,一经查实,本社区将立刻删除涉嫌侵权内容,举报邮箱: cloudbbs@huaweicloud.com
    • 点赞
    • 收藏
    • 关注作者

    评论(0

    0/1000
    抱歉,系统识别当前为高风险访问,暂不支持该操作

    全部回复

    上滑加载中

    设置昵称

    在此一键设置昵称,即可参与社区互动!

    *长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。

    *长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。