Mmad【免费下载链接】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 训练系列产品不支持功能说明头文件路径为tensor_api/tensor.h。Mmad接口用于完成L0A Buffer上左矩阵A和L0B Buffer上右矩阵B的矩阵乘加结果写入L0C Buffer上结果矩阵C。默认模式为普通矩阵计算。流水类型PIPE_M。数学表达式为C A * B带bias输入时bias作为结果矩阵C初始值参与计算等价于C A * B BiasMmad的矩阵乘加关系可参考下图图1Mmad矩阵乘加公式Mmad使用显式传入的MmadAtom配置矩阵计算功能。左矩阵A、右矩阵B、结果矩阵C的图示说明如下图2Mmad矩阵布局矩阵计算说明矩阵物理位置维度数据格式数据类型左矩阵AL0A Bufferm × kNZ参见数据类型右矩阵BL0B Bufferk × nZN参见数据类型结果矩阵CL0C Bufferm × nNZ参见数据类型函数原型执行不传bias的矩阵计算。template typename AtomType, typename DstTensor, typename FmTensor, typename FilterTensor __aicore__ inline void Mmad( const MmadAtomAtomType atomMmad, const DstTensor dst, const FmTensor fm, const FilterTensor filter)执行传入bias的矩阵计算。template typename AtomType, typename DstTensor, typename FmTensor, typename FilterTensor, typename BiasTensor, Std::enable_if_tIsAttrTensorVBiasTensor, int Enable __aicore__ inline void Mmad( const MmadAtomAtomType atomMmad, const DstTensor dst, const FmTensor fm, const FilterTensor filter, const BiasTensor bias)构造默认矩阵计算原子对象。template typename MmadOperationType __aicore__ inline constexpr auto MakeMmad(const MmadOperationType mmadOperation)构造指定Trait的矩阵计算原子对象。template typename MmadOperationType, typename MmadTraitType __aicore__ inline constexpr auto MakeMmad( const MmadOperationType mmadOperation, const MmadTraitType mmadTrait)参数说明表1Mmad接口参数说明参数名输入/输出描述atomMmad输入矩阵计算原子对象。通过MakeMmad(MmadOperation{})或MakeMmad(MmadOperation{}, MmadTraitDefault{})可以构造默认原子对象。dst输出结果矩阵C存储位置为Location::L0C数据格式为NZ。fm输入左矩阵A存储位置为Location::L0A数据格式为NZ。filter输入右矩阵B存储位置为Location::L0B数据格式为ZN。bias输入bias张量存储位置为Location::BIAS或Location::L0C数据格式为ND。表2MakeMmad接口参数说明参数名输入/输出描述mmadOperation输入矩阵计算操作对象通过MmadOperation{}构造。mmadTrait输入矩阵计算Trait对象用于指定矩阵计算静态特性。默认使用MmadTraitDefault{}。表3MmadParams参数说明参数名类型默认值描述muint16_t0左矩阵A的高度结果矩阵C的高度。nuint16_t0右矩阵B的宽度结果矩阵C的宽度。kuint16_t0左矩阵A的宽度右矩阵B的高度。unitFlaguint8_t0控制Mmad和后续矩阵数据搬出的细粒度并行。0表示不使能2表示使能且执行后不复位单元标记位3表示使能且执行后复位单元标记位。cmatrixInitValboolfalse不传bias时控制是否初始化结果矩阵C。true表示C矩阵默认初始化为0false表示C矩阵不进行默认操作通过设置cmatrixSource参数进行初始化。表4MmadTrait参数说明参数名类型默认值描述fmOffsetint32_t0左矩阵offset当前Tensor API实现中作为兼容参数保留。kDirectionAlignboolfalseK方向对齐控制当前Tensor API实现中作为兼容参数保留。cmatrixSourceboolfalse配置C矩阵初始值是否来源于BT Buffer带bias调用时该配置无效。false表示C矩阵不进行初始化操作true表示使用BT Buffer的数据对C矩阵进行初始化操作。disableGemvbooltrueM1场景下是否关闭GEMV模式。false表示开启GEMVtrue表示关闭GEMV。mmadTypeMmadTypeMmadType::NORMAL矩阵计算类型。默认使用MmadType::NORMAL。使用MmadType::MX表示MX场景Mmad计算。数据类型支持如下左矩阵A、右矩阵B、结果矩阵C的数据类型组合左矩阵A右矩阵B结果矩阵Cint8_tint8_tint32_thalfhalffloatfloatfloatfloatbfloat16_tbfloat16_tfloatfp8_e4m3fn_tfp8_e4m3fn_tfloatfp8_e4m3fn_tfp8_e5m2_tfloatfp8_e5m2_tfp8_e4m3fn_tfloatfp8_e5m2_tfp8_e5m2_tfloathifloat8_thifloat8_tfloat传入bias的场景中int8_t * int8_t - int32_t的组合要求bias数据类型为int32_t其余组合要求bias数据类型为float。返回值说明Mmad无返回值。MakeMmad返回MmadAtom对象。约束说明dst必须位于L0C Bufferfm必须位于L0A Bufferfilter必须位于L0B Buffer。结果矩阵C起始地址需要满足64Byte地址对齐要求。左矩阵A和右矩阵B起始地址需要满足512Byte地址对齐要求。上述地址对齐要求属于硬件访问约束编译器和运行时不对所有场景单独检查用户需要保证入参满足约束。m、n、k需要与实际参与Mmad计算的数据尺寸一致。开启unitFlag功能时需要结果矩阵C搬出接口的FixpipeParams::unitFlag与Mmad计算的MmadParams::unitFlag参数配合设置。连续两次Mmad沿K方向累加时需要关注结果矩阵C的写读依赖。同步优化的阈值关系可参考结果矩阵C初始化中的图示。当M、K、N不是16的倍数时硬件仍以16×16分形块组织数据尾块中的无效数据会占用分形块空间但不参与有效计算。有效数据与无效数据排布方式如下图所示图3尾块有效数据排布关键特性说明结果矩阵C初始化不传bias时cmatrixInitVal控制是否初始化结果矩阵C。通常第一次K方向累加时设置为true后续K分块累加设置为false。如果多K分块累加场景中首次Mmad如果设为false将导致C矩阵包含L0C Buffer残留数据计算结果错误。连续两次Mmad沿K方向累加时通常需要关注中间结果的写读依赖。同步优化的阈值关系可参考下图图4PipeBarrier阈值关系UnitFlagunitFlag用于控制Mmad与后续Fixpipe的细粒度并行。启用时MmadParams::unitFlag和结果搬出的FixpipeParams::unitFlag需要配合设置常见取值为2和3。GEMV模式当M1时可通过MmadTrait::disableGemv控制是否开启GEMV模式。默认值为true表示关闭GEMV。MX Mmad模式MX Mmad模式没有单独的函数接口而是通过MmadTrait::mmadType设置。执行MX Mmad时需要将MmadTrait::mmadType设置为MmadType::MX并通过MakeMmad(MmadOperation{}, MmadTraitMX{})构造矩阵计算原子对象。struct MmadTraitMX { using TraitType MmadTrait; static constexpr const TraitType value MmadTrait(0, false, false, true, MmadType::MX); };MX Mmad支持如下左矩阵A、右矩阵B、结果矩阵C的数据类型组合左矩阵A右矩阵B结果矩阵Cfp4x2_e2m1_tfp4x2_e2m1_tfloatfp4x2_e2m1_tfp4x2_e1m2_tfloatfp4x2_e1m2_tfp4x2_e2m1_tfloatfp4x2_e1m2_tfp4x2_e1m2_tfloatfp8_e4m3fn_tfp8_e4m3fn_tfloatfp8_e4m3fn_tfp8_e5m2_tfloatfp8_e5m2_tfp8_e4m3fn_tfloatfp8_e5m2_tfp8_e5m2_tfloatMX Mmad通常需要在调用Mmad前完成左矩阵缩放数据ScaleA和右矩阵缩放数据ScaleB的搬运。ScaleA搬运可参考L1到L0ScaleA数据搬运ScaleB搬运可参考L1到L0ScaleB数据搬运。矩阵缩放数据不作为Mmad函数参数传入而是通过Copy接口提前搬运到L0ScaleA Buffer和L0ScaleB Buffer。相关接口UnitFlagGEMVHF32L1到L0ScaleA数据搬运L1到L0ScaleB数据搬运L0C到GM数据搬运L0C到UB数据搬运调用示例Mmad调用示例#include tensor_api/tensor.h __aicore__ inline void NormalMmadExample() { using namespace AscendC::Te; constexpr uint16_t m 16; constexpr uint16_t n 16; constexpr uint16_t k 16; __ca__ half l0aBuf[m * k]; __cb__ half l0bBuf[k * n]; __cc__ float l0cBuf[m * n]; auto l0A MakeTensor(MakeMemPtr(l0aBuf), MakeFrameLayoutNZLayoutPtn, half(m, k)); auto l0B MakeTensor(MakeMemPtr(l0bBuf), MakeFrameLayoutZNLayoutPtn, half(k, n)); auto l0C MakeTensor(MakeMemPtr(l0cBuf), MakeFrameLayoutNZLayoutPtn, 16(m, n)); MmadParams params(m, n, k, 0, true); auto atom MakeMmad(MmadOperation{}, MmadTraitDefault{}).with(params); Mmad(atom, l0C, l0A, l0B); }MX场景Mmad调用示例#include tensor_api/tensor.h struct MmadTraitMX { using TraitType AscendC::Te::MmadTrait; static constexpr const TraitType value AscendC::Te::MmadTrait(0, false, false, true, AscendC::Te::MmadType::MX); }; __aicore__ inline void MxMmadExample() { using namespace AscendC::Te; constexpr uint16_t m 16; constexpr uint16_t n 16; constexpr uint16_t k 16; __ca__ fp8_e5m2_t l0aBuf[m * k]; __cb__ fp8_e5m2_t l0bBuf[k * n]; __cc__ float l0cBuf[m * n]; auto l0A MakeTensor(MakeMemPtr(l0aBuf), MakeFrameLayoutNZLayoutPtn, fp8_e5m2_t(m, k)); auto l0B MakeTensor(MakeMemPtr(l0bBuf), MakeFrameLayoutZNLayoutPtn, fp8_e5m2_t(k, n)); auto l0C MakeTensor(MakeMemPtr(l0cBuf), MakeFrameLayoutNZLayoutPtn, 16(m, n)); MmadParams params(m, n, k, 0, true); auto atom MakeMmad(MmadOperation{}, MmadTraitMX{}).with(params); Mmad(atom, l0C, l0A, l0B); }【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考