Add acceleration flag

Author: zhang677

include/taco/index_notation/index_notation.h

@@ -778,6 +778,16 @@ class IndexStmt : public util::IntrusivePtr<const IndexStmtNode> {
   IndexStmt assemble(TensorVar result, AssembleStrategy strategy, 
                      bool separately_schedulable = false) const;
 
+  /// The wsaccel primitive specifies the dimensions of a workspace that
+  /// will be accelerated. Acc controls whether acceleration will be applied.
+  /// If accels is empty it means all dimensions should be accelerated.
+  /// Currently, it only supports one-dimension acceleration. Acceleration is used
+  /// by default.
+  ///
+  /// Precondition:
+  /// Workspace can be accessed by the IndexVars in the accels.
+  IndexStmt wsaccel(TensorVar& ws, const std::vector<IndexVar>& accels, bool Acc = true);
+
   /// Casts index statement to specified subtype.
   template <typename SubType>
   SubType as() {
@@ -1156,6 +1166,15 @@ class TensorVar : public util::Comparable<TensorVar> {
   /// Gets the fill value of the tensor variable. May be left undefined.
   const Literal& getFill() const;
 
+  /// Gets the acceleration dimensions
+  const std::vector<IndexVar>& getAccels() const;
+
+  /// Gets the acceleration flag
+  bool getAcc() const;
+
+  /// Set the acceleration dimensions
+  void setAccels(const std::vector<IndexVar>& accels, bool Acc);
+
   /// Set the fill value of the tensor variable
   void setFill(const Literal& fill);
 

src/index_notation/index_notation.cpp

@@ -2048,6 +2048,30 @@ IndexStmt IndexStmt::assemble(TensorVar result, AssembleStrategy strategy,
   return transformed;
 }
 
+IndexStmt IndexStmt::wsaccel(TensorVar& ws, const std::vector<IndexVar>& accels, bool Acc) {
+    if (accels.size() == 0) {
+        ws.setAccels(accels, Acc);
+        return *this;
+    }
+    set<IndexVar> TempVars;
+    match(*this,
+          std::function<void(const WhereNode*)>([&](const WhereNode* where) {
+        auto Temp = getResultAccesses(where->producer).first[0];
+        if (Temp.getTensorVar() == ws) {
+            for (auto i :getIndexVars()){
+                TempVars.insert(i);
+            }
+        }
+    }));
+    for (auto i : accels) {
+        if (TempVars.find(i) == TempVars.end()) {
+            taco_uerror << "No matching indexVars in the Accel";
+        }
+    }
+    ws.setAccels(accels, Acc);
+    return *this;
+}
+
 std::ostream& operator<<(std::ostream& os, const IndexStmt& expr) {
   if (!expr.defined()) return os << "IndexStmt()";
   IndexNotationPrinter printer(os);
@@ -2520,6 +2544,8 @@ struct TensorVar::Content {
   Format format;
   Schedule schedule;
   Literal fill;
+  std::vector<IndexVar> accels;
+  bool Acc;
 };
 
 TensorVar::TensorVar() : content(nullptr) {
@@ -2552,6 +2578,8 @@ TensorVar::TensorVar(const int& id, const string& name, const Type& type, const
   content->type = type;
   content->format = format;
   content->fill = fill.defined()? fill : Literal::zero(type.getDataType());
+  content->accels = std::vector<IndexVar> {};
+  content->Acc = true;
 }
 
 int TensorVar::getId() const {
@@ -2595,6 +2623,19 @@ const Literal& TensorVar::getFill() const {
   return content->fill;
 }
 
+const std::vector<IndexVar>& TensorVar::getAccels() const {
+  return content->accels;
+}
+
+bool TensorVar::getAcc() const {
+  return content->Acc;
+}
+
+void TensorVar::setAccels(const std::vector<IndexVar>& accels, bool Acc) {
+  content->Acc = Acc;
+  content->accels = accels;
+}
+
 void TensorVar::setFill(const Literal &fill) {
   content->fill = fill;
 }

src/lower/lowerer_impl_imperative.cpp

@@ -2274,6 +2274,12 @@ std::pair<bool,bool> LowererImplImperative::canAccelerateDenseTemp(Where where)
   }
 
   TensorVar temporary = where.getTemporary();
+
+  // (0) Acceleration flag is true
+  if (!temporary.getAcc()) {
+    return std::make_pair(false, false);
+  }
+
   // (1) Temporary is dense vector
   if(!isDense(temporary.getFormat()) || temporary.getOrder() != 1) {
     return std::make_pair(false, false);
@@ -2302,9 +2308,6 @@ std::pair<bool,bool> LowererImplImperative::canAccelerateDenseTemp(Where where)
           return resultVar == tempVar[0] ||
                  provGraph.isDerivedFrom(tempVar[0], resultVar);
   });
-  if (resultVars.size() == 0){
-      return std::make_pair(false, false);
-  }
   if (it == resultVars.end()) {
     return std::make_pair(true, false);
   }

test/tests-workspaces.cpp

@@ -546,6 +546,7 @@ TEST(workspaces, tile_dotProduct_2) {
 
   stmt = stmt.concretize();
 
+  stmt = stmt.wsaccel(precomputed, {}, false);
   A.compile(stmt);
   A.assemble();
   A.compute();