Add comments

Author: zhang677

include/taco/index_notation/index_notation.h

@@ -71,9 +71,16 @@ struct SuchThatNode;
 class IndexExprVisitorStrict;
 class IndexStmtVisitorStrict;
 
+/// Describe the relation between indexVar sets of lhs and rhs in an Assignment node.
+/// equal: lhs = rhs
+/// none: lhs and rhs are mutually exclusive. And lhs and rhs are not empty sets.
+/// lcr: rhs is a proper subset of lhs. (lhs contains rhs)
+/// rcl: lhs is a proper subset of rhs. (rhs contains lhs)
+/// inter: lhs and rhs share common elements but are not equal or empty.
 enum IndexSetRel {
     equal, none, lcr, rcl, inter
 };
+
 /// Return true if the index statement is of the given subtype.  The subtypes
 /// are Assignment, Forall, Where, Sequence, and Multi.
 template <typename SubType> bool isa(IndexExpr);
@@ -823,7 +830,7 @@ class Assignment : public IndexStmt {
   /// Return the reduction index variables i nthe assign
   std::vector<IndexVar> getReductionVars() const;
 
-  /// Return the set relation of indexVars in lhs and rhd
+  /// Return the set relation of indexVars in lhs and rhs
   IndexSetRel getIndexSetRel() const;
 
   typedef AssignmentNode Node;

src/index_notation/transformations.cpp

@@ -433,13 +433,16 @@ IndexStmt Precompute::apply(IndexStmt stmt, std::string* reason) const {
             );
 
             IndexSetRel rel = a.getIndexSetRel();
+            /// The reduceOp depends on the relation between indexVar sets of rhs and lhs. For rcl and inter, reduceOp
+            /// must be +=. For lcr, reduceOp must be =. For none and equal, reduceOp can't be decided at this stage.
             switch (rel) {
-                case none: a = Assignment(a.getLhs(), a.getRhs());break; // =
-                case rcl:  a = Assignment(a.getLhs(), a.getRhs(), Add());break; // +=
-                case lcr: a = Assignment(a.getLhs(), a.getRhs());break; // =
-                case inter: a = Assignment(a.getLhs(), a.getRhs(), Add());break; // +=
-                case equal: a = Assignment(a.getLhs(), a.getRhs());break;// = OR +=
-            }return a;
+                case none: a = Assignment(a.getLhs(), a.getRhs());break;
+                case rcl:  a = Assignment(a.getLhs(), a.getRhs(), Add());break;
+                case lcr: a = Assignment(a.getLhs(), a.getRhs());break;
+                case inter: a = Assignment(a.getLhs(), a.getRhs(), Add());break;
+                case equal: a = Assignment(a.getLhs(), a.getRhs());break;
+            }
+            return a;
         }
 
         Assignment getProducerAssignment(TensorVar& ws,
@@ -450,14 +453,15 @@ IndexStmt Precompute::apply(IndexStmt stmt, std::string* reason) const {
 
             auto a = ws(iw_vars) = replace(e, substitutions);
             IndexSetRel rel = a.getIndexSetRel();
+            /// The reduceOp depends on the relation between indexVar sets of rhs and lhs. For rcl and inter, reduceOp
+            /// must be +=. For lcr, reduceOp must be =. For none and equal, reduceOp can't be decided at this stage.
             switch (rel) {
-                case none: a = Assignment(a.getLhs(), a.getRhs());break; // =
-                case rcl:  a = Assignment(a.getLhs(), a.getRhs(), Add());break; // +=
-                case lcr: a = Assignment(a.getLhs(), a.getRhs());break; // =
-                case inter: a = Assignment(a.getLhs(), a.getRhs(), Add());break; // +=
-                case equal: a = Assignment(a.getLhs(), a.getRhs());break;// = OR +=
+                case none: a = Assignment(a.getLhs(), a.getRhs());break;
+                case rcl:  a = Assignment(a.getLhs(), a.getRhs(), Add());break;
+                case lcr: a = Assignment(a.getLhs(), a.getRhs());break;
+                case inter: a = Assignment(a.getLhs(), a.getRhs(), Add());break;
+                case equal: a = Assignment(a.getLhs(), a.getRhs());break;
             }
-
             return a;
         }
 
@@ -565,6 +569,9 @@ IndexStmt Precompute::apply(IndexStmt stmt, std::string* reason) const {
         }
     };
 
+    /// RedundantVisitor uses Forall Context to determine reduceOp for none and equal.
+    /// We assume += is used if a workspace is accessed multiple times, otherwise =.
+    /// Forall Context describes the related indexVars of the given indexVar at a specific stage. `ctx_stack` implements such concept.
     struct RedundantVisitor: public IndexNotationVisitor {
         using IndexNotationVisitor::visit;
 
@@ -608,6 +615,9 @@ IndexStmt Precompute::apply(IndexStmt stmt, std::string* reason) const {
             Assignment a(node->lhs, node->rhs, node->op);
             vector<IndexVar> freeVars = a.getLhs().getIndexVars();
             set<IndexVar> seen(freeVars.begin(), freeVars.end());
+
+            /// For equal, if some indexVar in lhs has sibling in ctx stack, reduceOp will be +=.
+            bool is_equal = (a.getIndexSetRel() == equal);
             bool has_sibling = false;
             match(a.getRhs(),
                   std::function<void(const AccessNode*)>([&](const AccessNode* op) {
@@ -619,11 +629,12 @@ IndexStmt Precompute::apply(IndexStmt stmt, std::string* reason) const {
                           }
                       }
                   }));
-            bool is_equal = (a.getIndexSetRel() == equal);
-            bool is_none = (a.getIndexSetRel() == none);
             if (is_equal && has_sibling) {
                 to_change.push_back(a);
             }
+
+            /// For none, if ctx_stack except the top contains indexVars in lhs, reduceOp will be +=.
+            bool is_none = (a.getIndexSetRel() == none);
             bool has_outside = true;
             for (auto & var : seen) {
                 for (auto &svar: ctx_stack) {

test/tests-scheduling-eval.cpp

@@ -78,7 +78,6 @@ IndexStmt scheduleSpGEMMCPU(IndexStmt stmt, bool doPrecompute) {
     stmt = stmt.precompute(assign.getRhs(), j, j, w);
   }
   stmt = stmt.assemble(result, AssembleStrategy::Insert, true);
-  //stmt = stmt.assemble(result, AssembleStrategy::Append, true);
   auto qi_stmt = stmt.as<Assemble>().getQueries();
   IndexVar qi;
   if (isa<Where>(qi_stmt)) {

test/tests-workspaces.cpp

@@ -594,6 +594,7 @@ TEST(workspaces, tile_dotProduct_3) {
 
 
   stmt = stmt.concretize();
+
   A.compile(stmt);
   A.assemble();
   A.compute();