Merge pull request #369 from rohany/refactor-break

include,src: introduce a true break statement, rename current to continue

Author: stephenchouca

include/taco/ir/ir.h

@@ -65,8 +65,9 @@ enum class IRNodeType {
   BlankLine,
   Print,
   GetProperty,
-  Break,
-  Sort
+  Continue,
+  Sort,
+  Break
 };
 
 enum class TensorProperty {
@@ -719,7 +720,14 @@ struct BlankLine : public StmtNode<BlankLine> {
   static const IRNodeType _type_info = IRNodeType::BlankLine;
 };
 
-/** Breaks current loop */
+/** Continues past current iteration of current loop */
+struct Continue : public StmtNode<Continue> {
+  static Stmt make();
+
+  static const IRNodeType _type_info = IRNodeType::Continue;
+};
+
+/** Breaks out of the current loop */
 struct Break : public StmtNode<Break> {
   static Stmt make();
 

include/taco/ir/ir_printer.h

@@ -65,10 +65,11 @@ class IRPrinter : public IRVisitorStrict {
   virtual void visit(const Free*);
   virtual void visit(const Comment*);
   virtual void visit(const BlankLine*);
-  virtual void visit(const Break*);
+  virtual void visit(const Continue*);
   virtual void visit(const Print*);
   virtual void visit(const GetProperty*);
   virtual void visit(const Sort*);
+  virtual void visit(const Break*);
 
   std::ostream &stream;
   int indent;

include/taco/ir/ir_rewriter.h

@@ -65,10 +65,11 @@ class IRRewriter : public IRVisitorStrict {
   virtual void visit(const Free* op);
   virtual void visit(const Comment* op);
   virtual void visit(const BlankLine* op);
-  virtual void visit(const Break* op);
+  virtual void visit(const Continue* op);
   virtual void visit(const Print* op);
   virtual void visit(const GetProperty* op);
   virtual void visit(const Sort *op);
+  virtual void visit(const Break *op);
 };
 
 }}

include/taco/ir/ir_visitor.h

@@ -45,10 +45,11 @@ struct Allocate;
 struct Free;
 struct Comment;
 struct BlankLine;
-struct Break;
+struct Continue;
 struct Print;
 struct GetProperty;
 struct Sort;
+struct Break;
 
 /// Extend this class to visit every node in the IR.
 class IRVisitorStrict {
@@ -96,10 +97,11 @@ class IRVisitorStrict {
   virtual void visit(const Free*) = 0;
   virtual void visit(const Comment*) = 0;
   virtual void visit(const BlankLine*) = 0;
-  virtual void visit(const Break*) = 0;
+  virtual void visit(const Continue*) = 0;
   virtual void visit(const Print*) = 0;
   virtual void visit(const GetProperty*) = 0;
   virtual void visit(const Sort*) = 0;
+  virtual void visit(const Break*) = 0;
 };
 
 
@@ -150,10 +152,11 @@ class IRVisitor : public IRVisitorStrict {
   virtual void visit(const Free* op);
   virtual void visit(const Comment* op);
   virtual void visit(const BlankLine* op);
-  virtual void visit(const Break* op);
+  virtual void visit(const Continue* op);
   virtual void visit(const Print* op);
   virtual void visit(const GetProperty* op);
   virtual void visit(const Sort* op);
+  virtual void visit(const Break* op);
 };
 
 }}

src/codegen/codegen_cuda.cpp

@@ -1142,7 +1142,7 @@ void CodeGen_CUDA::visit(const Sqrt* op) {
   stream << ")";
 }
 
-void CodeGen_CUDA::visit(const Break*) {
+void CodeGen_CUDA::visit(const Continue*) {
   doIndent();
   if(!isHostFunction && deviceFunctionLoopDepth == 0) {
     // can't break out of kernel

src/codegen/codegen_cuda.h

@@ -46,7 +46,7 @@ class CodeGen_CUDA : public CodeGen {
   void visit(const Call*);
   void visit(const Store*);
   void visit(const Assign*);
-  void visit(const Break*);
+  void visit(const Continue*);
   void visit(const Free* op);
   std::string printDeviceFuncName(const std::vector<std::pair<std::string, Expr>> currentParameters, int index);
   void printDeviceFuncCall(const std::vector<std::pair<std::string, Expr>> currentParameters, Expr blockSize, int index, Expr gridSize);

src/ir/ir.cpp

@@ -786,6 +786,11 @@ Stmt BlankLine::make() {
   return new BlankLine;
 }
 
+// Continue
+Stmt Continue::make() {
+  return new Continue;
+}
+
 // Break
 Stmt Break::make() {
   return new Break;
@@ -954,14 +959,16 @@ template<> void StmtNode<Comment>::accept(IRVisitorStrict *v)
     const { v->visit((const Comment*)this); }
 template<> void StmtNode<BlankLine>::accept(IRVisitorStrict *v)
     const { v->visit((const BlankLine*)this); }
-template<> void StmtNode<Break>::accept(IRVisitorStrict *v)
-  const { v->visit((const Break*)this); }
+template<> void StmtNode<Continue>::accept(IRVisitorStrict *v)
+  const { v->visit((const Continue*)this); }
 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); }
+template<> void StmtNode<Break>::accept(IRVisitorStrict *v)
+  const { v->visit((const Break*)this); }
 
 // printing methods
 std::ostream& operator<<(std::ostream& os, const Stmt& stmt) {

src/ir/ir_printer.cpp

@@ -558,9 +558,14 @@ void IRPrinter::visit(const BlankLine*) {
   stream << endl;
 }
 
+void IRPrinter::visit(const Continue*) {
+  doIndent();
+  stream << "continue;" << endl;
+}
+
 void IRPrinter::visit(const Break*) {
   doIndent();
-  stream << "continue;" << endl; // TODO: add continue statement
+  stream << "break;" << endl;
 }
 
 void IRPrinter::visit(const Print* op) {

src/ir/ir_rewriter.cpp

@@ -447,6 +447,10 @@ void IRRewriter::visit(const BlankLine* op) {
   stmt = op;
 }
 
+void IRRewriter::visit(const Continue* op) {
+  stmt = op;
+}
+
 void IRRewriter::visit(const Break* op) {
   stmt = op;
 }

src/ir/ir_visitor.cpp

@@ -228,6 +228,9 @@ void IRVisitor::visit(const Comment*) {
 void IRVisitor::visit(const BlankLine*) {
 }
 
+void IRVisitor::visit(const Continue*) {
+}
+
 void IRVisitor::visit(const Break*) {
 }