Fix int32 overflow in max_pool3d forward kernel

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

Author: pbielak

src/ATen/native/xpu/sycl/DilatedMaxPool3d.cpp

@@ -33,7 +33,7 @@ DISABLE_RETURN_TYPE_WARNING_BEGIN
 namespace at::native::xpu {
 
 template <typename scalar_t, bool channels_last_>
-struct MaxPool3dKerenlFunctor {
+struct MaxPool3dKernelFunctor {
   void operator()(sycl::nd_item<1> item) const {
     auto outputIndex = item.get_global_id(0);
     if (outputIndex < OutputSize_) {
@@ -62,9 +62,12 @@ struct MaxPool3dKerenlFunctor {
       int tStart = oTime * dT_ - pT_;
       int hStart = oRow * dH_ - pH_;
       int wStart = oColumn * dW_ - pW_;
-      int tEnd = std::min(tStart + (kT_ - 1) * dilationT_ + 1, itime_);
-      int hEnd = std::min(hStart + (kH_ - 1) * dilationH_ + 1, iheight_);
-      int wEnd = std::min(wStart + (kW_ - 1) * dilationW_ + 1, iwidth_);
+      int tEnd = std::min(
+          tStart + (kT_ - 1) * dilationT_ + 1, static_cast<int>(itime_));
+      int hEnd = std::min(
+          hStart + (kH_ - 1) * dilationH_ + 1, static_cast<int>(iheight_));
+      int wEnd = std::min(
+          wStart + (kW_ - 1) * dilationW_ + 1, static_cast<int>(iwidth_));
 
       while (tStart < 0)
         tStart += dilationT_;
@@ -78,9 +81,9 @@ struct MaxPool3dKerenlFunctor {
       int64_t ioffset;
 
       if constexpr (!channels_last_) {
-        ioffset = (int64_t)slice * in_cf_c_stride_;
+        ioffset = slice * in_cf_c_stride_;
       } else {
-        ioffset = ((int64_t)batch * in_batch_stride_) + channel;
+        ioffset = batch * in_batch_stride_ + channel;
       }
 
       scalar_t max = at::numeric_limits<scalar_t>::lower_bound();
@@ -89,7 +92,7 @@ struct MaxPool3dKerenlFunctor {
         for (int h = hStart; h < hEnd; h += dilationH_) {
           for (int w = wStart; w < wEnd; w += dilationW_) {
             scalar_t val;
-            int index = t * in_hw_stride_ + h * iwidth_ + w;
+            int64_t index = t * in_hw_stride_ + h * iwidth_ + w;
             if constexpr (!channels_last_) {
               val = inputData_[ioffset + index];
             } else {
@@ -107,29 +110,28 @@ struct MaxPool3dKerenlFunctor {
 
       int64_t out_index;
       if constexpr (!channels_last_) {
-        out_index = (int64_t)slice * out_cf_c_stride_ +
-            oTime * out_cf_d_stride_ + oRow * owidth_ + oColumn;
+        out_index = slice * out_cf_c_stride_ + oTime * out_cf_d_stride_ +
+            oRow * owidth_ + oColumn;
       } else {
-        out_index = (int64_t)batch * out_batch_stride_ +
-            oTime * out_cl_d_stride_ + oRow * out_cl_h_stride_ +
-            oColumn * features_ + channel;
+        out_index = batch * out_batch_stride_ + oTime * out_cl_d_stride_ +
+            oRow * out_cl_h_stride_ + oColumn * features_ + channel;
       }
       outputData_[out_index] = max;
       indicesData_[out_index] = maxIndex;
     }
   }
-  MaxPool3dKerenlFunctor(
+  MaxPool3dKernelFunctor(
       const scalar_t* inputData,
       scalar_t* outputData,
       int64_t* indicesData,
-      int features,
-      int itime,
-      int iheight,
-      int iwidth,
-      int obatch,
-      int otime,
-      int oheight,
-      int owidth,
+      int64_t features,
+      int64_t itime,
+      int64_t iheight,
+      int64_t iwidth,
+      int64_t obatch,
+      int64_t otime,
+      int64_t oheight,
+      int64_t owidth,
       int kT,
       int kH,
       int kW,
@@ -143,17 +145,17 @@ struct MaxPool3dKerenlFunctor {
       int dilationH,
       int dilationW,
       int64_t OutputSize,
-      int out_cf_d_stride,
-      int out_cf_c_stride,
-      int in_cf_d_stride,
-      int in_cf_c_stride,
-      int out_cl_h_stride,
-      int out_cl_d_stride,
-      int in_cl_h_stride,
-      int in_cl_d_stride,
-      int in_batch_stride,
-      int out_batch_stride,
-      int in_hw_stride)
+      int64_t out_cf_d_stride,
+      int64_t out_cf_c_stride,
+      int64_t in_cf_d_stride,
+      int64_t in_cf_c_stride,
+      int64_t out_cl_h_stride,
+      int64_t out_cl_d_stride,
+      int64_t in_cl_h_stride,
+      int64_t in_cl_d_stride,
+      int64_t in_batch_stride,
+      int64_t out_batch_stride,
+      int64_t in_hw_stride)
       : inputData_(inputData),
         outputData_(outputData),
         indicesData_(indicesData),
@@ -194,14 +196,14 @@ struct MaxPool3dKerenlFunctor {
   const scalar_t* inputData_;
   scalar_t* outputData_;
   int64_t* indicesData_;
-  int features_;
-  int itime_;
-  int iheight_;
-  int iwidth_;
-  int obatch_;
-  int otime_;
-  int oheight_;
-  int owidth_;
+  int64_t features_;
+  int64_t itime_;
+  int64_t iheight_;
+  int64_t iwidth_;
+  int64_t obatch_;
+  int64_t otime_;
+  int64_t oheight_;
+  int64_t owidth_;
   int kT_;
   int kH_;
   int kW_;
@@ -215,32 +217,32 @@ struct MaxPool3dKerenlFunctor {
   int dilationH_;
   int dilationW_;
   int64_t OutputSize_;
-  int out_cf_d_stride_;
-  int out_cf_c_stride_;
-  int in_cf_d_stride_;
-  int in_cf_c_stride_;
-  int out_cl_h_stride_;
-  int out_cl_d_stride_;
-  int in_cl_h_stride_;
-  int in_cl_d_stride_;
-  int in_batch_stride_;
-  int out_batch_stride_;
-  int in_hw_stride_;
+  int64_t out_cf_d_stride_;
+  int64_t out_cf_c_stride_;
+  int64_t in_cf_d_stride_;
+  int64_t in_cf_c_stride_;
+  int64_t out_cl_h_stride_;
+  int64_t out_cl_d_stride_;
+  int64_t in_cl_h_stride_;
+  int64_t in_cl_d_stride_;
+  int64_t in_batch_stride_;
+  int64_t out_batch_stride_;
+  int64_t in_hw_stride_;
 };
 
 template <typename scalar_t, bool channels_last>
 void max_pool3d_with_indices_out_template(
     const scalar_t* inputData,
     scalar_t* outputData,
     int64_t* indicesData,
-    int features,
-    int itime,
-    int iheight,
-    int iwidth,
-    int obatch,
-    int otime,
-    int oheight,
-    int owidth,
+    int64_t features,
+    int64_t itime,
+    int64_t iheight,
+    int64_t iwidth,
+    int64_t obatch,
+    int64_t otime,
+    int64_t oheight,
+    int64_t owidth,
     int kT,
     int kH,
     int kW,
@@ -255,10 +257,10 @@ void max_pool3d_with_indices_out_template(
     int dilationW) {
   int64_t OutputSize = obatch * features * otime * oheight * owidth;
 
-  int out_cf_d_stride = 0, out_cf_c_stride = 0, in_cf_d_stride = 0,
-      in_cf_c_stride = 0;
-  int out_cl_h_stride = 0, out_cl_d_stride = 0, in_cl_h_stride = 0,
-      in_cl_d_stride = 0;
+  int64_t out_cf_d_stride = 0, out_cf_c_stride = 0, in_cf_d_stride = 0,
+          in_cf_c_stride = 0;
+  int64_t out_cl_h_stride = 0, out_cl_d_stride = 0, in_cl_h_stride = 0,
+          in_cl_d_stride = 0;
   if constexpr (!channels_last) {
     out_cf_d_stride = owidth * oheight;
     out_cf_c_stride = otime * out_cf_d_stride;
@@ -268,10 +270,10 @@ void max_pool3d_with_indices_out_template(
     out_cl_h_stride = owidth * features;
     out_cl_d_stride = oheight * out_cl_h_stride;
   }
-  auto in_batch_stride = itime * iheight * iwidth * features;
-  auto out_batch_stride = otime * oheight * owidth * features;
-  auto in_hw_stride = iwidth * iheight;
-  MaxPool3dKerenlFunctor<scalar_t, channels_last> kfn(
+  int64_t in_batch_stride = itime * iheight * iwidth * features;
+  int64_t out_batch_stride = otime * oheight * owidth * features;
+  int64_t in_hw_stride = iwidth * iheight;
+  MaxPool3dKernelFunctor<scalar_t, channels_last> kfn(
       inputData,
       outputData,
       indicesData,

test/regressions/test_max_pool3d_fwd_int64.py

@@ -0,0 +1,72 @@
+# Copyright 2020-2026 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Owner(s): ["module: intel"]
+"""
+Regression test for int32 overflow in max_pool3d forward kernel.
+
+The forward kernel computes strides (e.g. in_batch_stride = C * D * H * W)
+that can exceed INT32_MAX for large tensors. This test exercises shapes
+where batch strides overflow int32 in both channels_last_3d and contiguous
+memory formats.
+"""
+
+import torch
+from torch.testing._internal.common_device_type import largeTensorTest
+from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    parametrize,
+    run_tests,
+    subtest,
+    TestCase,
+)
+
+
+class TestMaxPool3dFwdInt64(TestCase):
+    """Test that the forward pass of max_pool3d handles int64 strides correctly.
+
+    Exercises two overflow scenarios:
+    - channels_last_3d: in_batch_stride = C * D * H * W > INT32_MAX
+      (shape [2, 2200, 100, 100, 100] -> stride 2.2B).
+    - contiguous: in_cf_c_stride = D * H * W > INT32_MAX
+      (shape [1, 2, 1300, 1300, 1300] -> stride 2.197B).
+    """
+
+    @largeTensorTest("20GB", device="cpu")
+    @largeTensorTest("12GB", device="xpu")
+    @parametrize(
+        "shape,memory_format",
+        [
+            subtest(
+                ((2, 2200, 100, 100, 100), torch.channels_last_3d),
+                name="channels_last_3d",
+            ),
+            subtest(
+                ((1, 2, 1300, 1300, 1300), torch.contiguous_format), name="contiguous"
+            ),
+        ],
+    )
+    def test_pool3d_fwd_large_size_int64(self, shape, memory_format):
+        x = torch.empty(
+            *shape, dtype=torch.half, device="xpu", memory_format=memory_format
+        ).normal_()
+        y = torch.nn.functional.max_pool3d(x, 5)
+        torch.xpu.synchronize()
+
+        ref_x = x.detach().cpu().float().contiguous()
+        ref_y = torch.nn.functional.max_pool3d(ref_x, 5)
+
+        self.assertEqual(y.cpu(), ref_y, exact_dtype=False)
+        del x, y, ref_x, ref_y
+
+
+instantiate_parametrized_tests(TestMaxPool3dFwdInt64)
+
+
+if __name__ == "__main__":
+    run_tests()