[pt2] metas and SymInt support for max_pool ops (pytorch#103951)

Pull Request resolved: pytorch#103951
Approved by: https://github.com/Chillee, https://github.com/kulinseth

Author: nkaretnikov

Diff for: aten/src/ATen/native/DilatedMaxPool2d.cpp

@@ -39,7 +39,7 @@ bool ceil_mode) {
                  stride.size() == 1 ? dH : safe_downcast<int, int64_t>(stride[1]);
 
   TORCH_CHECK(padding.size() == 1 || padding.size() == 2,
-    "max_pool2d: padding must be either be a single int, or a tuple of two ints");
+    "max_pool2d: padding must either be a single int, or a tuple of two ints");
   const int padH = safe_downcast<int, int64_t>(padding[0]);
   const int padW = padding.size() == 1 ? padH : safe_downcast<int, int64_t>(padding[1]);
 
@@ -112,7 +112,7 @@ const Tensor& indices) {
                  stride.size() == 1 ? dH : safe_downcast<int, int64_t>(stride[1]);
 
   TORCH_CHECK(padding.size() == 1 || padding.size() == 2,
-    "max_pool2d: padding must be either be a single int, or a tuple of two ints");
+    "max_pool2d: padding must either be a single int, or a tuple of two ints");
   const int padH = safe_downcast<int, int64_t>(padding[0]);
   const int padW = padding.size() == 1 ? padH : safe_downcast<int, int64_t>(padding[1]);
 

Diff for: aten/src/ATen/native/DilatedMaxPool3d.cpp

@@ -173,7 +173,7 @@ void max_pool3d_with_indices_out_cpu_template(
                  stride.size() == 1 ? dT : safe_downcast<int, int64_t>(stride[2]);
 
   TORCH_CHECK(padding.size() == 1 || padding.size() == 3,
-    "max_pool3d: padding must be either be a single int, or a tuple of three ints");
+    "max_pool3d: padding must either be a single int, or a tuple of three ints");
   const int pT = safe_downcast<int, int64_t>(padding[0]);
   const int pH = padding.size() == 1 ? pT : safe_downcast<int, int64_t>(padding[1]);
   const int pW = padding.size() == 1 ? pT : safe_downcast<int, int64_t>(padding[2]);
@@ -381,7 +381,7 @@ Tensor& max_pool3d_with_indices_backward_out_cpu_template(
                  stride.size() == 1 ? dT : safe_downcast<int, int64_t>(stride[2]);
 
   TORCH_CHECK(padding.size() == 1 || padding.size() == 3,
-    "max_pool3d: padding must be either be a single int, or a tuple of three ints");
+    "max_pool3d: padding must either be a single int, or a tuple of three ints");
   const int pT = safe_downcast<int, int64_t>(padding[0]);
   const int pH = padding.size() == 1 ? pT : safe_downcast<int, int64_t>(padding[1]);
   const int pW = padding.size() == 1 ? pT : safe_downcast<int, int64_t>(padding[2]);

Diff for: aten/src/ATen/native/MaxPooling.cpp

