[SCHEDULE] Refactor bound inference logic

apache · Feb 12, 2017 · a162619 · a162619
1 parent b8f0ec5
commit a162619
Show file tree

Hide file tree

Showing 4 changed files with 205 additions and 115 deletions.
diff --git a/include/tvm/schedule_pass.h b/include/tvm/schedule_pass.h
@@ -22,7 +22,7 @@ namespace schedule {
  * \param sch The root schedule to infer all the bounds.
  * \return the result bound of the iteration Variable
  */
-Map<IterVar, Range> InferBound(Schedule sch);
+Map<IterVar, Range> InferBound(const Schedule& sch);
 
 /*!
  * \brief Schedule s' dependent operations.

diff --git a/src/arithmetic/int_set.cc b/src/arithmetic/int_set.cc
@@ -432,7 +432,6 @@ TVM_STATIC_IR_FUNCTOR(IntSetEvaluator, vtable)
 .set_dispatch<And>(Binary<And>)
 .set_dispatch<Or>(Binary<Or>);
 
-
 IntSet EvalSet(Expr e,
                const std::unordered_map<const Variable*, IntSet>& dom_map) {
   return IntSetEvaluator(dom_map).Eval(e);
@@ -444,17 +443,12 @@ IntSet EvalSet(Expr e,
   for (auto kv : dom_map) {
     dmap[kv.first->var.as<Variable>()] = kv.second;
   }
-  IntSetEvaluator m(dmap);
-  return m.Eval(e);
+  return EvalSet(e, dmap);
 }
 
 IntSet EvalSet(Range r,
-               const Map<IterVar, IntSet>& dom_map) {
-  std::unordered_map<const Variable*, IntSet> dmap;
-  for (auto kv : dom_map) {
-    dmap[kv.first->var.as<Variable>()] = kv.second;
-  }
-  IntSetEvaluator m(dmap);
+               const std::unordered_map<const Variable*, IntSet>& dom_map) {
+  IntSetEvaluator m(dom_map);
   IntSet min_set = m.Eval(r->min);
   IntSet ext_set = m.Eval(r->extent).cover_interval();
   const Interval& ei = ext_set.as<IntervalSet>()->i;
@@ -463,13 +457,21 @@ IntSet EvalSet(Range r,
   return Combine<Add>(min_set, ext_set);
 }
 
+IntSet EvalSet(Range r,
+               const Map<IterVar, IntSet>& dom_map) {
+  std::unordered_map<const Variable*, IntSet> dmap;
+  for (auto kv : dom_map) {
+    dmap[kv.first->var.as<Variable>()] = kv.second;
+  }
+  return EvalSet(r, dmap);
+}
+
 TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable)
 .set_dispatch<IntervalSet>([](const IntervalSet *op, IRPrinter *p) {
     p->stream << "interval-set["
               << "[" << op->i.min << ", "
               << op->i.max << ']';
   });
 
-
 }  // namespace arith
 }  // namespace tvm
diff --git a/src/arithmetic/int_set.h b/src/arithmetic/int_set.h
@@ -103,6 +103,9 @@ IntSet EvalSet(Expr e,
  */
 IntSet EvalSet(Range r,
                const Map<IterVar, IntSet>& dom_map);
+IntSet EvalSet(Range r,
+               const std::unordered_map<const Variable*, IntSet>& dom_map);
+
 
 /*!
  * \brief Create an union set of all sets