AI实践哪家强?来 AICon, 解锁技术前沿,探寻产业新机! 了解详情
写点什么

CPU 推理性能提高数十倍,MegEngine 计算图、MatMul 优化解析 | 工程之道

  • 2020-08-14
  • 本文字数:4675 字

    阅读完需:约 15 分钟

CPU推理性能提高数十倍,MegEngine计算图、MatMul优化解析 | 工程之道

本文针对旷视天元深度学习框架在推理优化过程中所涉及的计算图优化与 MatMul 优化进行深度解读。

背景及引言

在深度学习大规模落地边缘端场景的今天,如何最大程度降本增效,是企业与开发者共同关注的话题。其中,模型的训练与推理是两个关键环节。


天元(MegEngine)深度学习框架凭借「训练与推理一体化」的独特范式,能够极大程度上(90%)节省模型从研发到部署的整体成本,降低转换难度,真正实现小时级转化;同时,天元(MegEngine)在 CPU 推理方面所做的大量优化工作,也使得开发者在推理时能够发挥出处理器的最佳性能。


在之前我们对天元的极致推理优化进行了综述《 工程之道,MegEngine 推理性能极致优化之综述篇》。本文则针对天元在推理优化过程中所涉及的计算图优化与 MatMul 优化进行深度解读,希望能够帮助广大开发者在利用天元 MegEngine「深度学习,简单开发」的同时,也能够了解 CPU 优化的相关知识。


从而帮助大家在模型部署的整体流程中更好地进行加速;在实际模型部署时能够评估模型在特定平台上运行所能达到的性能以及内存使用情况;以及在算法设计时可以设计出更利于 CPU 优化加速的卷积 Opr 等。

CPU 推理优化概览

对于产业应用而言,CPU 推理的性能优化至关重要,如下表所示,经过优化的推理性能可以较未经优化的原始性能提升数十倍。



在进行模型推理阶段的优化时,首先需要对模型的计算图进行优化,以避免冗余的计算与访存,确保计算图在推理时是最优的;其次,在大多 CV 相关模型中,卷积的计算比重最高,达到模型总计算量 90% 以上,因此对模型推理的优化主要聚焦在对卷积计算的优化上。


通常而言,卷积计算的实现方式有 3 种:direct 卷积,Im2col 卷积以及 Winograd 卷积。


Ÿ direct 卷积:根据卷积的计算公式直接对 FeatureMap 上进行滑窗计算;


Ÿ Im2col 卷积:根据卷积计算需要在输入通道上进行 reduce sum 的特点,将卷积运行转化为 MatMul 计算;


Ÿ winograd:在保证计算无误的前提下,使用加法替代乘法,达到优化卷积乘法计算量的目的,在中间过程需要使用 MatMul 进行计算。


由上文可知,以 Im2col 或 Winograd 方法进行的卷积计算会频繁使用到 MatMul,因此对 MatMul 这种基础算子的优化就显得尤为重要。


基于上述考量,本文将首先介绍模型优化中的图优化,然后介绍基础算子 MatMul 在 CPU 上的优化方法。

推理计算图优化

在训练阶段定义模型的计算图,主要是为了满足模型参数的训练需求。当训练结束,模型参数固定后,对计算图的进一步优化能够帮助模型推理更加高效。


推理和训练的计算图是一张有向无环图(DAG),在天元中,开发者能够以类似 LLVM 的方式对 DAG 计算图定义许多优化方法,这里简称 OptPass。OptPass 可以根据用户的配置有选择性的加入到图优化的 OptPass 列表中,从而帮助用户灵活地为图优化定义 OptPass。

计算图优化

天元定义了多个为推理进行计算图优化的 OptPass,开发者使用这些 OptPass 之后,将得到一张用于推理的最优计算图。下面以 MobileNetV1 中, Convolution+Batch Norm+Relu 这样的典型结构经过计算图优化之后 Fuse 为 ConvBias 的过程为例,介绍天元的计算图优化过程。


由于 Batch Norm 在推理阶段除了输入 Tensor 外其他都是常数,因此可以简化为多个 elemwise 的组合,天元实现了一个 ConvertBatchNormToElemwisePass 的 Optpass,这个 OptPass 将模型中的所有 Batch Norm 转化为 Elemwise,具体转化原理如下:



紧接这一过程,天元将运行 OptPass ParamRedistributePass,该 OptPass 会将上述 Batch Norm 转化而来的 Elemwise 中的 scale 融合到 convolution 的权重中,具体实现原理如下:



