功能说明

完成矩阵乘加操作。

定义原型

不含有矩阵乘偏置功能:

template <typename dst_T, typename src0_T, typename src1_T>
void Mmad(const LocalTensor<dst_T>& dstLocal, const LocalTensor<src0_T>& fmLocal,
const LocalTensor<src1_T>& filterLocal, const MmadParams& mmadParams)

参数说明

表1 不含有矩阵乘偏置接口参数说明
参数名称	输入/输出	含义
dstLocal	输出	目的操作数，结果矩阵，类型为LocalTensor。
fmLocal	输入	源操作数，左矩阵a，类型为LocalTensor。
filterLocal	输入	源操作数，右矩阵b，类型为LocalTensor。
mmadParams	输入	矩阵乘相关参数，类型为MmadParams，结构体具体定义为： struct MmadParams { uint16_t m; uint16_t n; uint16_t k; bool isBias; int32_t fmOffset; bool enSsparse; bool enWinogradA; bool enWinogradB; }; 参数说明请参考表2。

表2 MmadParams结构体内参数说明
参数名称	含义
m	左矩阵Height，取值范围：m∈[0, 4095] 。默认值为0。
n	右矩阵Width，取值范围：n∈[0, 4095] 。默认值为0。
k	左矩阵Width、右矩阵Height，取值范围：k∈[0, 4095] 。默认值为0。
isBias	配置是否需要累加初始矩阵，默认值为false，取值说明如下： false：矩阵乘，无需累加初始矩阵，C = A * B。 true：矩阵乘加，需要累加初始矩阵, C += A * B。
fmOffset	预留参数，用户无需关心，使用默认值0即可，当前暂不支持。
enSsparse	预留参数，用户无需关心，使用默认值false即可，当前暂不支持。
enWinogradA	预留参数，用户无需关心，使用默认值false即可, 当前暂不支持。
enWinogradB	预留参数，用户无需关心，使用默认值false即可, 当前暂不支持。

表3 dstLocal、fmLocal、filterLocal支持的精度类型组合
左矩阵fmLocal type	右矩阵filterLocal type	结果矩阵dstLocal type	支持的型号
uint8_t	uint8_t	uint32_t	Atlas 训练系列产品 Atlas推理系列产品AI Core
int8_t	int8_t	int32_t	Atlas 训练系列产品 Atlas推理系列产品AI Core
uint8_t	int8_t	int32_t	Atlas 训练系列产品 Atlas推理系列产品AI Core
uint8_t	uint8_t	int32_t	Atlas 训练系列产品 Atlas推理系列产品AI Core
half	half	half	Atlas 训练系列产品 Atlas推理系列产品AI Core
half	half	float	Atlas 训练系列产品 Atlas推理系列产品AI Core
float	float	float

注意事项

dstLocal只支持位于CO1，fmLocal只支持位于A2，filterLocal只支持位于B2。
该指令除了支持矩阵与矩阵乘/乘加，还支持向量与矩阵乘/乘加（m=1）、矩阵与向量乘/乘加（右矩阵可能需要进行填充）；向量与向量运算建议使用SIMD指令。
地址对齐要求请见：通用约束。

数据格式说明

Mmad 函数对于输入数据的格式要求和输出数据的要求如下图，矩阵 ABC 分别为 A2/B2/CO1 中的数据。下图中每个小方格代表一个 512Byte 的分形矩阵（对于 b8 来说是 16x32，对于 b16/b32 来说则是 16x16）。下图中Z字形的黑色线条代表对应位置数据在昇腾AI处理器上的排列顺序，起始点是左上角，终点是右下角。

矩阵A：每个分型矩阵内部是行主序，分型矩阵之间是行主序。简称小Z大Z格式。

矩阵B：每个分型矩阵内部是列主序，分型矩阵之间是行主序。简称小N大Z格式。

矩阵C：每个分型矩阵内部是行主序，分型矩阵之间是列主序。简称小Z大N格式。

以下是一个简单的例子，假设分型矩阵的大小是2x2，然后矩阵ABC的大小都是4x4

0	1	2	3
4	5	6	7
8	9	10	11
12	13	14	15

对于一个C风格的矩阵来说内部元素的排列顺序应该是0,1,2…15。

