CANN/asc-devkit:连续对齐搬出(StoreAlign)

CANN/asc-devkit:连续对齐搬出(StoreAlign) 连续对齐搬出StoreAlign【免费下载链接】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/reg_compute/kernel_reg_compute_datacopy_intf.h。StoreAlign能够实现数据从RegTensor连续搬运至Unified BufferUB支持单搬出模式和双搬出模式单搬出模式将一个RegTensorVL中的VL数据量搬出至UB。双搬出模式将两个RegTensor(2*VL)的数据交错搬出至UB偶数索引和奇数索引的元素分别来自两个RegTensor。连续对齐搬运有三类接口普通搬运接口、PostUpdate扩展搬运接口、使用AddrReg寄存器存储偏移量接口。RegTensor的模板参数regTrait支持RegTraitNumOne及RegTraitNumTwo具体支持情况如下RegTensor模板参数regTrait普通搬运接口PostUpdate扩展搬运接口使用AddrReg寄存器存储偏移量接口RegTraitNumOne单搬出完成一次搬运后UB地址不会自动更新每次迭代需要手动更新地址。• POST_MODE_NORMAL模式不支持。• POST_MODE_UPDATE模式完成一次搬运后UB地址会自动更新每次迭代不需要手动更新地址。在每次迭代中需要先调用CreateAddrReg手动设定地址偏移量再调用搬运指令。RegTraitNumOne双搬出完成一次搬运后UB地址不会自动更新每次迭代需要手动更新地址。不支持在每次迭代中需要先调用CreateAddrReg手动设定地址偏移量再调用搬运指令。RegTraitNumTwo单搬出完成一次搬运后UB地址不会自动更新每次迭代需要手动更新地址。• POST_MODE_NORMAL模式不支持。• POST_MODE_UPDATE模式完成一次搬运后UB地址会自动更新每次迭代不需要手动更新地址。不支持RegTraitNumTwo双搬出不支持不支持不支持数据搬出时可以通过模板参数配置搬运的数据分布模式能够实现压缩、只搬出第一个元素等功能分布模式具体介绍请参考表3 StoreDist参数说明单搬出模式、表6 StoreDist参数说明双搬出模式。图1展示了DIST_NORM、DIST_FIRST_ELEMENT_B16、DIST_PACK_B32分布模式。图 1数据搬出分布模式图示连续对齐搬入/搬出的三类搬运接口在迭代内调用方式请参考关键特性说明。函数原型单搬出模式普通搬运接口template typename T DefaultType, StoreDist dist StoreDist::DIST_NORM, typename U __simd_callee__ inline void StoreAlign(__ubuf__ T* dstAddr, U srcReg, MaskReg mask)单搬出模式PostUpdate扩展搬运接口template typename T DefaultType, PostLiteral postMode, StoreDist dist StoreDist::DIST_NORM, typename U __simd_callee__ inline void StoreAlign(__ubuf__ T* dstAddr, U srcReg, int32_t postUpdateStride, MaskReg mask)单搬出模式使用AddrReg存储偏移量需要配合CreateAddrReg接口使用template typename T DefaultType, StoreDist dist StoreDist::DIST_NORM, typename U __simd_callee__ inline void StoreAlign(__ubuf__ T* dstAddr, U srcReg, AddrReg offset, MaskReg mask)双搬出模式普通搬运接口template typename T DefaultType, StoreDist dist, typename U __simd_callee__ inline void StoreAlign(__ubuf__ T* dstAddr, U srcReg0, U srcReg1, MaskReg mask)双搬出模式使用AddrReg存储偏移量需要配合CreateAddrReg接口使用template typename T DefaultType, StoreDist dist, typename U __simd_callee__ inline void StoreAlign(__ubuf__ T* dstAddr, U srcReg0, U srcReg1, AddrReg offset, MaskReg mask)参数说明表 1模板参数说明单搬出模式参数名描述T目的操作数的数据类型。支持的数据类型请参考数据类型。postMode用于控制是否使能UB地址会自动更新功能post updatePostLiteral类型。具体参数说明请参考PostLiteral。dist搬运模式StoreDist类型。具体参数说明请参考表3 StoreDist参数说明单搬出模式。U源操作数的RegTensor类型。例如RegTensorhalf由编译器自动推导用户不需要手动填写。表 2参数说明单搬出模式参数名输入/输出描述dstAddr输入/输出目的操作数UB起始地址。srcReg输入源操作数类型为RegTensor。postUpdateStride输入地址更新步长int32_t类型单位元素个数。postUpdateStride * sizeof(T)的对齐约束请参考表3 StoreDist参数说明单搬出模式。根据postMode的取值有两种情况• POST_MODE_NORMAL不支持。• POST_MODE_UPDATE实际搬运UB起始地址为dstAddr搬运后执行地址更新dstAddr postUpdateStride。offset输入地址偏移寄存器AddrReg类型。需要在搬运前调用CreateAddrReg设定地址偏移量实际搬运UB地址为dstAddr 寄存器中存储的偏移量。AddrReg寄存器中存储的偏移量 * sizeof(T)的对齐约束请参考表3 StoreDist参数说明单搬出模式。mask输入源操作数元素操作的有效指示详细说明请参考MaskReg。表 3StoreDist参数说明单搬出模式StoreDist含义对齐约束ByteDIST_NORM_B8正常模式搬运数据量为VL数据类型为b8。32DIST_NORM_B16正常模式搬运数据量为VL数据类型为b16。32DIST_NORM_B32正常模式搬运数据量为VL数据类型为b32。32DIST_FIRST_ELEMENT_B8忽略mask向dst中搬运src第一个元素数据类型为b8。1DIST_FIRST_ELEMENT_B16忽略mask向dst中搬运src第一个元素数据类型为b16。可参考图1 数据搬出分布模式图示。2DIST_FIRST_ELEMENT_B32忽略mask向dst中搬运src第一个元素数据类型为b32。4DIST_PACK_B16压缩模式数据类型为b16根据mask将src中有效元素的低半部分bit数据连续存储于dst中。例数据类型为uint16_tsrc: [0x3210, 0x7654, 0xBA98, 0xFEDC, ..., 0xFEDC, 0xBA98, 0x7654, 0x3210]。dst: [0x5410, 0xDC98, ... 0x98DC, 0x1054]。min(32, VL/2)DIST_PACK_B32压缩模式数据类型为b32根据mask将src中有效元素的低半部分bit数据连续存储于dst中。可参考图1 数据搬出分布模式图示。min(32, VL/2)DIST_PACK_B64压缩模式数据类型为b64根据mask将src中有效元素的低半部分bit数据连续存储于dst中。min(32, VL/2)DIST_PACK4_B32压缩模式数据类型为b32根据mask将src中有效元素的低8bit四分之一数据连续存储于dst中。min(32, VL/4)DIST_NORM正常模式搬运VL数据支持数据类型b8/b16/b32系统会根据模板参数T自动选择DIST_NORM_B8、DIST_NORM_B16、DIST_NORM_B32。可参考图1 数据搬出分布模式图示。32表 4模板参数说明双搬出模式参数名描述T目的操作数的数据类型。支持的数据类型请参考数据类型。dist搬运模式StoreDist类型。具体参数说明请参考表6 StoreDist参数说明双搬出模式。U源操作数的RegTensor类型。例如RegTensorhalf由编译器自动推导用户不需要手动填写。表 5参数说明双搬出模式参数名输入/输出描述dstAddr输入/输出目的操作数UB起始地址。srcReg0输入第一个源操作数类型为RegTensor。srcReg1输入第二个源操作数类型为RegTensor。offset输入地址偏移量AddrReg类型。需要在搬运前调用CreateAddrReg设定地址偏移量实际搬运UB地址为dstAddr offset存储的偏移量搬运后地址不会更新。AddrReg寄存器中存储的偏移量*sizeof(T)需要32字节对齐。mask输入源操作数元素操作的有效指示详细说明请参考MaskReg。表 6StoreDist参数说明双搬出模式StoreDist含义对齐约束(Byte)DIST_INTLV_B8双搬出模式数据类型为b8忽略mask将src0src1中的元素交错存储于dst中dst长度需要为2*VL。例数据类型为uint8_tsrc0: [0, 2, 4, 6, ... 254, 0, 2, 4, ..., 252, 254]。src1: [1, 3, 5, 7, ..., 255, 1, 3, 5, ..., 253, 255]。dst: [0, 1, 2, 3, ..., 254, 255, 0, 1, 2, 3, ..., 253, 254, 255]。32DIST_INTLV_B16双搬出模式数据类型为b16忽略mask将src0src1中的元素交错存储于dst中dst长度需要为2*VL。32DIST_INTLV_B32双搬出模式数据类型为b32忽略mask将src0src1中的元素交错存储于dst中dst长度需要为2*VL。32数据类型目的操作数与源操作数的数据类型需要保持一致。当RegTensor模板参数regTrait为RegTraitNumOne时支持的数据类型为b8、b16、b32、b64。当RegTensor模板参数regTrait为RegTraitNumTwo时支持的数据类型为complex32、b64。返回值说明无约束说明位于Unified Buffer的地址约束、postUpdateStride * sizeof(T)的对齐约束、AddrReg寄存器中存储的偏移量 * sizeof(T)的对齐约束与分布模式StoreDist有关具体地址约束请参考表3 StoreDist参数说明单搬出模式、表6 StoreDist参数说明双搬出模式。单搬出模式b64数据类型只支持StoreDist中的DIST_NORM和DIST_PACK_B64模式。当RegTensor模板参数RegTrait为RegTraitNumOne和RegTraitNumTwo时支持情况如下 | RegTensor模板参数RegTrait取值 | 支持的接口 | 支持的数据类型 | |-----|-----|-----| | RegTraitNumOne | 所有单搬出接口和双搬出接口 | b8/b16/b32/b64 | | RegTraitNumTwo | • 单搬出 普通搬运接口• 单搬出 PostUpdate扩展搬运接口 | complex32/b64 |关键特性说明连续对齐搬入搬出三类搬运接口调用方式图2展示了连续对齐搬入/搬出的三类搬运模式在迭代内使用调用方式普通搬运接口完成一次搬运后UB地址不会自动更新每次迭代需要手动更新地址。PostUpdate扩展搬运接口完成一次搬运后UB地址会自动更新每次迭代不需要手动更新地址。使用AddrReg存储偏移量接口在每次迭代中需要先调用CreateAddrReg手动设定地址偏移量再调用搬运指令。图 2regTrait为RegTraitNumOne时的三类搬运模式调用方式调用示例单搬入/单搬出模式普通搬运接口template typename T __simd_vf__ inline void ComputeMode01(__ubuf__ T* dstAddr, __ubuf__ T* srcAddr, uint32_t dstSize, uint16_t oneRepeatSize, uint16_t repeatTimes) { AscendC::Reg::RegTensorT dstReg; AscendC::Reg::MaskReg mask; for (uint16_t i 0; i repeatTimes; i) { mask AscendC::Reg::UpdateMaskT(dstSize); AscendC::Reg::LoadAlign(dstReg, srcAddr i * oneRepeatSize); AscendC::Reg::StoreAlign(dstAddr i * oneRepeatSize, dstReg, mask); } }单搬入/单搬出模式PostUpdate扩展搬运接口template typename T __simd_vf__ inline void ComputeMode02(__ubuf__ T* dstAddr, __ubuf__ T* srcAddr, uint32_t dstSize, uint16_t oneRepeatSize, uint16_t repeatTimes) { AscendC::Reg::RegTensorT dstReg; AscendC::Reg::MaskReg mask; for (uint16_t i 0; i repeatTimes; i) { mask AscendC::Reg::UpdateMaskT(dstSize); AscendC::Reg::LoadAlignT, AscendC::Reg::PostLiteral::POST_MODE_UPDATE(dstReg, srcAddr, oneRepeatSize); AscendC::Reg::StoreAlignT, AscendC::Reg::PostLiteral::POST_MODE_UPDATE(dstAddr, dstReg, oneRepeatSize, mask); } }单搬入/单搬出模式使用AddrReg存储偏移量template typename T __simd_vf__ inline void ComputeMode03(__ubuf__ T* dstAddr, __ubuf__ T* srcAddr, uint16_t oneRepeatSize, uint16_t repeatTimes) { AscendC::Reg::RegTensorT dstReg; AscendC::Reg::MaskReg mask AscendC::Reg::CreateMaskT(); AscendC::Reg::AddrReg aReg; for (uint16_t i 0; i repeatTimes; i) { aReg AscendC::Reg::CreateAddrRegT(i, oneRepeatSize); AscendC::Reg::LoadAlign(dstReg, srcAddr, aReg); AscendC::Reg::StoreAlign(dstAddr, dstReg, aReg, mask); } }双搬入/双搬出模式普通搬运接口template typename T __simd_vf__ inline void ComputeMode04(__ubuf__ T* dstAddr, __ubuf__ T* srcAddr, uint16_t oneRepeatSize, uint16_t repeatTimes) { AscendC::Reg::RegTensorT srcReg0; AscendC::Reg::RegTensorT srcReg1; AscendC::Reg::MaskReg mask AscendC::Reg::CreateMaskuint8_t, AscendC::Reg::MaskPattern::ALL(); for (uint16_t i 0; i repeatTimes; i) { AscendC::Reg::LoadAlignT, AscendC::Reg::LoadDist::DIST_DINTLV_B8(srcReg0, srcReg1, srcAddr i * oneRepeatSize); AscendC::Reg::StoreAlignT, AscendC::Reg::StoreDist::DIST_INTLV_B8(dstAddr i * oneRepeatSize, srcReg0, srcReg1, mask); } }双搬入/双搬出模式使用AddrReg存储偏移量需要配合CreateAddrReg接口使用template typename T __simd_vf__ inline void ComputeMode05(__ubuf__ T* dstAddr, __ubuf__ T* srcAddr, uint16_t oneRepeatSize, uint16_t repeatTimes) { AscendC::Reg::RegTensorT srcReg0; AscendC::Reg::RegTensorT srcReg1; AscendC::Reg::MaskReg mask AscendC::Reg::CreateMaskT(); AscendC::Reg::AddrReg aReg; for (uint16_t i 0; i repeatTimes; i) { aReg AscendC::Reg::CreateAddrRegT(i, oneRepeatSize); AscendC::Reg::LoadAlignT, AscendC::Reg::LoadDist::DIST_DINTLV_B8(srcReg0, srcReg1, srcAddr, aReg); AscendC::Reg::StoreAlignT, AscendC::Reg::StoreDist::DIST_INTLV_B8(dstAddr, srcReg0, srcReg1, aReg, mask); } }【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考