CANN PyPTO 算子开发指南pypto.index_select 张量索引算子详解【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptopypto.index_select是 CANN PyPTOParallel Tensor/Tile Operation 编程范式提供的张量索引算子它沿指定维度dim依据索引张量index中的元素从输入张量input中挑选数据并重组为一个新张量常用于 embedding 查表、序列重排、数据采样等需要按索引收集数据的场景。本文以 pypto-index_select 接口文档 为主体结合仓库内 Python 前端封装、pybind11 绑定与 C 算子实现源码完整讲解该接口的功能语义、参数约束、TileShape 切分配置与调用示例帮助开发者在 Ascend 硬件上正确、高效地完成按索引取数。产品支持情况该接口在当前仓库所适配的昇腾平台上均有支持各产品线覆盖情况如下产品系列支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持不同型号所支持的输入数据类型有所差异具体以 约束说明 中的数据类型列表为准。功能说明pypto.index_select返回一个新的张量该张量使用索引index中的元素沿维度dim对输入张量进行索引。其行为可理解为把input沿dim轴展开成一串切片再用index里的下标去挑选这些切片并按顺序拼接。返回张量与原始输入张量具有相同的维度数其中第dim维度的大小与索引index的长度相同其余维度的大小与原始张量保持一致。形式化定义如下设输入input形状为 $(S_0, S_1, \ldots, S_{n-1})$dim d索引index形状为 $(I_0,)$$$ \begin{array}{l} \text{shape}(\mathbf{input}) (S_0, S_1, \ldots, S_{n-1}) \ dim d \ \text{shape}(\mathbf{index}) (I_0,) \ \text{shape}(\mathbf{result}) (S_0, \ldots, S_{d-1}, I_0, S_{d1}, \ldots, S_{n-1}) \ \mathbf{result}[s_0, \ldots, s_{d-1}, i, s_{d1}, \ldots, s_{n-1}] \mathbf{input}[s_0, \ldots, s_{d-1}, \mathbf{index}[i], s_{d1}, \ldots, s_{n-1}] \end{array} $$从结果形状可以看出dim轴被替换为索引的长度其余轴原样保留。仓库中 C 层的形状推导逻辑 gather.cpp 正是按此规则实现// result shape: params.shape[:aixs] indices.shape params.shape[axis1:] std::vectorint64_t resultShape params-shape; resultShape.erase(resultShape.begin() axis); resultShape.insert(resultShape.begin() axis, indices-shape.begin(), indices-shape.end());注意index本身支持 1~2 维因此上述公式可推广为结果第dim轴及紧随其后的若干轴被替换为index的完整形状详见下文返回值说明。函数原型index_select(input: Tensor, dim: int, index: Tensor) - Tensor参数说明参数名输入/输出说明input输入源操作数类型为 Tensor。不同型号支持的数据类型有所差异详见约束说明。不支持空 TensorShape 仅支持 1~4 维Shape Size 不大于 2147483647即 INT32_MAX。dim输入int 类型索引的维度支持任意不超过 input 维数的值取值范围-input.dim dim input.dim支持负数负值解释为dim input.dim详见约束说明。index输入索引操作数类型为 Tensor。支持的数据类型为 DT_INT32、DT_INT64不支持空 TensorShape 只支持 1~2 维Shape Size 不大于 2147483647即 INT32_MAX且元素必须为合法索引即不超过input在dim轴上的 Shape 大小。上述参数约束在 C 侧有对应的运行时校验逻辑见 gather.cppvoid CheckGatherParamsInvalid(const Tensor params, const Tensor indices, int axis, const std::string opName) { const auto supportedTypes ConfigManager::Instance().GetOpSupportedInputDtypes(Opcode::OP_GATHER); CheckTensorDataType(params.GetStorage(), supportedTypes, opName); CheckFp8ArchSupport(params, opName); std::unordered_setDataType indexSupportedTypes {DT_INT32, DT_INT64}; CheckTensorDataType(indices.GetStorage(), indexSupportedTypes, opName); CheckTensorDimRange(params.GetStorage(), 1, NUM_VALUE_4, opName); // input 维度 1~4 CheckTensorDimRange(indices.GetStorage(), 1, NUM_VALUE_2, opName); // index 维度 1~2 CheckTensorShapeSize(params.GetStorage(), opName); CheckTensorShapeSize(indices.GetStorage(), opName); CheckAxisRange(params, axis); CheckTensorsFormatConsistency(params.GetStorage(), indices.GetStorage(), opName); }返回值说明返回输出 Tensor输出 Tensor 的数据类型与input保持一致输出 Tensor 的 Shape 由input、dim以及index共同确定具体规则为input.shape[:dim] index.shape input.shape[dim1:]dim为负数时先做归一化处理。由于index支持 1~2 维当index为 2 维、形状为 $(I_0, I_1)$ 时结果形状为 $(S_0, \ldots, S_{d-1}, I_0, I_1, S_{d1}, \ldots, S_{n-1})$即input的dim轴被index的完整形状整体替换。仓库中 calc_torch.cpp 提供了基于 PyTorch 的参考计算实现CPU 侧数值参考可用来核对结果的正确性。约束说明使用该接口前必须满足以下约束否则可能导致结果错误或 AI Core Error索引类型与取值index必须是整数类型DT_INT32 或 DT_INT64且元素值为合法索引即不能超出input.shape[dim]。维度取值范围dim为 int 类型取值范围-input.dim dim input.dim。支持负数负值会被解释为dim input.dim。C 侧同样有归一化与越界检查见 gather.cpp。viewshape 约束索引轴不可切input.shape的dim轴 viewshape 不可切要求viewshape[dim] input.shape[dim]其余维度的 Shape 大小不做限制。该约束来自 index_select 的算子语义dim 轴作为索引源需要在当前 view 中整体可见而不是当前实现的额外限制。若 dim 轴按照小于input.shape[dim]的 viewshape 切分index可能引用当前 view 之外的数据导致结果精度错误或 AI Core Error。在 C 切分实现中可以看到对params[axis]轴的处理正是不切分、offset 置 0、shape 取整轴见 gather.cpp// params[axis]不切 paramsInput.tileInfo.offset[axis] 0; paramsInput.tileInfo.shape[axis] paramsInput.tensor.GetShape()[axis];Tensor 数据类型说明input支持的数据类型按产品线区分Ascend 950PR / Ascend 950DTDT_INT8DT_INT16DT_INT32DT_UINT8DT_UINT16DT_UINT32DT_FP16DT_FP32DT_BF16DT_BOOLDT_FP8E4M3DT_FP8E5M2DT_FP8E8M0Atlas A3 训练系列产品 / Atlas A3 推理系列产品DT_INT8DT_INT16DT_INT32DT_UINT8DT_UINT16DT_UINT32DT_FP16DT_FP32DT_BF16Atlas A2 训练系列产品 / Atlas A2 推理系列产品DT_INT8DT_INT16DT_INT32DT_UINT8DT_UINT16DT_UINT32DT_FP16DT_FP32DT_BF16其中 FP8 类型DT_FP8E4M3、DT_FP8E5M2、DT_FP8E8M0仅支持 DAV_3510 架构该检查在 gather.cpp 中通过CheckFp8ArchSupport完成。TileShape 约束TileShape 的维度与 result 相同用于切分 result。TileShape 设置需保证 result 不超过 UB 大小具体用法详见 TileShape 设置示例。调用示例TileShape 设置示例调用该 operation 接口前应通过set_vec_tile_shapes设置 TileShape。TileShape 的维度设置须与输出张量保持一致用于控制输出 Tile 块的大小。以输入 $input[B,S,D]$、索引 $index[T]$、轴 $\text{axis}-2$即 S 轴、输出 $output[B,T,D]$ 为例设 TileShape 为 $[b_1, t_1, d_1]$该配置直接作用于输出 output 的各维度同时映射至输入与索引$b_1$ 切分 input 的批次维 B$d_1$ 切分 input 的特征维 D输入的序列维 S即轴 -2不参与切分仅作为索引源需保证 viewshape 覆盖完整的 S 轴对应约束 3$t_1$ 作用于索引 index 的长度维 T。Tile 内存占用须满足约束$$ b_1 \cdot t_1 \cdot d_1 \cdot \text{sizeof}(\mathbf{output}) \text{UBSize} $$也就是说每个 Tile 块的数据量必须能被统一缓冲区UB容纳设置 TileShape 时需要根据输出数据类型和 UB 容量综合评估。仓库中的系统测试 test_gather_onboard.py 展示了该接口与set_vec_tile_shapes、pypto.loop、pypto.assemble组合使用的典型算子内核模式先在 kernel 入口通过set_vec_tile_shapes设置切分块在循环中切片输入与索引调用pypto.index_select得到当前块的结果再用pypto.assemble将结果写回输出张量pypto.frontend.jit(debug_options{runtime_debug_mode: 0, compile_debug_mode: 0}) def gather_onboard_2d_2input_2d_output_kernel( input0: pypto.Tensor(), input1: pypto.Tensor(), output0: pypto.Tensor(), config: GatherOnboardConfig ): pypto.set_vec_tile_shapes(*config.tile_shape) for index_0 in pypto.loop(config.loop_ranges[0]): for index_1 in pypto.loop(config.loop_ranges[1]): offsets [index_0 * config.execution_view_shape[0], index_1 * config.execution_view_shape[1]] input0_view input0[:] input1_view input1[:] result pypto.index_select(input0_view, config.axis, input1_view) output_offset [ 0 if config.output_offset_map[axis] 0 else offsets[config.output_offset_map[axis]] for axis in range(2) ] pypto.assemble(result, output_offset, output0)接口调用示例x pypto.tensor([3, 4], pypto.DT_FP32) indices pypto.tensor([2], pypto.DT_INT32) out1 pypto.index_select(x, 0, indices) out2 pypto.index_select(x, 1, indices)结果示例如下输入x: [[ 0.1427, 0.0231, -0.5414, -1.0009], [-0.4664, 0.2647, -0.1228, -1.1068], [-1.1734, -0.6571, 0.7230, -0.6004]] 输入index: [0, 2] 输出out1: [[ 0.1427, 0.0231, -0.5414, -1.0009], [-1.1734, -0.6571, 0.7230, -0.6004]] 输出out2: [[ 0.1427, -0.5414], [-0.4664, -0.1228], [-1.1734, 0.7230]]示例解读out1 pypto.index_select(x, 0, indices)沿 0 轴行方向取第 0 行和第 2 行结果形状为[2, 4]out2 pypto.index_select(x, 1, indices)沿 1 轴列方向取第 0 列和第 2 列结果形状为[3, 2]。该示例同时出现在 indexing.py 的 docstring 中作为接口的官方行为参考。调用链与底层实现解析pypto.index_select的完整调用链为Python 前端 API → pybind11 绑定 → C tile_fwk 算子实现 → 图内 OP_GATHER/OP_GATHER_FROM_UB 算子。理解这条链路有助于定位问题、评估性能。1. Python 前端封装接口在 indexing.py 中定义经op_wrapper装饰后转发至底层实现op_wrapper def index_select(input: Tensor, dim: int, index: Tensor) - Tensor: ... return pypto_impl.index_select(input, dim, index)2. pybind11 绑定Python 侧的pypto_impl.index_select由 operation.cpp 绑定到 C 实现底层实际复用 Gather 算子m.def( index_select, [](const Tensor params, int dim, const Tensor indices) { return npu::tile_fwk::Gather(params, indices, dim); }, Tensor index_select.);3. C 算子实现C 入口 gather.cpp 的Gather在完成格式要求 TILEOP_NZ与参数合法性校验后进入GatherOperation执行 tiling 切分形状推导GatherOperationResultShapegather.cpp按替换 dim 轴规则计算结果形状递归切分TiledGatherOperationgather.cpp按照 result 的 TileShape 逐维递归切分三个分支分别处理 gather 轴外层、index 轴本身与 gather 轴内层并在叶子节点为每个 Tile 块插入OP_GATHER_FROM_UB静态形状或OP_GATHER动态形状算子同时设置axis属性动态形状路径TensorGatherOperationgather.cpp针对动态 Shape 场景通过GraphUtils::AddDynOperation添加带动态 valid shape 的OP_GATHER算子形状推断 Pass图内OP_GATHER/OP_GATHER_FROM_UB算子的形状推断函数为InferFunc4Gather注册于 op_infer_shape_impl.cpp同样遵循替换 dim 轴的规则保证图中下游算子拿到的输出形状与接口语义一致。4. 配套参考实现与测试数值参考calc_torch.cpp 提供基于 PyTorch 的 CPU 参考实现用于算子结果的数值校验golden 对照并支持 FP8/FP4 等低比特类型的格式转换处理系统测试test_gather_onboard.py 覆盖 2D/3D 输入、2D/3D/4D 输出等组合场景测试中通过与torch.index_select的 CPU 结果对比来校验pypto.index_select的正确性见 test_gather_onboard.py并覆盖了轴、切分循环、TileShape 等参数的不同配置。总结pypto.index_select是 CANN PyPTO 中实现按索引取数的核心算子接口其语义与 PyTorch 的torch.index_select对齐但在昇腾场景下需要额外关注三点索引轴 viewshape 不可切dim轴作为索引源必须在当前 view 中整体可见这是算子语义决定的硬性约束TileShape 必须与输出维度一致并保证单个 Tile 的内存占用小于 UB 容量数据类型按产品线区分FP8 类型仅限 DAV_3510 架构Ascend 950 系列。正确设置set_vec_tile_shapes并理解其与输入/索引/输出的映射关系是编写高性能、无越界错误的 index_select 内核的关键。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
CANN PyPTO 算子开发指南pypto.index_select 张量索引算子详解【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptopypto.index_select是 CANN PyPTOParallel Tensor/Tile Operation 编程范式提供的张量索引算子它沿指定维度dim依据索引张量index中的元素从输入张量input中挑选数据并重组为一个新张量常用于 embedding 查表、序列重排、数据采样等需要按索引收集数据的场景。本文以 pypto-index_select 接口文档 为主体结合仓库内 Python 前端封装、pybind11 绑定与 C 算子实现源码完整讲解该接口的功能语义、参数约束、TileShape 切分配置与调用示例帮助开发者在 Ascend 硬件上正确、高效地完成按索引取数。产品支持情况该接口在当前仓库所适配的昇腾平台上均有支持各产品线覆盖情况如下产品系列支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持不同型号所支持的输入数据类型有所差异具体以 约束说明 中的数据类型列表为准。功能说明pypto.index_select返回一个新的张量该张量使用索引index中的元素沿维度dim对输入张量进行索引。其行为可理解为把input沿dim轴展开成一串切片再用index里的下标去挑选这些切片并按顺序拼接。返回张量与原始输入张量具有相同的维度数其中第dim维度的大小与索引index的长度相同其余维度的大小与原始张量保持一致。形式化定义如下设输入input形状为 $(S_0, S_1, \ldots, S_{n-1})$dim d索引index形状为 $(I_0,)$$$ \begin{array}{l} \text{shape}(\mathbf{input}) (S_0, S_1, \ldots, S_{n-1}) \ dim d \ \text{shape}(\mathbf{index}) (I_0,) \ \text{shape}(\mathbf{result}) (S_0, \ldots, S_{d-1}, I_0, S_{d1}, \ldots, S_{n-1}) \ \mathbf{result}[s_0, \ldots, s_{d-1}, i, s_{d1}, \ldots, s_{n-1}] \mathbf{input}[s_0, \ldots, s_{d-1}, \mathbf{index}[i], s_{d1}, \ldots, s_{n-1}] \end{array} $$从结果形状可以看出dim轴被替换为索引的长度其余轴原样保留。仓库中 C 层的形状推导逻辑 gather.cpp 正是按此规则实现// result shape: params.shape[:aixs] indices.shape params.shape[axis1:] std::vectorint64_t resultShape params-shape; resultShape.erase(resultShape.begin() axis); resultShape.insert(resultShape.begin() axis, indices-shape.begin(), indices-shape.end());注意index本身支持 1~2 维因此上述公式可推广为结果第dim轴及紧随其后的若干轴被替换为index的完整形状详见下文返回值说明。函数原型index_select(input: Tensor, dim: int, index: Tensor) - Tensor参数说明参数名输入/输出说明input输入源操作数类型为 Tensor。不同型号支持的数据类型有所差异详见约束说明。不支持空 TensorShape 仅支持 1~4 维Shape Size 不大于 2147483647即 INT32_MAX。dim输入int 类型索引的维度支持任意不超过 input 维数的值取值范围-input.dim dim input.dim支持负数负值解释为dim input.dim详见约束说明。index输入索引操作数类型为 Tensor。支持的数据类型为 DT_INT32、DT_INT64不支持空 TensorShape 只支持 1~2 维Shape Size 不大于 2147483647即 INT32_MAX且元素必须为合法索引即不超过input在dim轴上的 Shape 大小。上述参数约束在 C 侧有对应的运行时校验逻辑见 gather.cppvoid CheckGatherParamsInvalid(const Tensor params, const Tensor indices, int axis, const std::string opName) { const auto supportedTypes ConfigManager::Instance().GetOpSupportedInputDtypes(Opcode::OP_GATHER); CheckTensorDataType(params.GetStorage(), supportedTypes, opName); CheckFp8ArchSupport(params, opName); std::unordered_setDataType indexSupportedTypes {DT_INT32, DT_INT64}; CheckTensorDataType(indices.GetStorage(), indexSupportedTypes, opName); CheckTensorDimRange(params.GetStorage(), 1, NUM_VALUE_4, opName); // input 维度 1~4 CheckTensorDimRange(indices.GetStorage(), 1, NUM_VALUE_2, opName); // index 维度 1~2 CheckTensorShapeSize(params.GetStorage(), opName); CheckTensorShapeSize(indices.GetStorage(), opName); CheckAxisRange(params, axis); CheckTensorsFormatConsistency(params.GetStorage(), indices.GetStorage(), opName); }返回值说明返回输出 Tensor输出 Tensor 的数据类型与input保持一致输出 Tensor 的 Shape 由input、dim以及index共同确定具体规则为input.shape[:dim] index.shape input.shape[dim1:]dim为负数时先做归一化处理。由于index支持 1~2 维当index为 2 维、形状为 $(I_0, I_1)$ 时结果形状为 $(S_0, \ldots, S_{d-1}, I_0, I_1, S_{d1}, \ldots, S_{n-1})$即input的dim轴被index的完整形状整体替换。仓库中 calc_torch.cpp 提供了基于 PyTorch 的参考计算实现CPU 侧数值参考可用来核对结果的正确性。约束说明使用该接口前必须满足以下约束否则可能导致结果错误或 AI Core Error索引类型与取值index必须是整数类型DT_INT32 或 DT_INT64且元素值为合法索引即不能超出input.shape[dim]。维度取值范围dim为 int 类型取值范围-input.dim dim input.dim。支持负数负值会被解释为dim input.dim。C 侧同样有归一化与越界检查见 gather.cpp。viewshape 约束索引轴不可切input.shape的dim轴 viewshape 不可切要求viewshape[dim] input.shape[dim]其余维度的 Shape 大小不做限制。该约束来自 index_select 的算子语义dim 轴作为索引源需要在当前 view 中整体可见而不是当前实现的额外限制。若 dim 轴按照小于input.shape[dim]的 viewshape 切分index可能引用当前 view 之外的数据导致结果精度错误或 AI Core Error。在 C 切分实现中可以看到对params[axis]轴的处理正是不切分、offset 置 0、shape 取整轴见 gather.cpp// params[axis]不切 paramsInput.tileInfo.offset[axis] 0; paramsInput.tileInfo.shape[axis] paramsInput.tensor.GetShape()[axis];Tensor 数据类型说明input支持的数据类型按产品线区分Ascend 950PR / Ascend 950DTDT_INT8DT_INT16DT_INT32DT_UINT8DT_UINT16DT_UINT32DT_FP16DT_FP32DT_BF16DT_BOOLDT_FP8E4M3DT_FP8E5M2DT_FP8E8M0Atlas A3 训练系列产品 / Atlas A3 推理系列产品DT_INT8DT_INT16DT_INT32DT_UINT8DT_UINT16DT_UINT32DT_FP16DT_FP32DT_BF16Atlas A2 训练系列产品 / Atlas A2 推理系列产品DT_INT8DT_INT16DT_INT32DT_UINT8DT_UINT16DT_UINT32DT_FP16DT_FP32DT_BF16其中 FP8 类型DT_FP8E4M3、DT_FP8E5M2、DT_FP8E8M0仅支持 DAV_3510 架构该检查在 gather.cpp 中通过CheckFp8ArchSupport完成。TileShape 约束TileShape 的维度与 result 相同用于切分 result。TileShape 设置需保证 result 不超过 UB 大小具体用法详见 TileShape 设置示例。调用示例TileShape 设置示例调用该 operation 接口前应通过set_vec_tile_shapes设置 TileShape。TileShape 的维度设置须与输出张量保持一致用于控制输出 Tile 块的大小。以输入 $input[B,S,D]$、索引 $index[T]$、轴 $\text{axis}-2$即 S 轴、输出 $output[B,T,D]$ 为例设 TileShape 为 $[b_1, t_1, d_1]$该配置直接作用于输出 output 的各维度同时映射至输入与索引$b_1$ 切分 input 的批次维 B$d_1$ 切分 input 的特征维 D输入的序列维 S即轴 -2不参与切分仅作为索引源需保证 viewshape 覆盖完整的 S 轴对应约束 3$t_1$ 作用于索引 index 的长度维 T。Tile 内存占用须满足约束$$ b_1 \cdot t_1 \cdot d_1 \cdot \text{sizeof}(\mathbf{output}) \text{UBSize} $$也就是说每个 Tile 块的数据量必须能被统一缓冲区UB容纳设置 TileShape 时需要根据输出数据类型和 UB 容量综合评估。仓库中的系统测试 test_gather_onboard.py 展示了该接口与set_vec_tile_shapes、pypto.loop、pypto.assemble组合使用的典型算子内核模式先在 kernel 入口通过set_vec_tile_shapes设置切分块在循环中切片输入与索引调用pypto.index_select得到当前块的结果再用pypto.assemble将结果写回输出张量pypto.frontend.jit(debug_options{runtime_debug_mode: 0, compile_debug_mode: 0}) def gather_onboard_2d_2input_2d_output_kernel( input0: pypto.Tensor(), input1: pypto.Tensor(), output0: pypto.Tensor(), config: GatherOnboardConfig ): pypto.set_vec_tile_shapes(*config.tile_shape) for index_0 in pypto.loop(config.loop_ranges[0]): for index_1 in pypto.loop(config.loop_ranges[1]): offsets [index_0 * config.execution_view_shape[0], index_1 * config.execution_view_shape[1]] input0_view input0[:] input1_view input1[:] result pypto.index_select(input0_view, config.axis, input1_view) output_offset [ 0 if config.output_offset_map[axis] 0 else offsets[config.output_offset_map[axis]] for axis in range(2) ] pypto.assemble(result, output_offset, output0)接口调用示例x pypto.tensor([3, 4], pypto.DT_FP32) indices pypto.tensor([2], pypto.DT_INT32) out1 pypto.index_select(x, 0, indices) out2 pypto.index_select(x, 1, indices)结果示例如下输入x: [[ 0.1427, 0.0231, -0.5414, -1.0009], [-0.4664, 0.2647, -0.1228, -1.1068], [-1.1734, -0.6571, 0.7230, -0.6004]] 输入index: [0, 2] 输出out1: [[ 0.1427, 0.0231, -0.5414, -1.0009], [-1.1734, -0.6571, 0.7230, -0.6004]] 输出out2: [[ 0.1427, -0.5414], [-0.4664, -0.1228], [-1.1734, 0.7230]]示例解读out1 pypto.index_select(x, 0, indices)沿 0 轴行方向取第 0 行和第 2 行结果形状为[2, 4]out2 pypto.index_select(x, 1, indices)沿 1 轴列方向取第 0 列和第 2 列结果形状为[3, 2]。该示例同时出现在 indexing.py 的 docstring 中作为接口的官方行为参考。调用链与底层实现解析pypto.index_select的完整调用链为Python 前端 API → pybind11 绑定 → C tile_fwk 算子实现 → 图内 OP_GATHER/OP_GATHER_FROM_UB 算子。理解这条链路有助于定位问题、评估性能。1. Python 前端封装接口在 indexing.py 中定义经op_wrapper装饰后转发至底层实现op_wrapper def index_select(input: Tensor, dim: int, index: Tensor) - Tensor: ... return pypto_impl.index_select(input, dim, index)2. pybind11 绑定Python 侧的pypto_impl.index_select由 operation.cpp 绑定到 C 实现底层实际复用 Gather 算子m.def( index_select, [](const Tensor params, int dim, const Tensor indices) { return npu::tile_fwk::Gather(params, indices, dim); }, Tensor index_select.);3. C 算子实现C 入口 gather.cpp 的Gather在完成格式要求 TILEOP_NZ与参数合法性校验后进入GatherOperation执行 tiling 切分形状推导GatherOperationResultShapegather.cpp按替换 dim 轴规则计算结果形状递归切分TiledGatherOperationgather.cpp按照 result 的 TileShape 逐维递归切分三个分支分别处理 gather 轴外层、index 轴本身与 gather 轴内层并在叶子节点为每个 Tile 块插入OP_GATHER_FROM_UB静态形状或OP_GATHER动态形状算子同时设置axis属性动态形状路径TensorGatherOperationgather.cpp针对动态 Shape 场景通过GraphUtils::AddDynOperation添加带动态 valid shape 的OP_GATHER算子形状推断 Pass图内OP_GATHER/OP_GATHER_FROM_UB算子的形状推断函数为InferFunc4Gather注册于 op_infer_shape_impl.cpp同样遵循替换 dim 轴的规则保证图中下游算子拿到的输出形状与接口语义一致。4. 配套参考实现与测试数值参考calc_torch.cpp 提供基于 PyTorch 的 CPU 参考实现用于算子结果的数值校验golden 对照并支持 FP8/FP4 等低比特类型的格式转换处理系统测试test_gather_onboard.py 覆盖 2D/3D 输入、2D/3D/4D 输出等组合场景测试中通过与torch.index_select的 CPU 结果对比来校验pypto.index_select的正确性见 test_gather_onboard.py并覆盖了轴、切分循环、TileShape 等参数的不同配置。总结pypto.index_select是 CANN PyPTO 中实现按索引取数的核心算子接口其语义与 PyTorch 的torch.index_select对齐但在昇腾场景下需要额外关注三点索引轴 viewshape 不可切dim轴作为索引源必须在当前 view 中整体可见这是算子语义决定的硬性约束TileShape 必须与输出维度一致并保证单个 Tile 的内存占用小于 UB 容量数据类型按产品线区分FP8 类型仅限 DAV_3510 架构Ascend 950 系列。正确设置set_vec_tile_shapes并理解其与输入/索引/输出的映射关系是编写高性能、无越界错误的 index_select 内核的关键。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
