Made a += b for lists do an in place add

Chillee · Chillee · commit fab06a58c9a6 · 2019-06-17T22:00:54.000-07:00
diff --git a/aten/src/ATen/core/List.h b/aten/src/ATen/core/List.h
@@ -349,6 +349,12 @@ class ListPtr final {
    */
   void push_back(T&& value) const;
 
+  /**
+   * Appends the given list to the end of the container. Uses at most one memory allocation.
+   * May invalidate any references, pointers, or iterators referring to contained elements. Any past-the-end iterators may also be invalidated.
+   */
+  void append(ListPtr<T> lst) const;
+
   /**
    * Appends the given element value to the end of the container.
    * The new element is constructed with the given arguments.
diff --git a/aten/src/ATen/core/List_inl.h b/aten/src/ATen/core/List_inl.h
@@ -178,6 +178,21 @@ void ListPtr<T>::push_back(T&& value) const {
   impl_->list.push_back(detail::list_element_from<T, StorageT>(std::move(value)));
 }
 
+template<class T>
+void ListPtr<T>::append(ListPtr<T> b) const {
+  size_type neededSize = this->size() + b.size();
+  if (impl_->list.capacity() < neededSize) {
+    this->reserve(std::max(impl_->list.capacity() * 2, neededSize));
+  }
+  if (b.use_count() == 1) {
+    std::move(b.begin(), b.end(), this->begin());
+  } else {
+    for (const auto& el: b) {
+      this->push_back(el);
+    }
+  }
+}
+
 template<class T>
 template<class... Args>
 void ListPtr<T>::emplace_back(Args&&... args) const {
diff --git a/torch/csrc/jit/register_prim_ops.cpp b/torch/csrc/jit/register_prim_ops.cpp
@@ -1554,6 +1554,15 @@ int listAdd(Stack& stack) {
   push(stack, std::move(ret));
   return 0;
 }
+template <class T>
+int listInplaceAdd(Stack& stack) {
+  c10::ListPtr<T> a = c10::make_list<T>();
+  c10::ListPtr<T> b = c10::make_list<T>();
+  pop(stack, a, b);
+  a.append(b);
+  push(stack, std::move(a));
+  return 0;
+}
 
 template <class T>
 int listMulIntLeft(Stack& stack) {
@@ -1928,6 +1937,10 @@ RegisterOperators reg2({
           "aten::add(" decl_type "[] a, " decl_type "[] b) -> " decl_type           \
           "[]",                                                                     \
           listAdd<c_type::value_type>),                                             \
+      Operator(                                                                     \
+          "aten::add_(" decl_type "[](a!) self, " decl_type "[] b) -> " decl_type           \
+          "[]",                                                                     \
+          listInplaceAdd<c_type::value_type>),                                             \
       Operator(                                                                     \
           "aten::slice(" decl_type                                                  \
           "[] l, int start, int end=9223372036854775807, int step=1) -> " decl_type \
diff --git a/torch/csrc/jit/script/compiler.cpp b/torch/csrc/jit/script/compiler.cpp
@@ -1632,7 +1632,15 @@ struct to_ir {
   // Get the appropriate builtin op for this augmented assignment
   // 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, bool isTensor) {
+  Symbol getAugOp(const AugAssign& stmt, TypePtr type) {
+    if (type->cast<ListType>()) { // Lists also have in-place ops.
+      std::cout<<"generating list op"<<std::endl;
+      switch (stmt.aug_op()) {
+        case '+':
+          return aten::add_;
+      }
+    }
+    bool isTensor = type->isSubtypeOf(TensorType::get());
     switch (stmt.aug_op()) {
       case '+':
         return isTensor ? aten::add_ : aten::add;
@@ -1696,7 +1704,7 @@ struct to_ir {
       emitBuiltinCall(
           stmt.range(),
           *method.graph(),
-          getAugOp(stmt, /*isTensor=*/true),
+          getAugOp(stmt, lhsValue->type()),
           self,
           {rhs},
           {},
@@ -1713,14 +1721,15 @@ struct to_ir {
     const auto lhs = Var(stmt.lhs());
     const auto lhsValue = environment_stack->getSugaredVar(lhs.name())
                               ->asValue(lhs.range(), method);
-    if (lhsValue->type()->isSubtypeOf(TensorType::get())) {
+    auto lhsType = lhsValue->type();
+    if (lhsType->isSubtypeOf(TensorType::get()) || lhsType->cast<c10::ListType>()) {
       // for tensors, 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);
       const auto output = emitBuiltinCall(
           stmt.range(),
           *method.graph(),
-          getAugOp(stmt, /*isTensor=*/true),
+          getAugOp(stmt, lhsValue->type()),
           self,
           {rhs},
           {},
@@ -1761,7 +1770,7 @@ struct to_ir {
         emitBuiltinCall(
             stmt.range(),
             *method.graph(),
-            getAugOp(stmt, /*isTensor=*/true),
+            getAugOp(stmt, sliceable->type()),
             slicedArg,
             {rhs},
             {},
@@ -1778,7 +1787,7 @@ struct to_ir {
         const auto augmented = emitBuiltinCall(
             stmt.range(),
             *method.graph(),
-            getAugOp(stmt, /*isTensor=*/true),
+            getAugOp(stmt, sliceable->type()),
             indexed,
             {rhs},
             {},
@@ -1796,8 +1805,7 @@ struct to_ir {
       const auto listType = sliceable->type()->cast<ListType>();
       AT_ASSERT(listType != nullptr);
 
-      bool isTensorList =
-          listType->getElementType()->isSubtypeOf(TensorType::get());
+      auto elementType = listType->getElementType();
 
       // Get the idx to augment
       const auto subscriptExprs = lhs.subscript_exprs();
@@ -1817,7 +1825,7 @@ struct to_ir {
       const auto getItem =
           graph->insert(aten::select, {listArg, idxArg}, {}, stmt.range());
       const auto augmentedItem = graph->insert(
-          getAugOp(stmt, isTensorList), {getItem, valueArg}, {}, stmt.range());
+          getAugOp(stmt, elementType), {getItem, valueArg}, {}, stmt.range());
       graph->insert(
           aten::_set_item, {listArg, idxArg, augmentedItem}, {}, stmt.range());
     }
