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_pool3d_fwd_large_size_int64 with a shape that overflows
  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,9 @@ 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 +78,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 +89,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 +107,29 @@ struct MaxPool3dKerenlFunctor {
 
       int64_t out_index;
       if constexpr (!channels_last_) {
-        out_index = (int64_t)slice * out_cf_c_stride_ +
+        out_index = slice * out_cf_c_stride_ +
             oTime * out_cf_d_stride_ + oRow * owidth_ + oColumn;
       } else {
-        out_index = (int64_t)batch * out_batch_stride_ +
+        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 +143,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 +194,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 +215,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,9 +255,9 @@ 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,
+  int64_t 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,
+  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;
@@ -268,10 +268,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/xpu/nn/test_pooling_xpu.py

@@ -549,6 +549,44 @@ def _test_pool3d_large_size_int64(self, device):
 TestPoolingNNDeviceType.test_pool3d_large_size_int64 = _test_pool3d_large_size_int64
 
 
+@largeTensorTest("20GB", device=device_type)
+@largeTensorTest("40GB", "cpu")
+def _test_pool3d_fwd_large_size_int64(self, device):
+    """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).
+    The existing test_pool3d_large_size_int64 does not trigger these
+    because its per-stride values fit in int32.
+    """
+    cases = [
+        # (shape, memory_format)
+        ((2, 2200, 100, 100, 100), torch.channels_last_3d),
+        ((1, 2, 1300, 1300, 1300), torch.contiguous_format),
+    ]
+    for shape, memory_format in cases:
+        with self.subTest(memory_format=memory_format):
+            x = torch.randn(*shape, dtype=torch.half, device=device).to(
+                memory_format=memory_format
+            )
+            y = torch.nn.functional.max_pool3d(x, 5)
+            torch.get_device_module(device).synchronize()
+
+            ref_x = x.detach().cpu().float().contiguous()
+            ref_y = torch.nn.functional.max_pool3d(ref_x, 5)
+
+            self.assertEqual(y, ref_y, exact_dtype=False)
+            del x, y, ref_x, ref_y
+
+
+TestPoolingNNDeviceType.test_pool3d_fwd_large_size_int64 = (
+    _test_pool3d_fwd_large_size_int64
+)
+
+
 # Upstream test_pooling.py:test_pooling_large uses hard-coded device="cuda"
 # and torch.cuda.synchronize(); resolve both via the test-supplied `device`
 # (set by instantiate_device_type_tests) so the body is not bound to a