index_notation: fix a bug where windows would be dropped through +=

include/taco/index_notation/index_notation.h

@@ -714,9 +714,11 @@ class Assignment : public IndexStmt {
   Assignment(Access lhs, IndexExpr rhs, IndexExpr op = IndexExpr());
 
   /// Create an assignment. Can specify an optional operator `op` that turns the
-  /// assignment into a compound assignment, e.g. `+=`.
+  /// assignment into a compound assignment, e.g. `+=`. Additionally, specify
+  /// any modifers on reduction index variables (windows, index sets, etc.).
   Assignment(TensorVar tensor, std::vector<IndexVar> indices, IndexExpr rhs,
-             IndexExpr op = IndexExpr());
+             IndexExpr op = IndexExpr(),
+             const std::map<int, std::shared_ptr<IndexVarIterationModifier>>& modifiers = {});
 
   /// Return the assignment's left-hand side.
   Access getLhs() const;

src/index_notation/index_notation.cpp

@@ -899,7 +899,14 @@ Assignment Access::operator=(const TensorVar& var) {
 
 Assignment Access::operator+=(const IndexExpr& expr) {
   TensorVar result = getTensorVar();
-  Assignment assignment = Assignment(result, getIndexVars(), expr, Add());
+  Assignment assignment = Assignment(
+    result,
+    getIndexVars(),
+    expr,
+    Add(),
+    // Include any windows on LHS index vars.
+    getNode(*this)->packageModifiers()
+  );
   // check(assignment); TODO: fix check for precompute
   const_cast<AccessNode*>(getNode(*this))->setAssignment(assignment);
   return assignment;
@@ -1686,8 +1693,9 @@ Assignment::Assignment(Access lhs, IndexExpr rhs, IndexExpr op)
 }
 
 Assignment::Assignment(TensorVar tensor, vector<IndexVar> indices,
-                       IndexExpr rhs, IndexExpr op)
-    : Assignment(Access(tensor, indices), rhs, op) {
+                       IndexExpr rhs, IndexExpr op,
+                       const std::map<int, std::shared_ptr<IndexVarIterationModifier>>& modifiers)
+    : Assignment(Access(tensor, indices, modifiers), rhs, op) {
 }
 
 Access Assignment::getLhs() const {

src/tensor.cpp

@@ -760,8 +760,10 @@ vector<void*> packArguments(const TensorBase& tensor) {
 
   // Pack any index sets on the result tensor at the front of the arguments list.
   auto lhs = getNode(tensor.getAssignment().getLhs());
-  if (isa<AccessTensorNode>(lhs)) {
-    auto indexSetModes = to<AccessTensorNode>(lhs)->indexSetModes;
+  // We check isa<AccessNode> rather than isa<AccessTensorNode> to catch cases
+  // where the underlying access is represented with the base AccessNode class.
+  if (isa<AccessNode>(lhs)) {
+    auto indexSetModes = to<AccessNode>(lhs)->indexSetModes;
     for (auto& it : indexSetModes) {
       arguments.push_back(it.second.tensor.getStorage());
     }

test/tests-windowing.cpp

@@ -540,3 +540,15 @@ TEST(windowing, lhsIndexSet) {
   a.evaluate();
   ASSERT_TRUE(equals(a, expected)) << a << endl << expected << endl;
 }
+
+// See issue #451 for details about this test.
+TEST(windowing, compoundAssign) {
+  Tensor<double> A("A", {8}, Format{Dense});
+  Tensor<double> B("B", {3}, Format{Dense});
+  IndexVar i("i");
+  A(i(0,3)) += B(i);
+  // We explicitly call assemble here as the += makes evaluate not call it.
+  A.compile(); A.assemble(); A.compute();
+  A(i({1, 3, 6})) += B(i);
+  A.evaluate();
+}