矩阵A的排列顺序:0,1,4,5,2,3,6,7,8,9,12,13,10,11,14,15

矩阵B的排列顺序:0,4,1,5,2,6,3,7,8,12,9,13,10,14,11,15

矩阵C的排列顺序:0,1,4,5,8,9,12,13,2,3,6,7,10,11,14,15

以下是一个具体的例子，数据为half类型。

如图，当 M=30，K=70，N=40 的时候，A2 中应该有 2x5 个 16x16 矩阵，B2 中应该有 5x3 个16x16 矩阵，CO1 中应该有 2x3 个 16x16 矩阵。在这种场景下 M、K 和 N 都不是 16 的倍数，A2 中右下角的矩阵实际有效的数据只有 14x6 个，但是也需要占一个 16x16 矩阵的空间，其他无效数据在计算中会被忽略。

一个 16x16 分形的数据块中，无效数据与有效数据排布的方式示意如下：

调用示例

示例：src1数据类型是half, src2数据类型是half, dst的数据类型是float，mmad不含有矩阵乘偏置

/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2023. All rights reserved.
 *
 * Function : c = a * b (matrix multiplication)
 * This sample is a very basic sample that implements Matmul on Ascend plaform.
 * In this sample:
 * Shape of matrix a is [m, k]: [32, 32]
 * Shape of matrix b is [k, n]: [32, 32]
 * Shape of matrix c is [m, n]: [32, 32]
 */

#include "kernel_operator.h"
using namespace AscendC;

