CANN/Ascend C SIMD Select数据过滤API

Select【免费下载链接】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 训练系列产品不支持功能说明给定两个源操作数src0和src1根据maskTensor相应位置的值非bit位选取元素得到目的操作数dst。选择的规则为当Mask的值为0时从src0中选取否则从src1选取。该接口支持多维Shape需满足maskTensor和源操作数Tensor的前轴非尾轴元素个数相同且maskTensor尾轴元素个数大于等于源操作数尾轴元素个数maskTensor多余部分丢弃不参与计算。maskTensor尾轴需32字节对齐且元素个数为16的倍数。源操作数Tensor尾轴需32字节对齐。如下图样例源操作数src0为Tensorshape为(2,16)数据类型为half尾轴长度满足32字节对齐源操作数src1为scalar数据类型为halfmaskTensor的数据类型为bool为满足对齐要求shape为(2,32)仅有图中蓝色部分的mask掩码生效灰色部分不参与计算。输出目的操作数dstTensor如下图所示。实现原理以float类型ND格式shape为[m, k1]的source输入Tensorshape为[m, k2]的mask Tensor为例描述Select高阶API内部算法框图如下图所示。图1Select算法框图计算过程分为如下几步均在Vector上进行GatherMask步骤如果k1, k2不相等则根据src的shape[m, k1]对输入mask[m, k2]通过GatherMask进行reduce计算使得mask的k轴多余部分被舍去shape转换为[m, k1]Cast步骤将上一步的mask结果cast成half类型Compare步骤使用Compare接口将上一步的mask结果与0进行比较得到cmpmask结果Select步骤根据cmpmask的结果选择srcTensor相应位置的值或者scalar值输出Output。图2Select算法框图计算过程在Vector上进行循环m次每次对k1个元素进行如下操作Compare步骤使用Compare接口将mask值与0进行比较得到cmpmask结果Select步骤根据cmpmask的结果选择srcTensor相应位置的值或者scalar值输出Output。函数原型src0为srcTensortensor类型src1为srcScalarscalar类型template typename T, typename U, bool isReuseMask true __aicore__ inline void Select(const LocalTensorT dst, const LocalTensorT src0, T src1, const LocalTensorU mask, const LocalTensoruint8_t sharedTmpBuffer, const SelectWithBytesMaskShapeInfo info)src0为srcScalarscalar类型src1为srcTensortensor类型template typename T, typename U, bool isReuseMask true __aicore__ inline void Select(const LocalTensorT dst, T src0, const LocalTensorT src1, const LocalTensorU mask, const LocalTensoruint8_t sharedTmpBuffer, const SelectWithBytesMaskShapeInfo info)该接口需要额外的临时空间来存储计算过程中的中间变量。临时空间需要开发者申请并通过sharedTmpBuffer入参传入。临时空间大小BufferSize的获取方式如下通过GetSelectMaxMinTmpSize中提供的接口获取需要预留空间范围的大小。参数说明表1模板参数说明参数名描述T操作数的数据类型。支持的数据类型为half、float。U掩码Tensor mask的数据类型。支持的数据类型为bool、int8_t、uint8_t、int16_t、uint16_t、int32_t、uint32_t。isReuseMask是否允许修改maskTensor。默认为true。取值为true时仅在maskTensor尾轴元素个数和srcTensor尾轴元素个数不同的情况下maskTensor可能会被修改其余场景maskTensor不会修改。取值为false时任意场景下maskTensor均不会修改但可能会需要更多的临时空间。表2接口参数说明参数名称输入/输出含义dst输出目的操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。src0(srcTensor)src1(srcTensor)输入源操作数。源操作数Tensor尾轴需32字节对齐。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。src1(srcScalar)src0(srcScalar)输入源操作数。类型为scalar。mask输入掩码Tensor。用于描述如何选择srcTensor和srcScalar之间的值。maskTensor尾轴需32字节对齐且元素个数为16的倍数。src0为srcTensortensor类型src1为srcScalarscalar类型若mask的值为0选择srcTensor相应的值放入dstLocal否则选择srcScalar的值放入dstLocal。src0为srcScalarscalar类型src1为srcTensortensor类型若mask的值为0选择srcScalar的值放入dstLocal否则选择srcTensor相应的值放入dstLocal。sharedTmpBuffer输入该API用于计算的临时空间所需空间大小根据GetSelectMaxMinTmpSize获取。info输入描述SrcTensor和maskTensor的shape信息。SelectWithBytesMaskShapeInfo类型定义如下方代码所示其中参数的含义为firstAxissrcLocal/maskTensor的前轴元素个数。srcLastAxissrcLocal的尾轴元素个数。maskLastAxismaskTensor的尾轴元素个数。注意需要满足srcTensor和maskTensor的前轴元素个数相同均为firstAxis。需要满足firstAxis * srcLastAxis srcTensor.GetSize()firstAxis * maskLastAxis maskTensor.GetSize()。maskTensor尾轴的元素个数大于等于srcTensor尾轴的元素个数计算时会丢弃maskTensor多余部分不参与计算。struct SelectWithBytesMaskShapeInfo { __aicore__ SelectWithBytesMaskShapeInfo() {}; uint32_t firstAxis 0; uint32_t srcLastAxis 0; uint32_t maskLastAxis 0; };返回值说明无约束说明源操作数与目的操作数可以复用。操作数地址对齐要求请参见通用地址对齐约束。maskTensor尾轴元素个数和源操作数尾轴元素个数不同的情况下 maskTensor的数据有可能被接口改写。调用示例完整示例可以参考Select样例。// dstLocal结果数据 // srcLocal1输入张量 // scalar待比较标量 // maskLocal选择掩码 // tmpBuffer临时空间 // infosrcLocal和maskLocal的shape信息 AscendC::SelectWithBytesMaskShapeInfo info; AscendC::Select(dstLocal, srcLocal1, scalar, maskLocal, tmpBuffer, info);结果示例如下输入数据srcLocal1: [-84.6 -24.38 30.97 -30.25 22.28 -92.56 90.44 -58.72 -86.56 5.74 6.754 -86.3 -96.7 -37.38 -81.9 46.9 -99.4 94.2 -41.78 -60.3 -14.43 78.6 8.93 -65.2 79.94 -46.88 4.516 20.03 -25.56 24.73 0.3223 21.98 -87.4 -93.9 46.22 -69.9 90.8 -24.17 -96.2 -91. 90.44 9.766 68.25 -57.78 -75.44 -8.86 -91.56 21.6 76. 82.1 -78. -23.75 92. -66.44 75. 94.9 2.62 -90.9 15.945 38.16 50.84 96.94 -59.38 44.22 ] 输入数据scalar: [35.6] 输入数据maskLocal: [False True False False True True False True True False False True False True False True True False False False True True True True True False True False True True True True False False True False True False True False True False True False True True True False True False True False True False True True True False False False True False True True ] 输出数据dstLocal: [-84.6 35.6 30.97 -30.25 35.6 35.6 90.44 35.6 35.6 5.74 6.754 35.6 -96.7 35.6 -81.9 35.6 35.6 94.2 -41.78 -60.3 35.6 35.6 35.6 35.6 35.6 -46.88 35.6 20.03 35.6 35.6 35.6 35.6 -87.4 -93.9 35.6 -69.9 35.6 -24.17 35.6 -91. 35.6 9.766 35.6 -57.78 35.6 35.6 35.6 21.6 35.6 82.1 35.6 -23.75 35.6 -66.44 35.6 35.6 35.6 -90.9 15.945 38.16 35.6 96.94 35.6 35.6 ]【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考