Halide应用开发

Halide应用开发

1. 基本原理

1.1、介绍

随着人工智能的普及，深度学习网络的不断涌现，为了让各硬件(CPU, GPU, NPU,...)能够支持深度学习应用，各硬件芯片需要软件库去支持高性能的深度学习张量运算。目前，这些高性能计算库主要由资深HPC工程师(高性能计算优化工程师）进行开发，为了加快开发进程，缩短深度学习应用落地周期，自动化算子优化是一个趋势。

AutoKernel是由OPEN AI LAB提出的高性能算子自动优化工具，可以自动优化调度策略、生成底层优化代码，大幅减少各硬件芯片算子开发成本，提升算子优化效率，让工程师更快实现深度学习算法在各硬件芯片上的高性能部署。

1.2、 AutoKernel架构

AutoKernel分为3个模块：

算子生成器

该模块使用了开源项目Halide；Halide是业界广泛使用的自动代码生成项目，它首次提出将计算和调度分离。该模块的输入是和硬件无关的算子计算描述，输出是相应后端的优化汇编代码/目标文件；

自动搜索模块

该模块可以通过最优化算法/搜索算法/机器学习/强化学习搜索出相应后端的最优算子的调度策略参数（该模块仍在开发中）；

算子部署插件(AutoKernel Plugin)

Tengine是OPEN AILAB开源的深度学习推理框架，实现了AI算法在不同硬件的快速高效部署。该模块实现了将自动生成的优化算子代码以plugin的形式一键集成到Tengine中，实现自动优化算子的一键部署；

1.3、 AutoKernel特点

低门槛: 无需底层优化汇编的知识门槛

简单易用: 提供docker环境，无需安装环境，plugin一键集成到推理框架Tengine

高效率: 无需手写优化汇编，一键生成优化代码，一键部署

2 AutoKernel优化GEMM实践

本部分将带领大家一步步优化矩阵乘法GEMM。无需手工撸代码，编写繁杂冗长的底层汇编代码，只需十几行简洁的调度代码。

2.1、优化的本质

在详细讲解优化步骤前，先谈谈优化的本质。在谈”优化“的时候：

计算机底层做了什么？

优化的”瓶颈“是什么？

为什么通过一波"优化操作"，性能就能提升呢？

AutoKernel使用的Halide是如何实现自动优化的呢？

上图是典型的存储理器层次结构：主存容量大，访问速度慢，寄存器和缓存读取速度快，但容量有限。在寄存器的层级上，CPU可以在一个时钟周期内访问它们，如果CPU去访问外部的DDR的话，延迟是非常大的，大概是200个时钟周期左右。如果CPU去访问cache的话，一般需要6到12个cycle就够了。所以，第1个优化宗旨是：优化内存访问，充分利用寄存器和高速缓存去存数据。

第2个优化宗旨则是提高并行性：充分利用SIMD进行指令向量化和多核心并行。大部分现代CPU支持SIMD（Single Instruction Multiple Data，单指令流多数据流）。在同一个CPU循环中，SIMD可在多个值上同时执行相同的运算/指令。如果在4个数据点上进行向量化，一次计算四个数据，理论上就可以实现4倍的加速。

2.2、运行环境搭建

AutoKernel提供了docker镜像，docker里已经配置好运行环境，进入docker即可直接运行demo代码：

# 拉取镜像

docker pull openailab/autokernel

# 启动容器，进入开发环境

docker run -it openailab/autokernel /bin/bash

# 获取代码

git clone https://github.com/OAID/AutoKernel.git

cd AutoKernel/doc/tutorials/data/

目录下的build.sh是demo的执行脚本，运行需要指定优化步骤step，可选的step是从1 到7，其中step=1是默认不优化的，step=7是最极致优化的。

执行指令：

# 执行demo

./build.sh 1  ==> 默认不优化

./build.sh 7  ==> 最极致优化

下图展示了在Intel(R) Core(TM) i9-9900K CPU @ 3.60GHz的电脑上的优化效果，无需手工撸代码，无需编写繁杂冗长的底层汇编代码，只需十几行简洁的调度代码, 就能性能优化200+倍~

2.3、详细的优化步骤：

STEP-1

第1个步骤是不带任何优化的。用Halide语言直接描述GEMM的计算过程。

Var x,y;  

RDom k(0, K);  

Func gemm("gemm"); 

gemm(x, y) += A(k, y) * B(x, k);  

计算M=N=K=640的矩阵乘法。运行脚本第一个参数指定step=1。耗时结果如下：

root@xxxxxxxxxx:/AutoKernel/doc/tutorials/data# ./06_build.sh 1

step =  1

M N K = 640 640 640     err 0.00        [rep 50] autokernel | blas      240.8523 ms     1.1376 ms

STEP-2

这一步采用分块tile。分块的目的是为了充分利用缓存。如果原来的循环较大，tile分块改成小块数据去计算，可以使得每次计算的数据都比较舒适地呆在缓存里，不用经历重复的驱逐（在缓存中重复的添加和删除数据)。分块后进行reorder操作，交换两个嵌套循环的顺序，目的是最内层的内存访问友好。按照x,y维度划分成16x8的小分块去计算:

