[Operator Mechanism] Fix index_select empty-axis validation

[cangtianhuang]

PR Category

Operator Mechanism

PR Types

Bug fixes

Description

问题描述

对于 paddle.index_select(x, index, axis=0, name=None) 而言，输入应满足以下约束：

1. index 为空时，不读取 x 中的任何元素；
2. index 非空时，则应存在可选择的元素；

如果所选择的轴 x[axis] == 0，则合法索引区间为空，对于任意非空 index 均应报错。本次排查发现，CPU/GPU/XPU 在 0size 快速路径下，均未正确处理“选择轴为空且 Index 非空”输入，因此一并修改并补充单测。

1. CPU：output.numel() == 0 时 early return

import paddle

paddle.set_device("cpu")
x = paddle.empty([0, 0], dtype="float32")
index = paddle.to_tensor([0], dtype="int64")
out = paddle.index_select(x, index, axis=0)
print(out.shape)

# 输出
Tensor(shape=[1, 0], dtype=float32, place=Place(cpu), stop_gradient=True,
       [[]])

此时输出 shape 为 [1, 0]，output.numel() == 0：选择轴 axis=0 由 0 替换为 index.numel()==1，另一维仍为 0。虽然输出 shape 形状推导正确，但 axis=0 的输入轴长度为 0，而 index=[0] 非空，语义上不存在合法可选元素。

2. GPU：非空输出时报 CUDA 700

import paddle

paddle.set_device("gpu")
x = paddle.randn([1, 2048, 0], dtype="float32")
index = paddle.zeros([64], dtype="int64")
out = paddle.index_select(x, index, axis=-1)
paddle.device.synchronize()

# 输出
OSError: (External) CUDA error(700), an illegal memory access was encountered. 
  [Hint: 'cudaErrorIllegalAddress'. The device encountered a load or store instruction on an invalid memory address. This leaves the process in an inconsistentstate and any further CUDA work will return the same error. To continue using CUDA, the process must be terminated and relaunched. ] (at /paddle/paddle/fluid/pybind/cuda_streams_py.cc:164)

解决方案

原有实现根据 output.numel() == 0 判断 0size 并直接返回：

if (output && output->numel() == 0) {
  dev_ctx.template Alloc<T>(output);
  return;
}

其中，负 axis 的归一化、以及选择轴的相关处理均位于该早退之后。CPU 会因 numel==0 静默返回；GPU 用例则因 numel 非零，进入 CUDA kernel 后在空选择轴上计算非法内存访问。

本 PR 将 x.dims() 获取、负 axis 归一化及结构检查移到三个 kernel 的 early return 和后端数据处理前：

auto input_dim = x.dims();
dim = dim >= 0 ? dim : dim + input_dim.size();
if (input_dim[dim] == 0 && index.numel() > 0) {
  PADDLE_THROW(common::errors::InvalidArgument(
      "The dimension of Input(X) on the select axis in OP(index_select) "
      "must be greater than 0 when Input(Index) is not empty."));
}

修改后，相同输入会抛出明确报错：

ValueError: (InvalidArgument) The dimension of Input(X) on the select axis in OP(index_select)
must be greater than 0 when Input(Index) is not empty.

修改文件与含义

本 PR 共修改 5 个文件：

Kernel 实现

• paddle/phi/kernels/cpu/index_select_kernel.cc
  • 将 x.dims() 获取和负 axis 归一化移到 empty-output fast path 前。
  • 在早退前增加“选择轴大小为 0 且 Index 非空”的 host 元数据检查，防止非法输入因其他零维度导致输出为空而静默成功。
  • 删除后续重复的负 axis 处理，正常 int32/int64 分派保持不变。
• paddle/phi/kernels/gpu/index_select_kernel.cu
  • 将 x.dims() 获取和负 axis 归一化移到 empty-output fast path 前，并增加同一结构检查，非法输入不会启动 CUDA kernel。
  • output_dim、stride、size、delta、dtype 分派和设备指针处理继续保留在合法的非空正常路径，不新增一般 GPU Index 值域验证。
• paddle/phi/kernels/xpu/index_select_kernel.cc
  • 在输出早退、设备指针访问、临时 index buffer 分配和 XPU vendor 调用前完成 axis 归一化及结构检查。
  • 保留空 Index 的合法 fast path，并避免非法输入进入 vendor API。

回归测试

• test/legacy_test/test_index_select_op.py
  • 新增 TestIndexSelectInvalidIndexCPU，固定 CPU 覆盖 [1, 0], axis=-1 与 [0, 0], axis=0 两个非法案例，并验证 [0, 0] 配合空 Index 仍返回 [0, 0]。
  • 新增 TestIndexSelectInvalidIndex，在 GPU 可用时运行 GPU，否则运行 CPU；覆盖 [1, 2048, 0], axis=-1 和 [0, 0], axis=0，验证非法输入均抛出新增参数异常，同时保留空 Index 对照。
• test/xpu/test_index_select_op_xpu.py
  • 新增 TestIndexSelectInvalidInput，在 XPU dygraph 环境覆盖 [1, 0], axis=-1 与 [0, 0], axis=0 的异常路径。
  • 增加 [0, 0] 配合空 Index 的合法对照，验证输出 shape 为 [0, 0]。

是否引起精度变化

否

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[risemeup1111]

已复查，未发现需要阻塞合入的问题；本地窄测因环境缺少 numpy 未能执行。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

[codecov-commenter]

Codecov Report

❌ Patch coverage is 75.00000% with 1 line in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (develop@5390960). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
paddle/phi/kernels/cpu/index_select_kernel.cc	75.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             develop   #79509   +/-   ##
==========================================
  Coverage           ?   75.00%           
==========================================
  Files              ?        1           
  Lines              ?        4           
  Branches           ?        0           
==========================================
  Hits               ?        3           
  Misses             ?        1           
  Partials           ?        0           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[risemeup1111]

已复查当前提交，这次只是补充 index_select 的 int32 测试覆盖，没有看到新的回归点。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

[wanghuancoder]

LGTM

[cangtianhuang]

/re-run all-failed

[swgu98]

覆盖率未达标为cpu代码