[object Object]
Cube编程范式把算子的实现流程分为5个基本任务:CopyIn,Split,Compute,Aggregate,CopyOut。CopyIn负责搬入操作,Split负责数据切分操作,Compute负责矩阵指令计算操作,Aggregate负责数据汇聚操作,CopyOut负责搬出操作。
图 1 矩阵编程基本任务设计[object Object][object Object]
具体任务之间的交互流程和流程图如下。
Stage1:CopyIn任务。
Stage2:Split任务。
Stage3:Compute任务。
Stage4:Aggregate任务。
Stage5:CopyOut任务。
图 2 矩阵编程Queue队列[object Object][object Object]
基于Ascend C方式实现矩阵算子的流程如下图所示。
图 3 矩阵算子实现流程[object Object][object Object]
- 算子分析:分析算子的数学表达式、输入、输出以及计算逻辑的实现,明确需要调用的Ascend C接口。
- 核函数定义:定义Ascend C算子入口函数。
- 根据矩阵编程范式实现算子类:完成核函数的内部实现,调用私有成员函数CopyIn、SplitA、SplitB、Compute、Aggregate、CopyOut完成矩阵算子的五级流水操作。
下文将以Matmul算子为例对上述步骤进行详细介绍,Matmul算子的代码框架如下,完整代码请参见。
在开发算子代码之前需要分析算子的数学表达式、输入、输出以及计算逻辑的实现,明确需要调用的Ascend C接口。
明确算子的数学表达式及计算逻辑。
Matmul算子完成矩阵乘操作,其数学表达式如下,形状为[m, k]的矩阵a和形状为[k, n]的矩阵b相乘,得到形状为[m, n]的矩阵c。为了方便,令m=k=n=32。
[object Object]注意需要处理的数据过大时,需要对数据进行切分并分块搬运到A2、B2,分别计算后再进行汇聚。下文的计算逻辑为了展示Split和Aggregate阶段的样例,请您根据实际需要处理的数据大小决定是否需要切分和汇聚。
计算逻辑如下:
- 分别搬运输入数据矩阵a、b至Local Memory A1、B1。
- 将a矩阵从A1搬运至A2。为实现部分并行,将b矩阵切分为part1和part2,形状均为[k, n / 2],切分后再分块搬运至B2。
- a矩阵和b矩阵part1、part2分别做矩阵乘运算,获得矩阵c的part1和part2,形状均为[m, n / 2]。计算结果在CO1存储。
- 将矩阵c的part1和part2分别拷贝到CO2进行合并。
- 将合并后的输出数据从CO2搬出。
明确输入和输出。
- Matmul算子有两个输入:a与b,输出为c。
- 本样例中算子输入支持的数据类型为half(float16),算子输出的数据类型为float32。
- 矩阵a、b、c的形状均为[32, 32]。
- 算子输入输出支持的数据格式为:ND。
确定核函数名称和参数。
- 您可以自定义核函数名称,本样例中核函数命名为matmul_custom。
- 根据对算子输入输出的分析,确定核函数有3个参数a,b,c;a,b为输入在Global Memory上的内存地址,c为输出在Global Memory上的内存地址。
约束分析。
由于硬件架构对矩阵乘计算的输入输出有格式约束,需要在算子实现中增加格式转换的流程。
确定算子实现所需接口。
通过以上分析,得到Ascend C Matmul算子的计算流程图和设计规格如下:
图 4 Matmul算子的计算流程图[object Object][object Object]
表 1 Ascend C Matmul算子设计规格
[object Object][object Object]
[object Object]函数原型定义。
本样例中,函数名为matmul_custom(核函数名称可自定义);根据中对算子输入输出的分析,确定有3个参数a,b,c,其中a,b都为输入内存,c为输出内存。根据中核函数的规则介绍,函数原型定义如下所示:使用__global__函数类型限定符来标识它是一个核函数,可以被<<<>>>调用;使用__aicore__函数类型限定符来标识该核函数在设备端aicore上执行;为方便起见,统一使用GM_ADDR宏修饰入参,GM_ADDR宏定义请参考。
[object Object]调用算子类的Init和Process函数。
算子类的Init函数,完成内存初始化相关工作,Process函数完成算子实现的核心逻辑,具体介绍参见。
[object Object]对核函数进行封装,得到matmul_custom_do函数,便于主程序调用。#ifndef ASCENDC_CPU_DEBUG表示该封装函数仅在编译运行NPU侧的算子时会用到,编译运行CPU侧的算子时,可以直接调用matmul_custom函数。根据章节,调用核函数时,除了需要传入参数a,b,c,还需要传入numBlocks(核函数执行的核数),l2ctrl(保留参数,设置为nullptr),stream(应用程序中维护异步操作执行顺序的stream)来规定核函数的执行配置。
[object Object]
根据上一章节介绍,核函数中会调用算子类的Init和Process函数,本章具体讲解基于编程范式实现算子类。矩阵编程范式请参考。
算子类中主要包含对外开放的初始化Init函数和核心处理函数Process以及一些实现中会用到的私有成员。KernelMatmul算子类的定义如下:
KernelMatmul构造函数实现
构造函数中对私有成员变量进行初始化,具体代码如下:
矩阵a的形状为[m, k],矩阵b的形状为[k, n],矩阵c的形状为[m,n],此样例中m、n、k均设置为32。
aSize、bSize、cSize分别为矩阵a、b、c的数值个数。
mBlocks、 nBlocks、 kBlocks为m、n、k所占分形数量,half类型一个分形Shape为16 * 16,blocks计算公式为:
- mBlocks = m / 16
- nBlocks = n / 16
- kBlocks = k / 16
Init函数实现
Init函数主要完成以下内容:
设置输入输出Global Tensor的Global Memory内存地址。
以设置输入a在Global Memory上的内存偏移地址为例:
[object Object]注意,因为本样例中Init函数的入参统一设置为uint8_t*,这里需要强转成具体的数据类型(__gm__ half*),再进行偏移。
通过Pipe内存管理对象为输入输出Queue分配内存。
比如,为输入数据队列inQueueB2分配内存,可以通过如下代码段实现:
[object Object]此样例中将b矩阵切分为两个part,为inQueueB2分配内存时需要申请两块内存空间,每一块的大小为b矩阵大小的一半,outQueueCO1的内存初始化同理。
具体的初始化函数代码如下:
Process函数实现
基于矩阵编程范式,将核函数的实现分为5个基本阶段:CopyIn,Split,Compute,Aggregate,CopyOut。Split,Compute,Aggregate阶段需要区分a、b矩阵。Process函数中通过如下方式调用这几个函数。
两次循环内,SplitB需要从inQueueB1中分别搬运两个part的b矩阵,Compute需要分别计算a矩阵和两个part b矩阵的乘法,Aggregate要分别搬运两个part的c矩阵,具体五个阶段数据流通示意图如下:
图 5 数据流通示意图[object Object][object Object]
切分b矩阵,可以实现一部分的并行,本样例的流水并行示意图如下:
图 6 并行示意图[object Object][object Object]
Stage1:CopyIn函数实现。
使用接口将矩阵a、b搬运到Local Memory,同时将其数据格式从ND转换为NZ。
一次DataCopy指令搬运height*16个数,循环执行width/16次。DataCopy的参数设置如下:
- blockCount设置为height,共搬运height次。
- blockLen设置为1,一次搬运16个类型为half的数。
- srcStride设置为width/16 - 1,源矩阵每搬运一个block需要跳跃一行。
- dstStride设置为0,目的矩阵每个block在内存中连续存储。
- 每次循环迭代,源矩阵首地址移动16个数,目的矩阵首地址移动16*height个数。
格式转换示意图如下,第一次循环搬运蓝色部分,第二次循环搬运绿色部分;图中width为32,占两个分形,height为32,占两个分形,一共搬运4个16*16分形。
图 7 ND to NZ转换示意图[object Object][object Object]
注意:上述ND到NZ的格式转换仅作为举例说明,开发者可根据实际情况选择合适的转换方式。
具体代码如下:
[object Object]Stage2:SplitA函数实现。
使用将矩阵a搬运到A2,同时将a矩阵从NZ格式转换为ZZ格式。
搬运及格式转换示意图如下:图中k为32,占kBlocks(k/16=2)个分形,m为32,占mBlocks(m/16=2)个分形,一共搬运4个16*16分形。本示例中,调用一次LoadData接口完成两个16*16分形的搬运,循环调用两次LoadData。第一次循环搬运蓝色部分两个分形,第二次循环搬运绿色部分两个分形。
单次循环中LoadData(本样例中要完成2个分形的搬运,蓝色部分或者绿色部分)的参数设置如下:
- repeatTimes表示数据处理的迭代次数,因为LoadData每个迭代处理一个分形,所以也可以理解为待搬运分形的个数。本样例中即为k轴方向的分形个数,设置为kBlocks,表示搬运kBlocks个分形。
- srcStride表示,相邻迭代间源操作数分形首地址之间的间隔,以搬运蓝色部分分形为例:下图中左侧源操作数矩阵,第一个蓝色分形和第二个蓝色分形起始地址之间的间隔为mBlocks个分形,此处设置为mBlocks。
- dstGap使用默认值,目的矩阵两个分形连续存储。
- ifTranspose设置为false,每块分形搬运前搬运后都为Z格式,不使能转置。
- 每次循环迭代源矩阵首地址偏移16*16,目的矩阵首地址偏移16*k。
图 8 NZ to ZZ格式转换示意图[object Object][object Object]
-
具体代码如下:
[object Object]
Stage2:SplitB函数实现。
-
搬运及格式转换示意图如下:图中k为32,占kBlocks(k/16=2)个分形,n为32,占nBlocks(n/16=2)个分形,一共搬运4个16*16分形。本示例中,调用一次LoadData接口完成两个16*16分形的搬运,循环调用两次LoadData。第一次循环搬运蓝色部分两个分形,第二次循环搬运绿色部分两个分形。
单次循环中LoadData(本样例中要完成2个分形的搬运,蓝色部分或者绿色部分)的参数设置如下:
- repeatTimes表示数据处理的迭代次数,因为LoadData每个迭代处理一个分形,所以也可以理解为待搬运分形的个数。本样例中即为k轴方向的分形个数,设置为kBlocks,表示搬运kBlocks个分形。
- srcStride相邻迭代间源操作数分形首地址之间的间隔,以搬运蓝色部分分形为例:下图中左侧源操作数矩阵,第一个蓝色分形和第二个蓝色分形起始地址之间的间隔为1个分形,此处设置为1,源矩阵两个分形连续存储。
- dstGap使用默认值0,目的矩阵两个分形连续存储。
- ifTranspose设置为true,每块分形搬运前为Z格式,搬运后需要为N格式,需要使能转置。
- 每次循环迭代,源矩阵首地址需要偏移k*n/2。
图 9 NZ to ZN格式转换示意图[object Object][object Object]
具体代码如下:
[object Object]Stage3:Compute函数实现,完成核心的矩阵计算功能。
- Compute函数需要传入参数a2Local,a2Local从A2的Queue中使用取出。
- 使用从CO1的Queue中申请c1Local。
- 使用从B2中取出b2Local。
- 使用完成矩阵乘计算。
- 使用将计算结果c1Local放入到CO1的Queue中。
具体代码如下:
[object Object]Stage4:Aggregate函数实现,完成数据汇聚操作。
-
DataCopy参数设置如下:
- blockCount设置为1,blockLen设置为2,连续搬运两个分形,无需格式转换。
- blockMode设置为BlockMode::BLOCK_MODE_MATRIX,表示需要按分形搬运。
- c2Local首地址偏移量设置为index * cSize / 2。
具体代码如下:
[object Object]Stage5:CopyOut函数实现。
使用将结果矩阵从CO2搬运到Global Memory,同时需要将格式从NZ转换为ND。
每次循环移动一个分形,搬运m*16个数。DataCopy参数说明如下:
- blockCount设置为m,共搬运m次。
- blockLen设置为2,DataCopy指令一次搬运2个block,每个block16个数。
- srcStride设置为0,每两次搬运间没有间隙。
- dstStride设置为(nBlocks - 1) * 2,每两次搬运间隔2个block。
- 每次循环迭代,目的矩阵偏移16,源矩阵偏移m*16。
格式转换示意图如下,第一次循环搬运蓝色部分数据,第二次循环搬运绿色部分数据。
图 10 NZ to ND格式转换示意图[object Object][object Object]
具体代码如下:
[object Object]