gemm.update()  

  .tile(x, y, xo, yo, xi, yi, 16, 8)  

  .reorder(xi, yi, k, xo, yo);   

执行结果如下：

root@xxxxxxxxxx:/AutoKernel/doc/tutorials/data#./06_build.sh 2

step =  2

M N K = 640 640 640     err 0.00        [rep 50] halide | blas  81.8148 ms      1.1281 ms

性能从240ms优化到82ms，提升了近3倍。

STEP-3

在上一步的基础上增加向量化vectorize。向量化是把几个标量计算（scale)转换为一个向量计算（vector),充分利用SIMD向量指令。大部分现代CPU支持SIMD（Single Instruction Multiple Data，单指令流多数据流）。在同一个CPU循环中，SIMD可在多个值上同时执行相同的运算/指令。

gemm.update()  

  .tile(x, y, xo, yo, xi, yi, 16, 8)  

  .reorder(xi, yi, k, xo, yo)  

  .vectorize(xi, 8);  

执行结果

root@xxxxxxxxxx:/AutoKernel/doc/tutorials/data# ./06_build.sh 3

step =  3

M N K = 640 640 640     err 0.00        [rep 50] autokernel | blas      27.5433 ms      1.1445 ms

性能从82ms优化到27ms，又加速了接近3倍。可以看到，围绕前面提到的两条优化宗旨：优化内存访问和提高并行性，从step1到step3，性能已经提升了近9倍。

STEP-4

调度策略在step3的基础上增加并行化parallel。对一个循环并行化是把循环的每次迭代分给多个线程或者处理器去同时处理，每个线程处理通过代码段（loop body),但是处理不同的数据。

增加并行化后，build.sh默认指定四线程，性能直接翻了近4倍，从27ms到7.3ms.

STEP-5

调度策略在上一步的基础上增加unroll展开。如果循环体内的语句没有数据相关依赖，循环展开可以增加并发执行的机会，使得更充分利用寄存器，减少循环时每个操作内存加载和保存的次数。

unroll展开后，性能从7.3ms优化到4.8ms.

STEP-6

前面的分块成 16 x 8的小kernel,
这一步先划分成 16 x 32的分块，然后把每个分块再分成
16 x 8的子分块。把最外层的两层循环合并到一层，并对这一层进行并行化。这一步计算描述多了一个prod函数来定义子分块的计算，prod函数的计算公式和总的gemm是一样的，通过 compute_at指定在 yi维度之下计算prod，则prod计算的是 16x8的小kernel, 大致逻辑如下：

这一步距离STEP1性能已经优化了近80倍了，性能越来越接近OpenBlas了。

STEP-7

这一步添加的操作是对矩阵B进行数据重排，使得在计算小kernel 16x8时，内存读取更顺畅。因为小kernel的x维度是按照16划分的，因此重排数据B的x维度也是按照16重排。

至此，的每一步调优策略始终都围绕两条优化宗旨“优化内存访问”，“提高并行性”展开优化，到最后性能已经与OpenBlAS差不多了，距离STEP1已经加速了200+倍了。