[发明专利]面向张量计算单元卷积算子优化实现方法

说明书

一种面向张量计算单元卷积算子优化实现方法，通过深度学习编译器的DSL表示卷积算子，经对卷积计算进行坐标变换得到隐式通用矩阵乘法的计算表示；然后对卷积算子进行调度优化得到调度模板后，经搜索得到最优搜索参数并通过深度学习编译器的后端生成CUDA C代码，再将生成的CUDA C代码集成入神经网络，实现卷积神经网络在NVIDIA GPU平台上的推理速度提升。本发明能够提升半精度计算中卷积算子自动代码生成的性能，为神经网络推理计算中融合算子的自动代码生成的性能提供保证。

技术领域

本发明涉及的是一种人工智能领域的技术，具体是一种面向NVIDIA GPU的张量计算单元(Tensor Core)卷积算子优化实现方法。

背景技术

深度学习编译器技术被提出以提升深度学习算子的研发效率，其中TVM编译优化技术能够有效的解决融合算子的自动代码生成问题。然而TVM生成的算子性能，极度依赖于算子调度策略的开发以及调度空间的设计。实际深度学习的业务场景要求TVM生成的算子性能和手工优化一样发挥硬件的极致性能。

现有的通用矩阵乘(GEMM)运算加速技术应用领域较为狭窄，或无法进行卷积算子的计算或无法运行融合算子，现有的改进GEMM处理技术则基于汇编指令集的操作编程难度较大，也不具备跨平台的通用性。

发明内容

本发明针对现有技术存在的上述不足，提出一种面向张量计算单元卷积算子优化实现方法，通过领域领域特定语言(DSL)编写卷积算子的计算与面向张量计算单元的调度，然后通过自动调优技术生成卷积计算代码，本发明能够提升半精度计算中卷积算子自动代码生成的性能，为神经网络推理计算中融合算子的自动代码生成的性能提供保证。

本发明是通过以下技术方案实现的：

本发明涉及一种面向张量计算单元卷积算子优化实现方法，通过深度学习编译器的DSL表示卷积算子，经对卷积计算进行坐标变换得到隐式通用矩阵乘法的计算表示；然后对卷积算子进行调度优化得到调度模板后，经搜索得到最优搜索参数并通过深度学习编译器的后端生成CUDA C代码，再将生成的CUDA C代码集成入神经网络，实现卷积神经网络在NVIDIA GPU平台上的推理速度提升。

所述的对卷积算子进行面向Tensor Core的调度优化，得到调度模板，具体包括：

步骤1)对输入数据、权重数据以及输出结果分别进行局部存储器(sharedmemory)和张量寄存器(wmma fragment)的缓存。

步骤2)利用半精度浮点类型存储空间小的特点，对缓存读写步骤进行双缓冲调度优化。

步骤3)利用半精度浮点类型读写带宽高的优势，对缓存读写步骤进行向量化调度优化。

步骤4)对计算维度进行切分，即将GEMM_M维度切分为bm,tm,om,im四个维度；将GEMM_N维度切分为bn,tn,on,in四个维度；根据切分参数im、in将GEMM_K切分为ok,ik两个维度。

步骤5)进行GPU的线程块绑定和线程绑定，即分别将切分出的维度bm绑定至blockIdx.y、tm绑定至threadIdx.y、bn绑定至blockIdx.x、in绑定至threadIdx.x。

步骤6)将未绑定的计算维度，映射为wmma::mma_sync表示的GEMM计算；经绑定后实现每个线程块将计算om*on个GEMM，计算的结果尺寸为im*in。

所述的切分，其参数为可搜索的参数模板，通过深度学习编译器中的自动搜索方式得到最优的切分参数。

为了参数搜索的高效性，本发明针对计算尺寸和硬件架构的特点，对切分参数im、in的搜索空间与搜索方法进行了特殊的设计，具体包括：

①当GEMM_M为32的倍数时，设置im的搜索空间为[8,16,32]并转到步骤④；否则转到步骤②。