最后天元将运行 OptPass FuseConvBiasNonlinPass,该 OptPass 会将计算图中的 Convolution+Elemwise 转化为天元内部实现的 ConvBias Op 中,同时,它还会设置 ConvBias Op 中 NonelineaMode 参数。


如此便完成了从 Convolution+Batch Norm+Relu 到 ConvBias 的转换,整体转换过程如下图:


实验验证图优化之后的性能

在推理之前,如果要对完成训练的模型进行图优化,则需要在模型 dump 的期间进行,当然,也可以在 SDK 运行模型之前进行。下面是在模型 dump 时进行图优化的代码:


from megengine.jit import trace@trace(symbolic=True)def pred_fun(data, *, net):    net.eval()    pred = net(data)    pred_normalized = F.softmax(pred)    return pred_normalized    # 使用 trace 类的 trace 接口无需运行直接编译pred_fun.trace(data, net=xor_net)# 使用 trace 类的 dump 接口进行部署pred_fun.dump("xornet_deploy.mge", arg_names=["data"], optimize_for_inference=True, enable_fuse_conv_bias_nonlinearity=True)
复制代码


上面的 optimize_for_inference=True 将在 dump 模型时候针对 inference 进行优化, enable_fuse_conv_bias_nonlinearity=True 将在模型中进行 Op Fuse 的优化,此外天元还支持其他的优化参数,具体可见天元的文档。


下表展示的,是在实验过程中对模型进行图优化前、后,模型运行性能的测试对比。



可以看出图优化对模型性能的提升效果显著,具体提升比例由模型自身决定。

MatMul 优化

如前文所述,MatMul 作为卷积运算的基础算子,会频繁地被 Im2col、Winograd 以及 FullyConnect 使用。因此在天元中,MatMul 既被封装为单独的 Op,也可以作为单独的 algo 供卷积的实现使用。


此外,由于 MatMul 也是计算最为密集的算子,因此天元对它进行了极致的优化。


优化


MatMul 是线性代数中的矩阵乘,假设矩阵 A 大小为 MK,矩阵 B 大小为 KN,则得到矩阵 C 大小为 M*N,其中 C 的每个元素的计算公式如下:



可以发现,在 MatMul 的计算中乘法和加法的计算量为 2MNK (计算 C 中每个元素时,加法和乘法计算量分别为 K,C 的总元素个数为 MN),访存量为 2MNK (计算每个 C 中元素需要 2K 访存)+ 2MN(整个 C 矩阵读一次和写一次)。由于计算量固定(排除 Strassen),所以只能优化访存,使得乘法和加法运算达到处理器的极限性能,从而实现 MatMul 的最佳性能。


MatMul 分块


关于减少 MatMul 计算时的访存量,最有效的方法是对 MatMul 的计算进行分块,并将分块之后的数据保存在 CPU 的 Cache 中。


如下图所示,将 A 按照 mr 进行行分块,将 B 按照 nr 进行列分块,计算时将需要用到的分块保存在 CPU 的各级 Cache 中,从而极大的减少了内存的读写。



进一步细化上面的分块计算过程如下图所示,A 中的一个行块都要重复地和 B 中的每一个列块进行小分块的 MatMul,写入到 C 的小块中,为了使得分块尽量的大 (如上所述,能够减少内存访存量),Cache miss 率尽量低,因此需要根据 CPU 的 Cache 结构特点(速度:L1D>L2>L3,容量 L1D<L2<L3) 来分配 Cache 和数据块之间的存放关系,天元中进行的分配如下:


· 将下图中红色 (访问重复次数最多的 A 的行块,计算时需要的 B 的一个列块以及计算结果的 C 的小块) 部分都保存在 L1 中。


· 由于计算完每一个 C 的行块,都需要重复遍历一次整个 B 矩阵,因此将 B 存放在 L2 中使得每次读取 B 的一个列块都是从 L2 中读取,代价最小。


· 将重复访问率最高的 C 的累加中间结果保存在速度最快的 CPU 处理器的寄存器中。



通过上面的分配策略,并结合 CPU 中资源(寄存器数量,L1D 和 L2 的大小),便可以确定最佳的 MatMul 计算中的 Nr,Kr:


· 可以根据 CPU 处理器的寄存器数量得到 mr 和 nr 的具体大小,寄存器容量 > mr*nr


· 根据 L1D Cache 的大小结合 mr 和 nr 计算出 Kr,Kr=L1D/(mr+nr)


