BatchToSpaceND 算子深度解析原理、约束与 NPU 图模式调用实践CANN ops-math【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathBatchToSpaceND 是 CANN ops-math 仓库conversion目录下提供的一个张量重排类转换算子负责把批次维度的数据按指定块形状block_shape重新排列到空间维度并依据裁剪参数crops裁剪空间维度。本文将围绕该算子的产品支持情况、数学定义、参数语义、约束条件与图模式调用方式展开并结合 算子 IR 定义、Shape 推导实现 与 Tiling 实现 等仓库源码带读者从公式到 Kernel 完整理解该算子在 NPU 上的落地路径并能直接复用仓库中的示例代码进行图模式单算子验证。一、产品支持情况BatchToSpaceND 在 CANN 当前支持的产品系列中均可使用具体如下表所示产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品√Atlas 推理系列产品√Atlas 训练系列产品√从算子注册配置也可以印证这一点算子定义文件 中通过this-AICore().AddConfig(ascend950, aicore_config)与AddConfig(ascend350, aicore_config)为 ascend950、ascend350 两代昇腾 AI Core 架构注册了相同的OpAICoreConfig并在 ascend950 配置目录 和 ascend350 配置目录 下提供了对应的算子二进制描述文件。二、功能说明2.1 算子功能将批次维度的数据重新排列到空间维度并裁剪空间维度。2.2 功能描述BatchToSpaceND 是 SpaceToBatchND 的逆操作。它首先将输入张量的批次维度按照指定的块形状block_shape重新排列到空间维度中然后根据裁剪参数crops裁剪空间维度。具体来说它将批次维度中的数据分散到空间维度中从而增加空间维度的大小同时减少批次维度的大小。在 算子 IR 定义 的注释中官方将该操作描述为Permutes data from batch into blocks of spatial data and then prunes them即把批次维的值以空间块的形式移动到高度与宽度维度随后对高度和宽度维度进行裁剪同时该头文件明确指出算子兼容 TensorFlow 的BatchToSpaceND算子见 框架插件注册 中FrameworkType(TENSORFLOW)与OriginOpType(BatchToSpaceND)的映射因此在迁移 TensorFlow 网络时可以直接对齐语义。2.3 计算公式设输入张量 x 为 N 维形状为 $[x_0, x_1, \ldots, x_{N-1}]$block_shape 为 M 维 1D 张量 $[b_0, b_1, \ldots, b_{M-1}]$crops 为 $M \times 2$ 的 2D 张量 $[[c_{00}, c_{01}], [c_{10}, c_{11}], \ldots, [c_{M-1,0}, c_{M-1,1}]]$满足 $1 \leq M N \leq 8$。输出张量 y 形状为 $[y_0, y_1, \ldots, y_{N-1}]$计算方式如下$$ y_i \begin{cases} \frac{x_0}{\prod_{j0}^{M-1} b_j}, i 0 \ x_i \times b_{i-1} - c_{i-1,0} - c_{i-1,1}, 1 \leq i \leq M \ x_i, M1 \leq i \leq N-1 \end{cases} $$其中$x_0$ 必须能够被 $\prod_{j0}^{M-1} b_j$ 整除。以 README 中示例代码的参数为例见 test_geir_batch_to_space_nd.cpp输入 x 形状[4, 2, 3]block_shape[2]crops[1, 0]则输出形状按公式计算为$y_0 4 / 2 2$$y_1 2 \times 2 - 1 - 0 3$$y_2 3$即[2, 3, 3]与示例中声明的yShape完全一致。仓库中的 Shape 推导实现 batch_to_space_nd_infershape.cpp 正是按该公式逐段计算的AppendBatchDim()将xShape-GetDim(0)依次除以blockVec_的每个元素得到 batch 维infershape 源码AppendSpatialDims()对前 M 个空间维计算x.shape[i] * block_shape[i-1] - (crops[i-1][0] crops[i-1][1])infershape 源码AppendRemainDims()剩余维度原样保留infershape 源码。三、参数说明BatchToSpaceND 共 3 个输入x、block_shape、crops和 1 个输出y全部为必选参数参数名输入/输出/属性描述数据类型数据格式x输入表示输入张量支持多种数据类型INT8、UINT8、INT16、UINT16、INT32、UINT32、INT64、UINT64、BF16、FLOAT16、FLOAT、DOUBLE、BOOL、COMPLEX32、COMPLEX64NDblock_shape输入表示空间块的形状1D 张量形状为 [M]指定每个空间维度的块大小INT32、INT64-crops输入表示裁剪量2D 张量形状为 [M, 2]指定每个空间维度从顶部和底部或左侧和右侧裁剪的元素数量INT32、INT64-y输出表示输出张量与输入 x 具有相同的数据类型。输出形状根据 block_shape 和 crops 进行计算与 x 一致ND需要特别说明在 Atlas 训练系列产品、Atlas 推理系列产品、Atlas 200I/500 A2 推理产品、Atlas A2/A3 训练与推理系列产品上不支持 BOOL 类型。上述参数语义在仓库源码中有三重印证算子 IR 层batch_to_space_nd_proto.h 通过REG_OP(BatchToSpaceND)声明x为TensorType({BasicType(), DT_BOOL})、block_shape与crops为TensorType::IndexNumberType()输出y与 x 类型一致算子定义层batch_to_space_nd_def.cpp 显式枚举了 x 的 15 种支持类型VALUE_DATA_TYPE_ALL与 block_shape/crops 的索引类型{DT_INT32, DT_INT64}并声明 block_shape、crops 为ValueDepend(OPTIONAL)即 Shape 推导与 Tiling 计算依赖其常量值图融合层batch_to_space_nd_graph_infer.cpp 通过InferDataTypeOutputSameAsInput保证输出 y 的数据类型与输入 x 相同。四、约束说明使用 BatchToSpaceND 时必须满足以下约束输入张量 x 的维度 N 必须满足2 ≤ N ≤ 8block_shape 的维度 M 必须满足1 ≤ M Nblock_shape 的长度必须等于 crops 的第一维度长度crops 的形状必须为[M, 2]输入张量的第 0 维batch 维度必须能够被 block_shape 中所有元素的乘积整除block_shape 中的每个元素必须大于 0crops 中的每个元素必须是非负整数对于每个空间维度 ii 1, 2, ..., M裁剪后的维度大小必须大于等于 0即x.shape[i] × block_shape[i-1] - crops[i-1][0] - crops[i-1][1] ≥ 0。这些约束不仅是文档约定在 Tiling 阶段会被逐条强制校验batch_to_space_nd_tiling_arch35.cppCheckX()校验 x 的维度范围 [2, 8] 与 shape 溢出源码 L246-L281CheckBlockShape()校验 block_shape 为 1D、各元素为正数、乘积不溢出并用batch % block_size ! 0校验 batch 可整除性源码 L283-L342CheckCrops()校验 crops 形状为 [M, 2]、与 block_shape 第一维一致、各元素非负源码 L344-L387CheckY()校验裁剪后空间维不为负即cropedShape crops源码 L389-L418。维度上限的说明README 中约束为 2 ≤ N ≤ 8即 $1 \leq M N \leq 8$而 IR 定义注释中的示例以 N4、M2 的经典形态batch、height、width、depth演示了输出各维的计算关系。对于动态 shape 场景Shape 推导代码对未知维度用UNKNOWN_DIM -1标记且当输入为 unknown rank 或 block_shape 非常量时直接返回 unknown shapeinfershape 源码并在 注册表 中声明InputsDataDependency({block_shape, crops})要求这两个输入以常量形式参与推导。五、调用说明BatchToSpaceND 支持图模式调用即通过算子 IR 构图方式调用算子。README 中给出的调用方式与样例代码如下调用方式样例代码说明图模式调用test_geir_batch_to_space_nd.cpp通过算子IR构图方式调用 BatchToSpaceND 算子。5.1 图模式调用示例详解示例程序 test_geir_batch_to_space_nd.cpp 完整演示了初始化 GE → 构图 → 建 Session → 运行图 → 导出输入输出数据的全流程其核心构图逻辑在CreateOppInGraph()中// 创建单算子节点 auto node op::BatchToSpaceND(node); std::vectorint64_t xShape {4, 2, 3}; std::vectorint64_t yShape {2, 3, 3}; std::vectorint64_t blockShapeShape {1}; std::vectorint32_t blockShapeValue {2}; std::vectorint64_t cropsShape {1, 2}; std::vectorint32_t cropsValue {1, 0}; ADD_SEQUENCE_INPUT(1, x, inDtype, xShape); // 输入 x形状 [4,2,3]按序列 0,1,2,... 填充 ADD_INT_INPUT(2, block_shape, DT_INT32, blockShapeShape, blockShapeValue); // block_shape [2] ADD_INT_INPUT(3, crops, DT_INT32, cropsShape, cropsValue); // crops [[1, 0]] ADD_OUTPUT(1, y, inDtype, yShape); // 输出 y形状 [2,3,3]构建的图执行环境配置为设备ge.exec.deviceId 0、图运行模式ge.graphRunMode 1。运行成功后程序会把输入、输出张量分别落盘为tc_ge_irrun_test_0008_npu_input_*.bin与tc_ge_irrun_test_0008_npu_output_*.bin并逐元素打印输出结果便于核对公式计算结果。5.2 NPU 侧的执行链路从仓库源码可以梳理出 BatchToSpaceND 在图模式调用下的完整执行链路算子注册REG_OP(BatchToSpaceND)定义算子原型proto 头文件Shape 推导IMPL_OP_INFERSHAPE注册的Infershape4BatchToSpaceND按公式计算输出形状infershape 源码Tiling 计算IMPL_OP_OPTILING注册的Tiling4BatchToSpaceND在 Host 侧完成分块与分核决策tiling 源码Kernel 执行batch_to_space_nd入口按编译期 tiling key 分发到三种模板实现apt 源码。5.3 Tiling 分块策略源码级原理Tiling 阶段是 NPU 算子性能的关键。从 DoOpTiling() 可以看到BatchToSpaceND 会根据输入特征选择三种策略之一并通过 tiling key由 mode、blockShapeDimNum、isBigShape 编码GET_TPL_TILING_KEY 使用处告知 KernelTiling4SIMT当输出为空张量、合并后维度数 ≥ 6MIN_RANK_FOR_SIMT、或单字节类型dSize 1时使用按一维块划分块大小向下对齐 SIMT 线程数支持超过 UINT32 元素数的大 shapeisBigShape标志源码 L753-L797Tiling4LargeC当尾轴元素数 ≥ cache line 元素数时使用基于 UB 容量默认双缓冲、上限 64KB逐轴贪心塞入分块并对尾轴做 UB block 对齐crops 前缘按 block_shape 对齐源码 L566-L618Tiling4SmallC其余场景使用采用输入输出双切分的DualSideTiling方案UB tile 过小导致核利用率不足时会按减半因子迭代压缩 tile 尺寸以增加并行度源码 L692-L751。此外Tiling 前的MergeInput()会做合轴优化将 block_shape 为 1 且 crops 为 0 的空间维与相邻维合并把可合并的 remain 维度压缩为单轴从而降低实际参与分块的秩减小 Kernel 计算开销源码 L420-L485。对应地Kernel 入口 batch_to_space_nd_apt.cpp 通过if constexpr在编译期分别实例化BatchToSpaceND4Simt区分 32 位/64 位索引、BatchToSpaceND4LargeC与BatchToSpaceND4SmallCblockShapeDimNum三条执行路径并在 arch35 Kernel 头文件目录 下按模板拆分实现batch_to_space_nd_simt.h、batch_to_space_nd_large_c.h、batch_to_space_nd_small_c.h。六、测试与验证仓库为 BatchToSpaceND 提供了完备的 UT单元测试与 ST系统测试用例可作为验证算子行为与自行构造用例的参考Shape 推导 UTtest_batch_to_space_nd_infershape.cpp 覆盖了 4D 基本场景输入[4,1,1,1]、block_shape[2,2]、crops 全 0期望输出[1,2,2,1]、unknown rank、unknown shape 4D 等场景其中batch_to_space_nd_basic_4d用例同时验证了公式中 batch 维除以块乘积、空间维乘块后裁剪的完整推导Tiling UTarch35 目录 下分别针对 SIMT、LargeC、SmallC 三种模板提供独立测试test_batch_to_space_nd_tiling_simt.cpp、test_batch_to_space_nd_tiling_large_c.cpp、test_batch_to_space_nd_tiling_small_c.cpp系统级测试ttk_kernel_batch_to_space_nd_st.csv 以 CSV 用例表驱动 Kernel 级验证golden.py提供参考输出生成逻辑assets 目录。七、小结BatchToSpaceND 是 CANN ops-math 中一个语义清晰、工程实现完整的张量重排算子文档侧以分段公式精确刻画了 batch→space 的搬移与裁剪规则工程侧则由算子定义、Shape 推导、Host Tiling 与 AICore Kernel 四层实现闭环支撑并针对不同 shape 特征提供 SIMT / LargeC / SmallC 三套分块策略。无论是 TensorFlow 模型迁移中的算子对齐还是需要直接构图验证算子语义的场景都可以直接参考 示例程序 与 算子 IR 定义 快速上手。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 BatchToSpaceND 算子深度解析：原理、约束与 NPU 图模式调用实践（CANN ops-math）
BatchToSpaceND 算子深度解析原理、约束与 NPU 图模式调用实践CANN ops-math【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathBatchToSpaceND 是 CANN ops-math 仓库conversion目录下提供的一个张量重排类转换算子负责把批次维度的数据按指定块形状block_shape重新排列到空间维度并依据裁剪参数crops裁剪空间维度。本文将围绕该算子的产品支持情况、数学定义、参数语义、约束条件与图模式调用方式展开并结合 算子 IR 定义、Shape 推导实现 与 Tiling 实现 等仓库源码带读者从公式到 Kernel 完整理解该算子在 NPU 上的落地路径并能直接复用仓库中的示例代码进行图模式单算子验证。一、产品支持情况BatchToSpaceND 在 CANN 当前支持的产品系列中均可使用具体如下表所示产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品√Atlas 推理系列产品√Atlas 训练系列产品√从算子注册配置也可以印证这一点算子定义文件 中通过this-AICore().AddConfig(ascend950, aicore_config)与AddConfig(ascend350, aicore_config)为 ascend950、ascend350 两代昇腾 AI Core 架构注册了相同的OpAICoreConfig并在 ascend950 配置目录 和 ascend350 配置目录 下提供了对应的算子二进制描述文件。二、功能说明2.1 算子功能将批次维度的数据重新排列到空间维度并裁剪空间维度。2.2 功能描述BatchToSpaceND 是 SpaceToBatchND 的逆操作。它首先将输入张量的批次维度按照指定的块形状block_shape重新排列到空间维度中然后根据裁剪参数crops裁剪空间维度。具体来说它将批次维度中的数据分散到空间维度中从而增加空间维度的大小同时减少批次维度的大小。在 算子 IR 定义 的注释中官方将该操作描述为Permutes data from batch into blocks of spatial data and then prunes them即把批次维的值以空间块的形式移动到高度与宽度维度随后对高度和宽度维度进行裁剪同时该头文件明确指出算子兼容 TensorFlow 的BatchToSpaceND算子见 框架插件注册 中FrameworkType(TENSORFLOW)与OriginOpType(BatchToSpaceND)的映射因此在迁移 TensorFlow 网络时可以直接对齐语义。2.3 计算公式设输入张量 x 为 N 维形状为 $[x_0, x_1, \ldots, x_{N-1}]$block_shape 为 M 维 1D 张量 $[b_0, b_1, \ldots, b_{M-1}]$crops 为 $M \times 2$ 的 2D 张量 $[[c_{00}, c_{01}], [c_{10}, c_{11}], \ldots, [c_{M-1,0}, c_{M-1,1}]]$满足 $1 \leq M N \leq 8$。输出张量 y 形状为 $[y_0, y_1, \ldots, y_{N-1}]$计算方式如下$$ y_i \begin{cases} \frac{x_0}{\prod_{j0}^{M-1} b_j}, i 0 \ x_i \times b_{i-1} - c_{i-1,0} - c_{i-1,1}, 1 \leq i \leq M \ x_i, M1 \leq i \leq N-1 \end{cases} $$其中$x_0$ 必须能够被 $\prod_{j0}^{M-1} b_j$ 整除。以 README 中示例代码的参数为例见 test_geir_batch_to_space_nd.cpp输入 x 形状[4, 2, 3]block_shape[2]crops[1, 0]则输出形状按公式计算为$y_0 4 / 2 2$$y_1 2 \times 2 - 1 - 0 3$$y_2 3$即[2, 3, 3]与示例中声明的yShape完全一致。仓库中的 Shape 推导实现 batch_to_space_nd_infershape.cpp 正是按该公式逐段计算的AppendBatchDim()将xShape-GetDim(0)依次除以blockVec_的每个元素得到 batch 维infershape 源码AppendSpatialDims()对前 M 个空间维计算x.shape[i] * block_shape[i-1] - (crops[i-1][0] crops[i-1][1])infershape 源码AppendRemainDims()剩余维度原样保留infershape 源码。三、参数说明BatchToSpaceND 共 3 个输入x、block_shape、crops和 1 个输出y全部为必选参数参数名输入/输出/属性描述数据类型数据格式x输入表示输入张量支持多种数据类型INT8、UINT8、INT16、UINT16、INT32、UINT32、INT64、UINT64、BF16、FLOAT16、FLOAT、DOUBLE、BOOL、COMPLEX32、COMPLEX64NDblock_shape输入表示空间块的形状1D 张量形状为 [M]指定每个空间维度的块大小INT32、INT64-crops输入表示裁剪量2D 张量形状为 [M, 2]指定每个空间维度从顶部和底部或左侧和右侧裁剪的元素数量INT32、INT64-y输出表示输出张量与输入 x 具有相同的数据类型。输出形状根据 block_shape 和 crops 进行计算与 x 一致ND需要特别说明在 Atlas 训练系列产品、Atlas 推理系列产品、Atlas 200I/500 A2 推理产品、Atlas A2/A3 训练与推理系列产品上不支持 BOOL 类型。上述参数语义在仓库源码中有三重印证算子 IR 层batch_to_space_nd_proto.h 通过REG_OP(BatchToSpaceND)声明x为TensorType({BasicType(), DT_BOOL})、block_shape与crops为TensorType::IndexNumberType()输出y与 x 类型一致算子定义层batch_to_space_nd_def.cpp 显式枚举了 x 的 15 种支持类型VALUE_DATA_TYPE_ALL与 block_shape/crops 的索引类型{DT_INT32, DT_INT64}并声明 block_shape、crops 为ValueDepend(OPTIONAL)即 Shape 推导与 Tiling 计算依赖其常量值图融合层batch_to_space_nd_graph_infer.cpp 通过InferDataTypeOutputSameAsInput保证输出 y 的数据类型与输入 x 相同。四、约束说明使用 BatchToSpaceND 时必须满足以下约束输入张量 x 的维度 N 必须满足2 ≤ N ≤ 8block_shape 的维度 M 必须满足1 ≤ M Nblock_shape 的长度必须等于 crops 的第一维度长度crops 的形状必须为[M, 2]输入张量的第 0 维batch 维度必须能够被 block_shape 中所有元素的乘积整除block_shape 中的每个元素必须大于 0crops 中的每个元素必须是非负整数对于每个空间维度 ii 1, 2, ..., M裁剪后的维度大小必须大于等于 0即x.shape[i] × block_shape[i-1] - crops[i-1][0] - crops[i-1][1] ≥ 0。这些约束不仅是文档约定在 Tiling 阶段会被逐条强制校验batch_to_space_nd_tiling_arch35.cppCheckX()校验 x 的维度范围 [2, 8] 与 shape 溢出源码 L246-L281CheckBlockShape()校验 block_shape 为 1D、各元素为正数、乘积不溢出并用batch % block_size ! 0校验 batch 可整除性源码 L283-L342CheckCrops()校验 crops 形状为 [M, 2]、与 block_shape 第一维一致、各元素非负源码 L344-L387CheckY()校验裁剪后空间维不为负即cropedShape crops源码 L389-L418。维度上限的说明README 中约束为 2 ≤ N ≤ 8即 $1 \leq M N \leq 8$而 IR 定义注释中的示例以 N4、M2 的经典形态batch、height、width、depth演示了输出各维的计算关系。对于动态 shape 场景Shape 推导代码对未知维度用UNKNOWN_DIM -1标记且当输入为 unknown rank 或 block_shape 非常量时直接返回 unknown shapeinfershape 源码并在 注册表 中声明InputsDataDependency({block_shape, crops})要求这两个输入以常量形式参与推导。五、调用说明BatchToSpaceND 支持图模式调用即通过算子 IR 构图方式调用算子。README 中给出的调用方式与样例代码如下调用方式样例代码说明图模式调用test_geir_batch_to_space_nd.cpp通过算子IR构图方式调用 BatchToSpaceND 算子。5.1 图模式调用示例详解示例程序 test_geir_batch_to_space_nd.cpp 完整演示了初始化 GE → 构图 → 建 Session → 运行图 → 导出输入输出数据的全流程其核心构图逻辑在CreateOppInGraph()中// 创建单算子节点 auto node op::BatchToSpaceND(node); std::vectorint64_t xShape {4, 2, 3}; std::vectorint64_t yShape {2, 3, 3}; std::vectorint64_t blockShapeShape {1}; std::vectorint32_t blockShapeValue {2}; std::vectorint64_t cropsShape {1, 2}; std::vectorint32_t cropsValue {1, 0}; ADD_SEQUENCE_INPUT(1, x, inDtype, xShape); // 输入 x形状 [4,2,3]按序列 0,1,2,... 填充 ADD_INT_INPUT(2, block_shape, DT_INT32, blockShapeShape, blockShapeValue); // block_shape [2] ADD_INT_INPUT(3, crops, DT_INT32, cropsShape, cropsValue); // crops [[1, 0]] ADD_OUTPUT(1, y, inDtype, yShape); // 输出 y形状 [2,3,3]构建的图执行环境配置为设备ge.exec.deviceId 0、图运行模式ge.graphRunMode 1。运行成功后程序会把输入、输出张量分别落盘为tc_ge_irrun_test_0008_npu_input_*.bin与tc_ge_irrun_test_0008_npu_output_*.bin并逐元素打印输出结果便于核对公式计算结果。5.2 NPU 侧的执行链路从仓库源码可以梳理出 BatchToSpaceND 在图模式调用下的完整执行链路算子注册REG_OP(BatchToSpaceND)定义算子原型proto 头文件Shape 推导IMPL_OP_INFERSHAPE注册的Infershape4BatchToSpaceND按公式计算输出形状infershape 源码Tiling 计算IMPL_OP_OPTILING注册的Tiling4BatchToSpaceND在 Host 侧完成分块与分核决策tiling 源码Kernel 执行batch_to_space_nd入口按编译期 tiling key 分发到三种模板实现apt 源码。5.3 Tiling 分块策略源码级原理Tiling 阶段是 NPU 算子性能的关键。从 DoOpTiling() 可以看到BatchToSpaceND 会根据输入特征选择三种策略之一并通过 tiling key由 mode、blockShapeDimNum、isBigShape 编码GET_TPL_TILING_KEY 使用处告知 KernelTiling4SIMT当输出为空张量、合并后维度数 ≥ 6MIN_RANK_FOR_SIMT、或单字节类型dSize 1时使用按一维块划分块大小向下对齐 SIMT 线程数支持超过 UINT32 元素数的大 shapeisBigShape标志源码 L753-L797Tiling4LargeC当尾轴元素数 ≥ cache line 元素数时使用基于 UB 容量默认双缓冲、上限 64KB逐轴贪心塞入分块并对尾轴做 UB block 对齐crops 前缘按 block_shape 对齐源码 L566-L618Tiling4SmallC其余场景使用采用输入输出双切分的DualSideTiling方案UB tile 过小导致核利用率不足时会按减半因子迭代压缩 tile 尺寸以增加并行度源码 L692-L751。此外Tiling 前的MergeInput()会做合轴优化将 block_shape 为 1 且 crops 为 0 的空间维与相邻维合并把可合并的 remain 维度压缩为单轴从而降低实际参与分块的秩减小 Kernel 计算开销源码 L420-L485。对应地Kernel 入口 batch_to_space_nd_apt.cpp 通过if constexpr在编译期分别实例化BatchToSpaceND4Simt区分 32 位/64 位索引、BatchToSpaceND4LargeC与BatchToSpaceND4SmallCblockShapeDimNum三条执行路径并在 arch35 Kernel 头文件目录 下按模板拆分实现batch_to_space_nd_simt.h、batch_to_space_nd_large_c.h、batch_to_space_nd_small_c.h。六、测试与验证仓库为 BatchToSpaceND 提供了完备的 UT单元测试与 ST系统测试用例可作为验证算子行为与自行构造用例的参考Shape 推导 UTtest_batch_to_space_nd_infershape.cpp 覆盖了 4D 基本场景输入[4,1,1,1]、block_shape[2,2]、crops 全 0期望输出[1,2,2,1]、unknown rank、unknown shape 4D 等场景其中batch_to_space_nd_basic_4d用例同时验证了公式中 batch 维除以块乘积、空间维乘块后裁剪的完整推导Tiling UTarch35 目录 下分别针对 SIMT、LargeC、SmallC 三种模板提供独立测试test_batch_to_space_nd_tiling_simt.cpp、test_batch_to_space_nd_tiling_large_c.cpp、test_batch_to_space_nd_tiling_small_c.cpp系统级测试ttk_kernel_batch_to_space_nd_st.csv 以 CSV 用例表驱动 Kernel 级验证golden.py提供参考输出生成逻辑assets 目录。七、小结BatchToSpaceND 是 CANN ops-math 中一个语义清晰、工程实现完整的张量重排算子文档侧以分段公式精确刻画了 batch→space 的搬移与裁剪规则工程侧则由算子定义、Shape 推导、Host Tiling 与 AICore Kernel 四层实现闭环支撑并针对不同 shape 特征提供 SIMT / LargeC / SmallC 三套分块策略。无论是 TensorFlow 模型迁移中的算子对齐还是需要直接构图验证算子语义的场景都可以直接参考 示例程序 与 算子 IR 定义 快速上手。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考