Merge pull request tensor-compiler#339 from RawnH/workspace_reuse

Workspace reuse

include/taco/index_notation/index_notation.h

@@ -953,6 +953,10 @@ std::vector<TensorVar> getArguments(IndexStmt stmt);
 /// Returns the temporaries in the index statement, in the order they appear.
 std::vector<TensorVar> getTemporaries(IndexStmt stmt);
 
+// [Olivia]
+/// Returns the temporaries in the index statement, in the order they appear.
+std::map<Forall, Where> getTemporaryLocations(IndexStmt stmt);
+
 /// Returns the tensors in the index statement.
 std::vector<TensorVar> getTensorVars(IndexStmt stmt);
 

include/taco/ir/ir.h

@@ -65,7 +65,8 @@ enum class IRNodeType {
  BlankLine,
  Print,
  GetProperty,
- Break
+ Break,
+ Sort
 };
 
 enum class TensorProperty {
@@ -725,6 +726,13 @@ struct Break : public StmtNode<Break> {
  static const IRNodeType _type_info = IRNodeType::Break;
 };
 
+struct Sort : public StmtNode<Sort> {
+ std::vector<Expr> args;
+ static Stmt make(std::vector<Expr> args);
+
+ static const IRNodeType _type_info = IRNodeType::Sort;
+};
+
 /** A print statement.
  * Takes in a printf-style format string and Exprs to pass
  * for the values.

include/taco/ir/ir_printer.h

@@ -68,6 +68,7 @@ class IRPrinter : public IRVisitorStrict {
  virtual void visit(const Break*);
  virtual void visit(const Print*);
  virtual void visit(const GetProperty*);
+ virtual void visit(const Sort*);
 
  std::ostream &stream;
  int indent;

include/taco/ir/ir_rewriter.h

@@ -68,6 +68,7 @@ class IRRewriter : public IRVisitorStrict {
  virtual void visit(const Break* op);
  virtual void visit(const Print* op);
  virtual void visit(const GetProperty* op);
+ virtual void visit(const Sort *op);
 };
 
 }}

include/taco/ir/ir_visitor.h

@@ -48,6 +48,7 @@ struct BlankLine;
 struct Break;
 struct Print;
 struct GetProperty;
+struct Sort;
 
 /// Extend this class to visit every node in the IR.
 class IRVisitorStrict {
@@ -98,6 +99,7 @@ class IRVisitorStrict {
  virtual void visit(const Break*) = 0;
  virtual void visit(const Print*) = 0;
  virtual void visit(const GetProperty*) = 0;
+ virtual void visit(const Sort*) = 0;
 };
 
 
@@ -151,6 +153,7 @@ class IRVisitor : public IRVisitorStrict {
  virtual void visit(const Break* op);
  virtual void visit(const Print* op);
  virtual void visit(const GetProperty* op);
+ virtual void visit(const Sort* op);
 };
 
 }}

include/taco/lower/lowerer_impl.h

@@ -81,6 +81,16 @@ class LowererImpl : public util::Uncopyable {
  std::set<Access> reducedAccesses,
  ir::Stmt recoveryStmt);
 
+ /// Lower a forall that iterates over all the coordinates in the forall index
+ /// var's dimension, and locates tensor positions from the locate iterators.
+ virtual ir::Stmt lowerForallDenseAcceleration(Forall forall,
+ std::vector<Iterator> locaters,
+ std::vector<Iterator> inserters,
+ std::vector<Iterator> appenders,
+ std::set<Access> reducedAccesses,
+ ir::Stmt recoveryStmt);
+
+
  /// Lower a forall that iterates over the coordinates in the iterator, and
  /// locates tensor positions from the locate iterators.
  virtual ir::Stmt lowerForallCoordinate(Forall forall, Iterator iterator,
@@ -333,17 +343,29 @@ class LowererImpl : public util::Uncopyable {
  ir::Stmt codeToInitializeIteratorVars(std::vector<Iterator> iterators, std::vector<Iterator> rangers, std::vector<Iterator> mergers, ir::Expr coord, IndexVar coordinateVar);
  ir::Stmt codeToInitializeIteratorVar(Iterator iterator, std::vector<Iterator> iterators, std::vector<Iterator> rangers, std::vector<Iterator> mergers, ir::Expr coordinate, IndexVar coordinateVar);
 
+ /// Returns true iff the temporary used in the where statement is dense and sparse iteration over that
+ /// temporary can be automaticallty supported by the compiler.
+ bool canAccelerateDenseTemp(Where where);
+
+ /// Initializes a temporary workspace
+ std::vector<ir::Stmt> codeToInitializeTemporary(Where where);
+
+ /// Gets the size of a temporary tensorVar in the where statement
+ ir::Expr getTemporarySize(Where where);
+
+ /// Initializes helper arrays to give dense workspaces sparse acceleration
+ std::vector<ir::Stmt> codeToInitializeDenseAcceleratorArrays(Where where);
 
  /// Recovers a derived indexvar from an underived variable.
  ir::Stmt codeToRecoverDerivedIndexVar(IndexVar underived, IndexVar indexVar, bool emitVarDecl);
 
-  /// Conditionally increment iterator position variables.
+ /// Conditionally increment iterator position variables.
  ir::Stmt codeToIncIteratorVars(ir::Expr coordinate, IndexVar coordinateVar,
  std::vector<Iterator> iterators, std::vector<Iterator> mergers);
 
  ir::Stmt codeToLoadCoordinatesFromPosIterators(std::vector<Iterator> iterators, bool declVars);
 
-  /// Create statements to append coordinate to result modes.
+ /// Create statements to append coordinate to result modes.
  ir::Stmt appendCoordinate(std::vector<Iterator> appenders, ir::Expr coord);
 
  /// Create statements to append positions to result modes.
@@ -363,6 +385,9 @@ class LowererImpl : public util::Uncopyable {
  int markAssignsAtomicDepth = 0;
  ParallelUnit atomicParallelUnit;
 
+ /// Map used to hoist temporary workspace initialization
+ std::map<Forall, Where> temporaryInitialization;
+
  /// Map from tensor variables in index notation to variables in the IR
  std::map<TensorVar, ir::Expr> tensorVars;
 
@@ -371,6 +396,15 @@ class LowererImpl : public util::Uncopyable {
  };
  std::map<TensorVar, TemporaryArrays> temporaryArrays;
 
+ /// Map form temporary to indexList var if accelerating dense workspace
+ std::map<TensorVar, ir::Expr> tempToIndexList;
+
+ /// Map form temporary to indexListSize if accelerating dense workspace
+ std::map<TensorVar, ir::Expr> tempToIndexListSize;
+
+ /// Map form temporary to bitGuard var if accelerating dense workspace
+ std::map<TensorVar, ir::Expr> tempToBitGuard;
+
  /// Map from result tensors to variables tracking values array capacity.
  std::map<ir::Expr, ir::Expr> capacityVars;
 

src/codegen/codegen_c.cpp

@@ -182,14 +182,15 @@ class CodeGen_C::FindVars : public IRVisitor {
 
  virtual void visit(const Var *op) {
  if (varMap.count(op) == 0) {
- varMap[op] = codeGen->genUniqueName(op->name);
+ varMap[op] = op->is_ptr? op->name : codeGen->genUniqueName(op->name);
  }
  }
 
  virtual void visit(const VarDecl *op) {
  if (!util::contains(localVars, op->var)) {
  localVars.push_back(op->var);
  }
+ op->var.accept(this);
  op->rhs.accept(this);
  }
 

src/codegen/codegen_cuda.cpp

@@ -240,7 +240,7 @@ class CodeGen_CUDA::FindVars : public IRVisitor {
 
  virtual void visit(const Var *op) {
  if (varMap.count(op) == 0 && !inBlock) {
- varMap[op] = codeGen->genUniqueName(op->name);
+ varMap[op] = op->is_ptr? op->name : codeGen->genUniqueName(op->name);
  }
  }
 

src/index_notation/index_notation.cpp

@@ -2118,8 +2118,23 @@ bool isConcreteNotation(IndexStmt stmt, std::string* reason) {
  return;
  }
 
+ // Handles derived vars on RHS with underived vars on LHS.
+ Assignment assignPtrWrapper = Assignment(op);
+ std::vector<IndexVar> possibleReductionVars = assignPtrWrapper.getReductionVars();
+ std::vector<IndexVar> freeVars = assignPtrWrapper.getFreeVars();
+ std::set<IndexVar> freeVarsSet(freeVars.begin(), freeVars.end());
+
+ int numReductionVars = 0;
+ for(const auto& reductionVar : possibleReductionVars) {
+ std::vector<IndexVar> underivedParents = provGraph.getUnderivedAncestors(reductionVar);
+ for(const auto& parent : underivedParents) {
+ if(!util::contains(freeVarsSet, parent)) {
+ ++numReductionVars;
+ }
+ }
+ }
  // allow introducing precompute loops where we set a temporary to values instead of +=
- if (Assignment(op).getReductionVars().size() > 0 &&
+ if (numReductionVars > 0 &&
  op->op == IndexExpr() && !inWhereProducer) {
  *reason = "reduction variables in concrete notation must be dominated "
  "by compound assignments (such as +=)";
@@ -2342,6 +2357,22 @@ vector<TensorVar> getArguments(IndexStmt stmt) {
  return result;
 }
 
+std::map<Forall, Where> getTemporaryLocations(IndexStmt stmt) {
+ map<Forall, Where> temporaryLocs;
+ Forall f = Forall();
+ match(stmt,
+ function<void(const ForallNode*, Matcher*)>([&](const ForallNode* op, Matcher* ctx) {
+ f = op;
+ ctx->match(op->stmt);
+ }),
+ function<void(const WhereNode*, Matcher*)>([&](const WhereNode* w, Matcher* ctx) {
+ if (!(f == IndexStmt()))
+ temporaryLocs.insert({f, Where(w)});
+ })
+ );
+ return temporaryLocs;
+}
+
 std::vector<TensorVar> getTemporaries(IndexStmt stmt) {
  vector<TensorVar> temporaries;
  bool firstAssignment = true;

src/index_notation/index_notation_printer.cpp

@@ -81,7 +81,12 @@ void IndexNotationPrinter::visit(const NegNode* op) {
  Precedence precedence = Precedence::NEG;
  bool parenthesize = precedence > parentPrecedence;
  parentPrecedence = precedence;
- os << "-";
+ if(op->getDataType().isBool()) {
+ os << "!";
+ } else {
+ os << "-";
+ }
+
  if (parenthesize) {
  os << "(";
  }

src/index_notation/transformations.cpp

@@ -1114,17 +1114,19 @@ static IndexStmt optimizeSpMM(IndexStmt stmt) {
  return stmt;
  }
 
+ // I think we can to linear combination of rows as long as there are no permutations in the format and the
+ // level formats are ordered. The i -> k -> j loops should iterate over the data structures without issue.
  TensorVar B = Baccess.getTensorVar();
- if (B.getFormat().getModeFormats()[0].getName() != "dense" ||
- B.getFormat().getModeFormats()[1].getName() != "compressed" ||
+ if (!B.getFormat().getModeFormats()[0].isOrdered() ||
+ !B.getFormat().getModeFormats()[1].isOrdered() ||
  B.getFormat().getModeOrdering()[0] != 0 ||
  B.getFormat().getModeOrdering()[1] != 1) {
  return stmt;
  }
 
  TensorVar C = Caccess.getTensorVar();
- if (C.getFormat().getModeFormats()[0].getName() != "dense" ||
- C.getFormat().getModeFormats()[1].getName() != "compressed" ||
+ if (!C.getFormat().getModeFormats()[0].isOrdered() ||
+ !C.getFormat().getModeFormats()[1].isOrdered() ||
  C.getFormat().getModeOrdering()[0] != 0 ||
  C.getFormat().getModeOrdering()[1] != 1) {
  return stmt;

src/ir/ir.cpp

@@ -817,6 +817,13 @@ Expr GetProperty::make(Expr tensor, TensorProperty property, int mode,
  return gp;
 }
 
+// Sort
+Stmt Sort::make(std::vector<Expr> args) {
+ Sort* sort = new Sort;
+ sort->args = args;
+ return sort;
+}
+
 
 // GetProperty
 Expr GetProperty::make(Expr tensor, TensorProperty property, int mode) {
@@ -953,6 +960,8 @@ template<> void StmtNode<Print>::accept(IRVisitorStrict *v)
  const { v->visit((const Print*)this); }
 template<> void ExprNode<GetProperty>::accept(IRVisitorStrict *v)
  const { v->visit((const GetProperty*)this); }
+template<> void StmtNode<Sort>::accept(IRVisitorStrict *v)
+ const { v->visit((const Sort*)this); }
 
 // printing methods
 std::ostream& operator<<(std::ostream& os, const Stmt& stmt) {

src/ir/ir_printer.cpp

@@ -131,7 +131,11 @@ void IRPrinter::visit(const Var* op) {
 }
 
 void IRPrinter::visit(const Neg* op) {
- stream << "-";
+ if(op->type.isBool()) {
+ stream << "!";
+ } else {
+ stream << "-";
+ }
  parentPrecedence = Precedence::NEG;
  op->a.accept(this);
 }
@@ -575,6 +579,16 @@ void IRPrinter::visit(const GetProperty* op) {
  stream << op->name;
 }
 
+void IRPrinter::visit(const Sort* op) {
+ doIndent();
+ stream << "qsort(";
+ parentPrecedence = Precedence::CALL;
+ acceptJoin(this, stream, op->args, ", ");
+ stream << ");";
+ stream << endl;
+}
+
+
 void IRPrinter::resetNameCounters() {
  // seed the unique names with all C99 keywords
  // from: http:https://en.cppreference.com/w/c/keyword

src/ir/ir_rewriter.cpp

@@ -479,5 +479,23 @@ void IRRewriter::visit(const GetProperty* op) {
  }
 }
 
+void IRRewriter::visit(const Sort* op) {
+ std::vector<Expr> args;
+ bool rewritten = false;
+ for (auto& arg : op->args) {
+ Expr rewrittenArg = rewrite(arg);
+ args.push_back(rewrittenArg);
+ if (rewrittenArg != arg) {
+ rewritten = true;
+ }
+ }
+ if (rewritten) {
+ stmt = Sort::make(args);
+ }
+ else {
+ stmt = op;
+ }
+}
+
 
 }}

src/ir/ir_visitor.cpp

@@ -236,5 +236,10 @@ void IRVisitor::visit(const Print* op) {
  e.accept(this);
 }
 
+void IRVisitor::visit(const Sort* op) {
+ for (auto e: op->args)
+ e.accept(this);
+}
+
 } // namespace ir
 } // namespace taco

src/lower/lower.cpp

@@ -12,6 +12,7 @@
 #include "taco/ir/ir.h"
 #include "taco/ir/simplify.h"
 #include "ir/ir_generators.h"
+#include "taco/ir/ir_printer.h"
 
 #include "taco/lower/lowerer_impl.h"
 #include "taco/lower/iterator.h"