需求背景【免费下载链接】cann-ops-competitions本仓库用于 CANN 开源社区各类竞赛、开源课题、社区任务等课题发布、开发者作品提交和展示。项目地址: https://gitcode.com/cann/cann-ops-competitions需求来源基于 TBE MaskedScatter 算子历史版本使用 Ascend C 进行改造与适配。改造目标是保持 TBE 语义和关键调度逻辑一致并通过 aclnn 直调方式完成自定义 experimental 算子的编译、安装和 example 验证。MaskedScatter 的语义为按一维扁平顺序扫描x和mask当mask[i]为 false 时输出y[i] x[i]当mask[i]为 true 时按 true mask 出现顺序从updates中取值写入y[i]。等价伪代码如下int64_t updateIndex 0; for (int64_t i 0; i numElemX; i) { if (mask[i] ! 0) { y[i] updates[updateIndex]; updateIndex; } else { y[i] x[i]; } }TBE 源码分析通过对 TBE 内置 MaskedScatter 算子源码masked_scatter.py中MaskedScatter类、masked_scatter()顶层函数、task_schedule()、calc_updates_start()、compute()进行逐行分析当前支持的能力与核心逻辑如下。TBE 算子源码路径${ASCEND_INSTALL_PATH}/opp/built-in/op_impl/ai_core/tbe/impl/ops_legacy/dynamic/masked_scatter.py算子原型路径${ASCEND_INSTALL_PATH}/opp/built-in/op_proto/inc/算子信息库路径${ASCEND_INSTALL_PATH}/opp/built-in/op_impl/ai_core/tbe/kernel/config/ascend910b/ops_legacy/masked_scatter.json1. 支持的数据类型TBEmasked_scatter.py的 dtype 字节表支持如下类型数据类型字节大小说明float162正常处理float324正常处理int648TBE 字节表包含int324正常处理uint81正常处理int81正常处理bool1mask使用 boolTBE 内部按 int8 Tensor 读写int162正常处理bfloat162TBE 初始化时映射为 float16 处理当前 Ascend C experimental 算子注册支持数据类型字节大小说明float162支持float324支持uint81支持int81支持int162支持int324支持bool1支持bfloat162支持差异说明TBE 字节表包含int64当前 experimental 算子原型和 host dtype 校验未注册int64因此当前 Ascend C 版本不支持int64。2. 输入输出与 shape 约束TBEcheck_supported()的显式约束如下检查项规则动态未知 shapex_shape包含-1或-2时返回Unknownx与maskshape必须完全相同否则返回FalseTBE 顶层接口为def masked_scatter(x, mask, updates, y, kernel_namemasked_scatter)输入输出含义如下名称类别dtypeshape说明x输入dtype 表中的数据类型all原始输入mask输入bool与 x 相同控制是否替换updates输入与 x 相同all替换值来源y输出与 x 相同与 x 相同输出结果3. TBE 常量与 tiling 参数TBE 关键常量如下| 常量 | 值 | 说明 | | ---- | -- | ---- | |BYTES_PER_BLOCK| 32 | GM/UB 搬运按 32B 对齐 | |MAX_VEC_PROCESS_NUM| 64 | 向量处理基础长度也是 task 最小长度 | |TASK_ALIGN| 4096 | 大输入 task 长度也是calc_updates_start单次 reduce 的 mask 元素数 | |TILING_ARG_NUM| 4 | tiling GM 中 int64 参数数量 | |BOOL_DTYPE| int8 | bool mask 在 TIK Tensor 中按 int8 表示 |TBE 从tiling_gm读取 4 个 int64 参数tiling 下标字段说明0num_elem_xx总元素数1num_elem_maskmask总元素数TBE 中与num_elem_x一致2num_elem_updatesupdates总元素数3tiling_core_numhost 侧分配的 AI Core 数4. TBE 分段与分核策略TBE 先通过get_task_block_size()计算每个逻辑 task 的元素数量aligned_elem_per_coreblock_size num_elem_x / tiling_core_num block_size_aligned align_up(block_size, MAX_VEC_PROCESS_NUM) if num_elem_x tiling_core_num * TASK_ALIGN: block_size_aligned TASK_ALIGN if num_elem_x MAX_VEC_PROCESS_NUM: block_size_aligned MAX_VEC_PROCESS_NUM因此 task 长度明确分为三类输入规模aligned_elem_per_corenum_elem_x 646464 num_elem_x tiling_core_num * 4096num_elem_x / tiling_core_num向上对齐到 64num_elem_x tiling_core_num * 40964096随后计算逻辑 task 数并分配到 AI Corelogi_task_num ceil(num_elem_x / aligned_elem_per_core) logi_task_num_per_aicore logi_task_num / tiling_core_num logi_task_tail logi_task_num % tiling_core_num每个核的任务数为core_task_num logi_task_num_per_aicore if aicore_idx logi_task_tail: core_task_num core_task_num 1每个核开始前需要计算当前核之前已经处理的 mask true 数prefix_len pre_core_task_num * aligned_elem_per_core task_updates_start calc_updates_start(prefix_len)这个task_updates_start是该核第一个 task 从updates读取的全局起始下标。本核内多个 task 之间通过compute()返回的新updates_start串联不重复统计前缀。5. TBEcalc_updates_start逻辑calc_updates_start(num_elem_mask)统计mask[0 : num_elem_mask]中 true 元素数量。完整流程sum_res_int64 0。data_move_length TASK_ALIGN 4096。若0 num_elem_mask TASK_ALIGN则data_move_length num_elem_mask。iters ceil(num_elem_mask / data_move_length)。分配长度为data_move_length的bool_ub、fp16_ub、fp32_ub、work_tensor_ub、sum_fp32_ub。每轮按 32B 对齐从mask_gm搬入bool_ub。vec_conv(bool - fp16)。vec_conv(fp16 - fp32)。vec_reduce_add(fp32)得到本轮 mask true 数。scalar_conv(round)转为 int32。累加到sum_res_int64。关键对齐点TBE 每次 reduce 的 mask 元素分段长度是TASK_ALIGN 4096不是MAX_VEC_PROCESS_NUM 64。6. TBEcompute逻辑compute(input_offset, updates_start)处理一个逻辑 task。完整流程默认num_elem_per_input aligned_elem_per_core。若当前 task 越过num_elem_x尾部则num_elem_per_input num_elem_x - input_offset。计算burst_x ceil(num_elem_per_input / num_elem_per_block)。计算burst_maskmask 按 1 字节元素和 32B block 对齐搬运。从 GM 搬入x[input_offset]到x_ub。从 GM 搬入mask[input_offset]到mask_ub。计算num_remain_updates num_elem_updates - updates_start。若num_remain_updates num_elem_per_input则截断为num_elem_per_input。若num_remain_updates 0从updates_gm搬入本 task 最多需要的 updates。若对齐搬运会越过updates尾部则通过updates_back_offset回退 GM 读取起点避免越界。遍历 task 内 offsetmask_ub[offset] 0保留x_ub[offset]。mask_ub[offset] ! 0写入updates_ub[updates_offset_ub updates_back_offset]并递增updates_offset_ub。updates_start updates_start updates_offset_ub。将x_ub搬出到y_gm[input_offset]。返回新的updates_start。TBE 整体流程图流程图源码文件flowcharts/tbe/01_tbe_overall.mmdTBE task_schedule 流程图流程图源码文件flowcharts/tbe/02_tbe_task_schedule.mmdTBE calc_updates_start 流程图流程图源码文件flowcharts/tbe/03_tbe_calc_updates_start.mmdTBE compute 流程图流程图源码文件flowcharts/tbe/04_tbe_compute.mmd需求分析外部组件依赖不涉及新增外部组件依赖。本算子依赖 CANN 基础能力组件作用Ascend C kernel APIGM/UB Tensor、DataCopyPad、Cast、WholeReduceSum、PipeBarrier 等 kernel 侧能力op_host tiling APIhost 侧 dtype/shape 校验、tiling data 写入、block dim 设置GE op 注册框架算子原型、infer shape、AICore config 注册aclnn 调用框架example 中通过两段式 aclnn API 调用 custom 算子内部适配模块适配 Aclnn 接口调用并在 experimental 目录中提供算子原型、host tiling、kernel、example 和设计文档。模块路径作用算子原型op_graph/masked_scatter_proto.h声明MaskedScatter输入输出host 注册op_host/masked_scatter_def.cpp注册 dtype/format 和 AICore configinfer shapeop_host/masked_scatter_infershape.cpp输出 shape 等于xshapetilingop_host/masked_scatter_tiling.cpp校验 dtype/shape写 tiling data设置 block dimtiling dataop_kernel/masked_scatter_tiling_data.h定义 kernel 读取的 4 个 int64 tiling 字段tiling keyop_kernel/masked_scatter_tiling_key.h声明 Ascend C 模板参数kernel 入口op_kernel/masked_scatter.cpp读取 tiling调用Init和Processkernel 实现op_kernel/masked_scatter.h实现分核、mask 前缀计数、masked scatter 计算aclnn exampleexamples/test_aclnn_masked_scatter.cpp保留一条 eager custom 示例设计文档docs/masked_scatter_design.md记录 TBE 和 Ascend C 设计流程图源码docs/flowcharts/单独保存 Mermaid 流程图代码需求模块设计算子原型名称类别dtypeformatshape介绍x输入float16 / float32 / uint8 / int8 / int16 / int32 / bool / bfloat16NDall原始输入张量mask输入boolND与 x 相同控制替换位置的 bool maskupdates输入与 x 相同NDall替换值来源y输出与 x 相同ND与 x 相同masked scatter 后的输出张量属性MaskedScatter 当前无属性。shape 与 dtype 约束约束项规则实现位置mask dtype必须为 boolmasked_scatter_tiling.cppx/updates/y dtype必须完全一致masked_scatter_tiling.cppx/mask/y shape必须完全一致masked_scatter_tiling.cppy shape等于 x shapemasked_scatter_infershape.cpp算子支持型号Atlas A2 训练系列产品 / Atlas 800I A2 推理产品ascend910b。需求详细设计使能方式上层框架涉及的框架勾选TF训练/推理Pytorch训练/推理ATC推理Aclnn直调√OPAT调优SGAT子图切分需求总体设计host侧设计方案MaskedScatter host 侧负责输出 shape 推导、输入合法性校验、tiling data 写入和 block dim 设置。1) InferShape输出 shape 与输入x完全一致逐维度复制const gert::Shape* xShape context-GetInputShape(0); gert::Shape* yShape context-GetOutputShape(0); *yShape *xShape;该逻辑与 TBE 语义一致MaskedScatter 只改写元素值不改变张量 rank 和 shape。2) dtype 校验host tiling 中定义支持 dtype 集合const std::setge::DataType supportedDtype { ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_UINT8, ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_BOOL, ge::DT_BF16};校验规则xdtype 必须在支持集合中。maskdtype 必须为ge::DT_BOOL。x、updates、ydtype 必须完全一致。3) shape 校验host tiling 中读取x、mask、y的 storage shape并要求xShape maskShape xShape yShapeupdatesshape 不要求与x相同语义上只要求按 mask true 顺序提供替换值。4) tiling 参数Ascend C 使用与 TBEtiling_gm一致的 4 个字段struct MaskedScatterTilingData { int64_t numElemX; int64_t numElemMask; int64_t numElemUpdates; int64_t tilingCoreNum; };host 侧写入规则numElemX input x storage shape size numElemMask numElemX numElemUpdates input updates storage shape size tilingCoreNum PlatformAscendC.GetCoreNumAiv()5) 分核策略host 侧设置context-SetBlockDim(static_castuint32_t(coreNum));具体逻辑 task 长度、逻辑 task 数、每核 task 数在 kernelInit()和Process()中计算。该设计与 TBE 一致host 只传入tiling_core_numkernel 内按TASK_ALIGN和MAX_VEC_PROCESS_NUM计算分段。Ascend C Host Tiling 流程图流程图源码文件flowcharts/ascend/05_ascend_host_tiling.mmdkernel侧设计方案Kernel 侧执行Init和Process两个阶段。1) Kernel 入口masked_scatter.cpp中入口逻辑如下REGISTER_TILING_DEFAULT(MaskedScatterTilingData); GET_TILING_DATA_WITH_STRUCT(MaskedScatterTilingData, tilingData, tiling); NsMaskedScatter::MaskedScatterDTYPE_X op; op.Init(x, mask, updates, y, tilingData); op.Process();当前schMode仅作为模板参数声明存在实际 kernel 使用单一路径语义对应 TBE Python 版本的通用实现。2) InitInit完成 GM 绑定、tiling 读取、dtype 信息计算、task 参数计算和 UB buffer 初始化。GM 绑定成员GM 输入输出说明xGm_x原始输入maskGm_mask按 int8 访问用于逐元素判断maskGmUint8_mask按 uint8 访问用于前缀 reduce 统计updatesGm_updates替换值来源yGm_y输出task 参数计算dtypeBytesSize_ sizeof(DTYPE); numElemPerBlock_ BYTES_PER_BLOCK / dtypeBytesSize_; alignedElemPerCore_ GetTaskBlockSize(); logiTaskNum_ CeilDiv(numElemX_, alignedElemPerCore_); logiTaskNumPerAicore_ logiTaskNum_ / tilingCoreNum_; logiTaskTail_ logiTaskNum_ % tilingCoreNum_;UB buffer 初始化全部以TASK_ALIGN 4096为最大 task 长度buffer/queue大小inQueueX_TASK_ALIGN * dtypeBytesSize_inQueueMask_TASK_ALIGN * 1inQueueUpdates_TASK_ALIGN * dtypeBytesSize_outQueueY_TASK_ALIGN * dtypeBytesSize_boolBuf_TASK_ALIGN * 1fp16Buf_TASK_ALIGN * sizeof(float16)fp32Buf_TASK_ALIGN * sizeof(float32)sumBuf_TASK_ALIGN * sizeof(float32)3) GetTaskBlockSizeAscend C 与 TBE 保持一致int64_t blockSizeAligned MAX_VEC_PROCESS_NUM; if (tilingCoreNum_ ! 0) { int64_t blockSize numElemX_ / tilingCoreNum_; blockSizeAligned AlignDiv(blockSize, MAX_VEC_PROCESS_NUM); } if (numElemX_ tilingCoreNum_ * TASK_ALIGN) { blockSizeAligned TASK_ALIGN; } if (numElemX_ MAX_VEC_PROCESS_NUM) { blockSizeAligned MAX_VEC_PROCESS_NUM; } return blockSizeAligned;这里的TASK_ALIGN 4096、MAX_VEC_PROCESS_NUM 64与 TBE 常量保持一致。4) Process每个 AI Core 只处理分配给自己的逻辑 taskaicoreIdx GetBlockIdx()。计算本核coreTaskNum。计算本核之前已经分配的preCoreTaskNum。若coreTaskNum 0先调用taskUpdatesStart CalcUpdatesStart(preCoreTaskNum * alignedElemPerCore_);逐 task 调用taskUpdatesStart Compute((preCoreTaskNum taskIdx) * alignedElemPerCore_, taskUpdatesStart);该逻辑与 TBEtask_schedule()一致每核只在开头统计一次前缀 mask true 数本核内部 task 通过Compute返回值继续推进updatesStart。5) CalcUpdatesStartAscend C 的 mask 前缀计数流程与 TBE 对齐核心常量为constexpr int64_t TASK_ALIGN 4096; constexpr int64_t COUNT_REDUCE_LEN TASK_ALIGN;完整流程sumCount 0。dataMoveLength COUNT_REDUCE_LEN即 4096。若0 numElemMask COUNT_REDUCE_LEN则dataMoveLength numElemMask。iters CeilDiv(numElemMask, dataMoveLength)。每轮按 32B 对齐搬入 mask。计算curDataLen和alignedDataLen。调用CountMaskNonZero(maskUb, curDataLen, alignedDataLen)。累加每轮 true 数返回sumCount。CountMaskNonZero的逻辑对 pad 区域置 0避免尾部 32B 对齐补齐区域参与统计。Cast(uint8 mask - fp16)。Cast(fp16 - fp32)。Abs。Mins(..., 1.0)把任意非零 mask 归一化为 1。WholeReduceSum。通过V_S事件同步后读取 sum。round 后返回 int64。重点对齐项此前 Ascend C 使用COUNT_REDUCE_LEN 64会导致前缀 mask 计数分段与 TBE 不同当前已对齐为COUNT_REDUCE_LEN TASK_ALIGN 4096。6) ComputeCompute(inputOffset, updatesStart)处理一个逻辑 task。完整流程默认numElemPerInput alignedElemPerCore_。若 task 越过numElemX_尾部则缩短为尾部实际长度。计算burstX和burstMask。从xGm_搬入yUb相当于先默认y x。从maskGm_搬入maskUb。计算当前 task 可读取的 updates 数numRemainUpdates numElemUpdates_ - updatesStart numRemainUpdates min(numRemainUpdates, numElemPerInput)若还有 updates 可读则搬入updatesUb。若按 block 对齐读取会超过 updates 尾部则使用updatesBackOffset回退读起点。遍历 task 内 offsetmaskUb[offset] 0保留yUb[offset]即原始x值。maskUb[offset] ! 0 updatesOffsetUb numRemainUpdates写updatesUb[updatesOffsetUb updatesBackOffset]。updatesStart updatesOffsetUb。将yUb搬出到yGm_[inputOffset]。返回新的updatesStart。7) TBE 与 Ascend C 对齐关系设计点TBEAscend C结论算子输入输出x, mask, updates - yx, mask, updates - y一致tiling 字段4 个 int64同 4 个字段一致task 长度计算get_task_block_size()GetTaskBlockSize()一致大输入 task 长度TASK_ALIGN 4096TASK_ALIGN 4096一致前缀 mask reduce 分段每 4096 个 mask 元素COUNT_REDUCE_LEN TASK_ALIGN一致默认输出值先搬入 x再按 mask 覆盖先搬入 x 到 yUb再按 mask 覆盖一致updates 尾部处理updates_back_offsetupdatesBackOffset一致bfloat16映射为 float16BF16 注册元素字节数为 2task/搬运长度一致int64TBE 字节表包含当前未注册当前 experimental 不支持Ascend C Kernel 流程图1. Kernel 入口与 Init 流程图流程图源码文件flowcharts/ascend/06_ascend_kernel_entry_init.mmd2. Process 分核调度流程图流程图源码文件flowcharts/ascend/07_ascend_process_schedule.mmd3. CalcUpdatesStart 流程图流程图源码文件flowcharts/ascend/08_ascend_calc_updates_start.mmd4. Compute 流程图流程图源码文件flowcharts/ascend/09_ascend_compute.mmd5. Aclnn example 调用流程图流程图源码文件flowcharts/ascend/10_aclnn_example.mmdAclnn 示例设计示例文件examples/test_aclnn_masked_scatter.cpp当前 example 只保留一条测试与gather_elements_v3example 风格一致走 eager custom aclnn 两段式调用。输入x [1, 2, 3, 4, 5, 6, 7, 8] mask [1, 0, 1, 0, 1, 0, 1, 0] updates [10, 20, 30, 40]调用aclnnMaskedScatterGetWorkspaceSize(self, mask, source, out, workspaceSize, executor); aclnnMaskedScatter(workspaceAddr, workspaceSize, executor, stream);期望输出y [10, 2, 20, 4, 30, 6, 40, 8]编译安装与验证编译命令bash build.sh --pkg --socascend910b --opsmasked_scatter --experimental --vendor_namecustom安装命令./build_out/cann-ops-nn-custom_linux-aarch64.runexample 验证命令export LD_LIBRARY_PATH/home/developer/Ascend/cann-8.5.2/opp/vendors/custom_nn/op_api/lib/:${LD_LIBRARY_PATH} bash build.sh --run_example masked_scatter eager cust --vendor_namecustom --experimental当前验证输出result[0] is: 10.000000 result[1] is: 2.000000 result[2] is: 20.000000 result[3] is: 4.000000 result[4] is: 30.000000 result[5] is: 6.000000 result[6] is: 40.000000 result[7] is: 8.000000 Run test_aclnn_masked_scatter success.支持硬件支持的芯片版本涉及勾选香橙派OrangePi AIproAtlas 200I/500 A2推理产品Atlas 800I/T A2√Atlas A2训练系列产品√算子约束限制x、mask、yshape 必须完全相同。maskdtype 必须为 bool。x、updates、ydtype 必须完全相同。当前 experimental Ascend C 算子未注册int64即使 TBE dtype 字节表包含int64。当前 kernel 按扁平一维顺序处理所有元素不区分原始 rank。当前实现没有显式校验updates元素数量必须大于等于 mask true 数当当前 task 可读 updates 不足时updatesOffsetUb numRemainUpdates会阻止继续写入后续位置保留原始x值。特性交叉分析可维可测分析精度标准/性能标准验收标准描述(不涉及说明原因)标准来源精度标准与 TBE 版本按元素比对一致历史 TBE 对标分核标准task 长度、逻辑 task 分配、每核前缀 updates 起点与 TBE 一致TBE 源码对齐性能标准大输入 task 长度为 4096mask 前缀计数每 4096 个元素 reduce 一次TBE 源码对齐关联的 Issue暂无。文档更新本文档。类型标签Bug修复新特性性能优化文档更新其他请描述社区任务算子设计文档【免费下载链接】cann-ops-competitions本仓库用于 CANN 开源社区各类竞赛、开源课题、社区任务等课题发布、开发者作品提交和展示。项目地址: https://gitcode.com/cann/cann-ops-competitions创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
CANN算子设计：MaskedScatter实现方案
需求背景【免费下载链接】cann-ops-competitions本仓库用于 CANN 开源社区各类竞赛、开源课题、社区任务等课题发布、开发者作品提交和展示。项目地址: https://gitcode.com/cann/cann-ops-competitions需求来源基于 TBE MaskedScatter 算子历史版本使用 Ascend C 进行改造与适配。改造目标是保持 TBE 语义和关键调度逻辑一致并通过 aclnn 直调方式完成自定义 experimental 算子的编译、安装和 example 验证。MaskedScatter 的语义为按一维扁平顺序扫描x和mask当mask[i]为 false 时输出y[i] x[i]当mask[i]为 true 时按 true mask 出现顺序从updates中取值写入y[i]。等价伪代码如下int64_t updateIndex 0; for (int64_t i 0; i numElemX; i) { if (mask[i] ! 0) { y[i] updates[updateIndex]; updateIndex; } else { y[i] x[i]; } }TBE 源码分析通过对 TBE 内置 MaskedScatter 算子源码masked_scatter.py中MaskedScatter类、masked_scatter()顶层函数、task_schedule()、calc_updates_start()、compute()进行逐行分析当前支持的能力与核心逻辑如下。TBE 算子源码路径${ASCEND_INSTALL_PATH}/opp/built-in/op_impl/ai_core/tbe/impl/ops_legacy/dynamic/masked_scatter.py算子原型路径${ASCEND_INSTALL_PATH}/opp/built-in/op_proto/inc/算子信息库路径${ASCEND_INSTALL_PATH}/opp/built-in/op_impl/ai_core/tbe/kernel/config/ascend910b/ops_legacy/masked_scatter.json1. 支持的数据类型TBEmasked_scatter.py的 dtype 字节表支持如下类型数据类型字节大小说明float162正常处理float324正常处理int648TBE 字节表包含int324正常处理uint81正常处理int81正常处理bool1mask使用 boolTBE 内部按 int8 Tensor 读写int162正常处理bfloat162TBE 初始化时映射为 float16 处理当前 Ascend C experimental 算子注册支持数据类型字节大小说明float162支持float324支持uint81支持int81支持int162支持int324支持bool1支持bfloat162支持差异说明TBE 字节表包含int64当前 experimental 算子原型和 host dtype 校验未注册int64因此当前 Ascend C 版本不支持int64。2. 输入输出与 shape 约束TBEcheck_supported()的显式约束如下检查项规则动态未知 shapex_shape包含-1或-2时返回Unknownx与maskshape必须完全相同否则返回FalseTBE 顶层接口为def masked_scatter(x, mask, updates, y, kernel_namemasked_scatter)输入输出含义如下名称类别dtypeshape说明x输入dtype 表中的数据类型all原始输入mask输入bool与 x 相同控制是否替换updates输入与 x 相同all替换值来源y输出与 x 相同与 x 相同输出结果3. TBE 常量与 tiling 参数TBE 关键常量如下| 常量 | 值 | 说明 | | ---- | -- | ---- | |BYTES_PER_BLOCK| 32 | GM/UB 搬运按 32B 对齐 | |MAX_VEC_PROCESS_NUM| 64 | 向量处理基础长度也是 task 最小长度 | |TASK_ALIGN| 4096 | 大输入 task 长度也是calc_updates_start单次 reduce 的 mask 元素数 | |TILING_ARG_NUM| 4 | tiling GM 中 int64 参数数量 | |BOOL_DTYPE| int8 | bool mask 在 TIK Tensor 中按 int8 表示 |TBE 从tiling_gm读取 4 个 int64 参数tiling 下标字段说明0num_elem_xx总元素数1num_elem_maskmask总元素数TBE 中与num_elem_x一致2num_elem_updatesupdates总元素数3tiling_core_numhost 侧分配的 AI Core 数4. TBE 分段与分核策略TBE 先通过get_task_block_size()计算每个逻辑 task 的元素数量aligned_elem_per_coreblock_size num_elem_x / tiling_core_num block_size_aligned align_up(block_size, MAX_VEC_PROCESS_NUM) if num_elem_x tiling_core_num * TASK_ALIGN: block_size_aligned TASK_ALIGN if num_elem_x MAX_VEC_PROCESS_NUM: block_size_aligned MAX_VEC_PROCESS_NUM因此 task 长度明确分为三类输入规模aligned_elem_per_corenum_elem_x 646464 num_elem_x tiling_core_num * 4096num_elem_x / tiling_core_num向上对齐到 64num_elem_x tiling_core_num * 40964096随后计算逻辑 task 数并分配到 AI Corelogi_task_num ceil(num_elem_x / aligned_elem_per_core) logi_task_num_per_aicore logi_task_num / tiling_core_num logi_task_tail logi_task_num % tiling_core_num每个核的任务数为core_task_num logi_task_num_per_aicore if aicore_idx logi_task_tail: core_task_num core_task_num 1每个核开始前需要计算当前核之前已经处理的 mask true 数prefix_len pre_core_task_num * aligned_elem_per_core task_updates_start calc_updates_start(prefix_len)这个task_updates_start是该核第一个 task 从updates读取的全局起始下标。本核内多个 task 之间通过compute()返回的新updates_start串联不重复统计前缀。5. TBEcalc_updates_start逻辑calc_updates_start(num_elem_mask)统计mask[0 : num_elem_mask]中 true 元素数量。完整流程sum_res_int64 0。data_move_length TASK_ALIGN 4096。若0 num_elem_mask TASK_ALIGN则data_move_length num_elem_mask。iters ceil(num_elem_mask / data_move_length)。分配长度为data_move_length的bool_ub、fp16_ub、fp32_ub、work_tensor_ub、sum_fp32_ub。每轮按 32B 对齐从mask_gm搬入bool_ub。vec_conv(bool - fp16)。vec_conv(fp16 - fp32)。vec_reduce_add(fp32)得到本轮 mask true 数。scalar_conv(round)转为 int32。累加到sum_res_int64。关键对齐点TBE 每次 reduce 的 mask 元素分段长度是TASK_ALIGN 4096不是MAX_VEC_PROCESS_NUM 64。6. TBEcompute逻辑compute(input_offset, updates_start)处理一个逻辑 task。完整流程默认num_elem_per_input aligned_elem_per_core。若当前 task 越过num_elem_x尾部则num_elem_per_input num_elem_x - input_offset。计算burst_x ceil(num_elem_per_input / num_elem_per_block)。计算burst_maskmask 按 1 字节元素和 32B block 对齐搬运。从 GM 搬入x[input_offset]到x_ub。从 GM 搬入mask[input_offset]到mask_ub。计算num_remain_updates num_elem_updates - updates_start。若num_remain_updates num_elem_per_input则截断为num_elem_per_input。若num_remain_updates 0从updates_gm搬入本 task 最多需要的 updates。若对齐搬运会越过updates尾部则通过updates_back_offset回退 GM 读取起点避免越界。遍历 task 内 offsetmask_ub[offset] 0保留x_ub[offset]。mask_ub[offset] ! 0写入updates_ub[updates_offset_ub updates_back_offset]并递增updates_offset_ub。updates_start updates_start updates_offset_ub。将x_ub搬出到y_gm[input_offset]。返回新的updates_start。TBE 整体流程图流程图源码文件flowcharts/tbe/01_tbe_overall.mmdTBE task_schedule 流程图流程图源码文件flowcharts/tbe/02_tbe_task_schedule.mmdTBE calc_updates_start 流程图流程图源码文件flowcharts/tbe/03_tbe_calc_updates_start.mmdTBE compute 流程图流程图源码文件flowcharts/tbe/04_tbe_compute.mmd需求分析外部组件依赖不涉及新增外部组件依赖。本算子依赖 CANN 基础能力组件作用Ascend C kernel APIGM/UB Tensor、DataCopyPad、Cast、WholeReduceSum、PipeBarrier 等 kernel 侧能力op_host tiling APIhost 侧 dtype/shape 校验、tiling data 写入、block dim 设置GE op 注册框架算子原型、infer shape、AICore config 注册aclnn 调用框架example 中通过两段式 aclnn API 调用 custom 算子内部适配模块适配 Aclnn 接口调用并在 experimental 目录中提供算子原型、host tiling、kernel、example 和设计文档。模块路径作用算子原型op_graph/masked_scatter_proto.h声明MaskedScatter输入输出host 注册op_host/masked_scatter_def.cpp注册 dtype/format 和 AICore configinfer shapeop_host/masked_scatter_infershape.cpp输出 shape 等于xshapetilingop_host/masked_scatter_tiling.cpp校验 dtype/shape写 tiling data设置 block dimtiling dataop_kernel/masked_scatter_tiling_data.h定义 kernel 读取的 4 个 int64 tiling 字段tiling keyop_kernel/masked_scatter_tiling_key.h声明 Ascend C 模板参数kernel 入口op_kernel/masked_scatter.cpp读取 tiling调用Init和Processkernel 实现op_kernel/masked_scatter.h实现分核、mask 前缀计数、masked scatter 计算aclnn exampleexamples/test_aclnn_masked_scatter.cpp保留一条 eager custom 示例设计文档docs/masked_scatter_design.md记录 TBE 和 Ascend C 设计流程图源码docs/flowcharts/单独保存 Mermaid 流程图代码需求模块设计算子原型名称类别dtypeformatshape介绍x输入float16 / float32 / uint8 / int8 / int16 / int32 / bool / bfloat16NDall原始输入张量mask输入boolND与 x 相同控制替换位置的 bool maskupdates输入与 x 相同NDall替换值来源y输出与 x 相同ND与 x 相同masked scatter 后的输出张量属性MaskedScatter 当前无属性。shape 与 dtype 约束约束项规则实现位置mask dtype必须为 boolmasked_scatter_tiling.cppx/updates/y dtype必须完全一致masked_scatter_tiling.cppx/mask/y shape必须完全一致masked_scatter_tiling.cppy shape等于 x shapemasked_scatter_infershape.cpp算子支持型号Atlas A2 训练系列产品 / Atlas 800I A2 推理产品ascend910b。需求详细设计使能方式上层框架涉及的框架勾选TF训练/推理Pytorch训练/推理ATC推理Aclnn直调√OPAT调优SGAT子图切分需求总体设计host侧设计方案MaskedScatter host 侧负责输出 shape 推导、输入合法性校验、tiling data 写入和 block dim 设置。1) InferShape输出 shape 与输入x完全一致逐维度复制const gert::Shape* xShape context-GetInputShape(0); gert::Shape* yShape context-GetOutputShape(0); *yShape *xShape;该逻辑与 TBE 语义一致MaskedScatter 只改写元素值不改变张量 rank 和 shape。2) dtype 校验host tiling 中定义支持 dtype 集合const std::setge::DataType supportedDtype { ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_UINT8, ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_BOOL, ge::DT_BF16};校验规则xdtype 必须在支持集合中。maskdtype 必须为ge::DT_BOOL。x、updates、ydtype 必须完全一致。3) shape 校验host tiling 中读取x、mask、y的 storage shape并要求xShape maskShape xShape yShapeupdatesshape 不要求与x相同语义上只要求按 mask true 顺序提供替换值。4) tiling 参数Ascend C 使用与 TBEtiling_gm一致的 4 个字段struct MaskedScatterTilingData { int64_t numElemX; int64_t numElemMask; int64_t numElemUpdates; int64_t tilingCoreNum; };host 侧写入规则numElemX input x storage shape size numElemMask numElemX numElemUpdates input updates storage shape size tilingCoreNum PlatformAscendC.GetCoreNumAiv()5) 分核策略host 侧设置context-SetBlockDim(static_castuint32_t(coreNum));具体逻辑 task 长度、逻辑 task 数、每核 task 数在 kernelInit()和Process()中计算。该设计与 TBE 一致host 只传入tiling_core_numkernel 内按TASK_ALIGN和MAX_VEC_PROCESS_NUM计算分段。Ascend C Host Tiling 流程图流程图源码文件flowcharts/ascend/05_ascend_host_tiling.mmdkernel侧设计方案Kernel 侧执行Init和Process两个阶段。1) Kernel 入口masked_scatter.cpp中入口逻辑如下REGISTER_TILING_DEFAULT(MaskedScatterTilingData); GET_TILING_DATA_WITH_STRUCT(MaskedScatterTilingData, tilingData, tiling); NsMaskedScatter::MaskedScatterDTYPE_X op; op.Init(x, mask, updates, y, tilingData); op.Process();当前schMode仅作为模板参数声明存在实际 kernel 使用单一路径语义对应 TBE Python 版本的通用实现。2) InitInit完成 GM 绑定、tiling 读取、dtype 信息计算、task 参数计算和 UB buffer 初始化。GM 绑定成员GM 输入输出说明xGm_x原始输入maskGm_mask按 int8 访问用于逐元素判断maskGmUint8_mask按 uint8 访问用于前缀 reduce 统计updatesGm_updates替换值来源yGm_y输出task 参数计算dtypeBytesSize_ sizeof(DTYPE); numElemPerBlock_ BYTES_PER_BLOCK / dtypeBytesSize_; alignedElemPerCore_ GetTaskBlockSize(); logiTaskNum_ CeilDiv(numElemX_, alignedElemPerCore_); logiTaskNumPerAicore_ logiTaskNum_ / tilingCoreNum_; logiTaskTail_ logiTaskNum_ % tilingCoreNum_;UB buffer 初始化全部以TASK_ALIGN 4096为最大 task 长度buffer/queue大小inQueueX_TASK_ALIGN * dtypeBytesSize_inQueueMask_TASK_ALIGN * 1inQueueUpdates_TASK_ALIGN * dtypeBytesSize_outQueueY_TASK_ALIGN * dtypeBytesSize_boolBuf_TASK_ALIGN * 1fp16Buf_TASK_ALIGN * sizeof(float16)fp32Buf_TASK_ALIGN * sizeof(float32)sumBuf_TASK_ALIGN * sizeof(float32)3) GetTaskBlockSizeAscend C 与 TBE 保持一致int64_t blockSizeAligned MAX_VEC_PROCESS_NUM; if (tilingCoreNum_ ! 0) { int64_t blockSize numElemX_ / tilingCoreNum_; blockSizeAligned AlignDiv(blockSize, MAX_VEC_PROCESS_NUM); } if (numElemX_ tilingCoreNum_ * TASK_ALIGN) { blockSizeAligned TASK_ALIGN; } if (numElemX_ MAX_VEC_PROCESS_NUM) { blockSizeAligned MAX_VEC_PROCESS_NUM; } return blockSizeAligned;这里的TASK_ALIGN 4096、MAX_VEC_PROCESS_NUM 64与 TBE 常量保持一致。4) Process每个 AI Core 只处理分配给自己的逻辑 taskaicoreIdx GetBlockIdx()。计算本核coreTaskNum。计算本核之前已经分配的preCoreTaskNum。若coreTaskNum 0先调用taskUpdatesStart CalcUpdatesStart(preCoreTaskNum * alignedElemPerCore_);逐 task 调用taskUpdatesStart Compute((preCoreTaskNum taskIdx) * alignedElemPerCore_, taskUpdatesStart);该逻辑与 TBEtask_schedule()一致每核只在开头统计一次前缀 mask true 数本核内部 task 通过Compute返回值继续推进updatesStart。5) CalcUpdatesStartAscend C 的 mask 前缀计数流程与 TBE 对齐核心常量为constexpr int64_t TASK_ALIGN 4096; constexpr int64_t COUNT_REDUCE_LEN TASK_ALIGN;完整流程sumCount 0。dataMoveLength COUNT_REDUCE_LEN即 4096。若0 numElemMask COUNT_REDUCE_LEN则dataMoveLength numElemMask。iters CeilDiv(numElemMask, dataMoveLength)。每轮按 32B 对齐搬入 mask。计算curDataLen和alignedDataLen。调用CountMaskNonZero(maskUb, curDataLen, alignedDataLen)。累加每轮 true 数返回sumCount。CountMaskNonZero的逻辑对 pad 区域置 0避免尾部 32B 对齐补齐区域参与统计。Cast(uint8 mask - fp16)。Cast(fp16 - fp32)。Abs。Mins(..., 1.0)把任意非零 mask 归一化为 1。WholeReduceSum。通过V_S事件同步后读取 sum。round 后返回 int64。重点对齐项此前 Ascend C 使用COUNT_REDUCE_LEN 64会导致前缀 mask 计数分段与 TBE 不同当前已对齐为COUNT_REDUCE_LEN TASK_ALIGN 4096。6) ComputeCompute(inputOffset, updatesStart)处理一个逻辑 task。完整流程默认numElemPerInput alignedElemPerCore_。若 task 越过numElemX_尾部则缩短为尾部实际长度。计算burstX和burstMask。从xGm_搬入yUb相当于先默认y x。从maskGm_搬入maskUb。计算当前 task 可读取的 updates 数numRemainUpdates numElemUpdates_ - updatesStart numRemainUpdates min(numRemainUpdates, numElemPerInput)若还有 updates 可读则搬入updatesUb。若按 block 对齐读取会超过 updates 尾部则使用updatesBackOffset回退读起点。遍历 task 内 offsetmaskUb[offset] 0保留yUb[offset]即原始x值。maskUb[offset] ! 0 updatesOffsetUb numRemainUpdates写updatesUb[updatesOffsetUb updatesBackOffset]。updatesStart updatesOffsetUb。将yUb搬出到yGm_[inputOffset]。返回新的updatesStart。7) TBE 与 Ascend C 对齐关系设计点TBEAscend C结论算子输入输出x, mask, updates - yx, mask, updates - y一致tiling 字段4 个 int64同 4 个字段一致task 长度计算get_task_block_size()GetTaskBlockSize()一致大输入 task 长度TASK_ALIGN 4096TASK_ALIGN 4096一致前缀 mask reduce 分段每 4096 个 mask 元素COUNT_REDUCE_LEN TASK_ALIGN一致默认输出值先搬入 x再按 mask 覆盖先搬入 x 到 yUb再按 mask 覆盖一致updates 尾部处理updates_back_offsetupdatesBackOffset一致bfloat16映射为 float16BF16 注册元素字节数为 2task/搬运长度一致int64TBE 字节表包含当前未注册当前 experimental 不支持Ascend C Kernel 流程图1. Kernel 入口与 Init 流程图流程图源码文件flowcharts/ascend/06_ascend_kernel_entry_init.mmd2. Process 分核调度流程图流程图源码文件flowcharts/ascend/07_ascend_process_schedule.mmd3. CalcUpdatesStart 流程图流程图源码文件flowcharts/ascend/08_ascend_calc_updates_start.mmd4. Compute 流程图流程图源码文件flowcharts/ascend/09_ascend_compute.mmd5. Aclnn example 调用流程图流程图源码文件flowcharts/ascend/10_aclnn_example.mmdAclnn 示例设计示例文件examples/test_aclnn_masked_scatter.cpp当前 example 只保留一条测试与gather_elements_v3example 风格一致走 eager custom aclnn 两段式调用。输入x [1, 2, 3, 4, 5, 6, 7, 8] mask [1, 0, 1, 0, 1, 0, 1, 0] updates [10, 20, 30, 40]调用aclnnMaskedScatterGetWorkspaceSize(self, mask, source, out, workspaceSize, executor); aclnnMaskedScatter(workspaceAddr, workspaceSize, executor, stream);期望输出y [10, 2, 20, 4, 30, 6, 40, 8]编译安装与验证编译命令bash build.sh --pkg --socascend910b --opsmasked_scatter --experimental --vendor_namecustom安装命令./build_out/cann-ops-nn-custom_linux-aarch64.runexample 验证命令export LD_LIBRARY_PATH/home/developer/Ascend/cann-8.5.2/opp/vendors/custom_nn/op_api/lib/:${LD_LIBRARY_PATH} bash build.sh --run_example masked_scatter eager cust --vendor_namecustom --experimental当前验证输出result[0] is: 10.000000 result[1] is: 2.000000 result[2] is: 20.000000 result[3] is: 4.000000 result[4] is: 30.000000 result[5] is: 6.000000 result[6] is: 40.000000 result[7] is: 8.000000 Run test_aclnn_masked_scatter success.支持硬件支持的芯片版本涉及勾选香橙派OrangePi AIproAtlas 200I/500 A2推理产品Atlas 800I/T A2√Atlas A2训练系列产品√算子约束限制x、mask、yshape 必须完全相同。maskdtype 必须为 bool。x、updates、ydtype 必须完全相同。当前 experimental Ascend C 算子未注册int64即使 TBE dtype 字节表包含int64。当前 kernel 按扁平一维顺序处理所有元素不区分原始 rank。当前实现没有显式校验updates元素数量必须大于等于 mask true 数当当前 task 可读 updates 不足时updatesOffsetUb numRemainUpdates会阻止继续写入后续位置保留原始x值。特性交叉分析可维可测分析精度标准/性能标准验收标准描述(不涉及说明原因)标准来源精度标准与 TBE 版本按元素比对一致历史 TBE 对标分核标准task 长度、逻辑 task 分配、每核前缀 updates 起点与 TBE 一致TBE 源码对齐性能标准大输入 task 长度为 4096mask 前缀计数每 4096 个元素 reduce 一次TBE 源码对齐关联的 Issue暂无。文档更新本文档。类型标签Bug修复新特性性能优化文档更新其他请描述社区任务算子设计文档【免费下载链接】cann-ops-competitions本仓库用于 CANN 开源社区各类竞赛、开源课题、社区任务等课题发布、开发者作品提交和展示。项目地址: https://gitcode.com/cann/cann-ops-competitions创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
