Enable parallelization of forall statements with results assembled by ungrouped insertion

Author: stephenchouca

include/taco/index_notation/index_notation.h

@@ -1011,6 +1011,10 @@ std::vector<ir::Expr> createVars(const std::vector<TensorVar>& tensorVars,
                                  std::map<TensorVar, ir::Expr>* vars, 
                                  bool isParameter=false);
 
+/// Convert index notation tensor variables in the index statement to IR 
+/// pointer variables.
+std::map<TensorVar,ir::Expr> createIRTensorVars(IndexStmt stmt);
+
 
 /// Simplify an index expression by setting the zeroed Access expressions to
 /// zero and then propagating and removing zeroes.

include/taco/index_notation/index_notation_rewriter.h

@@ -128,5 +128,9 @@ IndexStmt replace(IndexStmt stmt,
 IndexStmt replace(IndexStmt stmt,
                   const std::map<TensorVar,TensorVar>& substitutions);
 
+/// Rewrites the statement to replace an index variable with a new variable.
+IndexStmt replace(IndexStmt stmt,
+                  const std::map<IndexVar,IndexVar>& substitutions);
+
 }
 #endif

src/index_notation/index_notation.cpp

@@ -2653,6 +2653,25 @@ vector<ir::Expr> createVars(const vector<TensorVar>& tensorVars,
   return irVars;
 }
 
+std::map<TensorVar,ir::Expr> createIRTensorVars(IndexStmt stmt)
+{
+  std::map<TensorVar,ir::Expr> tensorVars;
+
+  // Create result and parameter variables
+  vector<TensorVar> results = getResults(stmt);
+  vector<TensorVar> arguments = getArguments(stmt);
+  vector<TensorVar> temporaries = getTemporaries(stmt);
+
+  // Convert tensor results, arguments and temporaries to IR variables
+  map<TensorVar, ir::Expr> resultVars;
+  vector<ir::Expr> resultsIR = createVars(results, &resultVars);
+  tensorVars.insert(resultVars.begin(), resultVars.end());
+  vector<ir::Expr> argumentsIR = createVars(arguments, &tensorVars);
+  vector<ir::Expr> temporariesIR = createVars(temporaries, &tensorVars);
+
+  return tensorVars;
+}
+
 struct Zero : public IndexNotationRewriterStrict {
 public:
   Zero(const set<Access>& zeroed) : zeroed(zeroed) {}

src/index_notation/index_notation_rewriter.cpp

@@ -349,7 +349,29 @@ struct ReplaceIndexVars : public IndexNotationRewriter {
     }
   }
 
-  // TODO: Replace in assignments
+  void visit(const AssignmentNode* op) {
+    IndexExpr rhs = rewrite(op->rhs);
+    Access lhs = to<Access>(rewrite(op->lhs));
+    if (rhs == op->rhs && lhs == op->lhs) {
+      stmt = op;
+    }
+    else {
+      stmt = new AssignmentNode(lhs, rhs, op->op);
+    }
+  }
+
+  void visit(const ForallNode* op) {
+    IndexStmt s = rewrite(op->stmt);
+    IndexVar iv = util::contains(substitutions, op->indexVar) 
+                ? substitutions.at(op->indexVar) : op->indexVar;
+    if (s == op->stmt && iv == op->indexVar) {
+      stmt = op;
+    }
+    else {
+      stmt = new ForallNode(iv, s, op->parallel_unit, op->output_race_strategy, 
+                            op->unrollFactor);
+    }
+  }
 };
 
 struct ReplaceTensorVars : public IndexNotationRewriter {
@@ -404,4 +426,9 @@ IndexStmt replace(IndexStmt stmt,
   return ReplaceTensorVars(substitutions).rewrite(stmt);
 }
 
+IndexStmt replace(IndexStmt stmt,
+                  const std::map<IndexVar,IndexVar>& substitutions) {
+  return ReplaceIndexVars(substitutions).rewrite(stmt);
+}
+
 }

src/index_notation/transformations.cpp

