太极平台高效支持MKLDNN动态shape推理

| 导语  TIMO（Taiji Model Optimization）是太极机器学习平台旗下推出的一套推理加速方案，针对模型服务平台提供了从数据预处理到模型推理的一站式优化。为支持不同型号X86-CPU例如：AMD系列、Intel系列的高效推理，TIMO的CPU后端引擎基于MKLDNN进行了二次开发和优化，解决了MKLDNN在动态shape推理性能不佳的缺陷，并同时保留MKLDNN接口，不但兼容其原生计算库生态，而且性能相对于开源版本有明显提升。本文主要介绍TIMO选用MKLDNN作为计算库的原因，并阐述MKLDNN的动态shape推理优化方案，最后展示人脸识别等业务中的优化效果。

一、背景

1. 为什么采用MKLDNN

    太极机器学习平台模型服务场景和算法结构复杂多样，支撑CV、NLP、语音以及推荐等任务，推理加速计算的通用性受到重大挑战:（1）算法模型部署框架未统一，包括tensorflow、torch、kaldi、caffe和onnx-runtime等；（2）需要支持不同型号X86-CPU硬件如AMD系列、Intel系列。MKLDNN是Intel针对X86-CPU开发的高性能神经网络计算库, 在TF/Torch/Onnx-runtime/caffe等推理框架中都有对应版本集成，高效推理Convolution/Transormer等框架中的算子，并支持多种型号CPU(Intel-E5/Intel-Gold/Intel-8255C/AMD-EPYC等)，所以TIMO采用MKLDNN作为CPU后端计算库之一，可以较好节省开发成本以及实现多种框架适配。

2. 动态shape推理性能不佳

    其他神经网络计算库如Mblas, Openblas等仅针对Micro-kernel矩阵乘采用汇编优化，更为宏观粒度的计算优化则交给编译器。而MKLDNN则基于xybak(a JIT assembler)编写，除了在Micro-kernel循环层用汇编优化，在更高循环层也采用汇编优化，可以更好控制指令流水调节，寄存器分配等，因而优化的空间和上限更高。事实上，在输入为固定shape情况下测试卷积和GEMM等算子，MKLDNN相对比其他计算库一般性能更为优异。但由于计算资源如物理寄存器数目等是有限的，MKLDNN竭尽所能将分块尺寸或循环次数等kernel参数存为立即数，由输入shape计算所得的部分参数也会被固化为立即数。在静态shape的输入情况下，编译生成的kernel可执行二进制存到kernel cache，实际运行中会根据shape等查询对应的kernel，但对于动态shape数据，MKLDNN则需对每一种shape输入重新编译生成可执行代码，这样会引入的额外编译开销(create_primitive)，从而影响推理性能。

3. 高效支持MKLDNN的动态shape推理

  基于上述分析，TIMO在MKLDNN基础上进行扩展，对推理框架OpenVINO进行优化适配，并在底层开发一套支持动态shape推理 kernel库，无需重新编译就能即时推理，最后在基于TIMO自研的OpenVINO改进版本集成使用，已应用于多个业务场景中。

二、动态shape推理

1. 动态shape 推理框架

    TIMO的动态shape推理首先针对视觉计算场景优化，在静态shape的框架基础上，扩展出动态Batch和动态分辨率两种动态shape场景， 前者保证图片NCHW维度中N维度可变而H和W维度不可变，后者保证H、W维度可变而N维度不可变，整体框架流程示意图如下：

  其中静态shape的场景仅含离线流程，而动态shape场景包括在线和离线流程。离线流程只需执行一次，包含：模型解析，图优化，算子拓扑排序，静态内存分配，primitive kernel编译等步骤；在线流程需要针对每一个输入数据即时处理，主要包括：shape推断和primitive descriptor更新，kernel selection和reset kernel等步骤；以下重点介绍在线推理流程：

在线推理流程

（1）shape Inference ，即形状推断，作用是通过输入形状计算出输出形状，并初始化算子输出对应的memory_dims, 如下图卷积的shape Inference公式如下：

