每日视讯：深入浅出GPU优化系列：GEMM优化（二）

时间：2023-06-15 08:34:20 来源: 面包芯语

编辑：王佩

这篇文章的主要目的是让读者尽可能地理解GEMM优化的技巧，所以代码尽可能地通俗易懂。在兼顾易读性的同时，代码有不错的性能表现，在1024-4096的矩阵上平均能够达到cublas 91%的性能表现。

本篇文章主要分四部分进行讲述。第一部分是GEMM算法的概述，先让大家大概地了解代码在做什么事情。第二部分就是详细的代码解析，带着大家一行一行地来看。第三部分是实验结果，主要是和cublas的对比。第四部分是总结。

(资料图片)

一、GEMM算法概述

这个章节里主要来说一下GEMM的一个计算流程，其实这一点已经在GEMM优化（一）中提及。但上一篇文章主要说得是原理，关于具体计算逻辑，还是不太直观，所以我们在这里再提一下。然后这个具体的计算逻辑分为两个阶段介绍，分别是不采用数据预取和采用数据预取，这主要是考虑到直接说数据预取，有读者可能会看得云里雾里，比较难受，所以先把不采用数据预取这个内容说明白，然后再来讲这个数据预取。

1.1 不采用数据预取

随后再具体看看每一个大迭代中，block中的线程的计算逻辑。在进行一个大迭代时，shared memory中有128×8=1024个A矩阵元素和8×128=1024个B矩阵元素。随后，每个线程需要进行8次迭代，我们把这个迭代成为小迭代。bk=8，所以有8次小迭代。每一次小迭代中，每个线程需要从shared memory中拿到A矩阵的一小列和B矩阵的一小行，即8个A的元素和8个B的元素。线程将这8+8=16个元素放置在寄存器中。每个线程需要负责8×8=64个元素的计算，一共会产生64条FFMA指令。小迭代示意图如下：

1.2 采用数据预取

二、GEMM代码解析

2.1 参数说明

template <    const int BLOCK_SIZE_M,  // height of block of C that each  block calculate    const int BLOCK_SIZE_K,  // width of block of A that each  block load into shared memory    const int BLOCK_SIZE_N,  // width of block of C that each  block calculate    const int THREAD_SIZE_Y, // height of block of C that each thread calculate    const int THREAD_SIZE_X,  // width of block of C that each thread calculate    const bool ENABLE_DOUBLE_BUFFER // whether enable double buffering or not    >

// shared memory    __shared__ float As[2][BLOCK_SIZE_K][BLOCK_SIZE_M];    __shared__ float Bs[2][BLOCK_SIZE_K][BLOCK_SIZE_N];    // registers for C    float accum[THREAD_SIZE_Y][THREAD_SIZE_X] = {0};    // registers for A and B    float frag_a[2][THREAD_SIZE_Y];    float frag_b[2][THREAD_SIZE_X];    // registers load global memory    const int ldg_num_a = BLOCK_SIZE_M * BLOCK_SIZE_K / (THREAD_NUM_PER_BLOCK * 4);    const int ldg_num_b = BLOCK_SIZE_K * BLOCK_SIZE_N / (THREAD_NUM_PER_BLOCK * 4);    float ldg_a_reg[4*ldg_num_a];    float ldg_b_reg[4*ldg_num_b];

// threads number in one row    const int A_TILE_THREAD_PER_ROW = BLOCK_SIZE_K / 4;    const int B_TILE_THREAD_PER_ROW = BLOCK_SIZE_N / 4;    // row number and col number that needs to be loaded by this thread    const int A_TILE_ROW_START = tid / A_TILE_THREAD_PER_ROW;    const int B_TILE_ROW_START = tid / B_TILE_THREAD_PER_ROW;    const int A_TILE_COL = tid % A_TILE_THREAD_PER_ROW * 4;     const int B_TILE_COL = tid % B_TILE_THREAD_PER_ROW * 4;    // row stride that thread uses to load multiple rows of a tile    const int A_TILE_ROW_STRIDE = THREAD_NUM_PER_BLOCK / A_TILE_THREAD_PER_ROW;    const int B_TILE_ROW_STRIDE = THREAD_NUM_PER_BLOCK / B_TILE_THREAD_PER_ROW;

2.2 大迭代前预取数据

迭代前预取数据分为两个部分，第一个部分是将第一个大迭代的数据从global 预取到shared memroy中。第二个部分是将shared memory上的数据预取到寄存器中。先来看看第一个部分。这里面分别是将第一个大迭代中需要的A、B数据预取到shared memroy中。对于A矩阵而言，这个for循环代表着block中的线程需要搬运多少次才能将globa中的数据放到shared memory中。由于A需要先进行一次转置，所以先将数据先放置在寄存器中。数据按行取，然后按列存。对于B矩阵而言，数据不用转置，直接按行取，按行存。当然，这个过程中间也要经过寄存器，但是没有写出来的必要了。

然后就是第二个部分。将shared memory中的数据存到寄存器中。一共需要取THREAD_SIZE_Y个数，每次取4个数。这个倒没有什么好说的。

2.3 大迭代逻辑

2.4 大迭代详细解析

随后进入到小迭代的计算逻辑之中，load_stage_idx参数代表需要从As的哪个空间进行读数。然后是BLOCK_SIZE_K-1次小迭代。按照前面的参数配置，即需要在这里完成7次小迭代。由于在小迭代中也采用了双缓冲的方式，需要将下一轮小迭代的数据提前写入到寄存器中，这个过程需要对shared memory访存，会稍微慢点。与此同时，线程需要计算更新THREAD_SIZE_X x THREAD_SIZE_Y=8×8=64个C矩阵元素的结果。

而后需要将存储在临时寄存器的数据搬运到shared memory中。由于A矩阵需要经过一次转置，所以和B矩阵有一点不一样。

最后完成寄存器的预取，并将最后一个小迭代完成。

2.5 计算结果写回

三、实验

2.bm、bn、bk、rm、rn等相关参数对GEMM的性能表现有多大影响？

最后代码链接在下方：

《GPU优化教程系列》是澎峰科技收集、整理、创作的一个公益系列课程，已经完成的课程如下：

标签：

每日视讯：深入浅出GPU优化系列：GEMM优化（二）

一、GEMM算法概述

1.1 不采用数据预取

1.2 采用数据预取

二、GEMM代码解析

2.1 参数说明

2.2 大迭代前预取数据

2.3 大迭代逻辑

2.4 大迭代详细解析

2.5 计算结果写回

三、实验

精彩推送

韩国：2025年起将AI引入小中高课程到2028年实现全面覆盖

首台套50万千瓦冲击式水电机组项目开工

每日热点：希捷4TB机械硬盘史低仅售288元

九首歌180分钟未删减版百度网盘（九首歌180分钟未删）

【全球新视野】【夏收现场】知识落“地”助丰收耕地质量提升示范田小麦增产53.6%

高校食堂“鼠头鸭脖事件”背后：10万能拿下一个档口？

“文根”何以“深种”？看文博会如何寻找岭南新坐标

5月国内挖掘机销量降46％，出口增速重回两位数｜基建晴雨表_每日讯息

天天播报:边缘计算领域步入高速发展阶段 5月份以来超百家机构扎堆调研7家公司

新闻快讯

精品推荐

新闻快讯