人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载EmbeddingHashTableExport 是 CANN ops-nn 算子库hash 目录中负责将整张 embedding hash 表内容导出为可查询张量序列的算子。本文以 hash/embedding_hash_table_export/README.md 为主体结合其算子定义、InferShape、Tiling 与 kernel 源码系统讲解其产品支持范围、全部输入输出与属性参数、shape 推导规则、桶内数据结构与导出算法帮助读者在推荐系统/大规模稀疏特征场景中正确理解并调用该算子。产品支持情况当前算子仅支持以下产品形态其余训练/推理系列产品均不支持产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品✗Atlas A2 训练系列产品/Atlas A2 推理系列产品✗Atlas 200I/500 A2 推理产品✗Atlas 推理系列产品✗Atlas 训练系列产品✗该支持范围与源码中的 AICore 注册配置完全一致embedding_hash_table_export_def.cpp 中通过this-AICore().AddConfig(...)仅注册了ascend950、ascend960dt、ascend350三个平台配置与 README 表格中的 Ascend 950PR/Ascend 950DT 一一对应。对应平台还提供了独立的算子二进制配置清单ascend350、ascend950、ascend960dt因此使用前务必确认目标设备的 SoC 型号。功能说明算子功能导出整个 table 表。在基于 hash 结构的 embedding 参数存储方案中训练/推理过程中产生的 key-value 数据以桶bucket为单位存放在设备端 hash 表中。EmbeddingHashTableExport 的作用就是把一张或多张 hash 表中的全部桶数据按顺序取出整理为紧凑的 key、counter、filter_flag、value 张量序列输出供上层框架做检查点保存、参数序列化或后续离线处理使用。从源码结构看hash 目录下与它配套的还有 lookup/import、apply_adam_w 等一组 embedding hash 表算子它们共享table handle 桶数组这一底层内存契约本算子读取 handle 结构体中的表地址等元信息见 embedding_hash_table_export.h因此理解本算子对理解整套 hash 表算子体系同样有帮助。参数说明算子共包含 4 个输入、4 组动态输出和 2 个属性完整参数表如下与 README 一致并补充了源码中确认的细节参数名输入/输出/属性描述数据类型数据格式table_handles输入输入 hash 表 handle 句柄里面包含了 hash 表的表头地址等INT64NDtable_sizes输入table 数量INT64NDembedding_dims输入hash 表桶深度INT64NDbucket_sizes输入hash 表桶数量INT64NDkeys输出插入 key 序列INT64NDcounters输出导出 key 数量UINT64NDfilter_flags输出准入标志UINT8NDvalues输出导出 key 对应的 value 序列FLOATNDexport_mode输入属性导出模式默认 all全量导出STRING-filtered_export_flag输入属性是否按准入标志过滤导出默认 falseBOOL-输入输出在源码中的定义输入全部为必选REQUIRED且数据类型固定为DT_INT64、格式固定为FORMAT_ND见 embedding_hash_table_export_def.cpp。四组输出均为动态输出DYNAMIC即每个输出位置上实际包含的张量数量在编译期不固定由输入 table 数量决定keys 为DT_INT64counters 为DT_UINT64filter_flags 为DT_UINT8values 为DT_FLOATembedding_hash_table_export_def.cpp。属性默认值export_mode默认allfiltered_export_flag默认falseembedding_hash_table_export_def.cpp。在 graph 定义embedding_hash_table_export_proto.h中同样以all与false作为默认值且export_mode的合法取值在注释中明确为all或new两种。算子注册时开启了DynamicCompileStaticFlag(true)、DynamicShapeSupportFlag(true)、DynamicRankSupportFlag(true)等能力embedding_hash_table_export_def.cpp意味着该算子支持动态 shape 与动态编译上层可以传入运行期才确定形状的输入。二进制配置中的参数形态以 ascend950 的 binary 配置为例输入输出 shape 均为-2动态维度输出按[[...]]双层结构组织表示每个输出位置是一组张量属性区声明了export_modestringvalue 为 null即使用默认值与filtered_export_flagbool默认 false。这说明实际部署时四个输入都需要以 1 维动态张量传入。shape 推导规则InferShape 实现了关键推导逻辑embedding_hash_table_export_infershape.cpp输入形状一致性校验table_sizes、embedding_dims、table_handles、bucket_sizes四个输入的 shape size元素个数必须完全相等且table_sizes的 shape size 必须 ≥ 0否则报错并返回GRAPH_FAILEDCheckEmbeddingHashTableExportParams。输出张量数量校验每个动态输出列表中的张量个数 table 数量 × 4keys/counters/filter_flags/values 各一组总输出数必须等于numTable * 4。输出 shape 推导对第 i 张表keys、counters、filter_flags 输出 shape 均为[tableSizes[i]]values 输出 shape 为[tableSizes[i], embeddingDims[i]]InferShape4EmbeddingHashTableExport。若table_sizes为运行期值指针为空则各输出退化为[-1]/[-1, -1]的动态形状InferShape4EmbeddingHashTableExportNull。数据类型推导keys/counters/filter_flags/values 的输出类型分别固定为 INT64/UINT64/UINT8/FLOATInferDataType4EmbeddingHashTableExport。shape 推导依赖table_sizes与embedding_dims的具体取值InputsDataDependency({INPUT_TABLE_SIZE_IDX, INPUT_EMBEDDING_DIMS_IDX})这两个输入属于值依赖输入embedding_hash_table_export_def.cpp 中ValueDepend(OPTIONAL)。约束说明README 中约束说明为无但算子原型注释embedding_hash_table_export_proto.h与 Tiling 实现embedding_hash_table_export_tiling_arch35.cpp补充了两条使用约束可作为实操时的隐性要求table_handles、table_sizes、embedding_dims、bucket_sizes四个输入具有相同的 shape len元素个数相同每个动态输出张量列表keys/counters/filter_flags/values中的张量个数须等于table_handles的 shape lenvalues 数据类型在 kernel 侧仅支持 float32见 tiling 中的 SUPPORT_VALUES_DATA_TYPE。实现原理与内部细节桶内存布局kernel 源码开头的注释明确了每个桶bucket的数据结构current Bucket contains:int64_t key, uint64_t count, uint8 flag, int64_t value[embeddingDims]对应的字节布局常量embedding_hash_table_export.hflag 字节位于桶内偏移23 字节处KEY_FLAG_OFFSET_OF_BYTE 23value 区起始于偏移24 字节KEY_VALUE_OFFSET_OF_BYTE 24每个桶的行宽keyWidthByte 24 align8(sizeof(T) * embeddingDims)即 value 区按 8 字节向上对齐SingleTableCompute。flag 字节采用位图方式编码共使用 4 个 bit掩码含义VALID_FLAG_MASK 0b00000001桶位有效已被写入 keyFILTER_FLAG_MASK 0b00000010命中准入过滤条件EXPORT_FLAG_MASK 0b00000100已被导出过EVICTED_FLAG_MASK 0b00001000桶位已被驱逐导出判定条件kernel 在逐桶扫描时只有同时满足以下条件的桶才会被导出CountPerThread / ExportPerThreadflag VALID_FLAG_MASK为真桶位有效flag EVICTED_FLAG_MASK为假未被驱逐当export_mode new时要求flag EXPORT_FLAG_MASK为假从未导出过的新 key当export_mode all默认时不做该限制。导出流程对每张表Process()循环执行以下阶段embedding_hash_table_export.h统计阶段各 SIMT 线程扫描自己负责的桶区间统计待导出 key 数量并写入 UBCountPerThread随后用ReduceSum求整核导出数跨核同步每个 AICore 将本核导出数写入 workspace 中的核间同步区SaveToCoreSyncWorkspace再通过SyncAll()同步保证所有核的计数对彼此可见偏移计算各核根据前序核的导出数累加出本核输出起始偏移CalcOffset写出阶段ExportPerThread按偏移连续写出 keyINT64、counterUINT64与 filter_flagsUINT8命中FILTER_FLAG_MASK写 1否则写 0并将 value 区按float2/float4短向量访存合并方式拷贝到输出 values 张量对应行CopyExportValuesMergedembeddingDims%40 且行基址 16B 对齐时用 4 元素合并档否则回退 2/1 元素档刷新导出标志仅当桶此前未被导出时置位EXPORT_FLAG_MASK并统计首次导出数量原子更新计数通过asc_atomic_sub对 table handle 结构体索引 4 处的未导出数量字段做原子减操作AtomicSubToGm供export_modenew的下一次导出判定使用再SyncAll()保证跨核一致。值得注意的是 kernel 仅注册了一个入口模板EmbeddingHashTableExportfloatembedding_hash_table_export.cpp配合 tiling key 6BIT4WIDTH_TILING_KEY分发与 values 仅支持 float32 的约束一致。Tiling 策略Tiling 阶段embedding_hash_table_export_tiling_arch35.cpp做了以下决策解析属性export_mode为new时置exportMode1否则为 0全量导出filtered_export_flag直接透传计算可用核数/线程数核数取平台 AIV 核数与 64 的最小值MAX_CORE_NUM线程数取平台 SIMT 线程数与 256 的最小值MAX_THREAD_NUM平台能力在TilingPrepare阶段通过platform_ascendc::PlatformAscendC获取TilingPrepareForEmbeddingHashTableExport写入 TilingDatatableNum、exportMode、filteredExportFlag、bitWidth6同时作为 tiling key、maxCoreNum、maxThreadNum六个字段TilingData 结构定义见 embedding_hash_table_export_tiling_arch35.h设置调度模式为 batch 模式SetScheduleMode(1)注释说明因 kernel 使用SyncAll需要所有核同时启动并按maxThreadNum*6*8字节申请本地内存申请 workspace16MBASCENDC_TOOLS_WORKSPACE 核数×8×表数量其中核间同步区用于各核上报导出数workspace 计算。调用说明与验证README 的调用说明一节当前为无即该算子暂未在 README 中提供可直接粘贴的样例代码其标准用法是由上层图引擎按上述参数表构造算子节点输入为四个 1 维 INT64/ND 张量输出为四组动态张量列表。对调用方式与行为最直接的参考是仓库自带的单元测试test_embedding_hash_table_export_tiling.cpp 模拟了Ascend95064 AICore平台上的 TilingParse Tiling 全流程输入输出 shape 取{1}/{1024}属性取export_modeall、filtered_export_flagtrue验证 tiling 函数返回GRAPH_SUCCESS可直接作为调用参数形态的参考test_embedding_hash_table_export_infershape.cpp 在Ascend950平台注册信息下验证了 4 输入、4 组动态输出的 InferShape 与数据类型推导流程。小结EmbeddingHashTableExport 是 CANN ops-nn 中面向 embedding hash 表持久化的关键算子它通过table handle 桶数组的底层契约将整张或多张hash 表按桶序导出为 keys/counters/filter_flags/values 四组张量export_mode属性支持全量导出与仅导出新 key两种模式filtered_export_flag支持按准入标志过滤。在 950 系列设备上该算子以 SIMT 多核并行 跨核 workspace 同步 原子计数的方式实现高效导出。使用者只需保证四个输入 1 维张量长度一致、values 为 float32即可通过标准图接口接入该算子具体接入形式可参照上文提到的两个单元测试。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐G-Helper深度解析华硕笔记本性能调优的5大技术突破与实战应用G Helper深度解析华硕笔记本性能调优的5大技术突破与实战应用 G Helper作为一款专为华硕笔记本设计的轻量级开源控制工具在技术架构、性能优化和用户人工智能算子库深度学习CANNAscendCANN ops-nn 仓库 PReLU 算子aclnnPrelu详解功能、参数与 AICore 实现剖析CANN ops nn 仓库 PReLU 算子aclnnPrelu详解功能、参数与 AICore 实现剖析 导读 PReLUParametric Rec人工智能算子库深度学习CANNAscendCANN ops-nn 算子解读ApplyAdagradD 参数更新算子的原理、配置与 NPU 实现CANN ops nn 算子解读ApplyAdagradD 参数更新算子的原理、配置与 NPU 实现 ApplyAdagradD 是 CANN ops nn人工智能算子库深度学习CANNAscend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考