CANN ops-nn UpdateTensorDesc 算子深度解析：动态 Shape 场景下的 128×int64 描述缓冲区 RMW 更新机制

CANN ops-nn UpdateTensorDesc 算子深度解析动态 Shape 场景下的 128×int64 描述缓冲区 RMW 更新机制【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn导读UpdateTensorDesc 是 CANN ops-nn 算子库 control/update_tensor_desc 目录下的一个元数据控制类算子它不参与任何数值计算而是将一个占位输入x与必填属性shape翻译为输出y的 TensorDesc 更新结果——将目标维度个数与各维大小写入y对应的128×int64 描述缓冲区的固定槽位其余槽位原值透传。该算子专为动态 Shape 场景设计用于在算子图中显式刷新下游节点的输出描述信息。阅读本文后你将完整掌握该算子的功能语义、参数与约束、底层 Kernel 实现S1 CopyIn → S2 Compute → S3 CopyOut 的 RMW 流水、GE IR 图模式调用方式以及配套的单测与 ST 验证口径。产品支持情况产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品×Atlas A2 训练系列产品 / Atlas A2 推理系列产品×Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×当前仅 Ascend 950PR/DT系列产品支持该算子。在源码层面算子定义文件 update_tensor_desc_def.cpp 中通过this-AICore().AddConfig(ascend950, aicoreConfig)显式注册了 ascend950 平台的 AICore 配置与上表的产品支持矩阵一一对应。功能说明算子语义UpdateTensorDesc 接收一个占位输入xKernel 全程不读取其数据按必填属性shape对输出y进行RMW读-改-写将目标维度个数N len(shape)及各维大小shape[i]写入y对应的 128×int64 描述缓冲区的固定槽位其余槽位下标 0~2 与 4N~127保留执行前的原值实现读-改-写而非全量覆盖。该算子常用于动态 Shape 场景下输出 TensorDesc 的显式更新是算子图中用于校正输出描述信息的控制类节点。计算公式设N len(shape)输出y为 128×int64 描述缓冲区则$$ y[3] N,\quad y[4i] shape[i],\ i 0,\dots,N-1 $$其中y的其余槽位下标 0~2 与 4N~127保持原值不变。槽位布局由 update_tensor_desc_tiling_data.h 中的常量定义kDescSize 128描述缓冲区元素数128 × int64 1 KB即 Kernel RMW 的粒度kDimBaseIdx 3rank 写入下标基址即y[3] NkMaxRank 124rank(attr shape) 上限等于kDescSize - kDimBaseIdx - 1。示例输入x占位输入数据不参与计算 tensor([1., 1., 1., 1., 1., 1., 1., 1.], dtypetorch.float32) 属性shape [4, 8, 4] 输出yRMW后仅展示前8个元素 tensor([1, 1, 1, 3, 4, 8, 4, 1], dtypetorch.int64)可见y[3] 3即len([4,8,4])y[4]4、y[5]8、y[6]4依次写入 shape 各维而下标 0~2 保留执行前原值示例中为 1实际取决于执行通路对缓冲区的预填/分配初值。参数说明参数名输入/输出/属性描述数据类型数据格式x输入占位输入 Tensor数据不参与计算Kernel 不读取其数据。bool/float16/float/float64/int8/int16/int32/int64/uint8/uint16/uint32/uint64NDy输出描述缓冲区 Tensordtype 恒为 int64shape 由属性shape推导得出。INT64NDshape属性目标 shape必填ListInt 类型。ListInt-上述参数定义可在算子定义文件 update_tensor_desc_def.cpp 中逐条印证输入x通过.ParamType(REQUIRED)声明为必填DataType枚举了与上表一致的 12 种 dtypebool、FP16、FP32、FP64、INT8/16/32/64、UINT8/16/32/64Format与UnknownShapeFormat均为FORMAT_ND输出y的DataType固定为 12 份ge::DT_INT64与输入 dtype 一一对应无论输入是什么类型输出恒为 int64属性shape为.AttrType(REQUIRED).ListInt()即必填的整型列表。约束说明该算子有一个输入x、一个输出y、一个必填属性shape。输入xrank ∈ [0, 8]dtype 支持 bool/float16/float/float64/int8/int16/int32/int64/uint8/uint16/uint32/uint64 共 12 种format 仅支持 ND。输出ydtype 恒为 int64format 仅支持 NDshape 由属性shape推导得出。属性shaperank ∈ [1, 124]每个元素为非负整数且各元素乘积numel≥ 128。输出y的元素总个数numel≥ 128Kernel 固定按 128×int64 整块 RMWnumel 大于 128 时仅前 128 个元素被读写可观察效果等价于仅覆盖下标 [3, 3N) 区间。输出y的非写入槽位下标 0~2 与 4N~127保留执行前的原值透传。这些约束在 host 侧存在双重 fail-fast 校验形状/类型推导阶段与 tiling 阶段各自独立执行一次。以 update_tensor_desc_infershape.cpp 的 InferShape 为例它依次校验rank ∈ [1, kMaxRank]、每个元素非负、numel ≥ kDescSize任一失败即返回GRAPH_FAILEDupdate_tensor_desc_tiling_arch35.cpp 中的 TilingFunc 则按dtype → format → 维度 → attr → shape的顺序完成同样的校验链并额外校验y的实际 numel ≥ 128。调用说明调用方式调用样例说明图模式调用test_geir_update_tensor_desc参见 算子调用 完成算子编译和验证。GE IR 图模式调用样例解析test_geir_update_tensor_desc.cpp 给出了完整的 GE IR图模式调用框架核心流程如下构造 Data 占位节点通过op::Data(placeholder1).set_attr_index(0)创建输入节点TensorDesc使用FORMAT_ND与目标 dtypeGenPlaceholderData按 dtype 分配缓冲并填充固定字节模式示例注释明确说明 x 为占位输入Kernel 不读取其数据。构造 UpdateTensorDesc 节点op::UpdateTensorDesc(add1)调用set_input_x(placeholder1)与set_attr_shape(attrShape)输出描述通过update_output_desc_y显式声明为DT_INT64。静态/动态双模式RunMode枚举区分RUN_MODE_S静态 descgraph desc 直接使用真实 shape与RUN_MODE_D动态 descgraph desc 以 -1 占位由 InferShape 在运行时推导。输出校验VerifyOutput依次断言输出个数为 1、y的 dtype 为DT_INT64、y.shape attr(shape)、缓冲区字节数与 numel 匹配并逐槽位核对y[3] rank、y[4i] shape[i]。用例矩阵共 5 条用例覆盖 FP32/INT64/FP16/INT32/BOOL 五种 dtype、rank1 最小 numel{128}/rank3{4,8,4}/rank5{2,2,2,2,8}三种形状以及 S/D 两种运行模式最终打印逐用例报告并输出ALL CASES PASSED。GE 初始化参数样例在main中通过ge::GEInitialize配置了三个全局选项std::mapAscendString, AscendString global_options { {ge.exec.deviceId, 0}, {ge.graphRunMode, 1}, {ge.exec.precision_mode, must_keep_origin_dtype}};其中ge.graphRunMode 1表示图运行模式precision_mode must_keep_origin_dtype保证原始 dtype 不被精度优化改写对这类纯 int64 元数据搬运算子尤为重要。算子定义与 AICore 配置算子定义文件 update_tensor_desc_def.cpp 中的OpAICoreConfig集中体现了该算子的设计取向配置项值含义DynamicCompileStaticFlagtrueKernel 固定按 128×int64 RMW与 shape 无关可静态编译DynamicFormatFlagfalse固定 ND 格式DynamicRankSupportFlagtruex rank 0~8y rank 由 attr shape 决定DynamicShapeSupportFlagtrue面向动态 Shape 场景的算子NeedCheckSupportFlagfalsedtype/attr 约束由 host 侧 tiling 校验PrecisionReduceFlagfalse纯 int64 元数据搬运无浮点通路ExtendCfgInfo(opFile.value)update_tensor_desc_apt指定 Kernel 入口文件而 update_tensor_desc_proto.h 提供图模式 IR 注册REG_OP其 proto 内容与 canndev 保持一致供 GE 图编译链路使用示例中直接使用随 CANN 包安装的array_ops.h注册避免与本地 proto 头重复定义。Kernel 实现128×int64 的整块 RMW 流水单链三段式流水Kernel 实现位于 update_tensor_desc_kernel.h采用单 tile、单核、无 double buffer的 S1→S2→S3 单链结构Ascend 950 / arch35非模板化类S1 CopyInDataCopy将y_gmGlobal Memory 中的 128×int64 描述缓冲区以blockCount1, blockLen3232 × 32 B 1024 B整块读入 UBMTE2 搬运。先读入再覆写这是非写入槽位原值透传得以实现的物理基础。Sync1SetFlag/WaitFlagHardEvent::MTE2_SS 管线标量消费等待保证数据读入完成后才开始标量覆写。S2 Compute使用SetValue标量指令覆写yUb_[3] td_-rank并循环写入yUb_[4i] td_-shape[i]写入内容来自 host 侧 Tiling 数据。Sync2SetFlag/WaitFlagHardEvent::S_MTE3等待标量写完成后再触发搬出。S3 CopyOutDataCopy将 UB 中 1 KB 整块写回y_gmMTE3 搬运。入口与 Tiling 数据Kernel 入口 update_tensor_desc_apt.cpp 采用REGISTER_NONE_TILINGtilingKey 恒 0无 TPL 模板参数并显式声明KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY)纯标量/向量通路GET_TILING_DATA_WITH_STRUCT直接按同构布局取回 host 侧填充的 TilingData。TilingData 结构定义在 update_tensor_desc_tiling_data.hhost 与 kernel 共用同一布局仅承载 kernel 无法自取的 attr shape 数据struct UpdateTensorDescTilingData { int64_t rank; // N len(attr shape) ∈ [1, kMaxRank] int64_t shape[kMaxRank]; // attr shape 各维大小仅前 rank 个槽位有效 };Host 侧 Tiling极简单核配置update_tensor_desc_tiling_arch35.cpp 中的TilingFuncUpdateTensorDesc完成固定形状极简 tiling校验通过后填充tiling-rank与tiling-shape[]随后context-SetBlockDim(1)单核 RMW、workspaces[0] 0无 GM workspace。由于 RMW 粒度与 shape 无关编译期信息结构体UpdateTensorDescCompileInfo为空见 update_tensor_desc_tiling_arch35.h无需缓存任何编译期形状信息。测试与验证体系单测UTTiling 校验链全覆盖test_update_tensor_desc_tiling.cpp 通过gert::TilingContextFaker构造上下文模拟Ascend950 / NpuArch 3510平台共 14 条用例分为两类合法场景常规 rank3 FP32 输入校验blockDim1、workspace0、tilingData.rank3、shape[4,8,4]rank1 最小 numel{128} rank0 标量输入x rank8 上限 attr rank5动态 shapex desc 含 -1 未知维 / -2 未知秩——由于 x 是占位输入tiling 仅校验 rank(x)不影响结果非法场景fail-fastx dtype 为 bfloat16不在 12 种 dtype 集合内、y dtype 非 INT64、rank(x) 8、rank(attr shape) 0、attr shape 含负维度、numel(attr shape) 128、numel(y TensorDesc) 128、attr shape 含 -1 未知维、attr shape 为 -2 未知秩均断言返回GRAPH_FAILED。算子规格测试STgolden 与 third_party 双口径golden.py 定义了UpdateTensorDescSpec的完整验证口径golden以 numpy 构造平坦的 int64 缓冲flat[3] n、flat[4:4n] attr_shape其余保留区填 0 占位再 reshape 为attr_shapepre_compare由于保留区初值由执行通路决定kernel 通路 TTK 预填 1GEIR 通路 GE RunGraph 内部分配实测 0均无契约保证不属于算子契约因此比对前将 NPU 输出的保留区下标 3 与 ≥4N清零后整体返回窗口 [3, 4N) 不动交由默认binary_equal整块比对ThirdPartyImplPyTorch/TensorFlow 均无同名算子无对标竞品用 torch 小算子按同口径拼接flat[3] n、flat[4:4n] shape容差tolerance {int64: {standard: binary_equal}}——纯 int64 元数据搬运无浮点/累加因此采用逐位精确比对写入区必有值 ≠ 0y[3] N ≥ 1、shape[i] ≥ 1漏写、清零或错位必被检出。此外tests/st/arch35 目录下的ttk_kernel_update_tensor_desc_st.csv提供了 arch35 平台的 Kernel ST 用例清单用于在真实/仿真环境中验证 128×int64 RMW 的端到端行为。设计要点总结占位输入模式x仅用于保持算子图拓扑完整性Kernel 在Init中直接(void)x不绑定 GlobalTensor、不进 UB彻底避免无意义的数据搬运固定粒度 RMW128×int64 1 KB 的整块读-改-写与 shape 完全解耦从而支撑DynamicCompileStaticFlag(true)与 tilingKey 恒 0 的极简编译模型任何 rank ∈ [1, 124]、numel ≥ 128 的 shape 都只需同一份 Kernel 二进制双重 fail-fast 校验InferShape 与 Tiling 各自独立校验 rank、非负性、numel 下限非法输入在 host 侧即被拒绝不会进入 Kernel通路无关的输出契约算子唯一确定性输出是 RMW 窗口 [3, 4N)窗口外为保留区其值取决于执行通路TTK 预填 1 / GE 分配 0测试通过 pre_compare 抹平该差异将契约收敛到写入区本身。对于需要在动态 Shape 算子图中显式刷新输出描述信息的开发者UpdateTensorDesc 提供了一个零数值计算、纯元数据搬运的标准解法其 调用样例 与 完整测试体系 均可直接作为接入参考。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考