
CANN ops-math IsinPartV1 算子深度解析基于 CatSort 的 isin 核心部件实现与 aclnn 调用实战【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathIsinPartV1 是 CANN ops-math 仓库 experimental/math/isin_part_v1 目录下提供的一个实验性数学算子承担 isin 功能中判断成员关系的核心部件与 Cat、Sort 算子配合即可拼装出完整的 isin 语义。本文以该算子 README 为主体结合仓库内 op_host、op_kernel、config 与调用示例源码完整讲解 IsinPartV1 的算子语义、参数规格、算法流程、底层实现与 aclnn 单算子调用方法帮助读者在 Atlas A2 训练系列产品 / Atlas 800I A2 推理产品上独立完成该算子的理解、构建与验证。一、算子定位isin 的核心部件设计isin算子的功能是判断给定张量elements的每个元素是否出现在另一个张量test_elements中并返回一个与elements形状一致的布尔张量。在 CANN ops-math 中这一功能并未做成一个整体算子而是被拆解为Cat拼接→ Sort排序→ IsinPartV1成员判定三步IsinPartV1 正是其中负责核心判定的部分这也是其名字中 part 的由来。1.1 完整 isin 的数学表达README 给出了 isin 算子完整实现的公式化描述$$ elementsNum elements.length() \ combine_elements cat(elements, test_elements) \ value, index sort(combine_elements) \ z isin_part_v1(value, index, elementsNum) $$cat将elements与test_elements拼接为一个张量combine_elements其总长度记为totalLengthsort对拼接结果按值排序同时返回排序后的值value与每个值在拼接序列中的原始位置index其中下标0..elementsNum-1对应待判定元素elements下标elementsNum..totalLength-1对应test_elementsisin_part_v1基于排序后相等值必然相邻这一性质扫描value与index只要某个值既出现在elements段又出现在test_elements段就把输出z中对应位置置为true。由于sort后相等的值相邻判定只需要在相邻元素之间进行IsinPartV1 因此可以做到接近线性的复杂度这是该Cat Sort Part组合设计优于朴素两两比较的关键。1.2 从源码确认的输入输出关系在 experimental/math/isin_part_v1/examples/test_aclnn_isin_part_v1.cpp 的调用示例中注释明确给出了语义等价关系语义等价于isin(elements[1,1,4,1,1], test_elements[0,1,2,3])示例直接构造了三个输入value {0, 1, 1, 1, 1, 1, 2, 3, 4}排序后的值elements[1,1,4,1,1]与test_elements[0,1,2,3]拼接后排序index {5, 0, 1, 3, 4, 6, 7, 8, 2}对应原始位置0..4为待判定元素5..8为测试集元素elementsNum 5。期望输出为[1, 1, 0, 1, 1]即 1 出现在测试集中所以为 14 不在测试集中所以为 0——这正好对应了 isin 的成员判定语义。示例中直接构造三个输入、不依赖 Cat/Sort/Cast 等其它算子的做法也再次印证 IsinPartV1 是 isin 链路中可独立验证的核心部件。二、产品支持情况与约束说明根据 README 的产品支持矩阵产品是否支持Atlas A2 训练系列产品 / Atlas 800I A2 推理产品√约束说明无README 明确标注该算子当前无额外约束。从仓库构建配置看CMakeLists.txt 中通过add_all_modules_sources以OPTYPE isin_part_v1、ACLNNTYPE aclnn、COMPUTE_UNIT ascend950、TILING_DIR arch35方式登记该算子同时 isin_part_v1_def.cpp 中为 AICore 增加了ascend910b配置。可以推断该算子面向昇腾 A2 系列910B 体系NPU 编译具体以实际安装的 CANN 版本与硬件形态为准。需要说明的是该算子位于experimental/目录属于实验性算子使用前请关注对应版本的能力说明与配套工具链。三、参数说明README 给出了完整的参数表此处结合算子定义源码进一步补充每个参数在实现层面的细节参数名输入/输出/属性描述数据类型数据格式value输入elements 与 test_elements 拼接排序后的张量值FLOAT、INT32NDindex输入elements 与 test_elements 拼接排序后的张量索引INT32NDelementsNum输入elements 元素数量标量INT64、INT32NDz输出isin 操作后的输出张量形状由 elementsNum 参数决定BOOLND对参数表的补充说明依据源码确认value 与 index 的长度必须一致二者同为拼接后序列长度totalLength的排序结果。在 isin_part_v1_tiling.cpp 的GetShapeAttrsInfo中会校验输入 x、输入 y、输出 z 三者维度数一致且逐维检查inputX.dim outputZ.dim不满足则直接返回失败。elementsNum 是决定输出形状的关键标量它虽然以张量形式传入但语义上是一个标量。InferShape阶段isin_part_v1_infershape.cpp会直接从该输入张量中读出数值若为DT_INT32则取其首个元素将输出形状设置为[elementsNum]outShape-SetDim(0, elementsNum)即输出z是长度为elementsNum的一维 BOOL 张量。关于 elementsNum 的数据类型README 参数表标注支持INT64、INT32但当前仓库实现def.cpp中elementsNum仅注册ge::DT_INT32infershape的GetNumValue也仅处理DT_INT32示例同样以ACL_INT32传入实际只走 INT32 路径。使用 INT64 输入时请以实际版本的实现与校验结果为准。value 支持 FLOAT32 与 INT32 两种 dtype这在 isin_part_v1_binary.json 中体现为两个二进制条目IsinPartV1_a1532827238e1555db7b997c7bce2928int32 输入与IsinPartV1_1532827238e1555db7b997c7bce2928float32 输入index与elementsNum均为 int32、输出z为 bool全部使用 ND 格式format_match_mode为FormatAgnostic。所有输入均为required必选参数且开启AutoContiguousisin_part_v1_def.cpp框架侧会保证张量内存连续。四、算法原理如何用排序相邻性判定成员关系IsinPartV1 的 kernel 实现位于 op_kernel/isin_part_v1.h其核心思想是排序之后相等的值必然连续相邻。因此对于任意一个值 v只要存在某个下标i使得value[i] v且index[i] elementsNum说明 v 来自 elements同时存在某个下标j使得value[j] v且index[j] elementsNum说明 v 来自 test_elements那么 v 就在测试集中出现elements 中所有值为 v 的位置都应输出true。4.1 按 block 划分与边界修正kernel 按多核并行方式处理拼接后的序列按核划分区间Initisin_part_v1.h从 GM 中读取elementsNum根据totalLength与核数GetBlockNum()计算每个核负责的blockStart/blockEnd最后一个核处理到totalLength。边界修正保证等值不跨核排序后相等的值可能恰好被切分到两个核上。CopyInisin_part_v1.h会对区间做两处调整若blockStart处元素与前一个元素相等则向后移动blockStart把相同值完整收进本核区间若blockEnd处元素与后一个元素相等则向后扩展blockEnd。通过这种值边界对齐保证任意一个值的所有实例都落在同一个核内避免跨核比较产生遗漏或重复。4.2 单核内的扫描判定Computeisin_part_v1.h在区间[blockStart, blockEnd)上做单趟扫描维护指针pre指向当前连续相等值段的起点当发现index[j] elementsNum当前值来自测试集且index[pre] elementsNum段首值来自 elements且value[j] value[pre]时说明当前这个值同时出现在 elements 与 test_elements 中触发GMout(pre, j-1)GMoutisin_part_v1.h遍历段内所有index elementsNum的下标将输出z的对应位置写入 1当value[j] ! value[pre]时说明进入了新的值段重置pre j、清除触发标记扫描完成后若发生过写入调用zGmCache.Flush()将 UB 缓存回写到 GM。由于zGmCache使用了SetAtomicAddDataCopyPadisin_part_v1.h即使不同核有重叠回写也能保证最终结果正确输出缓存未命中时通过MTE3_V事件同步保证先刷旧块、再写新块的顺序isin_part_v1.h。4.3 缓存友好的 GM 访问排序后序列长度可能远超 UB 容量kernel 专门实现了两个缓存类GmInCacheisin_part_v1.h为value与index提供按需加载的读缓存以IN_ELEMENTS_SIZE64KB为粒度分块DataCopyPad到 UB索引不在缓存范围内时自动换块GmOutCacheisin_part_v1.h为输出z提供按需刷新的写缓存以OUT_ELEMENTS_SIZE32KB为粒度先Duplicate清零再逐位置 1换块前先 Flush。五、Tiling 与配置host 侧如何为 kernel 准备参数Tiling 是算子在 host 侧为 device kernel 计算切分参数的过程IsinPartV1 的 tiling 实现在 op_host/isin_part_v1_tiling.cpp主要流程如下获取平台信息GetPlatformInfo通过PlatformAscendC获取 AIV 核数coreNum与 UB 大小ubSize任一为 0 则报错退出校验 shape 与 dtypeGetShapeAttrsInfo校验elementsNum必须是 1 维、输入输出维度一致、value仅支持DT_FLOAT/DT_INT32、index与elementsNum仅支持DT_INT32totalIdx取第一个输入拼接后序列的总元素数即totalLength空输入处理totalIdx 0时 tiling 数据清零、SetBlockDim(1)直接返回避免 kernel 空跑申请 workspaceGetWorkspaceSize固定申请WS_SYS_SIZE 16MB的系统工作空间写入currentWorkspace[0]计算 tile 参数IsinPartV1TilingFunc以BLOCK_SIZE 32字节、BUFFER_NUM 2双缓冲为基准计算单 tile 元素数tileDataNum、总 block 数blocksTotal再按核数均分得到 small-core / big-core 两组数据量余数部分由前tailBlockNum个核多承担一个 block并把smallCoreDataNum、bigCoreDataNum、tileDataNum、smallTailDataNum、bigTailDataNum、finalSmallTileNum、finalBigTileNum、tailBlockNum、totalLength写入 IsinPartV1TilingData 结构体最后SetBlockDim(finalCoreNum)。Tiling 数据结构的 9 个字段定义在 isin_part_v1_tiling_data.hkernel 侧通过REGISTER_TILING_DEFAULT/GET_TILING_DATA_WITH_STRUCT读取isin_part_v1.cpp。调度模式tiling key在 isin_part_v1_tiling_key.h 中定义了ELEMENTWISE_TPL_SCH_MODE_0与ELEMENTWISE_TPL_SCH_MODE_1两种模板参数组合当前默认使用 MODE_0对应 isin_part_v1_simplified_key.ini 中的default0。算子注册侧的支撑信息算子原型isin_part_v1_def.cpp通过OpDef注册IsinPartV1声明 3 个输入 1 个输出AICore 配置开启动态编译静态标志、动态 shape/rank 支持并设置PrecisionReduceFlag(true)shape 推导isin_part_v1_infershape.cpp输出形状直接取自elementsNum的数值注册为IMPL_OP_INFERSHAPE(IsinPartV1)算子入口isin_part_v1.cpp__global__ __aicore__模板 kernelschMode为模板参数实例化NsIsinPartV1::IsinPartV1DTYPE_VALUE后依次执行Init与Process。六、aclnn 调用实战README 的调用说明指明通过 aclnn 接口方式调用 IsinPartV1 算子调用样例为 examples/test_aclnn_isin_part_v1.cpp。下面结合该示例讲解完整的 aclnn 单算子调用流程ACL 接口风格与 CANN 通用单算子调用一致。6.1 调用流程总览标准 aclnn 调用分为六步示例代码逐一步骤均有注释环境初始化aclInit→aclrtSetDevice→aclrtCreateContext→aclrtCreateStream示例固定使用deviceId 0构造输入输出 Tensor通过aclCreateTensor创建value、index、elementsNum三个输入 Tensor 与输出zTensor示例提供了CreateAclTensor与CreateEmptyTensor两个辅助函数查询 workspace 大小调用aclnnIsinPartV1GetWorkspaceSize(value, index, elementsNum, z, wsSize, executor)返回所需 workspace 字节数与算子执行器句柄申请 workspace 并执行aclrtMalloc申请 workspace 后调用aclnnIsinPartV1(wsAddr, wsSize, executor, stream)同步并取回结果aclrtSynchronizeStream后通过aclrtMemcpyACL_MEMCPY_DEVICE_TO_HOST拷贝输出示例的PrintOutResult按%u打印 BOOL 结果资源清理依次aclDestroyTensor、aclrtFree含 workspace、aclrtDestroyStream、aclrtDestroyContext、aclFinalize。6.2 示例中的关键数据示例构造的输入对应语义isin(elements[1,1,4,1,1], test_elements[0,1,2,3])参数值dtypevalue{0, 1, 1, 1, 1, 1, 2, 3, 4}ACL_FLOATindex{5, 0, 1, 3, 4, 6, 7, 8, 2}ACL_INT32elementsNum{5}ACL_INT32输出 z形状{5}ACL_BOOL期望输出[1, 1, 0, 1, 1]——elements 中的 1 出现在 test_elements 中index 为 0、1、3、4 的四个 1 都命中而 4 未出现在测试集中因此对应位置为 0。6.3 动手验证建议编译示例时需链接 CANN 的 acl 与 aclnn 运行库并包含acl/acl.h与aclnnop/aclnn_isin_part_v1.h头文件示例已在代码中引入示例中对每个 ACL 调用均用CHECK_RET宏检查返回值失败即打印错误码退出便于定位问题若想验证完整的 isin 语义可在调用 IsinPartV1 之前先用 aclnn 的 Cat 与 Sort 接口或框架层图编译完成拼接与排序再将value、index、elementsNum喂给aclnnIsinPartV1运行环境为 Atlas A2 训练系列 / Atlas 800I A2 推理产品构建登记为COMPUTE_UNIT ascend950CMakeLists.txt实验性算子请使用对应 CANN 版本配套的算子包与工具链。七、小结IsinPartV1 展示了 CANN ops-math 中一种典型的功能拆分 基础算子组合算子设计范式用cat与sort把 isin 问题规约到排序后相邻等值判定再用一个轻量核心部件完成 O(N) 的成员关系输出。其 host 侧 tiling 负责 shape/dtype 校验、workspace 申请与多核数据切分device 侧 kernel 则通过值边界对齐、UB 分块缓存与原子回写保证正确性与访存效率。对于需要在昇腾 NPU 上实现集合成员判定类算子或研究 CANN 单算子开发流程的开发者从 README 出发配合 tiling 实现、kernel 实现 与 aclnn 调用示例 三份材料对照阅读是一条完整且高效的学习路径。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考