[jit] Fix aug assign for non-tensor attributes (pytorch#32993)

Summary:
Instead of erroring out this de-sugars augmented assignments to class
members from `self.a += 1` to `self.a = self.a + 1`.

Fixes pytorch#32973
](https://our.intern.facebook.com/intern/diff/19737636/)
Pull Request resolved: pytorch#32993

Pulled By: driazati

Differential Revision: D19737636

fbshipit-source-id: 07307cde88d8c348a7affdafe26db21c74e28ec0

Author: davidriazati

test/test_jit.py

@@ -5294,6 +5294,88 @@ def foo(z):
             return y[0][1]
         self.checkScript(foo, ((1, [[1, 2], [3, 4]]),))
 
+    def test_nested_aug_assign(self):
+        @torch.jit.script
+        class SomeClass(object):
+            def __init__(self):
+                self.num = 99
+
+            def __iadd__(self, x):
+                # type: (int)
+                self.num += x
+                return self
+
+            def __eq__(self, other):
+                # type: (SomeClass) -> bool
+                return self.num == other.num
+
+        @torch.jit.script
+        class SomeOutOfPlaceClass(object):
+            def __init__(self):
+                self.num = 99
+
+            def __add__(self, x):
+                # type: (int)
+                self.num = x
+                return self
+
+            def __eq__(self, other):
+                # type: (SomeClass) -> bool
+                return self.num == other.num
+
+        class Child(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.x = 2
+                self.o = SomeClass()
+                self.oop = SomeOutOfPlaceClass()
+                self.list = [1, 2, 3]
+
+        class A(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.child = Child()
+
+            def forward(self):
+                self.child.x += 1
+                self.child.o += 5
+                self.child.oop += 5
+                some_list = [1, 2]
+                self.child.list += some_list
+                self.child.list *= 2
+                return self.child.x, self.child.o, self.child.list, self.child.oop
+
+        a = A()
+        sa = torch.jit.script(A())
+        eager_result = a()
+        script_result = sa()
+        self.assertEqual(eager_result, script_result)
+        self.assertEqual(a.child.x, sa.child.x)
+        self.assertEqual(a.child.o, sa.child.o)
+        self.assertEqual(a.child.list, sa.child.list)
+
+        @torch.jit.script
+        class SomeNonAddableClass(object):
+            def __init__(self):
+                self.num = 99
+
+            def __eq__(self, other):
+                # type: (SomeClass) -> bool
+                return self.num == other.num
+
+        # with self.assertRaisesRegex(RuntimeError, "")
+        class A(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.x = SomeNonAddableClass()
+
+            def forward(self):
+                self.x += SomeNonAddableClass()
+                return self.x
+
+        with self.assertRaisesRegex(RuntimeError, "Cannot emit inplace op"):
+            torch.jit.script(A())
+
     def test_nested_list_construct(self):
         def foo():
             return [[4]] + [[4, 5]]

torch/csrc/jit/register_prim_ops.cpp

@@ -1680,6 +1680,26 @@ int listInplaceAdd(Stack& stack) {
   return 0;
 }
 
+template <class T>
+int listMulIntLeftInPlace(Stack& stack) {
+  int64_t n = pop(stack).to<int64_t>();
+  c10::List<T> list = pop(stack).to<c10::List<T>>();
+
+  if (n <= 0) {
+    list.clear();
+  } else if (n > 1) {
+    size_t list_size = list.size();
+    for (auto i = 1; i < n; i++) {
+      for (size_t j = 0; j < list_size; j++) {
+        list.push_back(list.get(j));
+      }
+    }
+  }
+
+  push(stack, std::move(list));
+  return 0;
+}
+
 template <class T>
 int listMulIntLeft(Stack& stack) {
   int64_t n = pop(stack).to<int64_t>();
@@ -2295,6 +2315,10 @@ RegisterOperators reg2({
       Operator(                                                                     \
           "aten::mul(int n, " decl_type "[] l) -> " decl_type "[]",                 \
           listMulIntRight<c_type::value_type>,                                      \
+          aliasAnalysisFromSchema()),                                               \
+      Operator(                                                                     \
+          "aten::mul_(" decl_type "[](a!) l, int n) -> " decl_type "[](a!)",        \
+          listMulIntLeftInPlace<c_type::value_type>,                                \
           aliasAnalysisFromSchema())
 
     CREATE_LIST_OPS("int", c10::List<int64_t>),

torch/csrc/jit/script/ir_emitter.cpp

@@ -1780,22 +1780,37 @@ struct to_ir {
   // If the RHS is a tensor, return the corresponding ATen in-place op
   // If it's a list of scalars, then return the corresponding list augment op
   Symbol getAugOp(const AugAssign& stmt, const TypePtr& type) {
-    if (type->cast<ListType>()) { // Lists also have in-place ops.
-      switch (stmt.aug_op()) {
-        case '+':
-          return aten::add_;
-      }
+    bool use_inplace_op = type->isSubtypeOf(TensorType::get()) ||
+        type->kind() == TypeKind::ListType;
+    switch (stmt.aug_op()) {
+      case '+':
+        return use_inplace_op ? aten::add_ : aten::add;
+      case '-':
+        return use_inplace_op ? aten::sub_ : aten::sub;
+      case '/':
+        return use_inplace_op ? aten::div_ : aten::div;
+      case '*':
+        return use_inplace_op ? aten::mul_ : aten::mul;
+      default:
+        throw ErrorReport(stmt)
+            << "Unknown augmented assignment: " << kindToString(stmt.aug_op());
     }
-    bool isTensor = type->isSubtypeOf(TensorType::get());
+  }
+
+  // Get a pair of <in place magic method name, out of place magic method name>
+  // since the out of place method is called if the in place method is not
+  // present
+  std::pair<std::string, std::string> getAugMagicMethod(const AugAssign& stmt) {
     switch (stmt.aug_op()) {
       case '+':
-        return isTensor ? aten::add_ : aten::add;
+        return std::make_pair(std::string("__iadd__"), std::string("__add__"));
       case '-':
-        return isTensor ? aten::sub_ : aten::sub;
+        return std::make_pair(std::string("__isub__"), std::string("__sub__"));
       case '/':
-        return isTensor ? aten::div_ : aten::div;
+        return std::make_pair(
+            std::string("__itruediv__"), std::string("__truediv__"));
       case '*':
-        return isTensor ? aten::mul_ : aten::mul;
+        return std::make_pair(std::string("__imul__"), std::string("__mul__"));
       default:
         throw ErrorReport(stmt)
             << "Unknown augmented assignment: " << kindToString(stmt.aug_op());
@@ -1831,34 +1846,58 @@ struct to_ir {
   //
   //  def forward():
   //    self.num_batches += 1
-  //
-  // In this case we will only consider the scenario that the module
-  // buffer type is a tensor, and we emit the corresponding tensor
-  // in place op, and throw error for other unsupported types
   void emitAugAssignmentToSelectVar(const AugAssign& stmt) {
     const auto lhs = Select(stmt.lhs());
-    const auto lhsSugaredVar =
-        environment_stack->getSugaredVar(Var(lhs.value()).name());
+    auto lhsSugaredVar = emitSugaredExpr(lhs.value(), 1);
     const auto lhsValue =
         lhsSugaredVar->attr(lhs.range(), method, lhs.selector().name())
             ->asValue(lhs.range(), method);
-    if (lhsValue->type()->isSubtypeOf(TensorType::get())) {
-      // for module parameter/buffer assignment, only consider tensor types,
-      // emit the corresponding in-place op
-      const auto rhs = NamedValue(stmt.rhs().range(), emitExpr(stmt.rhs()));
-      const auto self = NamedValue(stmt.lhs().range(), "self", lhsValue);
-      emitBuiltinCall(
+  if (lhsValue->type()->kind() == TypeKind::ClassType) {
+      // Call `__iadd__` so updates happen in place on class types
+      // https://docs.python.org/3/reference/datamodel.html#object.__iadd__
+      std::string in_place_method_name;
+      std::string out_of_place_method_name;
+      std::tie(in_place_method_name, out_of_place_method_name) =
+          getAugMagicMethod(stmt);
+      const auto rhs = emitExpr(stmt.rhs());
+
+      // Determine whether to use __iadd__ or __add__ (use __add__ only if
+      // __iadd__ is not present)
+      auto type = lhsValue->type()->expect<ClassType>();
+      std::string magic_method_name;
+      if (type->getMethod(in_place_method_name)) {
+        magic_method_name = in_place_method_name;
+      } else if (type->getMethod(out_of_place_method_name)) {
+        magic_method_name = out_of_place_method_name;
+      } else {
+        throw ErrorReport(stmt.range())
+            << "Cannot emit inplace op on " << type->python_str()
+            << " since it does not define an " << in_place_method_name << " or "
+            << out_of_place_method_name << " method";
+      }
+
+      // Insert call to the magic method
+      MethodValue method_value(lhsValue, magic_method_name);
+      auto result = method_value.call(stmt.range(), method, {rhs}, {}, 0)
+                        ->asValue(stmt.range(), method);
+
+      // x += y is equivalent to x = x.__iadd__(y) or x = x.__add__(y) if
+      // __iadd__ is not present, so set the value to the function's return
+      // value
+      lhsSugaredVar->setAttr(
+          stmt.range(), method, lhs.selector().name(), result);
+    } else {
+      const auto rhs = NamedValue(stmt.rhs().range(), emitExpr(stmt.rhs()))
+                           .value(*method.graph());
+      auto rhsValue = emitBuiltinCall(
           stmt.range(),
           *method.graph(),
           getAugOp(stmt, lhsValue->type()),
-          {rhs},
+          {lhsValue, rhs},
           {},
-          self);
-
-    } else {
-      throw ErrorReport(stmt.lhs())
-          << "left-hand side of augmented assignment to module "
-          << "parameters/buffers can only be tensor types";
+          /*self=*/c10::nullopt);
+      lhsSugaredVar->setAttr(
+          stmt.range(), method, lhs.selector().name(), rhsValue);
     }
   }
 

torch/csrc/jit/script/sugared_value.cpp

@@ -525,6 +525,7 @@ std::shared_ptr<SugaredValue> MagicMethod::call(
           ->call(loc, m, inputs.slice(1), attributes, n_binders);
     }
   }
+  TORCH_INTERNAL_ASSERT(base_value_);
   return base_value_->call(loc, m, inputs, attributes, n_binders);
 }