CANN ops-math 算子解读AsStrided 张量视图算子as_strided功能、参数与 GE 图模式调用实战【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathAsStridedas_strided是 CANN ops-math 仓库中位于conversion/as_strided的转换类算子它允许用户通过指定新的形状size与步长stride在原张量的同一份底层存储上创建一个共享数据内存的视图view无需复制数据。本文以 conversion/as_strided/README.md 为核心结合算子原型as_strided_proto.h、算子定义as_strided_def.cpp、shape 推导、ACLNN API、Tiling 与 Kernel 实现以及测试用例系统讲解该算子的产品支持情况、数据语义、参数规格、源码级实现原理与 GE 图模式调用样例读完即可在 NPU 上正确使用并验证该算子。产品支持情况AsStrided 算子在 CANN ops-math 当前仓库中支持以下硬件产品“√”表示支持产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品√Atlas 推理系列产品√Atlas 训练系列产品√从仓库源码看算子定义as_strided_def.cpp为 AICore 注册了ascend950与ascend350两套配置均开启动态编译静态标志DynamicCompileStaticFlag(true)、动态 Rank 支持DynamicRankSupportFlag(true)与动态 shape 支持DynamicShapeSupportFlag(true)对应 CMake 构建脚本CMakeLists.txt也明确设置了SUPPORT_COMPUTE_UNIT ascend950 ascend350以及统一的arch35Tiling 目录。功能说明算子功能AsStrided 算子允许用户通过制定新的形状size和步长stride来创建一个与原张量共享相同数据内存的张量视图。也就是说输出张量 y 不会复制输入数据而是以新的逻辑形状、步长与存储偏移量重新“解读”输入张量 x 的底层存储。这是深度学习框架中常用的零拷贝视图操作例如 PyTorch 的torch.as_strided即与其语义兼容as_strided_proto.h 中明确标注了 “Compatible with the PyTorch operator as_strided”。计算公式设输出张量第 d 维的维度大小为size[d]、步长为stride[d]输出元素下标为i多维下标展开为i_0, i_1, ..., i_{D-1}D 为张量维度storage_offset为相对于输入张量底层存储的偏移量则输出元素与输入存储的映射关系为$$ out_iinput_{\text{storage_offset}\sum_{d0}^{D-1}(i_d\cdot \text{strided}[d])} $$即输出第 i 个元素取自输入底层存储的第storage_offset Σ(i_d × stride[d])个位置。以 test_geir_as_strided.cpp 中示例为例输入 x 形状为{3, 3}size 为{2, 2}stride 为{1, 2}storage_offset 为{1}则输出 y 为{2, 2}其元素依次取自输入扁平化存储的偏移1、3、2、4四个位置可见步长允许跨维跳跃读取也允许小于完整行的跨度从而形成重叠视图。参数说明AsStrided 共包含 4 个输入和 1 个输出均为 ND 格式。完整参数规格如下表参数名输入/输出/属性描述数据类型数据格式x输入公式中的 input_i即原输入张量。INT64、UINT64、INT32、UINT32、FLOAT、FLOAT16、INT8、UINT8、BF16、INT16、UINT16、BOOL、COMPLEX32、COMPLEX64、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FNNDsize输入输出张量的形状。INT32、INT64NDstride输入stride[d] 是输入张量在第 d 维的步幅。INT32、INT64NDstorage_offset输入是 out_i 中相对于原张量 input_i 存储的偏移量。INT32、INT64NDy输出公式中的 out_i即生成的视图张量。与 x 相同INT64、UINT64、INT32、UINT32、FLOAT、FLOAT16、INT8、UINT8、BF16、INT16、UINT16、BOOL、COMPLEX32、COMPLEX64、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FNND源码级佐证算子原型与定义算子原型 as_strided_proto.h 通过REG_OP(AsStrided)声明了x、size、stride、storage_offset四个输入与输出yx支持BasicType()、DT_HIFLOAT8、DT_FLOAT8_E5M2、DT_FLOAT8_E4M3FN、DT_BOOLsize/stride/storage_offset均使用IndexNumberType()即 INT32、INT64。原型注释进一步约束size 与 stride 的所有元素必须是非负整数storage_offset 必须为非负整数。算子定义 as_strided_def.cpp 中x与y枚举了全部支持的数据类型含 COMPLEX32/COMPLEX64、HIFLOAT8、FLOAT8 系列格式统一为FORMAT_NDsize、stride、storage_offset三个输入均标记为ValueDepend(OPTIONAL)表示这些输入的值参与后续 Tiling/Shape 推导其中storage_offset的ParamType为OPTIONAL即该输入在语义上可缺省缺省视为偏移 0而x、size、stride、y均为REQUIRED。约束说明原文档明确无约束。不过在实操中仍需遵循原型与定义中隐含的语义约束size、stride 各元素及 storage_offset 须为非负整数且各维步长应保证所访问的存储范围不越界。调用说明图模式调用GE IR 构图AsStrided 支持通过算子 IR 构图方式调用样例位于 test_geir_as_strided.cpp对应的算子 IR 定义见 as_strided_proto.h。调用方式与样例汇总调用方式调用样例说明图模式调用test_geir_as_strided.cpp通过算子IR构图方式调用 as_strided 算子图模式调用流程解析test_geir_as_strided.cpp完整演示了“初始化 GE → 构图 → 建 Session → 加图 → 运行 → 导出数据”的标准流程关键步骤与源码对应如下初始化 GE以ge.exec.deviceId0、ge.graphRunMode1作为全局选项调用ge::GEInitialize(global_options)失败则直接返回。创建算子节点通过op::AsStrided(as_strided1)创建算子节点算子名来自 as_strided_proto.h 中的REG_OP(AsStrided)注册。添加输入占位符示例中定义输入 x 形状{3, 3}float32全 2 填充size为{2, 2}、stride为{1, 2}、storage_offset为{1}三者均为 INT32 常量并分别通过node.set_input_x / set_input_size / set_input_stride / set_input_storage_offset绑定到算子输入。声明输出输出 y 形状{2, 2}与 x 同 dtype通过node.update_output_desc_y设置输出描述。建 Session 并运行创建ge::Sessionsession-AddGraph(graph_id, graph)添加计算图session-RunGraph(graph_id, input, output)执行并可通过aclgrphDumpGraph将图 dump 出来便于检查。结果落盘将输入输出按tc_ge_irrun_test_0008_npu_input_*.bin/tc_ge_irrun_test_0008_npu_output_*.bin命名导出便于后续与 golden 对比。ACLNN 单算子 API 调用除图模式外仓库还提供了 Level0 单算子 API。接口声明见 as_strided.hconst aclTensor* AsStrided(const aclTensor* x, const aclTensor* y, const aclTensor* size, const aclTensor* stride, const aclTensor* storageOffset, aclOpExecutor* executor);实现见 as_strided.cpp其调度逻辑值得关注依据当前 NPU 架构与数据类型判断是否走 AI Core 路径IsAiCoreSupport不同架构维护了各自的 dtype 支持表910B、910 与 regbase 系列其中 regbase 额外包含 HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN支持时进入AsStridedAiCore加入 launcher 列表执行。对于 AI Core 不支持的数据类型自动降级为复用 ViewCopy 算子conversion/view_copy目录下的视图拷贝算子完成等价的 strided 拷贝保证调用方无需感知内部差异。Shape 推导机制算子的输出形状由 host 侧 shape 推导逻辑完成见 as_strided_infershape.cpp推导函数读取size输入张量按DT_INT32/DT_INT64分别调用GetValueToShapeint32_t/GetValueToShapeint64_t将 size 的值直接转换为输出形状其他 dtype 报错并通过InputsDataDependency({IN_SIZE, IN_STRIDE, IN_OFFSET})声明 size、stride、storage_offset 三个输入参与值依赖即推导与 Tiling 都依赖这些输入的实际数值这正是它们在算子定义中被标记为ValueDepend(OPTIONAL)的原因。源码级实现原理多策略 Kernel 调度算子的 NPU Kernel 入口位于 as_strided_apt.cpp通过TILING_KEY_IS(...)按 tilingKey 分发到不同实现策略覆盖不同形状、步长模式与数据类型tilingKey 宏取值触发场景 / 实现策略AS_STRIDED_B8/B16/B32/B641/2/4/8按元素位宽1/2/4/8 字节分派的基础KernelAsStrided实现AS_STRIDED_MOVE_ALIGN_B*101~108可 32B 对齐搬运场景下的KernelAsStridedMoveAlignDataCopy 对齐优化AS_STRIDED_DUAL_CUT200双切分场景KernelAsStridedDualCut见 as_strided_dual_cut.hALL_STRIDEDS_ZERO_KEY300各维 stride 均为 0广播式视图的StridedIsZero快速路径SIMT_KEY400通用 SIMTvector 逐元素兜底路径AsStridedSimtAS_STRIDED_GATHER500基于 gather 的取数路径KernelAsStridedGatherEMPTY_TENSOR_KEY1000空张量尺寸为 0快速返回多级切分 Tiling 设计Tiling 逻辑位于 as_strided_tiling_arch35.cpp从源码结构可以看出其设计要点按 dtype 位宽映射 tilingKeytilingTypeKeyMap将 INT64/UINT64/COMPLEX64 等映射为 8 字节档位FLOAT/INT32/COMPLEX32 等映射为 4 字节档位FLOAT16/BF16/INT16 等映射为 2 字节档位INT8/UINT8/BOOL/HIFLOAT8/FLOAT8 系列映射为 1 字节档位与 Kernel 中的 B8/B16/B32/B64 一一对应。UB 容量感知针对不同位宽配置不同的 UB buffer 上限B8 为 126976 字节、B16 为 63488、B32 为 31744、B64 为 15872并将内轴/外轴切分因子、循环次数、32B 对齐标志等信息写入 tilingData见AsStridedSetTilingData与 as_strided_tiling_arch35.h。特殊场景独立 tiling全零 strideSetZeroStrideTilingData、SIMTSetSimtTilingData、gatherSetWithGatherUbParam等场景分别使用独立的 tiling 数据结构与 Kernel 侧分支一一对应此外源码中还体现了维度合并merge axis与双切分dual cut等优化策略文件as_strided_merge_axis_tiling_arch35.h、as_strided_dualcut_tiling_arch35.h。测试与 golden 验证仓库为 AsStrided 提供了完整的单测与 golden 基准可用于验证算子正确性Kernel goldengolden.py 中as_strided_golden使用numpy.lib.stride_tricks.as_strided(x[storage_offset:], size, stride * 元素宽度)生成参考结果complex32 场景则用torch.as_stridedview_as_real处理并依据dtype_width_map将 stride 从“元素个数”换算为“字节数”与算子“按元素计步长”的语义保持一致。host 侧单测test_as_strided_infershape.cpp 验证 shape 推导含动态 shapetest_as_strided_tiling.cpp 验证 arch35 Tiling 计算。kernel 侧单测test_as_strided.cpp 与数据生成脚本 gen_data.py、gen_tiling.py生成 tiling.bin配合可离线跑 kernel 并与 golden 对比tiling 数据中亦能看到{2,2}输出、{1,2}stride、offset1 等典型用例的 tiling 参数落盘。总结AsStrided 是 CANN ops-math 中实现“零拷贝张量视图”的关键转换算子它以out_i input[storage_offset Σ(i_d × stride[d])]为核心语义通过x / size / stride / storage_offset → y的接口在 ND 格式下支持 INT、UINT、FLOAT、BF16、BOOL、COMPLEX、FLOAT8 等十余种数据类型host 侧完成了基于 size 数值的动态 shape 推导与多策略 Tilingkernel 侧针对对齐搬运、全零 stride、gather、SIMT、空张量等场景分别优化ACLNN 接口还能在 AI Core 不支持的 dtype 上自动降级复用 ViewCopy 算子。开发者既可参照 test_geir_as_strided.cpp 走 GE 图模式构图调用也可使用 as_strided.h 声明的单算子 API并借助 golden.py 完成结果校验。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考