CANN/asc-devkit稀疏矩阵乘加API

# MmadWithSparse

【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言，原生支持C和C++标准规范，主要由类库和语言扩展层构成，提供多层级API，满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit

产品支持情况

• Ascend 950PR/Ascend 950DT：不支持

• Atlas A3 训练系列产品/Atlas A3 推理系列产品：支持

• Atlas A2 训练系列产品/Atlas A2 推理系列产品：支持

• Atlas 200I/500 A2 推理产品：不支持

• Atlas 推理系列产品AI Core：不支持

• Atlas 推理系列产品Vector Core：不支持

• Atlas 训练系列产品：不支持

功能说明

头文件路径为：basic_api/kernel_operator_mm_intf.h。

MmadWithSparse接口负责完成特殊稀疏矩阵乘加操作。稀疏矩阵是一种特殊类型的矩阵，即矩阵中包含较多的零元素。4：2结构化稀疏要求一个连续的4个权重或激活值的组（通常是张量中的一行或一列）中，最多只有2个值为非零，其余2个强制为零。

MmadWithSparse接口传入的左矩阵A为稀疏矩阵，右矩阵B为稠密矩阵。矩阵A是个全尺寸矩阵，在MmadWithSparse计算时完成稠密化；矩阵B是经过4：2结构化稀疏过滤掉零值之后的稠密矩阵，需要在计算执行前的输入数据准备时自行完成稠密化（按照下文中介绍的稠密算法进行稠密化）。B稠密矩阵需要通过调用LoadDataWithSparse载入，同时加载索引矩阵，索引矩阵在矩阵B稠密化的过程中生成，再用于A矩阵的稠密化。索引矩阵存储在内部缓冲区，该索引矩阵的布局和布局大小与矩阵B相同，用于在进行矩阵乘加操作之前进一步将矩阵A压缩。

跟Mmad接口实现昇腾NPU矩阵乘计算能力类似，MmadWithSparse接口的数学表达式为：

$$ C = A \times B + C $$

完整示例请参考：MmadWithSparse样例。

表1Sparse矩阵计算矩阵A、B、C解释说明

下面的图展示了Cube如何计算出其中一行和一列的内积：

