人工智能深度学习算子库CANNAscend【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址https://gitcode.com/cann/asc-devkit点击查看免费下载导读本文聚焦 CANN asc-devkit 开源仓库中 Reg 矢量计算模块的「比较与选择」接口族——Compare、Compares、Select与Squeeze系统讲解其功能语义、函数原型、参数与数据类型约束、掩码配合方式并结合仓库头文件实现与测试用例剖析底层指令映射。读完本文你将掌握在昇腾 SIMD 算子中完成逐元素比较、标量比较、掩码选值与压缩搬出Squeeze StoreUnAlign 连续非对齐搬出的完整编程范式。一、接口总览与适用场景「比较与选择」位于 Reg 矢量计算Reg Vector Compute范畴核心思路是用 MaskReg 掩码记录比较结果再用该掩码驱动数据选择与压缩。四个接口的分工如下表所示接口核心语义典型用途Compare逐元素比较两个 RegTensor结果写入 MaskReg 对应比特位求两向量逐元素大小关系产出布尔掩码Compares逐元素比较 RegTensor 与一个标量结果写入 MaskReg阈值过滤、边界判断如x thresholdSelect按 mask 比特位从两个源操作数中选取元素三目运算/条件选值常与 Compare 联动Squeeze将被 mask 选中的有效元素压缩、连续排列到目的寄存器稀疏数据紧凑化配合非对齐搬出写回从仓库的接口声明看Compare、Compares、Select三个接口集中声明于 include/basic_api/reg_compute/kernel_reg_compute_vec_cmpsel_intf.hSqueeze则由 include/basic_api/reg_compute/kernel_reg_compute_gather_mask_intf.h 提供文档中给出的头文件引用路径basic_api/reg_compute/kernel_reg_compute_vec_cmpsel_intf.h与basic_api/reg_compute/kernel_reg_compute_gather_mask_intf.h正是这两个公开接口头文件内部实现在impl/basic_api/reg_compute/下按芯片架构分目录存放。需要特别说明的是产品支持范围以仓库文档为准该接口族仅支持 Ascend 950PR950DT 系列、Kirin 9050 处理器系列以及 KirinDev0001/0002/0003 处理器系列Atlas A2/A3 系列、Atlas 推理系列AI Core/Vector Core、Atlas 训练系列与 Atlas 200I/500 A2 推理产品均不支持。编写算子前请务必核对目标芯片是否在支持清单内。二、Compare逐元素比较两个 RegTensor2.1 功能说明Compare逐元素比较两个RegTensor的大小根据模板参数mode指定的比较模式将比较结果写入目的操作数MaskReg的对应比特位比较结果为真对应比特位置 1否则置 0。比较模式支持EQ等于、GT大于、GE大于或等于、NE不等于、LE小于或等于、LT小于。2.2 函数原型与参数template typename T DefaultType, CMPMODE mode CMPMODE::EQ, typename U __simd_callee__ inline void Compare(MaskReg dst, U srcReg0, U srcReg1, MaskReg mask)模板参数与函数参数说明类别参数说明模板参数T操作数数据类型如half、float。模板参数mode比较模式取值LT/GT/GE/EQ/NE/LE。模板参数U源操作数的RegTensor类型由编译器自动推导无需填写。函数参数dst输出MaskReg类型目的操作数。函数参数srcReg0、srcReg1输入源操作数两者数据类型需一致。函数参数mask输入指示哪些元素参与比较元素级有效掩码。返回值无。2.3 数据类型组合文档给出的支持组合如下srcReg0/srcReg1数据类型 →dst数据类型源数据类型目的数据类型int8_t / uint8_t / int16_t / uint16_t / half / bfloat16_t / int32_t / uint32_t / float / int64_t / uint64_tMaskReg2.4 约束说明通过mask参数控制的未选中元素在目的操作数中被置零。操作数重叠约束srcReg0与srcReg1可以是同一个RegTensor即允许Compare(dst, x, x, mask)这类自比较写法。2.5 源码级实现印证从 impl/basic_api/reg_compute/dav_3510/kernel_reg_compute_vec_cmpsel_impl.h 可以看到CompareImpl的实现逻辑对应 3510 芯片系列分支对小于 8 字节的类型int8/16/32、uint8/16/32、half、float、bfloat16_t按mode直接映射为单条矢量比较指令vcmp_eq、vcmp_ne、vcmp_gt、vcmp_ge、vcmp_lt、vcmp_le见 CompareImpl。对 **int64_t/uint64_t8 字节**类型硬件没有直接的 64 位比较指令实现上先对高低 32 位分别比较如GT模式高字比较相等则取低字vcmp_gt结果、否则取高字vcmp_gt结果通过psel选出再组合出最终掩码——这正是文档中 int64_t/uint64_t 也受支持、但语义上需要拆字的底层原因见 CompareB64Impl 与CompareUint64Impl/CompareInt64Impl。同时入口头文件 include/basic_api/reg_compute/kernel_reg_compute_vec_cmpsel_intf.h 开头有一段编译期保护当启用--enable-simt编译时直接#error说明该接口族仅面向 SIMD 编译路径。2.6 调用示例比较后按掩码选值仓库文档示例完整展示了Compare Select StoreAlign的典型数据流掩码比较 → 条件选值 → 存回templatetypename T __simd_vf__ inline void CompareVF(__ubuf__ T* dstAddr, __ubuf__ T* src0Addr, __ubuf__ T* src1Addr, uint32_t count, uint32_t oneRepeatSize, uint16_t repeatTimes) { AscendC::Reg::RegTensorT srcReg0; // 源操作数0 AscendC::Reg::RegTensorT srcReg1; // 源操作数1 AscendC::Reg::RegTensorT dstReg; // 目的操作数存储Select选取的结果 AscendC::Reg::MaskReg mask; AscendC::Reg::MaskReg dstMask; // 比较结果掩码相等位置为1 for (uint16_t i 0; i repeatTimes; i) { mask AscendC::Reg::UpdateMaskT(count); // 根据实际元素个数更新mask AscendC::Reg::LoadAlign(srcReg0, src0Addr i * oneRepeatSize); // 加载源操作数0 AscendC::Reg::LoadAlign(srcReg1, src1Addr i * oneRepeatSize); // 加载源操作数1 AscendC::Reg::CompareT, AscendC::CMPMODE::EQ(dstMask, srcReg0, srcReg1, mask); // 逐元素比较是否相等 AscendC::Reg::Select(dstReg, srcReg0, srcReg1, dstMask); // 按比较结果掩码选取元素 AscendC::Reg::StoreAlign(dstAddr i * oneRepeatSize, dstReg, mask); // 存回dstAddr } }示例中half类型输入src0Addr [1.0, 4.0, 3.0, 2.0, 5.0, 6.0, 7.0, 8.0, ...]、src1Addr [1.0, 3.0, 3.0, 5.0, 2.0, 6.0, 9.0, 4.0, ...]EQ 模式下dstMask在相等位置pos0、pos2、pos5、pos10、pos14...为 1其余为 0随后Select将 mask 为 1 的位置取srcReg0的元素、为 0 的位置取srcReg1的元素。需要说明UpdateMask由当前有效元素个数count动态生成掩码LoadAlign/StoreAlign按oneRepeatSize对齐搬运repeatTimes控制重复次数这是 Reg 编程中标准的「分块循环」写法。三、ComparesRegTensor 与标量的逐元素比较3.1 功能说明Compares将源操作数srcReg中每个元素与一个标量scalarValue比较根据modeLT/GT/GE/EQ/NE/LE判断结果为真则输出MaskReg对应比特位置 1否则置 0。它是Compare的标量版特别适合阈值判断、边界过滤等场景无需先把标量复制成 RegTensor。3.2 函数原型与参数template typename T DefaultType, CMPMODE mode CMPMODE::EQ, typename U, typename S __simd_callee__ inline void Compares(MaskReg dst, U srcReg, S scalarValue, MaskReg mask)类别参数说明模板参数T矢量源操作数数据类型。模板参数mode比较模式取值LT/GT/GE/EQ/NE/LE。模板参数U矢量源操作数的RegTensor类型自动推导。模板参数S标量操作数数据类型。函数参数dst输出MaskReg类型。函数参数srcReg输入RegTensor类型源操作数。函数参数scalarValue输入标量源操作数。函数参数mask输入元素有效指示。返回值无。数据类型组合与Compare一致int8_t 至 int64_t/uint64_t 等 11 种类型 → MaskReg。3.3 约束与实现要点通过mask控制的未选中元素在目的操作数中被置零。从实现看CompareScalarImpl小于 8 字节的类型直接映射为vcmps_eq/vcmps_ne/vcmps_gt/vcmps_ge/vcmps_lt/vcmps_le标量比较指令8 字节类型int64_t/uint64_t则先经Duplicate将标量广播为 RegTensor再走CompareT, mode, U的矢量比较路径。实现同时用static_assert(Std::is_convertibleS, ActualT(), ...)保证标量类型可转换到 RegTensor 数据类型。3.4 调用示例templatetypename T __simd_vf__ inline void ComparesVF(__ubuf__ T* dstAddr, __ubuf__ T* src0Addr, __ubuf__ T* src1Addr, T scalarValue, uint32_t count, uint32_t oneRepeatSize, uint16_t repeatTimes) { AscendC::Reg::RegTensorT srcReg0; AscendC::Reg::RegTensorT srcReg1; AscendC::Reg::RegTensorT dstReg; AscendC::Reg::MaskReg mask; AscendC::Reg::MaskReg cmpMaskReg; for (uint16_t i 0; i repeatTimes; i) { mask AscendC::Reg::UpdateMaskT(count); AscendC::Reg::LoadAlign(srcReg0, src0Addr i * oneRepeatSize); AscendC::Reg::LoadAlign(srcReg1, src1Addr i * oneRepeatSize); AscendC::Reg::ComparesT, AscendC::CMPMODE::EQ(cmpMaskReg, srcReg0, scalarValue, mask); AscendC::Reg::Select(dstReg, srcReg0, srcReg1, cmpMaskReg); AscendC::Reg::StoreAlign(dstAddr i * oneRepeatSize, dstReg, mask); } }该示例以ComparesEQ 模式生成掩码再与Select联动实现「等于标量的位置取 src0否则取 src1」的条件写回。仓库另保留了历史接口CompareScalar与Compares语义相同其在 3510/5102/5162/3003/3113 等架构下可用但 Kirin v516 系列5101/5161/5163/5165不支持新代码建议统一使用Compares见 kernel_reg_compute_vec_cmpsel_intf_impl.h。四、Select按掩码从两个源操作数中选值4.1 功能说明给定srcReg0、srcReg1两个源操作数Select按mask比特位选取元素写入dstRegmask 比特位为 1 时取srcReg0对应位置元素为 0 时取srcReg1对应位置元素等价于 SIMD 化的三目运算符mask ? src0 : src1。4.2 函数原型与参数template typename T DefaultType, typename U __simd_callee__ inline void Select(U dstReg, U srcReg0, U srcReg1, MaskReg mask)类别参数说明模板参数T操作数数据类型。模板参数U源/目的操作数的RegTensor类型自动推导。函数参数dstReg输出目的操作数。函数参数srcReg0、srcReg1输入源操作数数据类型需与目的操作数一致。函数参数mask输入选择规则比特位为 1 取 srcReg0为 0 取 srcReg1。返回值无。4.3 数据类型Select支持的范围比 Compare 更宽包括bool、int8_t、uint8_t、int16_t、uint16_t、half、bfloat16_t、int32_t、uint32_t、float、int64_t、uint64_t且源/目的类型一一对应如 half→halfbool→bool不产生 MaskReg 输出。4.4 约束说明与 Compare 不同Select的约束是通过 mask 控制的未选中的元素保持源操作数的值而非清零即结果中未选中位置仍来自 srcReg1。4.5 源码级实现印证SelectImpl的实现见 SelectImpl按字节宽分派1/2/4 字节类型直接映射为单条vsel指令内部按uint8_t/uint16_t/uint32_t视图执行8 字节类型int64_t/uint64_t拆成两个 32 位视图各执行一次vsel若源 RegTensor 是RegTraitNumOne还需先用MaskPack扩展掩码、再经B64TraitOneToTraitTwo/B64TraitTwoToTraitOne完成单/双寄存器 trait 的转换——这也解释了为什么 64 位类型的 Select 语义与其他类型一致。4.6 调用示例文档给出的示例与 2.6 节基本一致Compare生成cmpRegSelect消费该掩码templatetypename T __simd_vf__ inline void SelectVF(__ubuf__ T* dstAddr, __ubuf__ T* src0Addr, __ubuf__ T* src1Addr, uint32_t count, uint32_t oneRepeatSize, uint16_t repeatTimes) { AscendC::Reg::RegTensorT srcReg0; AscendC::Reg::RegTensorT srcReg1; AscendC::Reg::RegTensorT dstReg; AscendC::Reg::MaskReg mask; AscendC::Reg::MaskReg cmpReg; for (uint16_t i 0; i repeatTimes; i) { mask AscendC::Reg::UpdateMaskT(count); AscendC::Reg::LoadAlign(srcReg0, src0Addr i * oneRepeatSize); AscendC::Reg::LoadAlign(srcReg1, src1Addr i * oneRepeatSize); AscendC::Reg::CompareT, AscendC::CMPMODE::EQ(cmpReg, srcReg0, srcReg1, mask); AscendC::Reg::Select(dstReg, srcReg0, srcReg1, cmpReg); AscendC::Reg::StoreAlign(dstAddr i * oneRepeatSize, dstReg, mask); } }五、Squeeze掩码驱动的元素压缩与连续搬出5.1 功能说明Squeeze将srcReg中被 mask 选中的有效元素依次复制到dstReg有效元素在目的寄存器中从低到高连续排列剩余位置元素置 0。这一行为示意可参考仓库文档配图 docs/zh/api/figures/reg_squeeze.pngSqueeze 计算示意图。该接口由 include/basic_api/reg_compute/kernel_reg_compute_gather_mask_intf.h 提供。5.2 函数原型与参数template typename T DefaultType, GatherMaskMode store GatherMaskMode::NO_STORE_REG, typename U __simd_callee__ inline void Squeeze(U dstReg, U srcReg, MaskReg mask)类别参数说明模板参数T目的/源操作数数据类型。模板参数storeGatherMaskMode枚举NO_STORE_REG有效元素总字节数不存入 AR 寄存器默认/STORE_REG有效元素总字节数存入 AR 寄存器。AR 寄存器说明见 GetSpr。模板参数U目的/源操作数RegTensor类型自动推导。函数参数dstReg输出目的操作数。函数参数srcReg输入源操作数。函数参数mask输入控制每次迭代内参与计算的元素。返回值无。5.3 数据类型支持的数据类型为int8_t、uint8_t、int16_t、uint16_t、half、int32_t、uint32_t、float注意与 Compare/Select 不同Squeeze 不支持 bfloat16_t、int64_t/uint64_t 与浮点 64 位类型。实现中的static_assert(SupportTypeActualT, int8_t, uint8_t, int16_t, uint16_t, int32_t, uint32_t, half, float(), ...)与文档完全一致见 GatherMaskImpl底层映射为vsqz指令并把GatherMaskMode转换为StoreMode作为指令模式参数。5.4 与 StoreUnAlign 的配合及硬件约束Squeeze最有价值的用法是与连续非对齐搬出接口配合实现「压缩后连续写回」当store GatherMaskMode::STORE_REG时Squeeze会把有效元素的总字节数存入AR 特殊寄存器SpecialPurposeReg::AR此时配合 StoreUnAlign_continuous 的场景 3使用 AR 寄存器存储偏移量、无需显式传入偏移量接口StoreUnAlign会自动从 AR 寄存器读取有效字节数作为地址偏移。硬件约束STORE_REG 模式下StoreUnAlign指令与Squeeze指令必须交替使用例如Squeeze(dstVreg, srcVreg, mask); StoreUnAlign(dstAddr, dstVreg, ureg); Squeeze(dstVreg, srcVreg, mask); StoreUnAlign(dstAddr, dstVreg, ureg);当store GatherMaskMode::NO_STORE_REG时不涉及 AR 寄存器Squeeze与StoreUnAlign不强制交替。此外按 GetSpr 的说明AR 寄存器的读取GetSprSpecialPurposeReg::AR()返回 int64_t单位字节只能在 VF 函数之外__aicore__上下文中调用在调用计算前可先用ClearSprSpecialPurposeReg::AR()将 AR 寄存器清零避免遗留脏数据影响偏移计算。5.5 调用示例STORE_REG 模式 StoreUnAlign 连续搬出仓库文档给出的典型示例SqueezeSTORE_REG 模式StoreUnAlign使用 AR 寄存器中的有效元素字节数作为偏移量templatetypename T __simd_vf__ inline void SqueezeVF(__ubuf__ T* dstAddr, __ubuf__ T* srcAddr, uint32_t oneRepeatSize, uint16_t repeatTimes) { AscendC::Reg::RegTensorT srcReg0; AscendC::Reg::RegTensorT srcReg1; AscendC::Reg::UnalignRegForStore ureg; AscendC::Reg::MaskReg sqzMask AscendC::Reg::CreateMaskT, AscendC::Reg::MaskPattern::H(); for (uint16_t i 0; i repeatTimes; i) { AscendC::Reg::LoadAlignT, AscendC::Reg::PostLiteral::POST_MODE_UPDATE(srcReg0, srcAddr, oneRepeatSize); AscendC::Reg::SqueezeT, AscendC::Reg::GatherMaskMode::STORE_REG(srcReg1, srcReg0, sqzMask); AscendC::Reg::StoreUnAlignT, AscendC::Reg::PostLiteral::POST_MODE_UPDATE(dstAddr, srcReg1, ureg); } AscendC::Reg::StoreUnAlignPost(dstAddr, ureg); }要点拆解CreateMaskT, MaskPattern::H()创建「最低一半元素有效」的掩码H是MaskPattern枚举ALL/VL1~VL128/M3/M4/H/Q/ALLF中的一种模式完整取值见 MaskReg 文档LoadAlign/StoreUnAlign均使用POST_MODE_UPDATE后置更新模式地址指针随每次迭代自动前进每次迭代内部严格按「Squeeze → StoreUnAlign」交替执行满足 STORE_REG 模式下的硬件约束循环结束后调用StoreUnAlignPost(dstAddr, ureg)收尾把未对齐余量写回。若需要验证 AR 寄存器中的有效字节数可在 VF 外通过AscendC::GetSprAscendC::SpecialPurposeReg::AR()读取示例可参考 GetSpr 的调用示例其中MaskPattern::M4掩码下 256 个 float 有效元素对应arNum 256字节。仓库中 tests/api/c_api/npu_arch_3510/vector_datamove/test_asc_squeeze_and_storeunalign.cpp 也通过 stub 桩验证了sprclr清 AR、vsqzSqueeze 底层指令断言MODE_STORED模式与vsturStoreUnAlign 底层指令的完整调用链覆盖 int8_t 至 bfloat16_t 等多种数据类型。六、综合实战阈值过滤 压缩输出的完整范式综合前述四个接口一个典型的「比较 → 选值/压缩 → 连续写回」算子内核可以组织为如下链路伪代码范式各步接口均可替换为上文接口// 阶段1生成比较掩码Compare 或 Compares AscendC::Reg::CompareT, AscendC::CMPMODE::GT(cmpMask, srcReg, thresholdReg, mask); // 或标量版 AscendC::Reg::ComparesT, AscendC::CMPMODE::GT(cmpMask, srcReg, threshold, mask); // 阶段2a条件选值Select1 取 src00 取 src1 AscendC::Reg::Select(dstReg, srcReg0, srcReg1, cmpMask); // 阶段2b压缩有效元素Squeeze STORE_REG 模式配合 StoreUnAlign 连续搬出 AscendC::Reg::SqueezeT, AscendC::Reg::GatherMaskMode::STORE_REG(sqzReg, srcReg, cmpMask); AscendC::Reg::StoreUnAlignT, AscendC::Reg::PostLiteral::POST_MODE_UPDATE(dstAddr, sqzReg, ureg);设计建议需要「逐位置保留三目结果」时选Select需要「只保留满足条件的元素并紧凑存放」时选Squeeze StoreUnAlign(场景3)。参与比较的元素范围由mask控制如UpdateMaskT(count)按实际元素个数更新比较结果再作为下一级 Select/Squeeze 的掩码实现掩码在指令间的级联复用。编写前先核对目标芯片支持情况Ascend 950PR950DT、Kirin 9050 及 KirinDev 系列并确保头文件引用basic_api/reg_compute/kernel_reg_compute_vec_cmpsel_intf.hCompare/Compares/Select与basic_api/reg_compute/kernel_reg_compute_gather_mask_intf.hSqueeze若要读取 AR 寄存器则需在 VF 外调用GetSprSpecialPurposeReg::AR()。七、相关文档与源码索引接口导航与概述compare_and_select、overview各接口详解Compare/Compares/Select/Squeeze 文档位于 compare_and_select 目录数据类型说明RegTensor、MaskReg关联功能连续非对齐搬出StoreUnAlign 场景3、AR 特殊寄存器读取 GetSpr公开接口头文件kernel_reg_compute_vec_cmpsel_intf.h、kernel_reg_compute_gather_mask_intf.h架构分派实现kernel_reg_compute_vec_cmpsel_intf_impl.h、kernel_reg_compute_gather_mask_intf_impl.h、dav_3510 实现测试用例test_asc_squeeze_and_storeunalign.cpp、test_operator_vec_compare.cpp赞分享人工智能深度学习算子库CANNAscend【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址https://gitcode.com/cann/asc-devkit点击查看免费下载相关推荐昇腾AscendC Reg矢量计算比较与选择API全解析Compare/Compares/Select/Squeeze原理与实战昇腾AscendC Reg矢量计算比较与选择API全解析Compare/Compares/Select/Squeeze原理与实战 本文基于 CANN 昇腾人工智能深度学习算子库CANNAscendCANN PyPTO vf 比较与选择指令详解eq/ne/lt/le/gt/ge/select/squeeze 编程指南CANN PyPTO vf 比较与选择指令详解eq/ne/lt/le/gt/ge/select/squeeze 编程指南 本指南围绕 PyPTOParall人工智能编译器模型编译深度学习高性能计算CANNAscendCANN Ascend C 比较与选择接口Compare / Select / GatherMask完全指南掩码驱动的向量级数据筛选与条件选取CANN Ascend C 比较与选择接口Compare / Select / GatherMask完全指南掩码驱动的向量级数据筛选与条件选取 本文是 C人工智能深度学习算子库CANNAscend上一篇Linux-NetSpeed 开源项目教程TCP加速算法全面解析与实践指南下一篇Drawer 开源项目教程创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
阅读完成 · 觉得有帮助?