Fix int32 overflow in max_pool3d backward kernel#3558
Conversation
pbielak
added
disable_e2e
There was a problem hiding this comment.
Pull request overview
This PR fixes a large-tensor indexing bug in the XPU max_pool3d backward SYCL kernel by widening stride/size arithmetic to int64_t, which aligns with the repository’s large-tensor safety requirements for backend kernels.
Changes:
- Widened backward-kernel stride members and constructor parameters from
inttoint64_t. - Made
gradOutputSizeand derived stride products explicitly compute in 64-bit arithmetic. - Preserved the existing kernel logic while removing the specific overflow path described in the PR for large 3D pooling workloads.
Read for this review: .github/skills/xpu-ops-pr-review/SKILL.md, .github/skills/xpu-ops-pr-review/references/torch-xpu-ops-review-notes.md, .github/skills/xpu-ops-pr-review/references/review-checklist.md, .github/skills/xpu-ops-pr-review/references/bc-guidelines.md.
76b5619 to
51b3e15
Compare
51b3e15 to
8f93b1c
Compare
8f93b1c to
f802d5d
Compare
BBBela
left a comment
BBBela
left a comment
There was a problem hiding this comment.
Looks good to me, thanks for addressing review comments!
f802d5d to
1e9981a
Compare
|
@pbielak Could you help check if the same issue persisted in the forward kernel? If yes, better to file another PR to fix it. |
1e9981a to
a1390df
Compare
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: - Fix typo: MaxPool3dKerenlFunctor -> MaxPool3dKernelFunctor - Widen dimension and stride params/members to int64_t (kernel, stride, padding, and dilation params remain int) - Remove now-redundant (int64_t) casts in operator() - Add test_pool3d_fwd_large_size_int64 with a shape that overflows int32 batch strides in the forward pass
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: - Fix typo: MaxPool3dKerenlFunctor -> MaxPool3dKernelFunctor - Widen dimension and stride params/members to int64_t (kernel, stride, padding, and dilation params remain int) - Remove now-redundant (int64_t) casts in operator() - Add test_pool3d_fwd_large_size_int64 with a shape that overflows int32 batch strides in the forward pass
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: - Fix typo: MaxPool3dKerenlFunctor -> MaxPool3dKernelFunctor - Widen dimension and stride params/members to int64_t (kernel, stride, padding, and dilation params remain int) - Remove now-redundant (int64_t) casts in operator() - Add test_pool3d_fwd_large_size_int64 with a shape that overflows int32 batch strides in the forward pass
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: - Fix typo: MaxPool3dKerenlFunctor -> MaxPool3dKernelFunctor - Widen dimension and stride params/members to int64_t (kernel, stride, padding, and dilation params remain int) - Remove now-redundant (int64_t) casts in operator() - Add test_pool3d_fwd_large_size_int64 with a shape that overflows int32 batch strides in the forward pass
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: - Fix typo: MaxPool3dKerenlFunctor -> MaxPool3dKernelFunctor - Widen dimension and stride params/members to int64_t (kernel, stride, padding, and dilation params remain int) - 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
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: - Fix typo: MaxPool3dKerenlFunctor -> MaxPool3dKernelFunctor - Widen dimension and stride params/members to int64_t (kernel, stride, padding, and dilation params remain int) - 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
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: - Fix typo: MaxPool3dKerenlFunctor -> MaxPool3dKernelFunctor - Widen dimension and stride params/members to int64_t (kernel, stride, padding, and dilation params remain int) - 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
Good catch! |
5dfcfca to
2f9c249
Compare
I think currently, the copilot does not have the skill for performance related issues. We need to manually give it skills. So when reviewing, it will aware such issues. As you could see there is a check list in |
@Stonepia I mean just add one line to ask copilot if this pr hurt perf， no extra skill needed |
Template `MaxPool3dBackwardKernelFunctor` and `max_pool3d_with_indices_backward_template` with an `index_t` type parameter, and dispatch between `int32_t` and `int64_t` at runtime using `canUse32BitIndexMath`. This avoids overflow when computing offsets for large tensors (e.g. shape `[70, 32, 100, 100, 100]`) while preserving int32 performance for smaller tensors. Fixes test case: `test_pool3d_large_size_int64`
2f9c249 to
8e6d572
Compare
Oh that's possible. Will add that. |
|
@jianyizh Added the template parameter of the index dtype - waiting for CI signal |
|
/merge -f "CI failure is unrelated (test_dynamic_lstm_retraceability_strict)" |
|
✅ 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 intel#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: - Fix typo: MaxPool3dKerenlFunctor -> MaxPool3dKernelFunctor - Widen dimension and stride params/members to int64_t (kernel, stride, padding, and dilation params remain int) - Remove now-redundant (int64_t) casts in operator() - Add test_pool3d_fwd_large_size_int64 with a shape that overflows int32 batch strides in the forward pass
Widen parameters and variables in
max_pool3d_with_indices_backward_templateandMaxPool3dBackwardKernelFunctorfrominttoint64_tto prevent overflow when computing offsets for large tensors (e.g. shape[70, 32, 100, 100, 100]).Fixes test case:
test_pool3d_large_size_int64