Generate predicates for non-root iteration variables as well (#2258)

apache · Dec 8, 2018 · 6fdc3e0 · 6fdc3e0
1 parent 3b3b8cb
commit 6fdc3e0
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/src/schedule/message_passing.cc b/src/schedule/message_passing.cc
@@ -475,16 +475,20 @@ std::vector<Expr> MakeBoundCheck(
     iset_dmap[kv.first->var.get()] = IntSet::range(kv.second);
   }
 
-  for (IterVar iv : stage->op->root_iter_vars()) {
+  for (const IterVar& iv : stage->all_iter_vars) {
     if (skip_iter.count(iv) || iv->iter_type == kOpaque) continue;
-    Range dom = dom_map.at(iv);
     if (bound_state.at(iv)) {
+      Range dom = dom_map.at(iv);
       Expr value = ComputeExpr<Sub>(value_map.at(iv), dom->min);
       Expr vmax = EvalSet(value, iset_dmap).max();
       if (vmax.type() != value.type() || !can_prove(vmax < dom->extent)) {
         preds.emplace_back(value < dom->extent);
       }
     }
+  }
+  for (const IterVar& iv : stage->op->root_iter_vars()) {
+    if (skip_iter.count(iv) || iv->iter_type == kOpaque) continue;
+    Range dom = dom_map.at(iv);
     CHECK(iv->dom.defined());
     if (!skip_ivar_domain && !iv->dom.same_as(dom)) {
       Expr value = ComputeExpr<Sub>(value_map.at(iv), iv->dom->min);