CANN ops-math SpaceToBatch 算子全解析：原理、公式、参数与 NPU 实现

CANN ops-math SpaceToBatch 算子全解析原理、公式、参数与 NPU 实现【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读SpaceToBatch 是 CANN ops-math 算子库中负责空间维度 → 批次维度数据重排的基础算子它将输入张量的空间维度按block_size × block_size的块进行切分把每个块搬移到批次维度并支持对空间维度先进行 zero-padding。本文以 conversion/space_to_batch/README.md 为核心骨架结合仓库中的算子定义、Infershape、Tiling、Kernel 与图模式示例源码系统讲解其功能语义、计算公式、参数约束、调用方式以及 NPU 上的切分与搬运实现帮助你理解该算子的完整落地链路并掌握图模式调用方法。产品支持情况SpaceToBatch 算子在不同昇腾产品上的支持情况如下表所示产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品√Atlas 训练系列产品√从源码结构看该算子的 Tiling 与 Kernel 实现存放于 conversion/space_to_batch/op_host/arch35/space_to_batch_tiling.cpp 与 conversion/space_to_batch/op_kernel/arch35/space_to_batch.h且算子定义中显式注册了ascend950的 AICore 配置见 space_to_batch_def.cpp与 README 中新一代昇腾产品950/A3/A2支持、Atlas 200I/500 A2 推理产品不支持的产品矩阵相互印证。功能说明算子功能将空间维度的数据按块重新排列到批次维度并对空间维度补零。功能描述算子首先根据paddings参数对输入的空间维度进行 zero-padding然后将 padded 空间划分为block_size × block_size的块每个块搬移到批次维度。输出批次维度变为输入的block_size × block_size倍空间维度相应缩小。算子定位在 space_to_batch_proto.h 的算子原型注释中明确说明SpaceToBatch 将空间数据划分为块并搬移到批次维度是 BatchToSpaceND 的逆操作且与 TensorFlow 框架的space_to_batch算子兼容。仓库中的 space_to_batch_tf_plugin.cpp 正是通过REGISTER_CUSTOM_OP(SpaceToBatch).FrameworkType(TENSORFLOW).OriginOpType(SpaceToBatch)完成对 TensorFlow 原始算子的映射接入方便 TF 模型在昇腾上进行图编译。计算公式设输入 x 为 4D NHWC 张量[N, H_in, W_in, C]block_size bspaddings [[pad_top, pad_bottom], [pad_left, pad_right]]H_padded H_in pad_top pad_bottom W_padded W_in pad_left pad_right H_out H_padded / bs W_out W_padded / bs要求H_padded和W_padded均能被bs整除。输出 y 形状为[N * bs * bs, H_out, W_out, C]。坐标映射输出 → 输入n n_out / (bs * bs) bh (n_out % (bs * bs)) / bs bw n_out % bs h_in h_out * bs bh - pad_top w_in w_out * bs bw - pad_left 若 0 ≤ h_in H_in 且 0 ≤ w_in W_in y[n_out, h_out, w_out, c] x[n, h_in, w_in, c] 否则 y[n_out, h_out, w_out, c] 0该坐标映射在 Kernel 侧有逐行对应的实现MapToInput与DecodeNOut见 space_to_batch.h将输出坐标解码为输入坐标若h_in/w_in越界落在 padding 区域则判定为无效位置输出补零。典型示例以图模式示例 test_geir_space_to_batch.cpp 中的用例为例输入x形状为[1, 4, 4, 2]block_size 2paddings [[0, 0], [0, 0]]无 padding则H_out (4 0 0) / 2 2W_out (4 0 0) / 2 2N_out 1 * 2 * 2 4输出y形状为[4, 2, 2, 2]。示例代码中对应的 shape 声明为yShape {4, 2, 2, 2}注释中也给出了完整的推导过程。参数说明参数名输入/输出/属性描述数据类型数据格式x输入表示输入张量4D NHWC 张量[N, H_in, W_in, C]支持多种数据类型INT8、UINT8、INT16、UINT16、INT32、INT64、FLOAT16、FLOAT、DOUBLENHWCpaddings输入表示空间维度 zero-padding 量2D 张量形状为[2, 2]值为[[pad_top, pad_bottom], [pad_left, pad_right]]INT32、INT64NDblock_size属性必需表示空间块的尺寸大小必须是大于 0 的整数。H_padded和W_padded必须能被block_size整除INT-y输出表示输出张量与输入 x 具有相同的数据类型。输出形状为[N*block_size*block_size, H_out, W_out, C]与 x 一致NHWC上述参数在算子定义源码中得到完整印证在 space_to_batch_proto.h 中REG_OP(SpaceToBatch)声明了输入xTensorType::BasicType()、输入paddingsTensorType::IndexNumberType()、输出yTensorType::BasicType()以及必需属性block_sizeInt。在 space_to_batch_def.cpp 中进一步细化了支持的数据类型组合x支持 9 种基本类型INT8/UINT8/INT16/UINT16/INT32/INT64/FLOAT16/FLOAT/DOUBLEpaddings支持 INT32/INT64 两种索引类型共组合出9 × 2 18种 dtype 组合同时定义了x为 NHWC 格式、paddings为 ND 格式输出y通过Follow(x)继承输入的类型与格式。约束说明输入张量 x 必须为 4D NHWC 格式。block_size必须大于 0。paddings形状为[2, 2]每个元素 0。H_padded H_in pad_top pad_bottom必须能被block_size整除。W_padded W_in pad_left pad_right必须能被block_size整除。block_size为编译期常量算子属性非运行时 tensor 输入。这些约束在 Infershape 与 Tiling 阶段均有硬性检查SpaceToBatchInferShapeHelper::Init中校验block_size 0会直接返回GRAPH_FAILED见 space_to_batch_infershape.cppTiling 的ParamCheck与ComputeOutShape同样会校验 x 必须为 4D、paddings必须含 4 个元素以及H_padded/W_padded对block_size的整除性不满足时记录错误日志并返回失败见 space_to_batch_tiling.cpp。此外Infershape 通过IMPL_OP_INFERSHAPE(SpaceToBatch)声明了.InputsDataDependency({INPUT_IDX_PADDINGS})即输出 shape 的计算依赖paddings的常量值当paddings为常量输入时可精确推导输出各维若paddings为动态输入或 x 存在未知维度则对应维度以-1UNKNOWN_DIM表示。输出数据类型则直接继承输入 x 的数据类型InferDataType4SpaceToBatch。调用说明调用方式样例代码说明图模式test_geir_space_to_batch通过算子 IR 构图方式调用 SpaceToBatch 算子。图模式示例解析示例CreateOppInGraph的核心构图流程如下见 test_geir_space_to_batch.cpp// SpaceToBatch: 4D NHWC [N, H_in, W_in, C], block_size2, no padding auto node op::SpaceToBatch(node); std::vectorint64_t xShape {1, 4, 4, 2}; std::vectorint64_t paddingsShape {2, 2}; std::vectorint32_t paddingsValue {0, 0, 0, 0}; // [[0,0],[0,0]], no padding int64_t blockSize 2;关键步骤创建算子节点通过op::SpaceToBatch(node)创建算子实例该符号由 space_to_batch_proto.h 中的REG_OP(SpaceToBatch)注册生成。添加输入使用宏ADD_SEQUENCE_INPUT为 x 生成 NHWC 格式、按 0 递增序列填充的输入数据形状[1, 4, 4, 2]使用宏ADD_INT_INPUT为 paddings 生成 ND 格式、值为{0, 0, 0, 0}的常量输入形状[2, 2]。设置属性node.SetAttr(block_size, blockSize)将block_size设为 2。声明输出通过宏ADD_OUTPUT声明输出 y 的数据类型与形状{4, 2, 2, 2}。构图运行main函数中依次执行ge::GEInitialize→ 创建Session→session-AddGraph→session-RunGraph→ge::GEFinalize并将输入输出张量落盘为tc_ge_irrun_test_0008_npu_input_*.bin/tc_ge_irrun_test_0008_npu_output_*.bin二进制文件同时打印输出结果供验证。该示例同时演示了 GEGraph Engine图模式下调用自定义算子的一般范式可用于本地验证算子的功能正确性。底层实现NPU 上的 Tiling 与 Kernel 链路Tiling面向多核并行的切分策略Tiling 阶段space_to_batch_tiling.cpp负责将算子任务切分为可在多个 AIV 核上并行的子任务核心流程为参数校验ParamCheck读取 x 的形状要求 4D、block_size属性要求 0、paddings常量要求恰好 4 个元素并换算元素字节数dSize。输出形状推导ComputeOutShape按[N*bs*bs, (Hpad)/bs, (Wpad)/bs, C]计算输出形状同时校验空间维整除性。获取 SoC 信息GetSocInfo通过PlatformAscendC读取 AIV 核数coreNum、UB 内存大小、CacheLine 大小、UB Block 大小换算每 CacheLine/UB Block 的元素数。UB 切分寻优DoOpTilingUB buffer 取ubSize / 2与64KB的较小值双缓冲BUFFER_NUM 2从 C 轴开始向 N 轴方向遍历寻找满足 CacheLine 对齐、UB 容量限制且能使核利用率realCoreNum / coreNum 0.8或核数更大最大的切分轴ubAxis与切分因子ubFactor若找不到满足条件的切分则回退为在 N 轴切分factor1。生成 TilingData将输入/输出形状、block_size、paddings、totalCount、perCoreCount、ubAxis、ubFactor、bufferSize写入SpaceToBatchTilingData结构定义见 space_to_batch_tiling_data.h并通过context_-SetTilingKey设置模板参数ubAxis、context_-SetBlockDim(realCoreNum_)设置实际使用的核数。Tiling 键的模板参数声明位于 space_to_batch_tiling_key.h支持ubAxis取 N/H/W/C 四种模板实例化算子的简化键配置见 space_to_batch_simplified_key.ini默认值为 0。Kernel坐标映射驱动的块搬运Kernel 侧space_to_batch.h以输出坐标 → 输入坐标的映射为核心实现数据搬运主要设计点包括Layout 计算根据输入/输出 shape 计算 N/H/W/C 四个轴的 inStride 与 outStrideC 轴 stride 为 1符合 NHWC 连续通道布局。块级并行Process每个 AIV 核根据blockIdx_与perCoreCount计算自己负责的块区间[startIdx, endIdx)逐块处理。Padding 快速路径BlockHasPadding对于不落在 padding 区域的块跳过Duplicate清零操作减少不必要的 UB 写入仅在涉及 padding 边界的块上先Duplicate补零。按轴分派的 CopyIn根据模板参数UbAxis选择CopyInAxisC / CopyInAxisW / CopyInAxisH / CopyInAxisN四种搬运模式其中 W 轴切分利用LoopModeParams循环模式源地址按blockSize * channels * sizeof(T)跳步搬运对应公式中w_in w_out * bs bw的映射关系。Ping-Pong 双缓冲UB 划分为两块BUFFER_NUM 2通过SetEvent/WaitFlag同步 MTE2拷贝入、MTE3拷贝出、V向量计算流水实现搬运与计算重叠。Kernel 入口函数space_to_batch见 space_to_batch_apt.cpp声明为KERNEL_TYPE_AIV_ONLY并依据DTYPE_X的字节宽度int8/int16/int32/int64实例化SpaceToBatchKernelT, UbAxis再按UbAxis模板分派到对应实现。测试佐证仓库为该算子的 Tiling 提供了单元测试 test_space_to_batch_tiling.cpp可对上述 Tiling 逻辑进行用例级验证测试资源目录 tests/assets/golden.py 中亦提供了 golden 数据生成脚本用于对比验证算子输出。读者可结合这些文件深入理解算子的期望行为。总结SpaceToBatch 在 CANN ops-math 中实现了空间分块 批次搬移 空间补零的完整语义算子原型与算子定义源码严格限定了 NHWC 4D 输入、[2,2]paddings 与编译期常量block_sizeInfershape 依赖 paddings 常量精确推导输出形状Tiling 阶段基于 SoC 信息在 N/H/W/C 四轴上寻优切分并以模板参数下发Kernel 阶段则以输出到输入的坐标映射为核心结合双缓冲流水与 padding 快速路径完成高效搬运。配合图模式示例与 TensorFlow 框架映射开发者可以在昇腾 NPU 上便捷地复用 TF 生态中的space_to_batch语义。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考