@@ -34,7 +34,7 @@ static void check_max_pool1d(
 
   TORCH_CHECK(
       self.dim() == 2 || self.dim() == 3,
-      "max_pool1d() Expected 2D or 3D input tensor, but got ", self.sizes());
+      "max_pool1d() Expected 2D or 3D input tensor, but got ", self.sym_sizes());
   TORCH_CHECK(
       kernel_size.size() == 1,
       "max_pool1d() kernel_size must be an int, list of ints or tuple of ints of size 1 but got size ",
@@ -74,7 +74,7 @@ static void check_max_pool1d(
   TORCH_CHECK(
       dilation[0] > 0, "max_pool1d() dilation must be greater than zero, but got ", dilation[0]);
 
-  const int64_t OW = pooling_output_shape(self.size(-1), kernel_size[0], padding[0], stride[0], dilation[0], ceil_mode);
+  const int64_t OW = pooling_output_shape(self.sym_size(-1).guard_int(__FILE__, __LINE__), kernel_size[0], padding[0], stride[0], dilation[0], ceil_mode);
   TORCH_CHECK(OW > 0, "max_pool1d() Invalid computed output size: ", OW);
 }
 
@@ -132,10 +132,10 @@ Tensor max_pool1d(
 
   auto ndim = self.ndimension();
    TORCH_CHECK(
-       (ndim == 2 && self.size(0) != 0 && self.size(1) != 0) ||
-           (ndim == 3 && self.size(1) != 0 && self.size(2) != 0),
+       (ndim == 2 && self.sym_size(0) != 0 && self.sym_size(1) != 0) ||
+           (ndim == 3 && self.sym_size(1) != 0 && self.sym_size(2) != 0),
        "max_pool1d: Expected 2D or 3D (batch mode) tensor with optional 0 dim batch size for input, but got:",
-       self.sizes());
+       self.sym_sizes());
 
   if (self.is_quantized()) {
     return at::quantized_max_pool1d(

Diff for: aten/src/ATen/native/cuda/DilatedMaxPool3d.cu

@@ -313,7 +313,7 @@ void max_pool3d_with_indices_out_cuda_template(
                  stride.size() == 1 ? dT : safe_downcast<int, int64_t>(stride[2]);
 
   TORCH_CHECK(padding.size() == 1 || padding.size() == 3,
-    "max_pool3d: padding must be either be a single int, or a tuple of three ints");
+    "max_pool3d: padding must either be a single int, or a tuple of three ints");
   const int pT = safe_downcast<int, int64_t>(padding[0]);
   const int pH = padding.size() == 1 ? pT : safe_downcast<int, int64_t>(padding[1]);
   const int pW = padding.size() == 1 ? pT : safe_downcast<int, int64_t>(padding[2]);
@@ -443,7 +443,7 @@ void max_pool3d_with_indices_backward_out_cuda_template(
                  stride.size() == 1 ? dT : safe_downcast<int, int64_t>(stride[2]);
 
   TORCH_CHECK(padding.size() == 1 || padding.size() == 3,
-    "max_pool3d: padding must be either be a single int, or a tuple of three ints");
+    "max_pool3d: padding must either be a single int, or a tuple of three ints");
   const int pT = safe_downcast<int, int64_t>(padding[0]);
   const int pH = padding.size() == 1 ? pT : safe_downcast<int, int64_t>(padding[1]);
   const int pW = padding.size() == 1 ? pT : safe_downcast<int, int64_t>(padding[2]);

Diff for: aten/src/ATen/native/mps/operations/Pooling.mm

@@ -70,7 +70,7 @@ static void pool2d_template(const Tensor& input,
               ": stride must either be omitted, a single int, or a tuple of two ints")
   TORCH_CHECK(padding.size() == 1 || padding.size() == 2,
               op_name,
-              ": padding must be either be a single int, or a tuple of two ints");
+              ": padding must either be a single int, or a tuple of two ints");
   TORCH_CHECK(dilation.size() == 1 || dilation.size() == 2,
               op_name,
               ": dilation must be either a single int, or a tuple of two ints");

Diff for: test/functorch/test_aotdispatch.py

@@ -2834,8 +2834,6 @@ def forward(self, x):
     xfail('nn.functional.interpolate', 'area'),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('nn.functional.interpolate', 'linear'),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('nn.functional.interpolate', 'trilinear'),  # Cannot call sizes() on tensor with symbolic sizes/st...
-    xfail('nn.functional.max_pool1d', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
-    xfail('nn.functional.max_pool3d', ''),  # aten.max_pool3d_with_indices.default - couldn't find symbolic m...
     xfail('nn.functional.multi_margin_loss', ''),  # could not find kernel
     xfail('nn.functional.multilabel_margin_loss', ''),  # could not find kernel
     xfail('nn.functional.nll_loss', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
@@ -2982,9 +2980,6 @@ def test_aot_autograd_symbolic_exhaustive(self, device, dtype, op):
                          # TypeError: unsupported operand type(s) for divmod(): 'SymInt' and 'int'
     torch.nn.FractionalMaxPool2d,  # int() argument must be a string, a bytes-like object or a number, not 'SymFloat'
     torch.nn.FractionalMaxPool3d,  # int() argument must be a string, a bytes-like object or a number, not 'SymFloat'
-    torch.nn.MaxPool1d,  # Cannot call sizes() on tensor with symbolic sizes/strides
-    torch.nn.MaxPool3d,  # torch._subclasses.fake_tensor.UnsupportedOperatorException:
-                         # aten.max_pool3d_with_indices.default
 }
 
 

Diff for: test/test_meta.py

@@ -631,8 +631,6 @@ def run_meta_crossref(
     torch.mode : {f64, i32, i64, f16, u8, i16, bf16, b8, i8, f32},
     torch.nn.functional.ctc_loss : {f64, f32},
     torch.nn.functional.gaussian_nll_loss : {f16, f64, bf16, f32},
-    torch.nn.functional.max_pool3d : {f64, f32},
-    torch.nn.functional.max_pool3d_with_indices : {f64, f32},
     torch.nn.functional.multi_margin_loss : {f64, f32},
     torch.nn.functional.multilabel_margin_loss : {f64, f32},
     torch.nn.functional.one_hot : {i64},
@@ -722,8 +720,6 @@ def run_meta_crossref(
     torch.histc: {i16, i32, i64, i8},  # aten::histc, aten::histc.out
     torch.kthvalue: {f16},  # aten::kthvalue.values
     torch.median: {f16},  # aten::median, aten::median.dim_values
-    torch.nn.functional.max_pool3d: {bf16, f16},  # aten::max_pool3d_with_indices
-    torch.nn.functional.max_pool3d_with_indices: {bf16, f16},  # aten::max_pool3d_with_indices
     torch.nn.functional.multi_margin_loss: {bf16, f16},  # aten::multi_margin_loss
     torch.nn.functional.multilabel_margin_loss: {bf16, f16},  # aten::multilabel_margin_loss_forward
     torch.ormqr: {f32, f64},  # aten::ormqr, aten::ormqr.out
@@ -847,7 +843,6 @@ def __torch_function__(self, func, types, args=(), kwargs=None):
     aten.histogram.bin_ct : {f32, f64},
     aten.histogram.bins_tensor : {f32, f64},
     aten.kthvalue.default : {i8, f64, i64, bf16, f32, i32, i16, u8},
-    aten.max_pool3d_with_indices.default : {f32, f64},
     aten.median.default : {i8, f64, i64, bf16, f32, i32, i16, u8},
     aten.median.dim : {i8, f64, i64, bf16, f32, i32, i16, u8},
     aten.mode.default : {f16, i8, f64, i64, bf16, f32, i32, b8, i16, u8},
@@ -907,7 +902,6 @@ def __torch_function__(self, func, types, args=(), kwargs=None):
     aten.linalg_eigvalsh.out: {f32, f64},  # aten::linalg_eigvalsh.out
     aten.log_sigmoid_forward.default: {bf16, f16, f64, f32},
     aten.log_sigmoid_forward.output : {bf16, f16, f64, f32},  # aten::log_sigmoid_forward.output
-    aten.max_pool3d_with_indices.default: {bf16, f16},  # aten::max_pool3d_with_indices
     aten.median.default: {f16},  # aten::median
     aten.median.dim: {f16},  # aten::median.dim_values
     aten.multi_margin_loss.default: {bf16, f16},  # aten::multi_margin_loss

Diff for: test/test_proxy_tensor.py

@@ -1541,8 +1541,6 @@ def f(a, b, c, d, e):
     xfail('nn.functional.grid_sample', ''),  # aten.grid_sampler_2d.default - couldn't find symbolic meta function/decompos...
     xfail('nn.functional.interpolate', 'linear'),  # aten.upsample_linear1d.vec - couldn't find symbolic meta function/dec...
     xfail('nn.functional.interpolate', 'trilinear'),  # aten.upsample_trilinear3d.vec - couldn't find symbolic meta functi...
-    xfail('nn.functional.max_pool1d', ''),  # Trying to call aten.size on a tensor with symbolic shapes.
-    xfail('nn.functional.max_pool3d', ''),  # aten.max_pool3d_with_indices.default - couldn't find symbolic meta function/d...
     xfail('nn.functional.multi_margin_loss', ''),  # Could not run 'aten::multi_margin_loss' with arguments from the...
     xfail('nn.functional.multilabel_margin_loss', ''),  # Could not run 'aten::multilabel_margin_loss_forward' with ...
     xfail('nn.functional.pixel_unshuffle', ''),  # aten.pixel_unshuffle.default - couldn't find symbolic meta function/deco...