@@ -483,20 +483,27 @@ IndexStmt Parallelize::apply(IndexStmt stmt, std::string* reason) const {
 
     Parallelize parallelize;
     ProvenanceGraph provGraph;
+    map<TensorVar,ir::Expr> tensorVars;
+    vector<ir::Expr> assembledByUngroupedInsert;
     set<IndexVar> definedIndexVars;
     set<ParallelUnit> parentParallelUnits;
     std::string reason = "";
 
     IndexStmt rewriteParallel(IndexStmt stmt) {
       provGraph = ProvenanceGraph(stmt);
+      tensorVars = createIRTensorVars(stmt);
+      assembledByUngroupedInsert.clear();
+      for (const auto& result : getAssembledByUngroupedInsertion(stmt)) {
+        assembledByUngroupedInsert.push_back(tensorVars[result]);
+      }
       return rewrite(stmt);
     }
 
     void visit(const ForallNode* node) {
       Forall foralli(node);
       IndexVar i = parallelize.geti();
 
-      Iterators iterators(foralli);
+      Iterators iterators(foralli, tensorVars);
       definedIndexVars.insert(foralli.getIndexVar());
       MergeLattice lattice = MergeLattice::make(foralli, iterators, provGraph, definedIndexVars);
       // Precondition 3: No parallelization of variables under a reduction
@@ -517,6 +524,9 @@ IndexStmt Parallelize::apply(IndexStmt stmt, std::string* reason) const {
 
         // Precondition 2: Every result iterator must have insert capability
         for (Iterator iterator : lattice.results()) {
+          if (util::contains(assembledByUngroupedInsert, iterator.getTensor())) {
+            continue;
+          }
           while (true) {
             if (!iterator.hasInsert()) {
               reason = "Precondition failed: The output tensor must allow inserts";
@@ -614,6 +624,34 @@ IndexStmt Parallelize::apply(IndexStmt stmt, std::string* reason) const {
       }
       IndexNotationRewriter::visit(node);
     }
+
+    void visit(const AssembleNode* op) {
+      IndexVar i = parallelize.geti();
+      IndexStmt queries = util::contains(op->queries.getIndexVars(), i) 
+                        ? rewrite(op->queries) : op->queries;
+      IndexStmt compute = util::contains(op->compute.getIndexVars(), i) 
+                        ? rewrite(op->compute) : op->compute;
+      if (queries == op->queries && compute == op->compute) {
+        stmt = op;
+      }
+      else {
+        stmt = new AssembleNode(queries, compute, op->results);
+      }
+    }
+
+    void visit(const WhereNode* op) {
+      IndexVar i = parallelize.geti();
+      IndexStmt producer = util::contains(op->producer.getIndexVars(), i) 
+                        ? rewrite(op->producer) : op->producer;
+      IndexStmt consumer = util::contains(op->consumer.getIndexVars(), i) 
+                        ? rewrite(op->consumer) : op->consumer;
+      if (producer == op->producer && consumer == op->consumer) {
+        stmt = op;
+      }
+      else {
+        stmt = new WhereNode(consumer, producer);
+      }
+    }
   };
 
   ParallelizeRewriter rewriter;
@@ -663,6 +701,12 @@ AssembleStrategy SetAssembleStrategy::getAssembleStrategy() const {
 IndexStmt SetAssembleStrategy::apply(IndexStmt stmt, string* reason) const {
   INIT_REASON(reason);
 
+  std::map<IndexVar,IndexVar> ivReplacements;
+  for (const auto& indexVar : getIndexVars(stmt)) {
+    ivReplacements[indexVar] = IndexVar("q" + indexVar.getName());
+  }
+  IndexStmt loweredQueries = replace(stmt, ivReplacements);
+
   // Tracks all tensors that correspond to attribute query results or that are 
   // used to compute attribute queries
   std::set<TensorVar> insertedResults;  
@@ -812,8 +856,8 @@ IndexStmt SetAssembleStrategy::apply(IndexStmt stmt, string* reason) const {
       expr = op;
     }
   };
-  IndexStmt loweredQueries = 
-      LowerAttrQuery(queryResults, insertedResults).lower(stmt);
+  loweredQueries = 
+      LowerAttrQuery(queryResults, insertedResults).lower(loweredQueries);
   std::cout << loweredQueries << std::endl;
 
   struct ReduceToAssign : public IndexNotationRewriter {
@@ -956,6 +1000,7 @@ IndexStmt SetAssembleStrategy::apply(IndexStmt stmt, string* reason) const {
   std::cout << loweredQueries << std::endl;
 
   //loweredQueries = parallelizeOuterLoop(loweredQueries);
+  //stmt = parallelizeOuterLoop(stmt);
   //std::cout << loweredQueries << std::endl;
 
   return Assemble(loweredQueries, stmt, queryResults);

src/lower/iterator.cpp

@@ -426,26 +426,6 @@ Iterators::Iterators()
 }
 
 
-static std::map<TensorVar, ir::Expr> createIRTensorVars(IndexStmt stmt)
-{
-  std::map<TensorVar, ir::Expr> tensorVars;
-
-  // Create result and parameter variables
-  vector<TensorVar> results = getResults(stmt);
-  vector<TensorVar> arguments = getArguments(stmt);
-  vector<TensorVar> temporaries = getTemporaries(stmt);
-
-  // Convert tensor results, arguments and temporaries to IR variables
-  map<TensorVar, Expr> resultVars;
-  vector<Expr> resultsIR = createVars(results, &resultVars);
-  tensorVars.insert(resultVars.begin(), resultVars.end());
-  vector<Expr> argumentsIR = createVars(arguments, &tensorVars);
-  vector<Expr> temporariesIR = createVars(temporaries, &tensorVars);
-
-  return tensorVars;
-}
-
-
 Iterators::Iterators(IndexStmt stmt) : Iterators(stmt, createIRTensorVars(stmt))
 {
 }