Fix int32 overflow in max_pool3d forward kernel#3632
Conversation
pbielak
added
disable_e2e
There was a problem hiding this comment.
Pull request overview
This PR addresses potential int32 overflow in the XPU SYCL max_pool3d forward implementation by widening shape/stride-related values to int64_t, and adds a large-tensor regression test intended to exercise the overflow scenario.
Skill files read: .github/skills/xpu-ops-pr-review/SKILL.md
Changes:
- Fixes a typo in the forward kernel functor name and widens forward kernel shape/stride parameters and members from
inttoint64_tto prevent overflow in stride math. - Updates forward indexing calculations to use widened stride/index types (removing now-redundant explicit casts).
- Adds a new large-tensor test intended to cover forward-pass stride overflow.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
src/ATen/native/xpu/sycl/DilatedMaxPool3d.cpp |
Widen forward max_pool3d functor/template parameters and stride computations to int64_t to avoid overflow. |
test/xpu/nn/test_pooling_xpu.py |
Adds a new large-tensor forward max_pool3d test targeting int32 stride overflow scenarios. |
cd613ce to
4086fba
Compare
4086fba to
35542e0
Compare
35542e0 to
9bb3a5c
Compare
9bb3a5c to
e4dd222
Compare
e4dd222 to
c5208d7
Compare
c5208d7 to
cd20c86
Compare
jianyizh
left a comment
jianyizh
left a comment
There was a problem hiding this comment.
Check bool canUse32BitIndexMath(const TensorBase& t, int64_t max_elem), you should set index dtype as a template parameter
e7874af to
86298d4
Compare
|
I discussed the results with @jianyizh and we compared our benchmark scripts. I re-ran the benchmark with parameters provided by @jianyizh i.e., kernel_size=3, stride=2, padding=1. With a larger pooling window (333 = 27 elements vs 222 = 8), the inner loop performs significantly more index arithmetic per output element, making the int32 savings meaningful. Ratios are int64-only/int32-dispatch (>1x = dispatch is faster):
The earlier regressions (0.6–0.9x) were specific to kernel_size=2, where the kernel is memory-bound and the dispatch overhead dominates. With kernel_size=3, the kernel is more compute-bound and int32 indexing provides a 14–32% speedup. Conclusion: the int32 dispatch is worth keeping for the forward kernel. |
cbfab56 to
6bef852
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (2)
src/ATen/native/xpu/sycl/DilatedMaxPool3d.cpp:567
- The index width is selected only from the input tensor, but this kernel indexes and launches over the output/indices tensor as well. Max-pool output can be larger than the input when padding is used, so an input that satisfies 32-bit index math can still produce an output whose numel/strides exceed INT_MAX; in that case
OutputSizeand output strides are truncated toint. Please include the output/indices (or at leastwork_output) in this decision, consistent with other kernels that dispatch based on the tensor being indexed.
AT_DISPATCH_INDEX_TYPES(
at::native::canUse32BitIndexMath(input, INT_MAX) ? ScalarType::Int
: ScalarType::Long,
src/ATen/native/xpu/sycl/DilatedMaxPool3d.cpp:259
OutputSizeshould not be narrowed toindex_tbefore computing the launch size. When 32-bit indexing is selected, a valid output size close to INT_MAX can overflow in the subsequent(OutputSize + work_group_size - 1)rounding expression, whereas the previousint64_tOutputSizeavoided that. Keep the launch-size arithmetic inint64_t/size_teven when per-element offsets use 32-bit indices.
index_t OutputSize = obatch * features * otime * oheight * owidth;
|
@copilot New UT failures detected in op_ut tests (2 total). Please check these new failures and analyze whether they are caused by the PR changes. Important: Do NOT update the plan in the PR description directly. Use reply comments to update the status. Job log: https://github.com/intel/torch-xpu-ops/actions/runs/26029952738 New failures:
Instructions:
|
Implemented in commit 6bef852: max_pool3d forward now templates index type (
I investigated run 26029952738 and downloaded the Conclusion: the 2 new failures are not caused by this PR’s max_pool3d changes.
Neither traceback/execution path touches Reproducer commands:
No code change was made on this PR branch for these failures. |
6bef852 to
d1022f7
Compare
|
@copilot New UT failures detected in op_ut tests (8 total). Please check these new failures and analyze whether they are caused by the PR changes. Important: Do NOT update the plan in the PR description directly. Use reply comments to update the status. Job log: https://github.com/intel/torch-xpu-ops/actions/runs/26099835430 New failures:
... and 3 more failures. See full list in artifact. Instructions:
|
I analyzed run 26099835430 and downloaded both Conclusion: these 8 new failures are not caused by this PR.
Pytest repro commands:
UI screenshot: N/A (no UI changes). |
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
d1022f7 to
433cae9
Compare
|
@copilot New UT failures detected in op_ut tests (7 total). Please check these new failures and analyze whether they are caused by the PR changes. Important: Do NOT update the plan in the PR description directly. Use reply comments to update the status. Job log: https://github.com/intel/torch-xpu-ops/actions/runs/26149912888 New failures:
... and 2 more failures. See full list in artifact. Instructions:
|
Investigated run 26149912888 and downloaded both Conclusion: the 7 new failures are not caused by this PR.
Pytest repro commands:
UI screenshot: N/A (no UI changes). |
|
/merge -f "All failing test cases are unrelated to the PR changes" |
|
✅ PR has been successfully merged by @pbielak (force merge). |
Widen parameters and variables in MaxPool3dKernelFunctor and max_pool3d_with_indices_out_template from int to int64_t to prevent overflow when computing strides for large tensors. This mirrors the backward kernel fix from #3558
The forward kernel was partially protected by int64_t local variables and explicit casts at usage sites, so the existing test shape [70, 32, 100, 100, 100] did not trigger failures. However, the stride values themselves (e.g. in_batch_stride = features * T * H * W) were computed as int, which silently overflows for shapes where a single stride exceeds INT32_MAX (e.g. [2, 2200, 100, 100, 100] gives in_batch_stride = 2.2B).
Changes: