Commit cbfab56
committed
Fix int32 overflow in max_pool3d forward kernel
Template MaxPool3dKernelFunctor and max_pool3d_with_indices_out_template
with an index_t type parameter, and dispatch between int32 and int64 at
runtime using canUse32BitIndexMath. This prevents overflow when computing
strides for large tensors while preserving int32 performance for smaller
ones.
The forward kernel computes stride values (e.g. in_batch_stride =
features * T * H * W) that can exceed INT32_MAX for shapes like
[2, 2200, 100, 100, 100] (in_batch_stride = 2.2B).
Changes:
- Fix typo: MaxPool3dKerenlFunctor -> MaxPool3dKernelFunctor
- Add index_t template parameter to functor and template function
- Dispatch index_t via AT_DISPATCH_INDEX_TYPES + canUse32BitIndexMath
- Remove now-redundant (int64_t) casts in operator()
- Add test/regressions/test_max_pool3d_fwd_int64.py with shapes
that overflow int32 batch strides in the forward pass1 parent 816a4cc commit cbfab56
2 files changed
Lines changed: 223 additions & 142 deletions
0 commit comments