class KernelMatmul {
public:
    __aicore__ inline KernelMatmul()
    {
        aSize = m * k;
        bSize = k * n;
        cSize = m * m;
        mBlocks = m / 16;
        nBlocks = n / 16;
        kBlocks = k / 16;
    }
    __aicore__ inline void Init(GM_ADDR a, GM_ADDR b, GM_ADDR c)
    {
        aGM.SetGlobalBuffer((__gm__ half*)a);
        bGM.SetGlobalBuffer((__gm__ half*)b);
        cGM.SetGlobalBuffer((__gm__ float*)c);
        pipe.InitBuffer(inQueueA1, 1, aSize * sizeof(half));
        pipe.InitBuffer(inQueueA2, 1, aSize * sizeof(half));
        pipe.InitBuffer(inQueueB1, 1, bSize * sizeof(half));
        pipe.InitBuffer(inQueueB2, 2, bSize * sizeof(half) / 2);
        pipe.InitBuffer(outQueueCO1, 2, cSize * sizeof(float) / 2);
        pipe.InitBuffer(outQueueCO2, 1, cSize * sizeof(float));
    }
    __aicore__ inline void Process()
    {
        CopyIn();
        SplitA();

        LocalTensor<half> b1Local = inQueueB1.DeQue<half>();
        LocalTensor<half> a2Local = inQueueA2.DeQue<half>();
        LocalTensor<float> c2Local = outQueueCO2.AllocTensor<float>();
        // split matrix b into 2 parts, [32, 16] and [32, 16]
        for (int i = 0; i < 2; ++i) {
            SplitB(b1Local, i);
            Compute(a2Local);
            Aggregate(c2Local, i);
        }
        inQueueB1.FreeTensor(b1Local);
        inQueueA2.FreeTensor(a2Local);
        outQueueCO2.EnQue<float>(c2Local);

        CopyOut();
    }

private:
    __aicore__ inline void CopyND2NZ(const LocalTensor<half>& dst, const GlobalTensor<half>& src, const uint16_t height,
        const uint16_t width)
    {
        for (int i = 0; i < width / 16; ++i) {
            int srcOffset = i * 16;
            int dstOffset = i * 16 * height;
            DataCopy(dst[dstOffset], src[srcOffset], { height, 1, uint16_t(width / 16 - 1), 0 });
        }
    }
    __aicore__ inline void CopyIn()
    {
        LocalTensor<half> a1Local = inQueueA1.AllocTensor<half>();
        LocalTensor<half> b1Local = inQueueB1.AllocTensor<half>();

        CopyND2NZ(a1Local, aGM, m, k);
        CopyND2NZ(b1Local, bGM, k, n);

        inQueueA1.EnQue(a1Local);
        inQueueB1.EnQue(b1Local);
    }
    __aicore__ inline void SplitA()
    {
        int srcOffset = 0;
        int dstOffset = 0;
        LocalTensor<half> a1Local = inQueueA1.DeQue<half>();
        LocalTensor<half> a2Local = inQueueA2.AllocTensor<half>();

        // transform nz to zz
        for (int i = 0; i < mBlocks; ++i) {
            LoadData2dParams loadDataParams;
            loadDataParams.repeatTimes = kBlocks;
            loadDataParams.srcStride = mBlocks;
            loadDataParams.ifTranspose = false;

            LoadData(a2Local[dstOffset], a1Local[srcOffset], loadDataParams);

            srcOffset += 16 * 16;
            dstOffset += k * 16;
        }

        inQueueA2.EnQue<half>(a2Local);
        inQueueA1.FreeTensor(a1Local);
    }
    __aicore__ inline void SplitB(const LocalTensor<half>& b1Local, const int bSplitIdx)
    {
        LocalTensor<half> b2Local = inQueueB2.AllocTensor<half>();

        // transform nz to zn
        LoadData2dParams loadDataParams;
        loadDataParams.repeatTimes = kBlocks;
        loadDataParams.srcStride = 1;
        loadDataParams.ifTranspose = true;

        LoadData(b2Local, b1Local[bSplitIdx * bSize / 2], loadDataParams);

        inQueueB2.EnQue<half>(b2Local);
    }
    __aicore__ inline void Compute(const LocalTensor<half>& a2Local)
    {
        LocalTensor<half> b2Local = inQueueB2.DeQue<half>();
        LocalTensor<float> c1Local = outQueueCO1.AllocTensor<float>();

        Mmad(c1Local, a2Local, b2Local, { m, uint16_t(n / 2), k, false, 0, false, false, false });

        outQueueCO1.EnQue<float>(c1Local);
        inQueueB2.FreeTensor(b2Local);
    }
    __aicore__ inline void Aggregate(const LocalTensor<float>& c2Local, const int bSplitIdx)
    {
        LocalTensor<float> c1Local = outQueueCO1.DeQue<float>();

        DataCopyParams dataCopyParams;
        dataCopyParams.blockCount = 1;
        dataCopyParams.blockLen = 2;
        DataCopyEnhancedParams enhancedParams;
        enhancedParams.blockMode = BlockMode::BLOCK_MODE_MATRIX;
        DataCopy(c2Local[bSplitIdx * cSize / 2], c1Local, dataCopyParams, enhancedParams);

        outQueueCO1.FreeTensor(c1Local);
    }
    __aicore__ inline void CopyOut()
    {
        LocalTensor<float> c2Local = outQueueCO2.DeQue<float>();

        // transform nz to nd
        for (int i = 0; i < nBlocks; ++i) {
            DataCopy(cGM[i * 16], c2Local[i * m * 16], { m, 2, 0, uint16_t((nBlocks - 1) * 2) });
        }

        outQueueCO2.FreeTensor(c2Local);
    }

private:
    TPipe pipe;

    TQue<QuePosition::A1, 1> inQueueA1;
    TQue<QuePosition::A2, 1> inQueueA2;
    TQue<QuePosition::B1, 1> inQueueB1;
    TQue<QuePosition::B2, 2> inQueueB2;
    // dst queue
    TQue<QuePosition::CO1, 2> outQueueCO1;
    TQue<QuePosition::CO2, 1> outQueueCO2;

    GlobalTensor<half> aGM, bGM;
    GlobalTensor<float> cGM;

    uint16_t m = 32;
    uint16_t n = 32;
    uint16_t k = 32;

    uint16_t aSize, bSize, cSize, mBlocks, nBlocks, kBlocks;
};

extern "C" __global__ __aicore__ void matmul_custom(GM_ADDR a, GM_ADDR b, GM_ADDR c)
{
    KernelMatmul op;
    op.Init(a, b, c);
    op.Process();
}

#ifndef __CCE_KT_TEST__
// call of kernel function
void matmul_custom_do(uint32_t blockDim, void* l2ctrl, void* stream, uint8_t* a, uint8_t* b, uint8_t* c)
{
    matmul_custom<<<blockDim, l2ctrl, stream>>>(a, b, c);
}
#endif