[Operator Mechanism]Check the guard in IndexType dispatch of big tensors (#79645)
* [Operator Mechanism] Fix int32 index overflow in allclose dispatch guard
Include `numel+stride` guard at int32_t/int64_t dispatch to avoid the grid-stride
loop's final increcent overflow in int32 path.
* [Operator Mechanism] Fix index overflow in ArgCudaKernel
- Add check pre*post*2 to avoid `height-1+gridDim.x` overflow
- Add check n+max_block_size-1 to avoid `width-1+blockDim.x` overflow
* fix the type of input parameters from int64 to IndexType in `ArgCUDAKernel`
* Fix index overflow of big tensor in min_max_with_index_kernel.cu
- add guard for outer grid-stride loop (peaking at `height-1+gridDim.x`)
- add guard for inner reduction loop (peaking at `width-1+blockDim.x`)
* [Operator Mechanism] Fix int32 index guard in pad3d forward/backward
In pad3d_kernel.cu and pad3d_grad_kernel.cu, the int32/int64 dispatch
was only decided by the output numel only, which misses two other
sources of magnitude inside the kernels:
- Every in_data / d_in offset is bounded by the input numel, not the output
one. Padding values may be negative (removing elements instead of adding
them), so the input can be far larger than the output.
- when model==reflect it computes `2 * in_dim - x - 2`, if a single spatial
dim(height / width / depth) > INT32_MAX / 2 it may overflow.
What we do(in both fwd and bwd):
- Add the guard of input->numel in condition
- Add the guard of three type of `in_dim` with `INT32_MAX / 2` when mode==reflect
- Drop the per-dim loop over the output dims, each dim is already bounded by the
the numel, so that loop can never reject a shape the numel check accepts.
- In the backward kernel the dim variables are hoisted above the guard so it can
inspect them; they are kept as int64_t and narrowed to IndexType at launch,
* [Operator Mechanism] Fix int32 index guard in temporal_shift forward/backward
- Add guard of `numel + blockDim.x * gridDim.x` in temporal_shift_kernel.cu
and temporal_shift_grad_kernel.cu
The grid-stride loop index `tid` peaks at `numel - 1 + blockDim.x * gridDim.x`
rather than `numel - 1`. The previous guard only checked `numel < INT32_MAX`,
leaving a window where `tid` overflows int32, wraps negative and indexes out
of bounds.
* [Operator Mechanism] Fix int32 index guard in roi_pool forward/backward
1.roi_pool_kernel.cu:
- Promote the scaled box coordinates (box_start/end_w/h) and the derived
box_width/box_height from int to IndexType: they are added to hstart/hend/wstart/wend
before the min/max clamp, so a narrower type wraps before being clamped.
- Promote maxidx and the h/w loop variables from int to IndexType:
maxidx caches h * width + w (bounded by height * width - 1).
- Add boxes.numel() to the int32 dispatch guard: n * kROISize overflows once
rois_num > INT32_MAX / kROISize, which output_size and x.numel() do not bound.
- Widen the host-side box_batch_id_data indices (start and the inner loop
variable) from int to int64_t and make the byte count size_t, since both are bounded by
rois_num rather than by any int32-guarded quantity.
2.roi_pool_grad_kernel.cu:
- Add last value increment of grid_stride loop in the int32 dispatch guard:
the grid-stride loop variable peak is nthreads - 1 + blockDim.x * gridDim.x.
- Widen the host-side box_batch_id_data indices from int to int64_t and make
the byte count size_t, mirroring the forward kernel.
* Fix int32 overflow of n * kROISize in roi_align GPU kernels
The IndexType guards only covered the x / out address spaces, leaving the
boxes address space unchecked.
In GPURoiAlignForward / GPURoiAlignBackward, the peak value of
`input_rois + n * kROISize` equals boxes.numel() - kROISize.
Since output_size = rois_num * channels * pooled_height * pooled_width can
degenerate to rois_num (channels = pooled_h = pooled_w = 1), the existing
output_size or x.numel() guards cannot bound 4 * rois_num, and the expression
wraps once rois_num > 536,870,912.
Add boxes.numel() to both guards. Since boxes is enforced to be
[rois_num, 4], boxes.numel() is exactly the upper bound of n * kROISize,
so the threshold is tight and no valid shape is needlessly downgraded.
* [Operator Mechanism] Fix int32 index overflow in kthvalue dispatch guard
Update the guard of Indextype dispatch in `kthvalue_kernel.cu`. Because
the strided loop in GatherKthValue, FindPattern, RadixCountUsingMask peaks
at `num_cols-1+blockDim.x`. And in int32 path `blockDim.x<=MAX_NUM_THREADS`
and `num_cols==input_width`.
The int32 loop counter overflows and wraps negative, which either reads
far out of bounds or deadlocks the block on `__syncthreads()`.
Repro before this change (hangs or reports CUDA 700):
x = paddle.rand([2147483647], dtype='float16')
paddle.kthvalue(x, k=1)
* [Operator Mechanism] Fix index overflow in nearest interpolate NCHW backward
1. Add the guard in IndexType dispatch on the `nearest` + NCHW path:
- The guard now covers the output size. The old IndexType guard in
`Interpolate2DCUDABwd` only measured the input size. The kernel indexes the
output grad with the same type, and upsampling makes that tensor the larger
one, so the int32 branch could be selected while `out_index` overflowed.
- The guard plus `nc - 1 + nc_stride` for the loop variable, whose final
increment feeds the loop condition and therefore must not rely on
undefined signed overflow.
2. Change type of num_img in `GetGpuLaunchConfig3D` from int to int64_t.
`num_img` (= n * c) is a product of two dimensions with no upstream
clamp, so it can exceed INT32_MAX and wrap negative.
* [Operator Mechanism] Fix int32 index overflow in Rope kernel when d>d2 and numel>INT_MAX
In both `FusedRopeKernelImpl` and `FusedRopeGradKernelImpl`:
- Change `offset_src` and `offser_dst` from `int` to `IndexT` type
- Change `d_id` and `h_id` from `int` to `IndexT` type
In`FusedRopeKernelLauncher`:
- Add conservative guard: the index in kernel's loop overshoot
the bound(d/h, both <= numel) up to one stride(block.x/y). Add
one stride of headroom on top of numel for int32 path.
* update comments
* Restore upstream roi_pool batch ID count O
omoYang committed
63b1f45a3c5eecaae7fc3c9132a2f02a6dde9686
Parent: 811fbbe
Committed by GitHub <noreply@github.com>
on 8/20/2026, 7:38:31 AM