· 再根据 L2 的大小计算出 B 矩阵中的 Nr,Nr=(L2-L1D)/Kr


在上面计算 N 时,用 L2-L1D 的原因是,由于当前 CPU 使用的 Cache 是 Inclusive 的,且 L2 是指令缓存和数据缓存合并的,另外上面 M 没有明确限制。


在得到上面最佳 Nr、Kr、mr 和 nr 之后,进一步便可以首先对 MatMul 计算中的 N、K 进行 Nr 和 Kr 分块,然后在 Nr、Kr 的基础上再进行 mr 和 nr 分块。如此,MatMul 计算中的计算访存比达到最高,且 CPU 处理器的资源也得到充分利用。


经过分块之后,由于在最内层的计算 Kernel 为 A(mrKr)B(Krnr)=C(mrnr)的分块矩阵乘,决定了整个 MatMul 的计算性能,因此需要极致地优化。

Kernel 计算

最核心的计算 Kernel 进行的是尺寸为 (mrKr)x(Krnr)–>(mr*nr) 的小尺寸 MatMul,其计算示意图如下:



如上图所示,Kernel 在计算时会读取 A 中一列, B 中一行,进行矩阵乘,得到 大小为 mr*nr 的 C,然后和原来 C 中的值相加,如此循环 Kr 次,完成该 Kernel 的计算。


在该 Kernel 的计算过程中乘法和加法的计算量为 2mrnrKr,访存量为 (mr+nr)Kr+2mrnr,可以根据处理器来判断该计算能否隐藏数据的访存。下面以 ARM cortex A76 为例进行分析,根据 A76 的数据手册得到:


· FP32 SIMD Load throughput=2,即单周期可以 load 8 个 float 数据


· FP32 SIMD FMLA throughput=2,即单周期可以进行 16 个乘加运算


因此,当 Kernel 的尺寸 mr=8、nr=12、Kr=256,计算量为 49152 次乘加运算,访存量为 5312 个 float 数据时,该计算访存量为 9.25,大于处理器的计算访存比 2。因此可以得出结论,如果 A 和 B 均在 L1 中,则该 Kernel 的计算不会因为数据的 Load 被阻塞,所以计算单元能够发挥出处理器的最佳性能。


虽然在对 MatMul 进行分块时,已经计划将 A 和 B 这样的小矩阵置于 L1D Cache 中,但是数据在真正运行时却不一定都在 L1D 中,原因在于,B 矩阵的列块在原来的大矩阵中内存并不连续,其次 A 中的一列由于内存不连续,也不能使用 SIMD 进行 load,为此需要对 A 和 B 进行数据 PACK。

MatMul 数据 PACK

上文 kernel 在计算过程中,需要同时读取 A 矩阵 mr 行数据,而每行数据之间内存不连续,因此这种访问对 Cache 不友好,同样,在读取矩阵 B 中 nr 列的时候也存在数据读取不连续的问题,加之 A 的所有行块和 B 中的所有列块将被读取多次,因此可以通过对 A 和 B 提前进行数据 PACK, 实现在后续计算中频繁多次访问数据时是连续的(只在第一次 PACK 时对 Cache 不友好),进而获得巨大收益。



对矩阵 A 进行数据 PACK 是将 A 中 mr 行数据的相同列拷贝到一起,如上图中将 A PACK 到 A’ 的步骤。重复完所有 A 中的行块便完成了 A 矩阵的数据 PACK。B 矩阵的 PACK 操作是,将 nr 列数据拷贝到连续的内存地址中,它对应上图 B PACK 到 B’ 的过程。

实 验

按照文介绍方式方式,天元在 X86 和 ARM 上分别对 MatMul 进行了优化。下表展示了 ARM64 上的性能测试结果,实验平台为 kirin 980。


首先,对该处理器进行分析可以看到,其主频为 2.6 GHz,每个周期能够进行 16 次乘加计算,因此其理论计算峰值为 16*2.6=41.6 Gflops。



可以看到,经天元优化的 MatMul 计算,发挥出了该处理器 90% 以上的计算性能。

总 结

本文以 Batch Norm 为例介绍了推理计算图的具体实现,以及 MatMul 在 CPU 上的优化细节。作为 CPU 推理优化的基石,最优的推理计算图是实现高性能 CPU 推理的前提条件,极致性能的 MatMul 计算基础算子将为实现卷积计算中的 Im2col 和 Winograd 提供性能保障。


在后面的文章中,我们将在详细介绍卷积计算中 Im2col 和 Winograd 的优化细节。


