Misc changes from compiled autograd branch (pytorch#104316)

This PR pulls out some standalone changes from pytorch#103822

Pull Request resolved: pytorch#104316
Approved by: https://github.com/ezyang

Author: jansel

aten/src/ATen/TensorGeometry.h

@@ -113,6 +113,26 @@ struct TORCH_API TensorGeometry {
     return r;
   }
 
+  std::vector<c10::SymInt>& mutable_sizes() {
+    return sizes_;
+  }
+  std::vector<c10::SymInt>& mutable_strides() {
+    return strides_;
+  }
+  c10::SymInt& mutable_storage_offset() {
+    return storage_offset_;
+  }
+  void recompute() {
+    // recalculate numel after a change
+    c10::SymInt numel = 1;
+    for (const auto& i : sizes_) {
+      numel = numel * i;
+    }
+    numel_ = std::move(numel);
+    has_symbolic_sizes_strides_ =
+        !c10::asIntArrayRefSlowOpt(sizes_).has_value();
+  }
+
  private:
   std::vector<c10::SymInt> sizes_;
   std::vector<c10::SymInt> strides_;

c10/core/SafePyObject.h

@@ -22,21 +22,33 @@ struct C10_API SafePyObject {
   // Steals a reference to data
   SafePyObject(PyObject* data, c10::impl::PyInterpreter* pyinterpreter)
       : data_(data), pyinterpreter_(pyinterpreter) {}
+  SafePyObject(SafePyObject&& other)
+      : data_(std::exchange(other.data_, nullptr)),
+        pyinterpreter_(other.pyinterpreter_) {}
 
   // In principle this could be copyable if we add an incref to PyInterpreter
   // but for now it's easier to just disallow it.
   SafePyObject(SafePyObject const&) = delete;
   SafePyObject& operator=(SafePyObject const&) = delete;
 
   ~SafePyObject() {
-    (*pyinterpreter_)->decref(data_, /*is_tensor*/ false);
+    if (data_ != nullptr) {
+      (*pyinterpreter_)->decref(data_, /*is_tensor*/ false);
+    }
   }
 
   c10::impl::PyInterpreter& pyinterpreter() const {
     return *pyinterpreter_;
   }
   PyObject* ptr(const c10::impl::PyInterpreter*) const;
 
+  // stop tracking the current object, and return it
+  PyObject* release() {
+    auto rv = data_;
+    data_ = nullptr;
+    return rv;
+  }
+
  private:
   PyObject* data_;
   c10::impl::PyInterpreter* pyinterpreter_;

tools/autograd/templates/Functions.h

@@ -52,7 +52,6 @@ struct TypeAndSize {
 
   Tensor zeros() { return at::zeros_symint(sym_sizes, options); }
 
-private:
   std::vector<c10::SymInt> sym_sizes;
   at::TensorOptions options;
 };

torch/_dynamo/variables/builder.py

@@ -1028,19 +1028,10 @@ def wrap_unspecialized_primitive(self, value):
                         guards=self.make_guards(GuardBuilder.CONSTANT_MATCH),
                     )
 
-                wrapped_value = shape_env.create_symintnode(
-                    # TODO: This is wrong wrong wrong, create_symbol will
-                    # generate something that is non-negative, but this is
-                    # not a sound assumption to make.
-                    # Not fixing as this was a preexisting condition.
-                    shape_env.create_symbol(
-                        value,
-                        source=self.source,
-                        dynamic_dim=dynamic_dim,
-                        constraint_dim=None,
-                    ),
-                    hint=value,
+                wrapped_value = shape_env.create_symint_and_symbol(
+                    value,
                     source=self.source,
+                    dynamic_dim=dynamic_dim,
                 )
                 self.tx.output.tracked_fakes.append(
                     TrackedFake(wrapped_value, self.source, None)

torch/_functorch/aot_autograd.py

@@ -2799,12 +2799,6 @@ def forward(ctx, *deduped_flat_tensor_args):
             )
 
             num_outputs = CompiledFunction.metadata.num_outputs
-            num_outputs_aliased_to_inputs = (
-                CompiledFunction.metadata.num_outputs_aliased_to_inputs
-            )
-            num_outputs_aliased_to_intermediates = (
-                CompiledFunction.metadata.num_outputs_aliased_to_intermediates
-            )
             num_outputs_aliased = CompiledFunction.metadata.num_outputs_aliased
             num_intermediate_bases = CompiledFunction.metadata.num_intermediate_bases
             num_symints_saved_for_bw = CompiledFunction.num_symints_saved_for_bw
@@ -2979,10 +2973,12 @@ def backward(ctx, *flat_args):
                 # Add the seed and offset to args
                 rng_args = CUDARngStateHelper.get_torch_state_as_tuple()
 
-            all_args = (
-                list(ctx.symints) + list(ctx.saved_tensors) + list(contiguous_args) + list(rng_args)
-            )
-
+            all_args = [
+                *ctx.symints,
+                *ctx.saved_tensors,
+                *contiguous_args,
+                *rng_args
+            ]
             del contiguous_args
 
             def call_compiled_backward():

torch/fx/experimental/symbolic_shapes.py

@@ -2245,6 +2245,21 @@ def create_symintnode(
             return int(sym)
         return SymInt(SymNode(sym, self, int, hint, fx_node=fx_node))
 
+    def create_symint_and_symbol(self, value, source, dynamic_dim):
+        # TODO: This is wrong wrong wrong, create_symbol will
+        # generate something that is non-negative, but this is
+        # not a sound assumption to make.
+        # Not fixing as this was a preexisting condition.
+        return self.create_symintnode(
+            self.create_symbol(
+                value,
+                source=source,
+                dynamic_dim=dynamic_dim,
+            ),
+            hint=value,
+            source=source,
+        )
+
     def create_symboolnode(self, sym: "sympy.Expr"):
         # This function is only being used in serialization, so we do not track it
         # for validation.