[Cherry-Pick][Operator Mechanism]Add CUDA small-stride index gradient kernel - #79731
[Cherry-Pick][Operator Mechanism]Add CUDA small-stride index gradient kernel#79731feixi139 wants to merge 10 commits into
Conversation
There was a problem hiding this comment.
Paddle-Bot Review Board （review完成)
| 序号 | 位置 | 优先级 | 规则来源 | 状态 |
|---|---|---|---|---|
| 1 | 非法偏移回退 |  |
默认规则 | ✅ |
| 2 | 负步长读取与归约 | |
默认规则 | ✅ |
| 3 | None 轴回归测试 | |
默认规则 | ✅ |
| 4 | 大索引张量位宽分派 |  |
默认规则 | ✅ |
| 5 | 非连续 value 位宽分派 | |
默认规则 | ✅ |
| 6 | CPU 大偏移回归覆盖 | |
默认规则 | ✅ |
| const size_t ndim = input_dims.size(); | ||
| const size_t nidx = index_strides.size(); | ||
| if (nidx == 0 || ndim == 0 || input_strides.size() != ndim || | ||
| index_dims.size() < nidx || slice_offset % elesize != 0) { |
There was a problem hiding this comment.
slice_offset 判为 false 后仍会继续执行 elementwise 回退，而入口只检查了 0 <= slice_offset < grad_bytes。例如 float32 的 slice_offset=1 会让回退路径构造未对齐的 T*，atomic 写可能越界或损坏相邻元素；index_strides 不能被 elesize 整除时同理。请在进入任一路径前直接拒绝这些非法字节属性，并让跨度校验覆盖回退路径。
| std::vector<std::pair<int64_t, int64_t>> axes; // (stride, extent) | ||
| for (size_t i = 0; i < dims.size(); ++i) { | ||
| if (dims[i] == 1) continue; // its stride is never used | ||
| if (strides[i] <= 0) return false; |
There was a problem hiding this comment.
x[::-1, idx] 会因此在后续回落到原来的 CudaAtomicAdd 路径，fp16/bf16 的重复下标仍会逐次舍入且结果不确定，未获得本 PR 的排序归约语义。请基于绝对 stride 判断重叠并正确规范化起始偏移，或提供安全的负 stride 路径，同时补充重复下标回归测试。
There was a problem hiding this comment.
IsNonOverlapping 虽使用绝对 stride，但前向 index_elementwise_get 仍在 paddle/phi/kernels/gpu/index_elementwise_get_kernel.cu 中以 make_offset_calculator_put<3, false, OffsetT> 构造 offset calculator；其 stride_t 是无符号类型，而 strided_slice 会产生负的 input_strides，因此 x[::-1, idx] 的读取偏移会被转换为巨大正数。排序反向只修正写回，无法挽救已经错误的前向值。请让前向 offset calculator 支持 signed strides 并规范化 base offset，或禁止该路径，并增加负步长 eager/反向回归测试。
| // here keeps `advanced_index_dim` in the same coordinate system, | ||
| // otherwise a `None` placed before an advanced index would bind the index | ||
| // to the wrong axis. | ||
| estimated_dim++; |
There was a problem hiding this comment.
None 位于高级索引前时轴绑定错误的核心改动，但新增测试只覆盖 CUDA 反向，仓库中也没有 x[None, :, i]、x[:, None, i] 等回归。该逻辑同时影响 eager/static 的 getitem 和 setitem，轴号偏差会直接产生静默错误结果；请补充两种模式下与 NumPy 对照的 getitem/setitem 测试，并覆盖多个 None 及被 None 分隔的高级索引。
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## release/3.4 #79731 +/- ##
===============================================
Coverage ? 100.00%
===============================================
Files ? 2
Lines ? 8
Branches ? 0
===============================================
Hits ? 8
Misses ? 0
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Paddle-Bot
left a comment
Paddle-Bot
left a comment
There was a problem hiding this comment.
本轮复核后仍有 1 条 P1 未解决：负步长切片进入前向 index_elementwise_get 时，offset calculator 使用无符号 stride，负 stride 会被解释为巨大正偏移；当前反向排序路径的修正不能弥补前向读取错误。请修复该读取路径或明确禁止此组合并补充回归测试。
Advanced indexing on a reversed view (`x[::-1, idx]`) raised CUDA error 700. `strided_slice` returns a view whose base offset points at the highest address of the reversed axis, so the gather/scatter kernels have to reach backwards from `slice_offset`, but their offset calculators were instantiated with an unsigned stride type that wrapped those offsets into huge positive values. - Instantiate the offset calculators of index_elementwise get/put, forward and backward, CPU and GPU, with signed strides. Add `funcs::IsInInt32Range` for the 32/64-bit dispatch, since the unsigned bound no longer applies, and guard the new `CheckOffsetRange` with `if constexpr (signed_strides)` so the unsigned instantiations under `stride/` keep their previous behaviour. The forward gather kernel now validates both ends of the touched span and reports a diagnosable error instead of faulting. - Order the iteration dimensions in `reorder_dimensions` by |stride| instead of by signed value. A negative stride always compared smallest and was sorted innermost, and the contiguous strides that `allocate_or_resize_outputs()` assigns afterwards then described a permuted layout while the output tensor's meta stayed C-contiguous, silently transposing the result for rank >= 3. `coalesce_dimensions` keeps its signed comparison on purpose, otherwise a reversed axis would be merged into its neighbour.
Failed CI看板
日志分析报告失败的测试 case: 根本原因分析:
修复建议:
Powered by Nyanpasu with gpt-5.6-luna 默认推理级别, please check the suggestions carefully. |
Paddle-Bot
left a comment
Paddle-Bot
left a comment
There was a problem hiding this comment.
本轮复核确认上一轮负步长读取问题已修复，但发现 1 条新的 P2：大索引张量的位宽分派未纳入第三个 operand，详见内联评论。
| if (out->numel() == 0) return; | ||
|
|
||
| if (funcs::IsInUint32Range(x.numel() * sizeof(T), out->numel() * sizeof(T))) { | ||
| if (funcs::IsInInt32Range(x.numel() * sizeof(T), out->numel() * sizeof(T))) { |
There was a problem hiding this comment.
OffsetT 使用有符号类型，CheckOffsetRange 会检查第三个 operand（索引张量）的字节跨度，但这里的 32/64 位分派只看 x 和 out，没有覆盖索引张量。例如 x 为 bool[1]、索引为 int64[300000000] 时，x 和输出（300 MB）都在 int32 范围，仍会选择 OffsetT=uint32_t（其 stride_t 为 int32_t），随后第三路 2.4 GB 的 stride 会在构造 calculator 时直接触发溢出；此前无符号 32 位路径仍可表示该范围。请把索引张量的实际 span 纳入分派，或在 32 位检查失败时回退到 64 位，并同步 get/put/grad 的入口。
The indexing offset calculators are instantiated with signed_strides=true, so their 32 bit fast path is bounded by int32_t rather than uint32_t. The dispatch, however, only measured x and out. The third operand walks the index tensors with byte strides scaled by sizeof(int64_t), so it can reach past int32_t while x and out stay small: a bool x[1] gathered by an int64 index of 3e8 elements produces a 2.4 GB span and CheckOffsetRange rejected it, even though the previous unsigned path handled that shape. Add IndexOperandByteSpan(), which derives the broadcast index shape with cal_shape_stride and returns (elements - 1) * sizeof(int64_t), and feed it into the dispatch at all six entries that build a signed calculator (index_elementwise get / get_grad / put / put_grad). The index tensors may be broadcast, so index[0]->numel() would understate the span and is not used.
Paddle-Bot
left a comment
Paddle-Bot
left a comment
There was a problem hiding this comment.
本轮复核确认上一条大索引张量位宽分派问题已修复，但发现 1 条新的 P2：index_elementwise_put_with_tensor 的 32 位分派仍未覆盖非连续 value 的实际 stride span，详见内联评论。
| if (funcs::IsInUint32Range(x.numel() * sizeof(T), out->numel() * sizeof(T))) { | ||
| if (funcs::IsInInt32Range(x.numel() * sizeof(T), | ||
| out->numel() * sizeof(T), | ||
| funcs::IndexOperandByteSpan(index_dims))) { |
There was a problem hiding this comment.
put_with_tensor 的 offset calculator 还会通过 value.strides() 计算第二路偏移，这里仍只检查 x、out 和索引。比如 value 是形状为 [2]、字节 stride 大于 INT32_MAX 的非连续视图时，x/out/index 都可以很小，当前仍选择 32 位，随后 CheckOffsetRange 会因 operand 1 溢出而抛错，尽管 64 位路径能够处理。请按 value 的 dims/strides 计算最大可达偏移并纳入分派，或在 32 位范围检查失败时回退到 64 位。
index_elementwise_put_with_tensor hands `value` to the offset calculator with its own dims and strides, but the dispatch only weighed x, out and the index operand. Their extents do not bound value's: a complex128 value of 2e8 elements reaches 3.2e9 bytes while x, out and the int64 index operand (1.6e9 bytes) all stay inside int32_t, so the 32 bit path was chosen and CheckOffsetRange then rejected a shape the 64 bit path handles.
Paddle-Bot
left a comment
Paddle-Bot
left a comment
There was a problem hiding this comment.
本轮复核确认上一条非连续 value 位宽分派问题已修复，但发现 1 条新的 P2：CPU signed offset calculator 没有 64 位回退，会拒绝原先可处理的 2–4 GiB 连续张量，详见内联评论。
| strides_vec); | ||
| auto offset_calc = | ||
| funcs::CPUmake_offset_calculator_put<3>(desired_shape, strides_array); | ||
| auto offset_calc = funcs::CPUmake_offset_calculator_put<3, true>( |
There was a problem hiding this comment.
CPUOffsetCalculator<..., uint32_t, true>，其 stride_t 是 int32_t，但 CPU 没有 64 位 offset 回退。对于连续 float32 张量，只要可达字节偏移超过 INT32_MAX（约 2 GiB、但仍小于 4 GiB），这里的 CheckOffsetRange 就会直接抛错；改动前的 unsigned calculator 仍可处理这段范围。请保留非负 stride 的 unsigned fast path，或为 CPU 增加 64 位分派。
There was a problem hiding this comment.
当前实现已通过统一使用有符号 int64_t offset 修复原问题。不过，上一版提交中用于直接覆盖 INT32_MAX/UINT32_MAX 以上 offset 的轻量 C++ 测试在当前 head 中已被移除，没有对应的边界回归覆盖。请恢复该测试，避免这一位宽修复后续无保护地回退。
CPUOffsetCalculator took a signed_strides flag mirroring the GPU one, and with it the CPU kernels inherited the GPU's tradeoff: signed offsets make reversed views work but halve the reach, and the GPU pays for that by dispatching between a 32 bit and a 64 bit calculator. The CPU has no such dispatch, so a contiguous float32 [2, 600000000] -- 2.4 GB of reachable byte offsets, past int32_t but inside uint32_t -- was rejected by CheckOffsetRange, a shape the previous unsigned instantiation handled. The tradeoff does not exist on the CPU: a 64 bit add costs the same as a 32 bit one, and only the divmod over linear_idx benefits from index_t staying 32 bit, which is independent of the offset width. So accumulate offsets in int64_t unconditionally. That leaves signed_strides with no false user, so drop the flag from the calculator and its three factories and revert the five call sites; the reach check goes with it, since int64 bounds cannot be exceeded by a real allocation.
Three places needed the same quantity: the lowest and highest offset that walking a set of dims with a set of strides can produce. The bounds check in GPUIndexElementwiseGetKernel and the layout derivation in DeriveSortedPathLayout each carried their own loop, and StridedOperandByteSpan summed |(dim - 1) * stride| for the dispatch, which is the same hi - lo written differently. Each copy has to get the same two details right: a size-1 axis never uses its stride, and a negative stride pulls the low end below the base instead of raising the high end. Collect that in OperandReach plus AccumulateReach and have all three use it. AccumulateReach adds into an existing reach so that operands sharing a base -- x and the index tensor are both offset from slice_offset -- combine, and takes a scale so a caller can pass element strides and get bytes. Also correct the comment above IsNonOverlapping. It claimed a reversed view cannot reach the grad kernel because the forward kernel truncates negative strides to unsigned, but the forward kernel instantiates its offset calculator with signed_strides = true, so reversed views do arrive here and the negative case is live.
a8b7ce2 to
d81337d
Compare
|
/re-run all-failed |
4 participants
PR Category
Operator Mechanism
PR Types
Bug fixes
Description
本 PR 修复 GPU 高级索引（advanced indexing）反向的三个问题：
x[:, idx]这类表达式的梯度精度明显低于x[idx]；None（newaxis）出现在高级索引之前时，索引被绑定到错误的轴，前向 shape 就是错的；index_elementwise_get_grad没有校验slice_offset和被索引视图的跨度，非法值会变成越界写。1. 高级索引反向的低精度累加
torch 的
index_put_(accumulate=True)CUDA 实现完全没有 atomicAdd，只有排序路径，所以它的结果确定、可复现，且重复下标只舍入一次。本节的三处改动合起来把 Paddle 的所有 accumulate 反向也导向这条路径：1.1 放开路由，1.2 补上排序路径缺失的 kernel，1.3 让排序路径能处理放开后新进来的跨步视图。1.1 路由判据用错了维度，
x[:, idx]落到逐项舍入的 atomicAdd问题
index_elementwise_get_grad原来的路由条件是：is_combined来自 pybind：eager_method.cc里index_size = PyTuple_GET_SIZE(index_ptr)是__getitem__索引元组的元素个数（slice/None/Ellipsis都算），slice_utils.h里is_combined = (index_size == 1) ? false : true。因此真实判据是「索引元组只有一个元素」：x[idx]→ 走排序 kernel（IndexingBackwardKernelStride1/IndexingBackwardKernel），重复下标在opmath_t（fp32）寄存器里归约，只在写回时舍入一次；x[:, idx]、x[..., idx]、x[idx, None]→is_combined=true，落到GPUIndexElementwiseGetGrad，IndexEleGetGradAccKernel内是CudaAtomicAdd(reinterpret_cast<T*>(...), T value)，即 fp16/bf16 原生 atomicAdd，每加一次就舍回低精度，误差随重复次数增长，且累加顺序不确定、run-to-run 不可复现。触发条件是同一下标重复 ≥3 次：N=2 时低精度逐步累加恒等于正确舍入（
0+a精确，a+b只舍一次），所以必须重复三次以上才会暴露。[16, 8]配x[:, idx]（索引长 5）200 组随机数据统计：改动
路由门从
accumulate && index.size() == 1 && !is_combined改为if (accumulate)（index_elementwise_get_grad_kernel.cu:724）。放开之后有两类输入是排序路径原来处理不了的，分别由 1.2、1.3 补上。1.2 排序路径缺
1 < sliceSize <= 32的 kernel问题
sliceSize（被索引轴之后的元素个数）落在1 < sliceSize <= 32时，即使走到排序路径，也只有IndexingBackwardKernelStride1（仅sliceSize == 1）和按 feature 展开的通用IndexingBackwardKernel两个选择，后者对每个重复项都做一次grad_weight的 read-modify-write，同样是每步舍入。torch 在这里有专门的indexing_backward_kernel_small_stride，且它真正的分派维度就是sliceSize与 warp size 的比较，而不是索引表达式的写法。改动
新增
IndexingBackwardKernelSmallStride（:344），覆盖1 < sliceSize <= WARP_SIZE：一个线程独占一个 feature 列，该下标的所有重复项在opmath_t（fp32）寄存器内归约，写回时才舍入一次。分派处（:638）由此按sliceSize三分：== 1走IndexingBackwardKernelStride1，<= WARP_SIZE走新 kernel，其余走通用IndexingBackwardKernel。1.3 排序 kernel 无法寻址跨步视图
问题
排序 kernel 把
grad_weight当作一块平坦 buffer 寻址；原IndexPutWithSortKernel自己做 transpose / inverse-perm 并用funcs::makeLinearIndex构造线性索引，只在「被索引对象就是整块连续x_grad」时成立。1.1 放开路由后，x[1::2, idx]这类先切片再索引也会进来，此时被索引对象是x_grad内的一块跨步子区域，直接跑排序 kernel 会写到错误位置。改动
IndexPutWithSortKernel改为接收一个描述被索引视图的SortedPathLayout（:401），并直接用funcs::computeLinearIndex在该视图上构造线性索引。被索引轴的位置由 layout 给出，且必然是连续的一段，因此transposeToFrontAndInvPerm那一步在此路径上不再需要，也省掉了 permute 后的Contiguous与写回时的反向 transpose。视图信息的恢复方式（
DeriveSortedPathLayout，:437）：kernel 参数描述的是 restride 之后的视图 ——restride_src()（paddle/fluid/pybind/slice_utils.h）把被索引轴换成广播后的 index shape 并给这些轴零 stride，所以input_strides里连续的一段 0 就标出了被索引块，起点即dims_before；而index_dims（即AdvancedIndex::indexed_sizes）开头正是这些轴的 extent，于是视图在索引之前的 shape 及其在x_grad内的 stride 都可以还原。据此分三种情况：x_grad时，直接在x_grad上跑排序 kernel；x[1::2, idx]），先分配一个与该视图同形的连续 scratch，在其上归约，再用StridedCopyKernel按view_strides/view_offset搬回x_grad—— 这与 torch 用slice_backward组合index_backward的做法一致。x_grad已经清零，所以搬回时无需累加。注意phi::StridedCopyKernel会覆写out的 meta，调用方需自行 save / restore；IsNonOverlapping（:411）拒绝自重叠的as_strided式视图，这类布局回落到原来的 elementwise 路径（精度更低、run-to-run 不确定），与上游行为保持一致。slice_offset/index_strides是字节偏移（slice_utils.h里stride * element_size_bytes），派生 layout 时统一按elesize换算回元素单位，不能整除即视为非法并回落。2.
None出现在高级索引之前导致轴绑定错位问题
None位于高级索引之前时，前向 shape 就与 numpy / torch 不一致：x为[16, 8]，i长 3）x[None, :, i](1, 16, 3)(1, 3, 8)x[:, None, i](16, 1, 3)(16, 3, 8)x[None, i](1, 3, 8)(3, 16, 8)x[i, None](3, 1, 8)(3, 1, 8)（正确）ParseIndex里的estimated_dim是基础索引输出张量的轴号（slice保留 →++，整数 → 不++（会 decrease），Ellipsis→+= rank - specified，bool tensor →+= rank），advanced_index_dim[estimated_dim] = estimated_dim用它记录高级索引落在哪个轴。但Py_None分支只 pushnone_axes而不推进estimated_dim，而getTensorWithBasicIndexing结尾的unsqueeze_ad_func(out, *none_axes)确实在输出里插入了这个轴（none_axes已按decrease_axis修正，就是基础输出的坐标系）。两个坐标系因此错开none_count个位置，索引被绑到错误的轴上。同分支的单 bool（
True/ 0-D bool tensor）处理是对的：它 pushnone_axes的同时也estimated_dim++，正好印证None分支缺这一步。静态图
parse_index有同构的 bug，且更早暴露：replace_none先把None从indices里剥掉，循环根本看不到None，修改前静态图x[None, i]直接在gather_nd报 index out of range。改动
slice_utils.h的ParseIndexPy_None分支补estimated_dim++；eager_method.cc两处advanced_index_dim容量从rank * 2改为rank * 2 + index_size（setitem 侧为+ size）。否则 rank 2 张量上x[None, None, None, None, None, i]会让estimated_dim到 5，发生越界写；variable_index.py的parse_index改成_, none_axes = replace_none(indices)保留原indices，循环里遇到None时estimated_dim += 1; continue，并把advanced_index扩容len(indices)。3.
slice_offset与视图跨度的越界校验问题
slice_offset是前向记录的「被索引视图在x_gradbuffer 内的字节偏移」。排序路径和 elementwise 回落路径都会把它加到x_grad的基址上，因此一个非法值（负数，或超出 buffer）会直接变成越界写，而不是报错。同样地，跨步视图是稀疏的，它触达的最远位置由跨度决定而不是元素个数决定，只检查
numel <= grad_numel不足以保证不越界。改动
IndexElementwiseGetGradKernel入口校验0 <= slice_offset < x_grad 字节数；DeriveSortedPathLayout里校验视图触达的最后一个位置slice_offset / elesize + Σ (view_dims[k] - 1) * view_strides[k] < x_grad->numel()，报错信息带上 offset、view shape 和 strides，便于定位。修改后的行为
x[:, idx]等重复下标反向1 < sliceSize <= 32的排序路径x[1::2, idx]）x[None, :, i]等None在前__setitem__多数直接报 broadcast 失败、静态图gather_nd越界Noneadvanced_index_dim可能越界写index_size扩展slice_offset/ 越界视图自重叠的
as_strided式视图仍回落到原 elementwise 路径，行为与之前相同。devPR:#79716
是否引起精度变化
是