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接口负责完成特殊稀疏矩阵乘加操作。稀疏矩阵是一种特殊类型的矩阵即矩阵中包含较多的零元素。42结构化稀疏要求一个连续的4个权重或激活值的组通常是张量中的一行或一列中最多只有2个值为非零其余2个强制为零。MmadWithSparse接口传入的左矩阵A为稀疏矩阵右矩阵B为稠密矩阵。矩阵A是个全尺寸矩阵在MmadWithSparse计算时完成稠密化矩阵B是经过42结构化稀疏过滤掉零值之后的稠密矩阵需要在计算执行前的输入数据准备时自行完成稠密化按照下文中介绍的稠密算法进行稠密化。B稠密矩阵需要通过调用LoadDataWithSparse载入同时加载索引矩阵索引矩阵在矩阵B稠密化的过程中生成再用于A矩阵的稠密化。索引矩阵存储在内部缓冲区该索引矩阵的布局和布局大小与矩阵B相同用于在进行矩阵乘加操作之前进一步将矩阵A压缩。跟Mmad接口实现昇腾NPU矩阵乘计算能力类似MmadWithSparse接口的数学表达式为$$ C A \times B C $$完整示例请参考MmadWithSparse样例。表1Sparse矩阵计算矩阵A、B、C解释说明矩阵计算逻辑矩阵计算物理位置维度输入/输出数据格式数据类型AL0A BufferM x KZzint8_tBL0B BufferK/2 x NZnint8_tCL0C BufferM x NNzint32_t下面的图展示了Cube如何计算出其中一行和一列的内积图1MmadWithSparse接口计算流程示意图 其中矩阵A原始分形为(16, 2*C0)索引矩阵Index分形为(C016)每一行矩阵A的数据会基于索引矩阵Index中对应的一列数据进行4选2索引矩阵分形格式及生成方式请参考4选2稀疏索引矩阵选择算法参考矩阵A稀疏选择算法说明经过选择处理后的矩阵A分形变成(16, C0)矩阵B原始分形为(C0, 16)接下来会执行普通Mmad运算即矩阵A中一行和矩阵B中一列完成内积运算得到结果矩阵C中对应一个元素。矩阵A稀疏选择算法说明索引矩阵经过LoadDataWithSparse指令后存储于Cube上内置的专用buffer空间数据类型为uint8分形格式为小n大Z对应上图中的分形大小为3216。每一个uint8类型的索引元素由4个uint2的原始数据组成每两个2位索引数据可对应4位原始矩阵A。针对每一组2个索引数据A矩阵的4个元素的选择过滤规则示例如下表第一个索引数据0用于指示前3个元素中第1个非零元素的相对位置。第二个索引数据1用于指示第2个非零元素在后3个元素中的相对位置。其中“-”表示不关心该位置上的值即会被过滤。表2矩阵A选择过滤规则表索引数据0索引数据1元素0元素1元素2元素32’b102’b10--XY2’b012’b10-X-Y2’b002’b10X--Y2’b012’b01-XY-2’b002’b01X-Y-2’b002’b00XY--2’b102’b00--X/X-2’b012’b00-X/X--图2矩阵A 42选择算法模型图2展示了一个uint8类型的索引元素对应选择8个原始矩阵A元素的算法模型最后输出4个选择后的矩阵A元素。在正常使用情况下软件应确保最多存在两个非零元素。如果发生错误即存在三个或更多非零元素时只会使用最低有效位LSB位置的前两个非零元素。上表中使用的“-”表示“不关心该位置上的值”即暗示可能存在三个或更多非零元素的情况。函数原型template typename T int32_t, typename U int8_t, typename Std::enable_ifStd::is_samePrimTT, int32_t::value, bool::type true, typename Std::enable_ifStd::is_samePrimTU, int8_t::value, bool::type true __aicore__ inline void MmadWithSparse(const LocalTensorT dst, const LocalTensorU fm, const LocalTensorU filter, const MmadParams mmadParams)参数说明表3模板参数说明参数名描述Tdst的数据类型。Ufm、filter的数据类型。当dst、fm、filter为基础数据类型时T必须为int32_t类型U必须为int8_t类型否则编译失败。表4参数说明参数名称输入/输出含义dst输出目的操作数结果矩阵C类型为LocalTensor支持的物理存储位置为L0C BufferTPosition:CO1。LocalTensor的起始地址需要256个元素1024字节对齐。fm输入源操作数左矩阵A类型为LocalTensor支持的物理存储位置为L0A BufferTPosition: A2。LocalTensor的起始地址需要512字节对齐。filter输入源操作数右矩阵B类型为LocalTensor支持的物理存储位置为L0B BufferTPosition:B2。LocalTensor的起始地址需要512字节对齐。mmadParams输入矩阵乘相关参数类型为MmadParams。具体定义请参考${INSTALL_DIR}/asc/include/basic_api/interface/kernel_struct_mm.h${INSTALL_DIR}请替换为CANN软件安装后文件存储路径。参数说明请参考表5。表5MmadParams结构体内参数说明Sparse场景参数名称含义m左矩阵Height取值范围m∈[0, 4095]。默认值为0。n右矩阵Width取值范围n∈[0, 4095]。默认值为0。k左矩阵Width、右矩阵Height取值范围k∈[0, 4095]。默认值为0。cmatrixInitVal是否开启C矩阵默认初始化清零操作。默认值为 true。• trueC矩阵默认初始化为0• falseC矩阵不进行默认操作通过设置cmatrixSource参数进行初始化。cmatrixSource配置C矩阵初始值是否来源于BT Buffer。默认值为false。• false不对L0C Buffer进行初始化操作• true使用BT BufferTPosition:C2的数据对L0C Buffer进行初始化操作。Atlas A2 训练系列产品/Atlas A2 推理系列产品支持配置为true/false。Atlas A3 训练系列产品/Atlas A3 推理系列产品支持配置为true/false。Atlas 200I/500 A2 推理产品支持配置为true/false。注意带Bias输入的接口配置该参数无效会根据Bias输入的位置来判断C矩阵初始值是否来源于BT Buffer。isBias该参数废弃新开发内容不要使用该参数。如果需要累加初始矩阵请使用带Bias的接口来实现也可以通过cmatrixInitVal和cmatrixSource参数配置C矩阵的初始值来源来实现。推荐使用带Bias的接口相比于配置cmatrixInitVal和cmatrixSource参数更加简单方便。配置是否需要累加初始矩阵默认值为false取值说明如下• false矩阵乘无需累加初始矩阵C A * B。• true矩阵乘加需要累加初始矩阵C A * B。unitFlagunitFlag是一种Mmad指令和Fixpipe指令细粒度的并行开启该功能后硬件每计算完一个分形计算结果就会被搬出。取值说明如下• 02b00不开启unitFlag• 22b10开启unitFlag硬件执行完指令之后不复位单元标记位• 32b11开启unitFlag硬件执行完指令之后复位单元标记位。开启该功能时须将Mmad指令和Fixpipe指令的unitFlag值设置为2或3。该参数仅支持如下型号Atlas A2 训练系列产品/Atlas A2 推理系列产品Atlas A3 训练系列产品/Atlas A3 推理系列产品。参数设置方案和特性细节可参考UnitFlag。kDirectionAlignSparse场景本开关默认为false不支持配置为true。K方向对齐的核心功能是通过kDirectionAlign参数控制在使用float数据类型时L0A Buffer和L0B Buffer矩阵在K方向上的对齐方式。fmOffset左矩阵offset整个左矩阵对应一个值支持Scalar应与src_fm.dtype一致/立即数默认0。注未使用兼容旧款产品接口传入Atlas A2 训练系列产品/Atlas A2 推理系列产品及往后产品不做处理。enSsparse开启结构化稀疏特性默认false注未使用兼容旧款产品接口传入Atlas A2 训练系列产品/Atlas A2 推理系列产品及往后产品不做处理。enWinogradA指示矩阵a是否通过winograd_feature_map_transform()生成用于支持winograd特性bool类型默认false注未使用兼容旧款产品接口传入Atlas A2 训练系列产品/Atlas A2 推理系列产品及往后产品不做处理。enWinogradB指示矩阵b是否通过winograd_weight_transform()生成用于支持winograd特性bool类型默认false注未使用兼容旧款产品接口传入Atlas A2 训练系列产品/Atlas A2 推理系列产品及往后产品不做处理。数据类型表6A、B、C支持的精度类型组合Atlas 200I/500 A2 推理产品Atlas A2 训练系列产品/Atlas A2 推理系列产品Atlas A3 训练系列产品/Atlas A3 推理系列产品左矩阵A右矩阵B结果矩阵Cint8_tint8_tint32_t返回值说明无约束说明不同矩阵对于存储位置的约束结果矩阵C只支持位于物理存储位置为L0C BufferTPosition:CO1左矩阵A只支持位于物理存储位置为L0A BufferTPosition:A2右矩阵B只支持位于物理存储位置为L0B BufferTPosition:B2原始稀疏矩阵B每4个元素中应保证最多2个非零元素如果存在3个或更多非零元素则仅使用前2个非零元素。当M、K、N中的任意一个值为0时表示指令不会执行该接口将被视为NOP空操作。MmadWithSparse接口不支持Gemv模式。其他特殊场景约束可参考Mmad接口约束说明。调用示例完整使用样例请参见MmadWithSparse样例。AscendC::LocalTensorint8_t a1Local(AscendC::TPosition::A1, a1Addr, aSize); AscendC::LocalTensorint8_t a2Local(AscendC::TPosition::A2, a2Addr, aSize); AscendC::LocalTensorint8_t b1Local(AscendC::TPosition::B1, b1Addr, bSize); AscendC::LocalTensoruint8_t idxB1Local(AscendC::TPosition::B1, idxB1Addr, bSize / 4); AscendC::LocalTensorint8_t b2Local(AscendC::TPosition::B2, b2Addr, bSize); AscendC::LocalTensorint32_t cLocal(AscendC::TPosition::CO1, cAddr, cSize); // GM-L1 Buffer将原始矩阵a稠密化矩阵b与对应idx矩阵搬运至L1 Buffer CopyIn(a1Local, b1Local, idxB1Local); AscendC::SetFlagAscendC::HardEvent::MTE2_MTE1(EVENT_ID0); AscendC::WaitFlagAscendC::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::SetFlagAscendC::HardEvent::MTE1_M(EVENT_ID0); AscendC::WaitFlagAscendC::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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考