Skip to content

Commit cbfab56

Browse files
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 pass
1 parent 816a4cc commit cbfab56

2 files changed

Lines changed: 223 additions & 142 deletions

File tree

0 commit comments

Comments
 (0)