参考文献


  1. Anatomy of High-Performance Matrix Multiplication

  2. Fusing batch normalization and convolution in runtime

  3. Arm® Cortex®-A76 Software Optimization Guide


作者简介:


陈其友,旷视研究院高级深度学习架构师。从事深度学习框架开发,深度学习推理优化加速,在 C++,计算机体系结构和高性能计算方面有深度研究。


了解更多信息可访问:


· MegEngine WebSite: https://megengine.org.cn


· MegEngine GitHub(欢迎 Star): https://github.com/MegEngine


2020-08-14 09:102428
用户头像
刘燕 InfoQ高级技术编辑

发布了 1112 篇内容, 共 570.9 次阅读, 收获喜欢 1979 次。

关注

评论

发布
暂无评论
发现更多内容

【IT运维小知识】安全组是什么意思?

行云管家

网络安全 安全 数据安全 安全组 IT运维

三轮融资笔笔过亿,来画如何实现弯道超车?

ToB行业头条

阿里最新春招面经,腾讯/美团/字节1万道Java中高级面试题

收到请回复

Java 程序员 面试 职场 后端

如何赋能APaaS平台应用管理员

明道云

第 24 章 -《Linux 一学就会》- shell中色彩处理和awk使用技巧

学神来啦

Linux linux运维 linux云计算 linux一学就会

给JSONObject添加自定义遍历方法

FunTester

单元测试 Fastjson JSON库 spock FunTester

一文详解从云计算到 Serverless

Serverless Devs

云计算 Serverless

IOS技术分享| ARCall视频通话重构

anyRTC开发者

ios 音视频 语音通话 视频通话 视频呼叫

业内首款云原生技术中台产品云原生 Stack 来了!

阿里巴巴云原生

阿里云 云原生 技术中台

官方线索 | 阿里云1024程序员创造营

穿过生命散发芬芳

1024我在现场

网易云音乐音视频算法的 Serverless 探索之路

阿里巴巴中间件

云计算 阿里云 Serverless 云原生 中间件

用建造者模式实现一个防SQL注入的ORM框架

Tom弹架构

华为云网络覆盖全球2500+站点,打造高品质、低成本接入体验

华为云开发者联盟

RTC 华为云 媒体服务 LiveVideoStack

阿里二面:为什么要分库分表?

Java MySQL 数据库 架构 面试

英特尔全面升级开发者专区(Developer Zone),携手开发者共塑科技未来

科技新消息

短短 146 天就成为 Apache APISIX Committer,我是怎么做到的?

API7.ai 技术团队

成长笔记 API网关 社区 Apache APISIX

安全漏洞的复现与总结

网络安全学海

网络安全 信息安全 渗透测试 WEB安全 漏洞挖掘

英特尔On技术创新峰会公布开发者重点投入计划,全面赋能开发者

科技新消息

直面Z级计算挑战,英特尔携手生态合作伙伴加速开发者工具和软件部署

科技新消息

[2021年新鲜出炉]K8s工程师资料合辑,书籍推荐,面试题,精选文章,开源项目,PPT,视频,大厂资料

奔着腾讯去

golang Docker Kubernetes 容器 云原生

递推算法与递推套路(手撕算法篇)

有道技术团队

学习 编程 算法

瑜伽CRM系统使老师再也不用教学分心

低代码小观

CRM 管理工具 CRM系统 瑜伽 瑜伽馆管理

医药研发企业数据中台的选择

鲸品堂

产业互联网

Linux指令日志分析(一)

正向成长

日志分析

iOS 崩溃日志在线符号化实践

百度Geek说

企业级开发工具,J2PaaS低代码平台核心能力解析!

J2PaaS低代码平台

软件开发 低代码 低代码开发 低代码平台

堪称神级的Spring Boot手册,从基础入门到实战进阶

Java 程序员 Spring Boot

网易云音乐音视频算法的 Serverless 探索之路

Serverless Devs

阿里云 Serverless 算法 音视频

先行一步,7 大技术创新和突破,阿里云把 Serverless 领域的这些难题都给解了

Serverless Devs

阿里云 gpu VPC 函数计算

这一篇 K8S(Kubernetes)我觉得可以了解一下!!!

牧小农

Docker 容器 kubenetes

Redis核心原理与实践--Redis启动过程源码分析

binecy

redis 源码阅读

CPU推理性能提高数十倍,MegEngine计算图、MatMul优化解析 | 工程之道_AI&大模型_陈其友_InfoQ精选文章