（2）Reset  Primitive Descriptor，即MLKLDNN的核函数修饰器更新，通过Batch/OH/OW/IH/IW等计算kernel计算的相关参数，如图所示关于OH和OW相关的分块和子分块参数以及计算循环次数被重新计算，并存入指定内存变量中。

1x1 conv分块图

 (3) Select Optimal Kernel，即核函数选取，作用是在编译器阶段会缓存同一个算子相关的多种kernel算法，例如3x3卷积算子的算法实现有三种分别为：winograd kernel，direct kernel以及im2col kernel，其中每种kernel算法的实现在不同shape的参数下性能表现不同。编译过程中，遍历编译每一种kernel，而再实时计算阶段，则会预测目标开销(如乘法次数)，选择性能最优的kernel。

(4) Reset Primitive Kernel，即核函数重置，在(2)步骤中计算所得的primitive descriptor作为参数传递入选取的primitive kernel，将OH/OW/IH/IW/Batch推导出kernel参数推到的算法参数传递给kernel的寄存器，无需重新编译。

2. 动态shape Kernel

 动态shape kernel，指对不同shape输入的数据不需重新编译实现精度性能无损推理，CV场景中动态shape计算包括动态batch和动态分辨率两种场景，调用逻辑和相关指定参数如下所示：

   动态shape kernel核心思想是：使用寄存器替代立即数存储shape变化相关的算法参数，从而不需要重新编译，一般采用两种方法：

（1）使用寄存器装载，在实际运行中不断将相关参数拷贝到固定内存地址中，在相关作用域范围内再从该固定地址将参数加载到寄存器中计算；

（2）在一定范围内遍历整数，对比该整数和寄存器存的值，然后将该整数直接作为参数传入函数中。

寄存器替代计算

遍历整数法

    从代码重构工作量和性能考虑：前者适用于数值较大，且参数传入函数调用栈较浅的情况；后者更适用于可选数值较少，函数调用栈较深的情况。使用寄存器替代立即数或者遍历整数法情况下理论上会增加访存或逻辑计算的开销，但实际增加计算开销极小，可以忽略不计。

3. 快速上手

   TIMO在MKLDNN的原生API接口基础上扩展reset_primitive_desc和reset_primitive两个函数：

 以1x1 conv卷积示例，实际运行中各函数运行时间开销结果如下：

   从表格中函数统计结果来看，使用动态shape优化前后，扩展接口reset_primitive_desc和reset_primitive时间开销相对于原接口primitive_desc和primitive分别减少至25%和5%左右，这样jit编译时间开销极大减小，而卷积推理执行的时间和精度保持不变。

4. 业务测试效果

  将MKLDNN动态shape推理库集成到内部改进OpenVINO动态shape版本中(我们在2021.11.1已经支持OpenVINO的动态shape推理，而官方直到OpenVINO-2022.1.0版本才在框架层面支持动态shape，但是目前官方对于动态shape输入数据仍需即时编译MKLDNN)，调用代码如下

    人脸识别等业务中方案已上线验证。人脸识别算法主要分为两个步骤：分别为人脸检测+特征提取 ；人脸检测模型为MTCNN，需要动态shape的图片输入，由PNET、RNET、ONET三个CNN网络串联组成，所以算子数目较多，网络较深，编译开销较大。整体pipeline优化前后(AMD EPYC 7K62 48-Core Processor)的结果如下：

Mtcnn + vino_face inference  时间对比

  动态shape相对于原版性能提升明显，另外动态shape的MKLDNN对比Onnx-runtime在1080p图片测试也有50%的性能提升。

三、结语

    这篇文章主要讲述TIMO对MKLDNN的动态shape推理功能开发和优化，在单个算子和整个模型上性能表现优异。当前已经在人脸识别和视频相似度分析等业务中上线验证。团队在服务框架、推理框架、GPU/CPU计算和编译优化方面做了大量工作，当前在CNN、NLP和语音等多种场景已落地，并取得良好的效果。有需要看机会的可发送简历至waynexrli@tencent.com。

关注【腾讯太极机器学习平台】公众号

获取更多内容~