Revert "Relandx3 "SymIntify cat and narrow" (pytorch#86289)"

This reverts commit a00f848.

Reverted pytorch#86289 on behalf of https://github.com/malfet due to @seemether  unlanded the rest of the stack and it will fail intern import anyway

Author: pytorchmergebot

aten/src/ATen/BatchingRegistrations.cpp

@@ -1096,9 +1096,9 @@ TORCH_LIBRARY_IMPL(aten, Batched, m) {
   m.impl("expand_as", native::expand_as); // composite wrt autograd
   m.impl("movedim.intlist", movedim_batching_rule);
   m.impl("movedim.int", static_cast<Tensor(*)(const Tensor&,int64_t,int64_t)>(native::movedim)); // composite wrt autograd
-  // There is another variant of narrow.  However, we don't
+  // NB: static_cast because there's another variant of narrow. However, we don't
   // want to support the other variant yet bc it isn't documented...
-  m.impl("narrow", native::narrow_symint); // composite wrt autograd
+  m.impl("narrow", static_cast<Tensor(*)(const Tensor&,int64_t,int64_t,int64_t)>(native::narrow)); // composite wrt autograd
   m.impl("numpy_T", native::numpy_T);   // composite wrt autograd
   m.impl("matrix_H", native::matrix_H); // composite wrt autograd
   m.impl("mT", native::mT);             // composite wrt autograd

aten/src/ATen/WrapDimUtils.h

@@ -8,17 +8,22 @@
 
 namespace at {
 
-// if dim_post_expr is 0 and wrap_scalar is true, then dim must be in the
-// range [-1, 0]. This is a special case for scalar tensors and manifests in
-// e.g. torch.sum(scalar_tensor, 0) Otherwise, dim should be in the range
-// [-dim_post_expr, dim_post_expr-1].
-using c10::maybe_wrap_dim;
+static inline int64_t maybe_wrap_dim(
+    int64_t dim,
+    int64_t dim_post_expr,
+    bool wrap_scalar = true) {
+  // if dim_post_expr is 0 and wrap_scalar is true, then dim must be in the
+  // range [-1, 0]. This is a special case for scalar tensors and manifests in
+  // e.g. torch.sum(scalar_tensor, 0) Otherwise, dim should be in the range
+  // [-dim_post_expr, dim_post_expr-1].
+  return c10::maybe_wrap_dim(dim, dim_post_expr, wrap_scalar);
+}
 
-inline int64_t maybe_wrap_dim(int64_t dim, TensorImpl* tensor) {
+static inline int64_t maybe_wrap_dim(int64_t dim, TensorImpl* tensor) {
   return maybe_wrap_dim(dim, tensor->dim());
 }
 
-inline int64_t maybe_wrap_dim(int64_t dim, TensorList tensors) {
+static inline int64_t maybe_wrap_dim(int64_t dim, TensorList tensors) {
   if (tensors.size() == 0) {
     // can't wrap empty TensorList; rely on underlying implementation to throw
     // error if necessary.
@@ -27,7 +32,7 @@ inline int64_t maybe_wrap_dim(int64_t dim, TensorList tensors) {
   return maybe_wrap_dim(dim, tensors[0].dim());
 }
 
-inline int64_t maybe_wrap_dim(
+static inline int64_t maybe_wrap_dim(
     int64_t dim,
     const std::vector<std::vector<int64_t>>& tensor_sizes) {
   if (tensor_sizes.size() == 0) {
@@ -40,7 +45,7 @@ inline int64_t maybe_wrap_dim(
 
 // wrap each dim in the dims array, taking dim_post_expr as the true number of
 // dimensions
-inline void maybe_wrap_dims_n(
+static inline void maybe_wrap_dims_n(
     int64_t* dims,
     int64_t ndims,
     int64_t dim_post_expr) {
@@ -80,32 +85,19 @@ inline void maybe_wrap_dims(Container& dims, int64_t dim_post_expr) {
 // dimension behavior and dimension size checking). We maintain this behavior
 // for backwards compatibility, but only for this specific size (i.e. other
 // empty sizes are not skipped).
-template <typename T>
-inline int64_t _legacy_cat_wrap_dim(
+static inline int64_t legacy_cat_wrap_dim(
     int64_t dim,
-    const std::vector<std::vector<T>>& tensor_sizes) {
+    const std::vector<std::vector<int64_t>>& tensor_sizes) {
   for (auto& sizes : tensor_sizes) {
-    if (sizes == std::vector<T>({0})) {
+    if (sizes == std::vector<int64_t>({0})) {
       continue;
     }
     return maybe_wrap_dim(dim, sizes.size());
   }
   return dim;
 }
 
-inline int64_t legacy_cat_wrap_dim(
-    int64_t dim,
-    const std::vector<std::vector<int64_t>>& tensor_sizes) {
-  return _legacy_cat_wrap_dim<int64_t>(dim, tensor_sizes);
-}
-
-inline int64_t legacy_cat_wrap_dim_symint(
-    int64_t dim,
-    const std::vector<std::vector<c10::SymInt>>& tensor_sizes) {
-  return _legacy_cat_wrap_dim<c10::SymInt>(dim, tensor_sizes);
-}
-
-inline int64_t legacy_cat_wrap_dim(
+static inline int64_t legacy_cat_wrap_dim(
     int64_t dim,
     const MaterializedITensorListRef& tensors) {
   for (const Tensor& tensor : tensors) {
@@ -118,7 +110,7 @@ inline int64_t legacy_cat_wrap_dim(
 }
 
 // wrap negative dims in a vector
-inline void wrap_all_dims(
+static inline void wrap_all_dims(
     std::vector<int64_t>& dims_to_wrap,
     int64_t tensor_total_dims) {
   for (const auto i : c10::irange(dims_to_wrap.size())) {

aten/src/ATen/functorch/BatchRulesDecompositions.cpp

@@ -168,7 +168,7 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchBatched, m) {
   OP_DECOMPOSE2(movedim, int);
   OP_DECOMPOSE(msort);
   OP_DECOMPOSE(mT);
-  m.impl("narrow", native::narrow_symint);
+  OP_DECOMPOSE(narrow);
   OP_DECOMPOSE(negative);
   OP_DECOMPOSE2(frobenius_norm, dim);
   OP_DECOMPOSE2(nuclear_norm, dim);

aten/src/ATen/native/NonSymbolicBC.h

@@ -9,5 +9,4 @@ namespace native {
 // However, in certain cases (such as static runtime), we call the native versions of the ops directly.
 // In those cases, we will duplicate the signature here with non-symbolic ints, and also duplicate the C++ implementation.
 TORCH_API at::Tensor reshape(const at::Tensor& self, at::IntArrayRef proposed_shape);
-TORCH_API at::Tensor narrow(const at::Tensor& self, int64_t dim, int64_t start, int64_t length);
 }}

aten/src/ATen/native/TensorShape.cpp

@@ -11,7 +11,6 @@
 #include <ATen/core/DimVector.h>
 #include <ATen/core/IListRef.h>
 #include <ATen/native/Copy.h>
-#include <ATen/native/NonSymbolicBC.h>
 #include <ATen/native/Resize.h>
 #include <ATen/native/TensorIterator.h>
 #include <ATen/native/TensorShape.h>
@@ -1137,24 +1136,11 @@ Tensor narrow(const Tensor& self, int64_t dim, int64_t start, int64_t length) {
   return at::slice(self, dim, start, start + length, 1);
 }
 
-Tensor narrow_symint(const Tensor& self, int64_t dim, SymInt start, SymInt length) {
-  TORCH_CHECK(self.dim() > 0, "narrow() cannot be applied to a 0-dim tensor.");
-  auto cur_size = self.sym_size(dim);
-  if (start != cur_size) {  // start being the end is valid, but not a valid dim specification.
-    start = maybe_wrap_dim(start, cur_size);
-  }
-  TORCH_CHECK(length >= 0 && start <= cur_size - length,
-           "start (", start, ") + length (", length, ") exceeds dimension size (", cur_size, ").");
-  return at::slice_symint(self, dim, start, start + length, 1);
-}
-
-// This overload exists purely for XLA, because they wanted to pass in "symbolic"
-// start via Tensor.
-Tensor narrow_tensor_symint(const Tensor& self, int64_t dim, const Tensor& start, SymInt length) {
+Tensor narrow(const Tensor& self, int64_t dim, const Tensor& start, int64_t length) {
   TORCH_CHECK(start.dim() == 0 && isIntegralType(start.scalar_type(), /*includeBool=*/false),
               "start must be an 0-dim integral Tensor.");
   int64_t st = start.item<int64_t>();
-  return at::narrow_symint(self, dim, c10::SymInt(st), length);
+  return at::narrow(self, dim, st, length);
 }
 
 std::tuple<DimVector, DimVector, std::vector<int64_t>>

aten/src/ATen/native/native_functions.yaml

@@ -3710,19 +3710,15 @@
   dispatch:
     CPU: narrow_copy_dense_cpu_out
 
-- func: narrow(Tensor(a) self, int dim, SymInt start, SymInt length) -> Tensor(a)
+- func: narrow(Tensor(a) self, int dim, int start, int length) -> Tensor(a)
   variants: function, method
   device_check: NoCheck
   device_guard: False
-  dispatch:
-    CompositeImplicitAutograd: narrow_symint
 
-- func: narrow.Tensor(Tensor(a) self, int dim, Tensor start, SymInt length) -> Tensor(a)
+- func: narrow.Tensor(Tensor(a) self, int dim, Tensor start, int length) -> Tensor(a)
   variants: function, method
   device_check: NoCheck
   device_guard: False
-  dispatch:
-    CompositeImplicitAutograd: narrow_tensor_symint
 
 - func: native_batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps) -> (Tensor, Tensor, Tensor)
   dispatch:

aten/src/ATen/native/nested/cuda/NestedTensorTransformerFunctions.cpp

@@ -17,7 +17,6 @@
 #include <ATen/native/transformers/cuda/flash_attn/fmha_api.h>
 #endif
 
-#include <ATen/native/NonSymbolicBC.h>
 #include <ATen/native/nested/NestedTensorTransformerFunctions.h>
 #include <ATen/native/nested/NestedTensorMath.h>
 #include <ATen/native/nested/NestedTensorUtils.h>

aten/src/ATen/native/transformers/cuda/attention.cu

@@ -16,7 +16,6 @@
 
 #include <c10/cuda/CUDAMathCompat.h>
 
-#include <ATen/native/NonSymbolicBC.h>
 #include <ATen/native/nested/NestedTensorUtils.h>
 #include <ATen/native/nested/NestedTensorTransformerFunctions.h>
 

c10/core/WrapDimMinimal.cpp

@@ -3,8 +3,10 @@
 namespace c10 {
 namespace detail {
 
-template <typename T>
-T maybe_wrap_dim_slow(T dim, T dim_post_expr, bool wrap_scalar) {
+int64_t maybe_wrap_dim_slow(
+    int64_t dim,
+    int64_t dim_post_expr,
+    bool wrap_scalar) {
   TORCH_CHECK_INDEX(
       dim_post_expr >= 0, "Rank cannot be negative but got ", dim_post_expr);
 
@@ -17,8 +19,8 @@ T maybe_wrap_dim_slow(T dim, T dim_post_expr, bool wrap_scalar) {
     return c10::maybe_wrap_dim(dim, /*dim_post_expr=*/1, /*wrap_scalar=*/false);
   }
 
-  T min = dim_post_expr * -1;
-  T max = dim_post_expr - 1;
+  int64_t min = -dim_post_expr;
+  int64_t max = dim_post_expr - 1;
   TORCH_CHECK_INDEX(
       min <= dim && dim <= max,
       "Dimension out of range (expected to be in range of [",
@@ -33,11 +35,5 @@ T maybe_wrap_dim_slow(T dim, T dim_post_expr, bool wrap_scalar) {
       false, "should never reach here as dim should be out-of-bounds");
 }
 
-// Explicitly instantiate the template at the two types it will be used
-template C10_API int64_t
-maybe_wrap_dim_slow(int64_t dim, int64_t dim_post_expr, bool wrap_scalar);
-template C10_API SymInt
-maybe_wrap_dim_slow(SymInt dim, SymInt dim_post_expr, bool wrap_scalar);
-
 } // namespace detail
 } // namespace c10

c10/core/WrapDimMinimal.h

@@ -1,44 +1,25 @@
 #pragma once
 
-#include <c10/core/SymInt.h>
 #include <c10/util/Exception.h>
 
 namespace c10 {
 
 namespace detail {
-// This template can only be specialized at int64_t and c10::SymInt;
-// you'll get linker errors otherwise
-template <typename T>
-C10_API T maybe_wrap_dim_slow(T dim, T dim_post_expr, bool wrap_scalar);
-} // namespace detail
-
-template <typename T>
-T _maybe_wrap_dim(T dim, T dim_post_expr, bool wrap_scalar = true) {
-  // Inline the fast paths
-  if (C10_LIKELY(dim_post_expr * -1 <= dim && dim < dim_post_expr)) {
-    // For SymInts, we want an explicit control flow to trigger a guard, so we
-    // may as well branch too.
-    if (dim < 0) {
-      return dim + dim_post_expr;
-    }
-    return dim;
-  }
-  // Check edge-cases out-of-line (wrapping scalars and out-of-bounds errors)
-  return c10::detail::maybe_wrap_dim_slow<T>(dim, dim_post_expr, wrap_scalar);
+C10_API int64_t
+maybe_wrap_dim_slow(int64_t dim, int64_t dim_post_expr, bool wrap_scalar);
 }
 
-inline int64_t maybe_wrap_dim(
+static inline int64_t maybe_wrap_dim(
     int64_t dim,
     int64_t dim_post_expr,
     bool wrap_scalar = true) {
-  return _maybe_wrap_dim(dim, dim_post_expr, wrap_scalar);
-}
-
-inline c10::SymInt maybe_wrap_dim(
-    c10::SymInt dim,
-    c10::SymInt dim_post_expr,
-    bool wrap_scalar = true) {
-  return _maybe_wrap_dim(dim, dim_post_expr, wrap_scalar);
+  // Inline the fast paths
+  if (C10_LIKELY(-dim_post_expr <= dim && dim < dim_post_expr)) {
+    // Branch-less version of dim + (dim < 0 ? dim_post_expr : 0)
+    return dim + dim_post_expr * (dim < 0);
+  }
+  // Check edge-cases out-of-line (wrapping scalars and out-of-bounds errors)
+  return c10::detail::maybe_wrap_dim_slow(dim, dim_post_expr, wrap_scalar);
 }
 
 } // namespace c10

functorch/test/test_aotdispatch.py

@@ -738,13 +738,15 @@ def assert_compiler(gm: torch.fx.GraphModule, _):
     xfail('block_diag', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('broadcast_tensors', ''),  # 'int' and 'torch._C.SymIntNode'
     xfail('cartesian_prod', ''),  # Cannot call numel() on tensor with symbolic sizes/strides
+    xfail('cat', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('cdist', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('cholesky_inverse', ''),  # could not find kernel
     xfail('cholesky_solve', ''),  # could not find kernel
     xfail('chunk', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('column_stack', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('combinations', ''),  # aten.masked_select.default
     xfail('complex', ''),  # aten.view_as_real.default - couldn't find symbolic meta function/decomposition
+    xfail('constant_pad_nd', ''),  # aten.fill.Scalar - couldn't find symbolic meta function/decomposition
     xfail('cross', ''),  # aten.linalg_cross.default - couldn't find symbolic meta function/decomposition
     xfail('cummax', ''),  # aten.cummax.default - couldn't find symbolic meta function/decomposition
     xfail('cummin', ''),  # aten.cummin.default - couldn't find symbolic meta function/decomposition
@@ -759,6 +761,7 @@ def assert_compiler(gm: torch.fx.GraphModule, _):
     xfail('digamma', ''),  # aten.polygamma.default - couldn't find symbolic meta function/decomposition
     xfail('dist', ''),  # aten.dist.default - couldn't find symbolic meta function/decomposition
     xfail('dsplit', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
+    xfail('dstack', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('einsum', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('expand_as', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('fft.fft2', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
@@ -788,6 +791,7 @@ def assert_compiler(gm: torch.fx.GraphModule, _):
     xfail('gather', ''),  # aten.gather.default - couldn't find symbolic meta function/decomposition
     xfail('gradient', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('hsplit', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
+    xfail('hstack', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('i0', ''),  # aten.i0.default - couldn't find symbolic meta function/decomposition
     xfail('index_copy', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('index_fill', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
@@ -938,6 +942,7 @@ def assert_compiler(gm: torch.fx.GraphModule, _):
     xfail('nn.functional.nll_loss', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('nn.functional.normalize', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('nn.functional.pad', 'circular'),  # Cannot call sizes() on tensor with symbolic sizes/strides
+    xfail('nn.functional.pad', 'constant'),  # aten.fill.Scalar - couldn't find symbolic meta function/decom...
     xfail('nn.functional.pad', 'reflect'),  # aten.reflection_pad1d.default - couldn't find symbolic meta fu...
     xfail('nn.functional.pad', 'replicate'),  # aten.replication_pad1d.default - couldn't find symbolic meta...
     xfail('nn.functional.pairwise_distance', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
@@ -1029,6 +1034,7 @@ def assert_compiler(gm: torch.fx.GraphModule, _):
     xfail('view_as_complex', ''),  # aten.view_as_complex.default - couldn't find symbolic meta function/deco...
     xfail('view_as', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('vsplit', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
+    xfail('vstack', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
 }
 
 def _test_aot_autograd_helper(self, device, dtype, op):

test/test_proxy_tensor.py

@@ -1066,6 +1066,7 @@ def f(a, b, c, d, e):
     xfail('cholesky_solve', ''),  # Could not run 'aten::_cholesky_solve_helper' with arguments from the 'Meta' back...
     xfail('chunk', ''),  # aten.size.default - couldn't find symbolic meta function/decomposition
     xfail('column_stack', ''),  # Tensors of type TensorImpl do not have numel
+    xfail('constant_pad_nd', ''),  # aten.fill.Scalar - couldn't find symbolic meta function/decomposition
     xfail('count_nonzero', ''),  # Could not run 'aten::count_nonzero.dim_IntList' with arguments from the 'Meta' ba...
     xfail('cross', ''),  # aten.linalg_cross.default - couldn't find symbolic meta function/decomposition
     xfail('cummax', ''),  # aten.cummax.default - couldn't find symbolic meta function/decomposition
@@ -1206,6 +1207,7 @@ def f(a, b, c, d, e):
     xfail('nn.functional.interpolate', 'linear'),  # aten.upsample_linear1d.vec - couldn't find symbolic meta function/dec...
     xfail('nn.functional.interpolate', 'nearest'),  # aten.upsample_nearest1d.vec - couldn't find symbolic meta function/d...
     xfail('nn.functional.interpolate', 'trilinear'),  # aten.upsample_trilinear3d.vec - couldn't find symbolic meta functi...
+    xfail('nn.functional.local_response_norm', ''),  # Tensors of type TensorImpl do not have numel
     xfail('nn.functional.margin_ranking_loss', ''),  # The underlying op of 'aten.stride' has no overload name '_schema'
     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...
@@ -1215,6 +1217,7 @@ def f(a, b, c, d, e):
     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.pad', 'circular'),  # aten.size.default - couldn't find symbolic meta function/decomposition
+    xfail('nn.functional.pad', 'constant'),  # aten.fill.Scalar - couldn't find symbolic meta function/decomposition
     xfail('nn.functional.pad', 'reflect'),  # aten.reflection_pad1d.default - couldn't find symbolic meta function/decompo...
     xfail('nn.functional.pad', 'replicate'),  # aten.replication_pad1d.default - couldn't find symbolic meta function/deco...
     xfail('nn.functional.pdist', ''),  # Could not run 'aten::_pdist_forward' with arguments from the 'Meta' backend...
@@ -1284,6 +1287,8 @@ def f(a, b, c, d, e):
     xfail('special.scaled_modified_bessel_k1', ''),  # aten.special_scaled_modified_bessel_k1.default - couldn't find symbo...
     xfail('special.xlog1py', ''),  # aten.special_xlog1py.default - couldn't find symbolic meta function/decomposition
     xfail('split', ''),  # 'torch._C.SymIntNode' and 'int'
+    xfail('split', 'list_args'),  # aten.size.default - couldn't find symbolic meta function/decomposition
+    xfail('split_with_sizes', ''),  # aten.size.default - couldn't find symbolic meta function/decomposition
     xfail('stft', ''),  # argument 'size' must be tuple of ints, but found element of type torch._C.SymIntNode at...
     xfail('sum_to_size', ''),  # aten.size.default - couldn't find symbolic meta function/decomposition
     xfail('svd', ''),  # aten._linalg_svd.default - couldn't find symbolic meta function/decomposition

tools/autograd/derivatives.yaml

@@ -354,7 +354,7 @@
   self, other: matmul_backward(grad, self, other, grad_input_mask)
 
 - name: cat(Tensor[] tensors, int dim=0) -> Tensor
-  tensors: cat_tensors_backward(grad, to_args_sizes_symint(tensors), to_args_scalartypes(tensors), dim)
+  tensors: cat_tensors_backward(grad, to_args_sizes(tensors), to_args_scalartypes(tensors), dim)
   result: cat_jvp(tensors, dim)
 
 - name: cauchy_(Tensor(a!) self, float median=0, float sigma=1, *, Generator? generator=None) -> Tensor(a!)

tools/autograd/gen_python_functions.py

@@ -1160,11 +1160,6 @@ def is_arg_smaller(t1: Type, t2: Type) -> bool:
             # Prioritize IntArrayRef overload over SymIntArrayRef
             str(t1) == "SymInt[]"
             and str(t2) == "int[]"
-            or
-            # Make sure both in, SymInt are sorted consistently w.r.t. Tensor since Tensor can be implicitly
-            # converted to either int or SymInt.  Prioritize the Tensor overload since it otherwise gets shadowed.
-            (str(t1) == "SymInt" or str(t1) == "int")
-            and str(t2) == "Tensor"
         )
 
     def is_smaller(s1: PythonSignature, s2: PythonSignature) -> bool: