CANN ops-math 算子解读：AsStrided 张量视图算子（as_strided）功能、参数与 GE 图模式调用实战

CANN ops-math 算子解读：AsStrided 张量视图算子（as_strided）功能、参数与 GE 图模式调用实战

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math

AsStrided（as_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 当前仓库中支持以下硬件产品（“√”表示支持）：

从仓库源码看，算子定义（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_i=input_{\text{storage_offset}+\sum_{d=0}^{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 格式。完整参数规格如下表：

源码级佐证：算子原型与定义

• 算子原型 as_strided_proto.h 通过REG_OP(AsStrided)声明了x、size、stride、storage_offset四个输入与输出y：x支持BasicType()、DT_HIFLOAT8、DT_FLOAT8_E5M2、DT_FLOAT8_E4M3FN、DT_BOOL；size/stride/storage_offset均使用IndexNumberType()（即 INT32、INT64）。原型注释进一步约束：size 与 stride 的所有元素必须是非负整数，storage_offset 必须为非负整数。
• 算子定义 as_strided_def.cpp 中，x与y枚举了全部支持的数据类型（含 COMPLEX32/COMPLEX64、HIFLOAT8、FLOAT8 系列），格式统一为FORMAT_ND；size、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完整演示了“初始化 GE → 构图 → 建 Session → 加图 → 运行 → 导出数据”的标准流程，关键步骤与源码对应如下：

1. 初始化 GE：以ge.exec.deviceId=0、ge.graphRunMode=1作为全局选项调用ge::GEInitialize(global_options)，失败则直接返回。
2. 创建算子节点：通过op::AsStrided("as_strided1")创建算子节点（算子名来自 as_strided_proto.h 中的REG_OP(AsStrided)注册）。
3. 添加输入占位符：示例中定义输入 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绑定到算子输入。
4. 声明输出：输出 y 形状{2, 2}，与 x 同 dtype，通过node.update_output_desc_y设置输出描述。
5. 建 Session 并运行：创建ge::Session，session->AddGraph(graph_id, graph)添加计算图，session->RunGraph(graph_id, input, output)执行，并可通过aclgrphDumpGraph将图 dump 出来便于检查。
6. 结果落盘：将输入输出按tc_ge_irrun_test_0008_npu_input_*.bin/tc_ge_irrun_test_0008_npu_output_*.bin命名导出，便于后续与 golden 对比。

ACLNN 单算子 API 调用

除图模式外，仓库还提供了 Level0 单算子 API。接口声明见 as_strided.h：

const 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分别调用GetValueToShape<int32_t>/GetValueToShape<int64_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 分发到不同实现策略，覆盖不同形状、步长模式与数据类型：

多级切分 Tiling 设计

Tiling 逻辑位于 as_strided_tiling_arch35.cpp，从源码结构可以看出其设计要点：

• 按 dtype 位宽映射 tilingKey：tilingTypeKeyMap将 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：全零 stride（SetZeroStrideTilingData）、SIMT（SetSimtTilingData）、gather（SetWithGatherUbParam）等场景分别使用独立的 tiling 数据结构，与 Kernel 侧分支一一对应；此外源码中还体现了维度合并（merge axis）与双切分（dual cut）等优化策略文件（as_strided_merge_axis_tiling_arch35.h、as_strided_dualcut_tiling_arch35.h）。

测试与 golden 验证

仓库为 AsStrided 提供了完整的单测与 golden 基准，可用于验证算子正确性：

• Kernel golden：golden.py 中as_strided_golden使用numpy.lib.stride_tricks.as_strided(x[storage_offset:], size, stride * 元素宽度)生成参考结果（complex32 场景则用torch.as_strided+view_as_real处理），并依据dtype_width_map将 stride 从“元素个数”换算为“字节数”，与算子“按元素计步长”的语义保持一致。
• host 侧单测：test_as_strided_infershape.cpp 验证 shape 推导（含动态 shape），test_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、offset=1 等典型用例的 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 推导与多策略 Tiling，kernel 侧针对对齐搬运、全零 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