Annotate default 'if's

Change-Id: I0098bd1bf6788dd6366810dcefa84f1ebbffaab0
apache · Apr 1, 2020 · b35023b · b35023b
1 parent c7a51ec
commit b35023b
Showing 1 changed file with 20 additions and 9 deletions.
diff --git a/src/relay/transforms/merge_compiler_regions.cc b/src/relay/transforms/merge_compiler_regions.cc
@@ -182,6 +182,26 @@ class AnnotateRestDefault : public ExprMutator {
     return updated_get;
   }
 
+  Expr VisitExpr_(const IfNode *op) {
+    auto region = regions_->GetRegion(GetRef<TupleGetItem>(op));
+    auto new_e = ExprMutator::VisitExpr_(op);
+    auto iff = Downcast<If>(new_e);
+
+    if (!region.defined()) {
+      return If(
+        InsertBegin(InsertEnd(iff->cond)),
+        InsertBegin(InsertEnd(iff->true_branch)),
+        InsertBegin(InsertEnd(iff->false_branch)));
+    } else {
+      Expr updated_iff = If(
+        InsertEnd(iff->cond),
+        InsertEnd(iff->true_branch),
+        InsertEnd(iff->false_branch));
+      annotated_nodes_.insert(updated_iff);
+      return updated_iff;
+    }
+  }
+
   Expr VisitExpr_(const LetNode *op) {
     auto new_e = ExprMutator::VisitExpr_(op);
     auto let = Downcast<Let>(new_e);
@@ -191,15 +211,6 @@ class AnnotateRestDefault : public ExprMutator {
       InsertEnd(let->body));
   }
 
-  Expr VisitExpr_(const IfNode *op) {
-    auto new_e = ExprMutator::VisitExpr_(op);
-    auto iff = Downcast<If>(new_e);
-    return If(
-      InsertEnd(iff->cond),
-      InsertEnd(iff->true_branch),
-      InsertEnd(iff->false_branch));
-  }
-
   Expr VisitExpr_(const RefCreateNode *op) {
     auto new_e = ExprMutator::VisitExpr_(op);
     auto create = Downcast<RefCreate>(new_e);