Merge pull request tensor-compiler#391 from tensor-compiler/assembly-v2

Partial support for parallel assembly with compressed formats

Author: stephenchouca

apps/tensor_times_vector/CMakeLists.txt

@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 2.8)
+cmake_minimum_required(VERSION 2.8.12)
 project(tensor_times_vector)
 
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11")

include/taco/codegen/module.h

@@ -5,6 +5,7 @@
 #include <vector>
 #include <string>
 #include <utility>
+#include <random>
 
 #include "taco/target.h"
 #include "taco/ir/ir.h"
@@ -80,6 +81,10 @@ class Module {
 
   void setJITLibname();
   void setJITTmpdir();
+
+  static std::string chars;
+  static std::default_random_engine gen;
+  static std::uniform_int_distribution<int> randint;
 };
 
 } // namespace ir

include/taco/format.h

@@ -12,6 +12,8 @@ namespace taco {
 class ModeFormat;
 class ModeFormatPack;
 class ModeFormatImpl;
+class AttrQuery;
+class IndexVar;
 
 
 /// A Format describes the data layout of a tensor, and the sparse index data
@@ -135,6 +137,16 @@ class ModeFormat {
   bool hasInsert() const;
   bool hasAppend() const;
 
+  /// Returns true if a mode format has ungrouped insertion functions with 
+  /// specific attributes, false otherwise
+  bool hasSeqInsertEdge() const;
+  bool hasInsertCoord() const;
+  bool isYieldPosPure() const;
+
+  std::vector<AttrQuery> getAttrQueries(
+      std::vector<IndexVar> parentCoords, 
+      std::vector<IndexVar> childCoords) const;
+
   /// Returns true if mode format is defined, false otherwise. An undefined mode
   /// type can be used to indicate a mode whose format is not (yet) known.
   bool defined() const;

include/taco/index_notation/index_notation.h

@@ -20,7 +20,6 @@
 #include "taco/index_notation/intrinsic.h"
 #include "taco/index_notation/index_notation_nodes_abstract.h"
 #include "taco/ir_tags.h"
-#include "taco/lower/iterator.h"
 #include "taco/index_notation/provenance_graph.h"
 
 namespace taco {
@@ -57,6 +56,7 @@ struct YieldNode;
 struct ForallNode;
 struct WhereNode;
 struct SequenceNode;
+struct AssembleNode;
 struct MultiNode;
 struct SuchThatNode;
 
@@ -224,16 +224,22 @@ class Access : public IndexExpr {
   Access() = default;
   Access(const Access&) = default;
   Access(const AccessNode*);
-  Access(const TensorVar& tensorVar,
-         const std::vector<IndexVar>& indices={},
-         const std::map<int, std::shared_ptr<IndexVarIterationModifier>>& modifiers={});
+  Access(const TensorVar& tensorVar, const std::vector<IndexVar>& indices={}, 
+         const std::map<int, std::shared_ptr<IndexVarIterationModifier>>& modifiers={},
+         bool isAccessingStructure=false);
 
   /// Return the Access expression's TensorVar.
   const TensorVar &getTensorVar() const;
 
   /// Returns the index variables used to index into the Access's TensorVar.
   const std::vector<IndexVar>& getIndexVars() const;
 
+  /// Returns whether access expression returns sparsity pattern of tensor.
+  /// If true, the access expression returns 1 for every physically stored 
+  /// component. If false, the access expression returns the value that is  
+  /// stored for each corresponding component.
+  bool isAccessingStructure() const;
+
   /// hasWindowedModes returns true if any accessed modes are windowed.
   bool hasWindowedModes() const;
 
@@ -675,6 +681,8 @@ class IndexStmt : public util::IntrusivePtr<const IndexStmtNode> {
   /// integer number of iterations
   /// Preconditions: unrollFactor is a positive nonzero integer
   IndexStmt unroll(IndexVar i, size_t unrollFactor) const;
+
+  IndexStmt assemble(TensorVar result, AssembleStrategy strategy) const;
 };
 
 /// Check if two index statements are isomorphic.
@@ -818,6 +826,27 @@ class Sequence : public IndexStmt {
 Sequence sequence(IndexStmt definition, IndexStmt mutation);
 
 
+class Assemble : public IndexStmt {
+public:
+  typedef std::map<TensorVar,std::vector<std::vector<TensorVar>>> AttrQueryResults;
+
+  Assemble() = default;
+  Assemble(const AssembleNode*);
+  Assemble(IndexStmt queries, IndexStmt compute, AttrQueryResults results);
+
+  IndexStmt getQueries() const;
+  IndexStmt getCompute() const;
+
+  const AttrQueryResults& getAttrQueryResults() const;
+
+  typedef AssembleNode Node;
+};
+
+/// Create an assemble index statement.
+Assemble assemble(IndexStmt queries, IndexStmt compute, 
+                  Assemble::AttrQueryResults results);
+
+
 /// A multi statement has two statements that are executed separately, and let
 /// us compute more than one tensor in a concrete index notation statement.
 class Multi : public IndexStmt {
@@ -1095,10 +1124,18 @@ std::vector<TensorVar> getArguments(IndexStmt stmt);
 /// Returns the temporaries in the index statement, in the order they appear.
 std::vector<TensorVar> getTemporaries(IndexStmt stmt);
 
+/// Returns the attribute query results in the index statement, in the order 
+/// they appear.
+std::vector<TensorVar> getAttrQueryResults(IndexStmt stmt);
+
 // [Olivia]
 /// Returns the temporaries in the index statement, in the order they appear.
 std::map<Forall, Where> getTemporaryLocations(IndexStmt stmt);
 
+/// Returns the results in the index statement that should be assembled by 
+/// ungrouped insertion.
+std::vector<TensorVar> getAssembledByUngroupedInsertion(IndexStmt stmt);
+
 /// Returns the tensors in the index statement.
 std::vector<TensorVar> getTensorVars(IndexStmt stmt);
 
@@ -1123,6 +1160,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_nodes.h

@@ -71,10 +71,12 @@ struct IndexSet : IndexVarIterationModifier {
 };
 
 struct AccessNode : public IndexExprNode {
-  AccessNode(TensorVar tensorVar,
-             const std::vector<IndexVar> &indices,
-             const std::map<int, std::shared_ptr<IndexVarIterationModifier>> &modifiers = {})
-      : IndexExprNode(tensorVar.getType().getDataType()), tensorVar(tensorVar), indexVars(indices) {
+  AccessNode(TensorVar tensorVar, const std::vector<IndexVar>& indices, 
+             const std::map<int, std::shared_ptr<IndexVarIterationModifier>> &modifiers,
+             bool isAccessingStructure)
+      : IndexExprNode(isAccessingStructure ? Bool : tensorVar.getType().getDataType()), 
+        tensorVar(tensorVar), indexVars(indices), 
+        isAccessingStructure(isAccessingStructure) {
     // Unpack the input modifiers into the appropriate maps for each mode.
     for (auto &it : modifiers) {
       IndexVarIterationModifier::match(it.second, [&](std::shared_ptr<AccessWindow> w) {
@@ -108,11 +110,14 @@ struct AccessNode : public IndexExprNode {
   std::vector<IndexVar> indexVars;
   std::map<int, AccessWindow> windowedModes;
   std::map<int, IndexSet> indexSetModes;
+  bool isAccessingStructure;
 
 protected:
   /// Initialize an AccessNode with just a TensorVar. If this constructor is used,
   /// then indexVars must be set afterwards.
-  explicit AccessNode(TensorVar tensorVar) : IndexExprNode(tensorVar.getType().getDataType()), tensorVar(tensorVar) {}
+  explicit AccessNode(TensorVar tensorVar) : 
+      IndexExprNode(tensorVar.getType().getDataType()), 
+      tensorVar(tensorVar), isAccessingStructure(false) {}
 };
 
 struct LiteralNode : public IndexExprNode {
@@ -360,6 +365,20 @@ struct SequenceNode : public IndexStmtNode {
   IndexStmt mutation;
 };
 
+struct AssembleNode : public IndexStmtNode {
+  AssembleNode(IndexStmt queries, IndexStmt compute, 
+               Assemble::AttrQueryResults results)
+      : queries(queries), compute(compute), results(results) {}
+
+  void accept(IndexStmtVisitorStrict* v) const {
+    v->visit(this);
+  }
+
+  IndexStmt queries;
+  IndexStmt compute;
+  Assemble::AttrQueryResults results;
+};
+
 
 /// Returns true if expression e is of type E.
 template <typename E>

include/taco/index_notation/index_notation_printer.h

@@ -35,6 +35,7 @@ class IndexNotationPrinter : public IndexNotationVisitorStrict {
   void visit(const WhereNode*);
   void visit(const MultiNode*);
   void visit(const SequenceNode*);
+  void visit(const AssembleNode*);
   void visit(const SuchThatNode*);
 
 private:

include/taco/index_notation/index_notation_rewriter.h

@@ -56,6 +56,7 @@ class IndexStmtRewriterStrict : public IndexStmtVisitorStrict {
   virtual void visit(const ForallNode* op) = 0;
   virtual void visit(const WhereNode* op) = 0;
   virtual void visit(const SequenceNode* op) = 0;
+  virtual void visit(const AssembleNode* op) = 0;
   virtual void visit(const MultiNode* op) = 0;
   virtual void visit(const SuchThatNode* op) = 0;
 };
@@ -101,6 +102,7 @@ class IndexNotationRewriter : public IndexNotationRewriterStrict {
   virtual void visit(const ForallNode* op);
   virtual void visit(const WhereNode* op);
   virtual void visit(const SequenceNode* op);
+  virtual void visit(const AssembleNode* op);
   virtual void visit(const MultiNode* op);
   virtual void visit(const SuchThatNode* op);
 };
@@ -126,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

include/taco/index_notation/index_notation_visitor.h

@@ -31,6 +31,7 @@ struct ForallNode;
 struct WhereNode;
 struct MultiNode;
 struct SequenceNode;
+struct AssembleNode;
 struct SuchThatNode;
 
 /// Visit the nodes in an expression.  This visitor provides some type safety
@@ -65,6 +66,7 @@ class IndexStmtVisitorStrict {
   virtual void visit(const ForallNode*) = 0;
   virtual void visit(const WhereNode*) = 0;
   virtual void visit(const SequenceNode*) = 0;
+  virtual void visit(const AssembleNode*) = 0;
   virtual void visit(const MultiNode*) = 0;
   virtual void visit(const SuchThatNode*) = 0;
 };
@@ -107,6 +109,7 @@ class IndexNotationVisitor : public IndexNotationVisitorStrict {
   virtual void visit(const ForallNode* node);
   virtual void visit(const WhereNode* node);
   virtual void visit(const SequenceNode* node);
+  virtual void visit(const AssembleNode* node);
   virtual void visit(const MultiNode* node);
   virtual void visit(const SuchThatNode* node);
 };
@@ -176,6 +179,7 @@ class Matcher : public IndexNotationVisitor {
   RULE(WhereNode)
   RULE(MultiNode)
   RULE(SequenceNode)
+  RULE(AssembleNode)
   RULE(SuchThatNode)
 };
 

include/taco/index_notation/provenance_graph.h

@@ -1,6 +1,8 @@
 #ifndef TACO_PROVENANCE_GRAPH_H
 #define TACO_PROVENANCE_GRAPH_H
 
+#include "taco/lower/iterator.h"
+
 namespace taco {
 struct IndexVarRelNode;
 enum IndexVarRelType {UNDEFINED, SPLIT, DIVIDE, POS, FUSE, BOUND, PRECOMPUTE};

include/taco/index_notation/transformations.h

@@ -21,6 +21,7 @@ class ForAllReplace;
 class AddSuchThatPredicates;
 class Parallelize;
 class TopoReorder;
+class SetAssembleStrategy;
 
 /// A transformation is an optimization that transforms a statement in the
 /// concrete index notation into a new statement that computes the same result
@@ -34,6 +35,7 @@ class Transformation {
   Transformation(Parallelize);
   Transformation(TopoReorder);
   Transformation(AddSuchThatPredicates);
+  Transformation(SetAssembleStrategy);
 
   IndexStmt apply(IndexStmt stmt, std::string *reason = nullptr) const;
 
@@ -108,6 +110,7 @@ class Precompute : public TransformationInterface {
 /// Print a precompute command.
 std::ostream &operator<<(std::ostream &, const Precompute &);
 
+
 /// Replaces all occurrences of directly nested forall nodes of pattern with
 /// directly nested loops of replacement
 class ForAllReplace : public TransformationInterface {
@@ -129,6 +132,10 @@ class ForAllReplace : public TransformationInterface {
   std::shared_ptr<Content> content;
 };
 
+/// Print a ForAllReplace command.
+std::ostream &operator<<(std::ostream &, const ForAllReplace &);
+
+
 /// Adds a SuchThat node if it does not exist and adds the given IndexVarRels
 class AddSuchThatPredicates : public TransformationInterface {
 public:
@@ -147,6 +154,9 @@ class AddSuchThatPredicates : public TransformationInterface {
   std::shared_ptr<Content> content;
 };
 
+std::ostream& operator<<(std::ostream&, const AddSuchThatPredicates&);
+
+
 /// The parallelize optimization tags a Forall as parallelized
 /// after checking for preconditions
 class Parallelize : public TransformationInterface {
@@ -169,13 +179,28 @@ class Parallelize : public TransformationInterface {
   std::shared_ptr<Content> content;
 };
 
-/// Print a ForAllReplace command.
-std::ostream &operator<<(std::ostream &, const ForAllReplace &);
-
 /// Print a parallelize command.
 std::ostream& operator<<(std::ostream&, const Parallelize&);
 
-std::ostream& operator<<(std::ostream&, const AddSuchThatPredicates&);
+
+class SetAssembleStrategy : public TransformationInterface {
+public:
+  SetAssembleStrategy(TensorVar result, AssembleStrategy strategy);
+
+  TensorVar getResult() const;
+  AssembleStrategy getAssembleStrategy() const;
+
+  IndexStmt apply(IndexStmt stmt, std::string *reason = nullptr) const;
+
+  void print(std::ostream &os) const;
+
+private:
+  struct Content;
+  std::shared_ptr<Content> content;
+};
+
+/// Print a SetAssembleStrategy command.
+std::ostream &operator<<(std::ostream &, const SetAssembleStrategy&);
 
 // Autoscheduling functions
 
@@ -207,5 +232,6 @@ IndexStmt scalarPromote(IndexStmt stmt);
  * 1. The result is a is scattered into but does not support random insert.
  */
 IndexStmt insertTemporaries(IndexStmt stmt);
+
 }
 #endif