图1MmadWithSparse接口计算流程示意图 ![MmadWithSparse接口计算流程示意图](https://raw.gitcode.com/cann/asc-devkit/raw/1cd1d1192e01ae4a8684cd07552c35e13e5c4bdf/docs/zh/api/figures/mmadwithsparse_workflow_demo.png "MmadWithSparse接口计算流程示意图"?utm_source=gitcode_repo_files)

其中矩阵A原始分形为(16, 2*C0)，索引矩阵Index分形为(C0，16)，每一行矩阵A的数据会基于索引矩阵Index中对应的一列数据进行4选2，索引矩阵分形格式及生成方式请参考4选2稀疏索引矩阵，选择算法参考矩阵A稀疏选择算法说明；经过选择处理后的矩阵A分形变成(16, C0)，矩阵B原始分形为(C0, 16)，接下来会执行普通Mmad运算，即矩阵A中一行和矩阵B中一列完成内积运算得到结果矩阵C中对应一个元素。

• 矩阵A稀疏选择算法说明

  索引矩阵经过LoadDataWithSparse指令后存储于Cube上内置的专用buffer空间，数据类型为uint8，分形格式为小n大Z，对应上图中的分形大小为（32，16）。每一个uint8类型的索引元素由4个uint2的原始数据组成，每两个2位索引数据可对应4位原始矩阵A。针对每一组2个索引数据，A矩阵的4个元素的选择过滤规则示例如下表：

  • 第一个索引数据0用于指示前3个元素中第1个非零元素的相对位置。
  • 第二个索引数据1用于指示第2个非零元素在后3个元素中的相对位置。
  • 其中，“-”表示不关心该位置上的值，即会被过滤。

  表2矩阵A选择过滤规则表

  索引数据0	索引数据1	元素0	元素1	元素2	元素3
  2’b10	2’b10	-	-	X	Y
  2’b01	2’b10	-	X	-	Y
  2’b00	2’b10	X	-	-	Y
  2’b01	2’b01	-	X	Y	-
  2’b00	2’b01	X	-	Y	-
  2’b00	2’b00	X	Y	-	-
  2’b10	2’b00	-	-	X/X	-
  2’b01	2’b00	-	X/X	-	-

  图2矩阵A 4：2选择算法模型![矩阵A-4-2选择算法模型](https://raw.gitcode.com/cann/asc-devkit/raw/1cd1d1192e01ae4a8684cd07552c35e13e5c4bdf/docs/zh/api/figures/sparse_mmad_4select2.png "矩阵A-4-2选择算法模型"?utm_source=gitcode_repo_files)

  图2展示了一个uint8类型的索引元素对应选择8个原始矩阵A元素的算法模型，最后输出4个选择后的矩阵A元素。

  1. 在正常使用情况下，软件应确保最多存在两个非零元素。如果发生错误，即存在三个或更多非零元素时，只会使用最低有效位（LSB）位置的前两个非零元素。
  2. 上表中使用的“-”表示“不关心该位置上的值”，即暗示可能存在三个或更多非零元素的情况。

函数原型

template <typename T = int32_t, typename U = int8_t, typename Std::enable_if<Std::is_same<PrimT<T>, int32_t>::value, bool>::type = true, typename Std::enable_if<Std::is_same<PrimT<U>, int8_t>::value, bool>::type = true> __aicore__ inline void MmadWithSparse(const LocalTensor<T>& dst, const LocalTensor<U>& fm, const LocalTensor<U>& filter, const MmadParams& mmadParams)

参数说明

表3模板参数说明

表4参数说明

表5MmadParams结构体内参数说明（Sparse场景）

数据类型

表6A、B、C支持的精度类型组合（Atlas 200I/500 A2 推理产品）（Atlas A2 训练系列产品/Atlas A2 推理系列产品）（Atlas A3 训练系列产品/Atlas A3 推理系列产品）

返回值说明

无

约束说明

• 不同矩阵对于存储位置的约束：

  • 结果矩阵C只支持位于物理存储位置为L0C Buffer（TPosition:CO1）
  • 左矩阵A只支持位于物理存储位置为L0A Buffer（TPosition:A2）
  • 右矩阵B只支持位于物理存储位置为L0B Buffer（TPosition:B2）

• 原始稀疏矩阵B每4个元素中应保证最多2个非零元素，如果存在3个或更多非零元素，则仅使用前2个非零元素。

• 当M、K、N中的任意一个值为0时，表示指令不会执行，该接口将被视为NOP（空操作）。

• MmadWithSparse接口不支持Gemv模式。

• 其他特殊场景约束可参考Mmad接口约束说明。

调用示例

完整使用样例请参见MmadWithSparse样例。

AscendC::LocalTensor<int8_t> a1Local(AscendC::TPosition::A1, a1Addr, aSize); AscendC::LocalTensor<int8_t> a2Local(AscendC::TPosition::A2, a2Addr, aSize); AscendC::LocalTensor<int8_t> b1Local(AscendC::TPosition::B1, b1Addr, bSize); AscendC::LocalTensor<uint8_t> idxB1Local(AscendC::TPosition::B1, idxB1Addr, bSize / 4); AscendC::LocalTensor<int8_t> b2Local(AscendC::TPosition::B2, b2Addr, bSize); AscendC::LocalTensor<int32_t> cLocal(AscendC::TPosition::CO1, cAddr, cSize); // GM->L1 Buffer，将原始矩阵a，稠密化矩阵b与对应idx矩阵搬运至L1 Buffer CopyIn(a1Local, b1Local, idxB1Local); AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0); AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0); // L1 Buffer->L0A Buffer/L0B Buffer，将原始矩阵a，稠密化矩阵b与对应idx矩阵搬运至L0A Buffer/L0B Buffer SplitA(a1Local, a2Local); SplitB(b2Local, b1Local, idxB1Local); AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0); AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0); // mmad需要指定矩阵的维度进行计算 uint32_t m = 128; uint32_t k = 64; uint32_t n = 128; AscendC::MmadWithSparse(c1Local, a2Local, b2Local, { m, n, k, false, 0, false, false, false });

【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言，原生支持C和C++标准规范，主要由类库和语言扩展层构成，提供多层级API，满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit