[SCHEDULE] Improve bound inference, support reduce codegen. (#30)

include/tvm/expr.h

@@ -32,6 +32,9 @@ using Halide::Internal::IRPrinter;
 using Halide::Internal::Variable;
 
 using Halide::Internal::make_const;
+using Halide::Internal::make_zero;
+using Halide::Internal::as_const_int;
+using Halide::Internal::as_const_uint;
 
 
 inline Type TVMType2Type(TVMType t) {
@@ -126,25 +129,25 @@ using Halide::abs;
 using Halide::select;
 
 /*!
- * \brief sum of of source expression over rdom
+ * \brief sum of of source expression over axis
  * \param source The source expression.
- * \param rdom List of iteration variables that will be used for reduction.
+ * \param axis List of iteration variables that will be used for reduction.
  */
-Expr sum(Expr source, Array<IterVar> rdom);
+Expr sum(Expr source, Array<IterVar> axis);
 
 /*!
- * \brief max of of source expression over rdom
+ * \brief max of of source expression over axis
  * \param source The source expression.
- * \param rdom List of iteration variables that will be used for reduction.
+ * \param axis List of iteration variables that will be used for reduction.
  */
-Expr max(Expr source, Array<IterVar> rdom);
+Expr max(Expr source, Array<IterVar> axis);
 
 /*!
- * \brief max of of source expression over rdom
+ * \brief max of of source expression over axis
  * \param source The source expression.
- * \param rdom List of iteration variables that will be used for reduction.
+ * \param axis List of iteration variables that will be used for reduction.
  */
-Expr min(Expr source, Array<IterVar> rdom);
+Expr min(Expr source, Array<IterVar> axis);
 
 
 // print functions for expr

include/tvm/ir.h

@@ -30,8 +30,8 @@ struct Reduce : public ExprNode<Reduce> {
   std::string op;
   /*! \brief The source operand */
   Expr source;
-  /*! \brief The reduction domains */
-  Array<IterVar> rdom;
+  /*! \brief The reduction axis */
+  Array<IterVar> axis;
 
   /*! \brief construct expr from op and rdom */
   static Expr make(std::string op, Expr src, Array<IterVar> rdom);
@@ -40,7 +40,7 @@ struct Reduce : public ExprNode<Reduce> {
     v->Visit("dtype", &type);
     v->Visit("op", &op);
     v->Visit("source", &source);
-    v->Visit("rdom", &rdom);
+    v->Visit("axis", &axis);
   }
   static const IRNodeType _type_info = IRNodeType::ExtensionExpr;
   static constexpr const char* _type_key = "Reduce";

include/tvm/ir_pass.h

@@ -3,8 +3,8 @@
  * \file ir_pass.h
  * \brief Collection of IR pass functions
  *
- *  All the pass functions in this file are for Stmt,
- *  We can use PassFunction(Evaluate(expr)) to apply it to Expr
+ *  When the pass functions in this file are for Stmt,
+ *  we can use PassFunction(Evaluate(expr)) to apply it to Expr
  */
 #ifndef TVM_IR_PASS_H_
 #define TVM_IR_PASS_H_
@@ -38,15 +38,6 @@ inline Stmt Simplify(Stmt a) {
 }
 
 /*!
- * \brief Schedule s' dependent operations.
- *
- * \param s The schedule to be realized
- * \param dom_map The domain of each iter vars.
- * \return the result Stmt
- */
-Stmt ScheduleOps(Schedule s, Map<IterVar, Range> dom_map);
-
-/*!
  * \brief verifies whether the IR stmt or Expr is in SSA form.
  *  That is: each VarExpr is defined and assigned once(in Let/For)
  *
@@ -70,6 +61,14 @@ bool HasSideEffect(const Expr& e);
 Stmt ConvertSSA(Stmt stmt);
 
 /*!
+ * \brief Substitute the var specified in key->var to be value.
+ * \param stmt The source statement to be substituted
+ * \param value_map The map of new values.
+ * \return The converted form.
+ */
+Stmt Substitute(Stmt stmt, const Map<IterVar, Expr>& value_map);
+
+/*!
  * \brief inline all calls of f in stmt.
  *
  * \param f The function reference to be inlined

include/tvm/operation.h

@@ -49,6 +49,8 @@ class ComputeOpNode : public OperationNode {
  public:
   /*! \brief IterVar on each axis */
   Array<IterVar> axis;
+  /*! \brief IterVar on each reduction axis, if the body is a Reduce */
+  Array<IterVar> reduce_axis;
   /*! \brief the compute expression */
   Expr body;
   /*! \brief constructor */
@@ -64,6 +66,7 @@ class ComputeOpNode : public OperationNode {
   void VisitAttrs(AttrVisitor* v) final {
     v->Visit("name", &name);
     v->Visit("axis", &axis);
+    v->Visit("reduce_axis", &reduce_axis);
     v->Visit("body", &body);
   }
   static Operation make(std::string name,

include/tvm/schedule.h

@@ -123,6 +123,8 @@ class Stage : public NodeRef {
               IterVar* p_x_outer, IterVar* p_y_outer,
               IterVar* p_x_inner, IterVar* p_y_inner,
               Expr x_factor, Expr y_factor);
+  // declare container type
+  using ContainerType = StageNode;
 };
 
 /*!
@@ -153,10 +155,21 @@ class Schedule : public NodeRef {
     return this->operator[](tensor->op);
   }
   /*!
+   * \brief Normalize the schedule.
+   *  This is needed before bound inference.
+   *  Insert necessary RebaseNode to make sure all leaf_iter_vars
+   *  are in form [0, extent)
+   *
+   * \return A normalized schedule, can be same as current one.
+   */
+  void normalize();
+  /*!
    * \brief access the internal node container
    * \return the pointer to the internal node container
    */
   inline const ScheduleNode* operator->() const;
+  // declare container type
+  using ContainerType = ScheduleNode;
 };
 
 /*!
@@ -308,6 +321,30 @@ class FuseNode : public IterVarRelationNode {
   TVM_DECLARE_NODE_TYPE_INFO(FuseNode);
 };
 
+/*!
+ * \brief Rebase the iteration to make min to be 0.
+ *  This is useful to normalize the Schedule
+ *  to make every leaf variable's min to be 0.
+ */
+class RebaseNode : public IterVarRelationNode {
+ public:
+  /*! \brief The parent domain */
+  IterVar parent;
+  /*! \brief The inner domain */
+  IterVar rebased;
+
+  void VisitAttrs(AttrVisitor* v) final {
+    v->Visit("parent", &parent);
+    v->Visit("rebased", &rebased);
+  }
+
+  static IterVarRelation make(IterVar parent, IterVar rebased);
+
+  static constexpr const char* _type_key = "Rebase";
+  TVM_DECLARE_NODE_TYPE_INFO(RebaseNode);
+};
+
+
 // implementations
 inline const StageNode* Stage::operator->() const {
   return static_cast<const StageNode*>(node_.get());

include/tvm/schedule_pass.h

@@ -24,6 +24,15 @@ namespace schedule {
  */
 Map<IterVar, Range> InferBound(Schedule sch);
 
+/*!
+ * \brief Schedule s' dependent operations.
+ *
+ * \param s The schedule to be realized
+ * \param dom_map The domain of each iter vars.
+ * \return the result Stmt
+ */
+Stmt ScheduleOps(Schedule s, Map<IterVar, Range> dom_map);
+
 }  // namespace schedule
 }  // namespace tvm
 #endif  // TVM_SCHEDULE_PASS_H_

python/tvm/api.py

@@ -212,51 +212,51 @@ def IterVar(dom=None, name=None, thread_tag=''):
     return _api_internal._IterVar(dom, name, thread_tag)
 
 
-def sum(expr, rdom):
-    """Create a sum expression over rdom
+def sum(expr, axis):
+    """Create a sum expression over axis
 
     Parameters
     ----------
     expr : Expr
         The source expression.
 
-    rdom : RDomain
-        The reduction domainx
+    axis : IterVar
+        The reduction IterVar axis
     """
-    rdom = rdom if isinstance(rdom, list) else [rdom]
-    x = _make.Reduce("Add", expr, rdom)
+    axis = axis if isinstance(axis, list) else [axis]
+    x = _make.Reduce("Add", expr, axis)
     return x
 
 
-def min(expr, rdom):
-    """Create a min expression over rdom
+def min(expr, axis):
+    """Create a min expression over axis
 
     Parameters
     ----------
     expr : Expr
         The source expression.
 
-    rdom : RDomain
-        The reduction domainx
+    axis : IterVar
+        The reduction IterVar axis
     """
-    rdom = rdom if isinstance(rdom, list) else [rdom]
-    x = _make.Reduce("Min", expr, rdom)
+    axis = axis if isinstance(axis, list) else [axis]
+    x = _make.Reduce("Min", expr, axis)
     return x
 
 
-def max(expr, rdom):
-    """Create a min expression over rdom
+def max(expr, axis):
+    """Create a min expression over axis
 
     Parameters
     ----------
     expr : Expr
         The source expression.
 
-    rdom : RDomain
-        The reduction domainx
+    axis : IterVar
+        The reduction IterVar axis
     """
-    rdom = rdom if isinstance(rdom, list) else [rdom]
-    x = _make.Reduce("Max", expr, rdom)
+    axis = axis if isinstance(axis, list) else [axis]
+    x = _make.Reduce("Max", expr, axis)
     return x
 
 

python/tvm/build.py

@@ -62,9 +62,10 @@ def build(sch,
 
     # lowering
     bounds = schedule.InferBound(sch)
-    stmt = ir_pass.ScheduleOps(sch, bounds)
+    stmt = schedule.ScheduleOps(sch, bounds)
     stmt = ir_pass.StorageFlatten(stmt, binds)
     stmt = ir_pass.Simplify(stmt)
+    print(stmt)
     fapi = codegen.MakeAPI(stmt, name, arg_list, len(arg_list))
     fsplits = codegen.SplitHostDevice(fapi)
 
@@ -73,7 +74,8 @@ def build(sch,
         for i, f in enumerate(fsplits):
             t = target if i >= 1 else "c"
             record_codes.append(codegen.CompileToC(f, output_ssa, t))
-
+    for c in record_codes:
+        print(c)
     if target == "cuda":
         ret = codegen.BuildNVRTC(fsplits, "stackvm")
     elif target == "opencl":

python/tvm/schedule.py

@@ -33,6 +33,14 @@ class Schedule(NodeBase):
             raise ValueError("Cannot find the operation %s in schedule" % (str(k)))
         return self.stage_map[k]
 
+    def normalize(self):
+        """Build a normalized schedule.
+
+        Insert necessary rebase to make certain iter var to start from 0.
+        This is needed before bound inference and followup step.
+        """
+        _api_internal._ScheduleNormalize(self)
+
 @register_node
 class Stage(NodeBase):
     """A Stage represents schedule for one operation."""

src/api/api_lang.cc

@@ -253,4 +253,10 @@ TVM_REGISTER_API(_StageTile)
     *ret = Array<IterVar>({x_outer, y_outer, x_inner, y_inner});
   });
 
+TVM_REGISTER_API(_ScheduleNormalize)
+.set_body([](TVMArgs args, TVMRetValue* ret) {
+    args[0].operator Schedule()
+        .normalize();
+  });
+
 }  // namespace tvm

src/api/api_pass.cc

@@ -51,7 +51,6 @@ TVM_REGISTER_API(_pass_Equal)
 REGISTER_PASS1(ConvertSSA);
 REGISTER_PASS1(VerifySSA);
 REGISTER_PASS4(Inline);
-REGISTER_PASS2(ScheduleOps);
 REGISTER_PASS2(StorageFlatten);
 
 }  // namespace ir

src/api/api_schedule.cc

@@ -29,6 +29,7 @@ namespace schedule {
 REGISTER_SCHEDULE_PASS1(InferBound);
 REGISTER_SCHEDULE_PASS1(CreateReadGraph);
 REGISTER_SCHEDULE_PASS2(PostDFSOrder);
+REGISTER_SCHEDULE_PASS2(ScheduleOps);
 
 }  // namespace schedule
 }  // namespace tvm

src/codegen/codegen_c.cc

@@ -2,6 +2,7 @@
  *  Copyright (c) 2017 by Contributors
  * \file codegen_c.cc
  */
+#include <iomanip>
 #include "./codegen_c.h"
 
 namespace tvm {
@@ -216,7 +217,7 @@ inline void PrintConst(const FloatImm* op, std::ostream& os, CodeGenC* p) { // N
   switch (op->type.bits()) {
     case 64: case 32: {
       std::ostringstream temp;
-      temp << op->value;
+      temp << std::scientific << op->value;
       if (op->type.bits() == 32) temp << 'f';
       p->MarkConst(temp.str());
       os << temp.str();
@@ -225,7 +226,7 @@ inline void PrintConst(const FloatImm* op, std::ostream& os, CodeGenC* p) { // N
     case 16: {
       os << '(';
       p->PrintType(op->type, os);
-      os << ')' << op->value << 'f';
+      os << ')' << std::scientific <<op->value << 'f';
       break;
     }
     default: LOG(FATAL) << "Bad bit-width for float: " << op->type << "\n";

src/lang/ir.cc

@@ -26,7 +26,7 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable)
             << op->op
             << ", ";
   p->print(op->source);
-  p->stream << ", rdom=" << op->rdom << ")";
+  p->stream << ", axis=" << op->axis << ")";
 });
 
 }  // namespace Internal
@@ -35,16 +35,16 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable)
 namespace tvm {
 namespace ir {
 
-Expr Reduce::make(std::string op, Expr source, Array<IterVar> rdom) {
+Expr Reduce::make(std::string op, Expr source, Array<IterVar> axis) {
   auto n = std::make_shared<Reduce>();
   CHECK(source.defined());
-  for (size_t i = 0; i < rdom.size(); ++i) {
-    CHECK(rdom[i].defined());
+  for (size_t i = 0; i < axis.size(); ++i) {
+    CHECK(axis[i].defined());
   }
   n->type = source.type();
   n->source = source;
   n->op = op;
-  n->rdom = rdom;
+  n->axis = axis;
   return Expr(n);
 }
 

src/lang/operation.cc

@@ -4,6 +4,7 @@
  */
 #include <tvm/operation.h>
 #include <tvm/tensor.h>
+#include <tvm/ir.h>
 #include <memory>
 
 namespace tvm {
@@ -57,7 +58,12 @@ Tensor Placeholder(Array<Expr> shape, Type dtype, std::string name) {
 
 // ComputeOpNode
 Array<IterVar> ComputeOpNode::root_iter_vars() const {
-  return axis;
+  if (reduce_axis.size() == 0) return axis;
+  Array<IterVar> ret = axis;
+  for (IterVar iv : reduce_axis) {
+    ret.push_back(iv);
+  }
+  return ret;
 }
 
 Type ComputeOpNode::output_dtype(size_t i) const {
@@ -101,6 +107,9 @@ Operation ComputeOpNode::make(std::string name,
   n->name = name;
   n->axis = axis;
   n->body = body;
+  if (n->body->is_type<ir::Reduce>()) {
+    n->reduce_axis = n->body.as<ir::Reduce>()->axis;
+  }
   return Operation(n);
 }
 

src/pass/ir_mutator.cc

@@ -37,7 +37,7 @@ inline Array<Expr> MutateArray(Array<Expr> arr, IRMutator *m) {
   }
 }
 
-inline Array<IterVar> MutateRDom(Array<IterVar> rdom, IRMutator *m) {
+inline Array<IterVar> MutateIterVarArr(Array<IterVar> rdom, IRMutator *m) {
   std::vector<IterVar> new_dom(rdom.size());
   bool changed = false;
   for (size_t i = 0; i < rdom.size(); i++) {
@@ -237,13 +237,13 @@ Expr IRMutator::Mutate_(const Let *op, const Expr& e) {
 
 TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_expr)
 .set_dispatch<Reduce>([](const Reduce* op, const Expr& e, IRMutator* m) {
-    Array<IterVar> new_rdom  = MutateRDom(op->rdom, m);
+    Array<IterVar> new_axis  = MutateIterVarArr(op->axis, m);
     Expr new_source  = m->Mutate(op->source);
-    if (op->rdom.same_as(new_rdom) &&
+    if (op->axis.same_as(new_axis) &&
         op->source.same_as(new_source)) {
       return e;
     } else {
-      return Reduce::make(op->op, new_source, new_rdom);
+      return Reduce::make(op->op, new_source, new_axis);
     }
   });
 

src/pass/ir_visitor.cc

@@ -120,7 +120,7 @@ void IRVisitor::Visit_(const Call *op) {
 
 TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
 .set_dispatch<Reduce>([](const Reduce* op, IRVisitor* v) {
-    VisitRDom(op->rdom, v);
+    VisitRDom(op->axis, v);
     v->Visit(op->source);
   })
 .set_dispatch<IntImm>(NoOp)

src/pass/simple_passes.cc

@@ -5,6 +5,7 @@
  */
 #include <tvm/ir.h>
 #include <tvm/ir_visitor.h>
+#include <tvm/ir_mutator.h>
 #include <tvm/ir_pass.h>
 
 namespace tvm {
@@ -32,5 +33,26 @@ bool HasSideEffect(const Expr& e) {
   v.Visit(e);
   return v.has_side_effect_;
 }
+
+class IRSubstitue : public IRMutator {
+ public:
+  Expr Mutate_(const Variable* op, const Expr& e) final {
+    auto it = smap.find(op);
+    if (it != smap.end()) {
+      return it->second;
+    } else {
+      return e;
+    }
+  }
+  std::unordered_map<const Variable*, Expr> smap;
+};
+
+Stmt Substitute(Stmt stmt, const Map<IterVar, Expr>& value_map) {
+  IRSubstitue m;
+  for (auto kv : value_map) {
+    m.smap[kv.first->var.get()] = kv.second;
+  }
+  return m.Mutate(stmt);
+}
 }  // namespace ir
 }  // namespace tvm

src/schedule/bound.cc

@@ -54,6 +54,11 @@ void PassDown(const Stage& s,
       const Range& range_inner = state.at(r->inner);
       state[r->fused] = Range::make_with_min_extent(
           0, range_outer->extent * range_inner->extent);
+    } else if (rel.as<RebaseNode>()) {
+      const RebaseNode* r = rel.as<RebaseNode>();
+      CHECK(state.count(r->parent));
+      state[r->rebased] = Range::make_with_min_extent(
+          0, state.at(r->parent)->extent);
     } else {
       LOG(FATAL) << "unknown relation type";
     }
@@ -85,6 +90,13 @@ void PassUp(const Stage& s,
              &outer, &inner);
       state[r->outer] = outer;
       state[r->inner] = inner;
+    } else if (rel.as<RebaseNode>()) {
+      IntSet parent;
+      const RebaseNode* r = rel.as<RebaseNode>();
+      PassUp(r, dom_map,
+             state.at(r->rebased),
+             &parent);
+      state[r->parent] = parent;
     } else {
       LOG(FATAL) << "unknown relation type";
     }
@@ -109,9 +121,15 @@ void PassToOperation(
   // Eventually, we need to change the inference to be a Pull style inference
   if (tensor->op.as<ComputeOpNode>()) {
     auto root_iter_vars = tensor->op->root_iter_vars();
-    CHECK_EQ(tensor.ndim(), root_iter_vars.size());
-    for (size_t i = 0; i < tensor.ndim(); ++i) {
-      (*result)[root_iter_vars[i]].push_back(dim_bounds[i]);
+    const ComputeOpNode* op = tensor->op.as<ComputeOpNode>();
+    CHECK_EQ(op->axis.size() + op->reduce_axis.size(), root_iter_vars.size());
+    for (size_t i = 0; i < op->axis.size(); ++i) {
+      (*result)[op->axis[i]].push_back(dim_bounds[i]);
+    }
+    // reduction.
+    for (size_t i = 0; i < op->reduce_axis.size(); ++i) {
+      (*result)[op->reduce_axis[i]].push_back(
+          IntSet::range(op->reduce_axis[i]->dom));
     }
   } else {
     LOG(FATAL) << "unknown operation mode " << tensor->op->type_key();
@@ -173,9 +191,9 @@ bool ScopeRelax(const IterVar& iv, const std::string& scope) {
     {"local", 2}
   };
   static std::unordered_map<std::string, int> thread_tag_rank{
-    {"gridIdx.x", 0},
-    {"gridIdx.y", 0},
-    {"gridIdx.z", 0},
+    {"blockIdx.x", 0},
+    {"blockIdx.y", 0},
+    {"blockIdx.z", 0},
     {"threadIdx.x", 1},
     {"threadIdx.y", 1},
     {"threadIdx.z", 1}
@@ -194,8 +212,6 @@ void InferBound(const Stage& stage,
       (*rmap)[iv] = iv->dom;
     }
   }
-  // get range of all child iter vars.
-  PassDown(stage, rmap);
 
   if (stage->attach_type == kScope) {
     Stage parent = stage->attach_stage;
@@ -206,10 +222,18 @@ void InferBound(const Stage& stage,
 
     bool fix_value = true;
     for (auto iv : parent->leaf_iter_vars) {
+      Range vrange = rmap->at(iv);
+      CHECK(is_zero(vrange->min))
+          << "InferBound requires every leaf iter var's min equals 0, "
+          << "call schedule.normalize to achieve this.";
+      // special optimization to remove trivial loop
+      if (is_one(vrange->extent)) {
+        up_state[iv] = IntSet::single_point(vrange->min);
+      }
       if (fix_value && !ScopeRelax(iv, stage->scope)) {
-        up_state[iv] = IntSet::make_point(iv->var);
+        up_state[iv] = IntSet::single_point(iv->var);
       } else {
-        up_state[iv] = IntSet::make_range(rmap->at(iv));
+        up_state[iv] = IntSet::range(vrange);
       }
       if (stage->attach_ivar == iv) {
         fix_value = false;
@@ -223,12 +247,30 @@ void InferBound(const Stage& stage,
       bp_state[iv] = {up_state.at(iv)};
     }
     auto result = BoundProp(post_order, &bp_state);
+
+    // Set relaxation
+    Map<IterVar, IntSet> relax_set;
+    Stage s = stage;
+    while (s->attach_type == kScope) {
+      s = s->attach_stage;
+      for (auto iv : s->leaf_iter_vars) {
+        if (ScopeRelax(iv, stage->scope)) {
+          relax_set.Set(iv, IntSet::range(rmap->at(iv)));
+        }
+      }
+    }
     for (auto iv : stage->op->root_iter_vars()) {
       CHECK(result.count(iv));
       CHECK(!rmap->count(iv));
-      (*rmap)[iv] = result.at(iv).GetCoverRange();
+      Range r = result.at(iv).cover_range(iv->dom);
+      if (relax_set.size() != 0) {
+        r = EvalSet(r, relax_set).cover_range(iv->dom);
+      }
+      (*rmap)[iv] = r;
     }
   }
+  // get range of all child iter vars.
+  PassDown(stage, rmap);
 }
 
 

src/schedule/compute_expr.h

+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file compute_expr.h
+ * \brief Utility integer expression with quick eager simplification.
+ *  This is weaker than Simplify but can be done Eagerly.
+ */
+#ifndef TVM_SCHEDULE_COMPUTE_EXPR_H_
+#define TVM_SCHEDULE_COMPUTE_EXPR_H_
+
+#include <tvm/ir.h>
+#include <pass/Interval.h>
+
+namespace tvm {
+namespace schedule {
+
+using Halide::Internal::add_would_overflow;
+using Halide::Internal::sub_would_overflow;
+using Halide::Internal::mul_would_overflow;
+
+/*!
+ * \brief Compute the expression with the given binary op.
+ * \param lhs The left operand
+ * \param rhs The right operand
+ * \return The result.
+ */
+template<typename OP>
+inline Expr ComputeExpr(Expr lhs, Expr rhs) {
+  return OP::make(lhs, rhs);
+}
+
+template<typename T>
+inline bool GetConst(Expr e, T* out);
+
+template<>
+bool GetConst<int64_t>(Expr e, int64_t *out) {
+  if (e.type().is_vector()) return false;
+  const int64_t *v = as_const_int(e);
+  if (v) {
+    *out = *v; return true;
+  } else {
+    return false;
+  }
+}
+template<>
+bool GetConst<uint64_t>(Expr e, uint64_t *out) {
+  if (e.type().is_vector()) return false;
+  const uint64_t *v = as_const_uint(e);
+  if (v) {
+    *out = *v; return true;
+  } else {
+    return false;
+  }
+}
+
+#define TVM_CONST_PROPAGATION(OP_NAME, OP)                       \
+  int64_t ia = 0, ib = 0;                                        \
+  if (GetConst(a, &ia) && GetConst(b, &ib)) {                    \
+    if (OP_NAME ## _would_overflow(a.type().bits(), ia, ib)) {   \
+      LOG(FATAL) << "signed int overflow";                       \
+    }                                                            \
+    return ir::IntImm::make(a.type(), ia OP ib);                 \
+  }                                                              \
+  uint64_t ua = 0, ub = 0;                                       \
+  if (GetConst(a, &ua) && GetConst(b, &ub)) {                    \
+    return ir::UIntImm::make(a.type(), ua + ub);                 \
+  }                                                              \
+
+template<>
+inline Expr ComputeExpr<ir::Add>(Expr a, Expr b) {
+  if (is_zero(a)) return b;
+  if (is_zero(b)) return a;
+  TVM_CONST_PROPAGATION(add, +);
+  return ir::Add::make(a, b);
+}
+
+template<>
+inline Expr ComputeExpr<ir::Sub>(Expr a, Expr b) {
+  if (is_zero(b)) return a;
+  TVM_CONST_PROPAGATION(sub, -);
+  return ir::Add::make(a, b);
+}
+
+template<>
+inline Expr ComputeExpr<ir::Mul>(Expr a, Expr b) {
+  if (is_one(a)) return b;
+  if (is_one(b)) return a;
+  TVM_CONST_PROPAGATION(mul, *);
+  return ir::Mul::make(a, b);
+}
+
+template<>
+inline Expr ComputeExpr<ir::Div>(Expr a, Expr b) {
+  if (is_one(b)) return a;
+  return ir::Mul::make(a, b);
+}
+
+template<>
+inline Expr ComputeExpr<ir::Max>(Expr a, Expr b) {
+  return Halide::Internal::Interval::make_max(a, b);
+}
+
+template<>
+inline Expr ComputeExpr<ir::Min>(Expr a, Expr b) {
+  return Halide::Internal::Interval::make_min(a, b);
+}
+
+}  // namespace schedule
+}  // namespace tvm
+#endif   // TVM_SCHEDULE_COMPUTE_EXPR_H_

src/schedule/int_set.cc

 /*!
  *  Copyright (c) 2016 by Contributors
- * \file int_set.cc
+ * \file int_set_impl.cc
  * \brief The integer set functions
  */
 #include <tvm/ir.h>
+#include <tvm/ir_pass.h>
+#include <pass/Interval.h>
 #include "./int_set.h"
+#include "./compute_expr.h"
 
 namespace tvm {
 namespace schedule {
 
+using Halide::Internal::Interval;
+
 using namespace ir;
 
+/*! \brief Set of continuous interval */
+struct IntervalSet : public IntSetNode {
+  /*! \brief the internal interval*/
+  Interval i;
+
+  static IntSet make(Interval i) {
+    std::shared_ptr<IntervalSet> n =
+        std::make_shared<IntervalSet>();
+    n->i = i;
+    return IntSet(n);
+  }
+  static IntSet make(Expr min, Expr max) {
+    std::shared_ptr<IntervalSet> n =
+        std::make_shared<IntervalSet>();
+    n->i.min = min;
+    n->i.max = max;
+    return IntSet(n);
+  }
+
+  static constexpr const char* _type_key = "IntervalSet";
+  TVM_DECLARE_NODE_TYPE_INFO(IntervalSet);
+};
+
 /*!
- * \brief Internal node container of int set.
+ * \brief set represented by strided integers
+ *  Reserved for cases where strided access is supported.
  */
-class IntSetNode : public Node {
- public:
-  /*! \brief The base range scope */
-  Range base;
-  /*! \brief additional strided domain */
-  Array<Range> domain;
-  /*! \brief The stride of each strided domain */
-  Array<Expr> stride;
-  /*!
-   * \brief The concrete set,
-   *  used when concrete execution is enabled.
-   */
-  std::vector<int32_t> concrete;
-
-  void VisitAttrs(AttrVisitor* v) final {
-    v->Visit("base", &base);
-    v->Visit("domain", &domain);
-    v->Visit("stride", &stride);
-  }
-
-  static constexpr const char* _type_key = "IntSet";
-  TVM_DECLARE_NODE_TYPE_INFO(IntSetNode);
+struct StrideSet : public IntSetNode {
+  /*! \brief the base inetrval */
+  Interval base;
+  /*! \brief additional extents in positive number */
+  Array<Expr> extents;
+  /*! \brief additional strides in positive number */
+  Array<Expr> strides;
+
+  static constexpr const char* _type_key = "StrideSet";
+  TVM_DECLARE_NODE_TYPE_INFO(StrideSet);
 };
 
-TVM_REGISTER_NODE_TYPE(IntSetNode);
+inline IntSet IntSet::cover_interval() const {
+  if ((*this).as<IntervalSet>()) return *this;
+  const StrideSet* s =  (*this).as<StrideSet>();
+  if (s) {
+    CHECK_NE(s->extents.size(), 0U);
+    Expr max = s->base.max;
+    for (size_t i = 0; i < s->extents.size(); ++i) {
+      max = max + s->extents[i] * s->strides[i] - s->strides[i];
+    }
+    return IntervalSet::make(s->base.min, max);
+  }
+  LOG(FATAL) << "cannot convert set " << (*this)->type_key() << " to interval";
+  return IntSet::everything();
+}
+
+Range IntSet::cover_range(Range max_range) const {
+  IntSet temp;
+  const IntervalSet* s_int = (*this).as<IntervalSet>();
+  if (s_int == nullptr) {
+    temp = this->cover_interval();
+    s_int = temp.as<IntervalSet>();
+  }
+  if (s_int->i.is_bounded()) {
+    return Range::make_with_min_extent(
+        s_int->i.min, Simplify(s_int->i.max + 1 - s_int->i.min));
+  }
+  return max_range;
+}
 
-namespace {
+bool IntSet::is_everything() const {
+  const IntervalSet* s_int = (*this).as<IntervalSet>();
+  return (s_int && s_int->i.is_everything());
+}
 
-inline bool Match(const Expr& e, int64_t value) {
-  const ir::IntImm* v = e.as<ir::IntImm>();
-  return v != nullptr && v->value;
+bool IntSet::is_single_point() const {
+  const IntervalSet* s_int = (*this).as<IntervalSet>();
+  return (s_int && s_int->i.is_single_point());
 }
 
-// whether a exactly matches b.
-inline bool Match(const IntSet& a,
-                  const Range& b) {
-  if (a->base == b &&
-      a->domain.size() == 0 &&
-      a->concrete.size() == 0) {
-    return true;
-  } else {
-    return false;
-  }
+IntSet IntSet::everything() {
+  return IntervalSet::make(Interval::everything());
 }
 
-// whether a exactly matches b.
-inline bool Match(const IntSet& a,
-                  const Expr& b) {
-  if (a->domain.size() == 0 &&
-      a->concrete.size() == 0) {
-    return Match(a->base->extent, 1) && a->base->min.same_as(b);
-  } else {
-    return false;
-  }
+IntSet IntSet::single_point(Expr x) {
+  return IntervalSet::make(Interval::single_point(x));
 }
 
-inline bool IsNumber(const IntSet& s) {
-  if (s->domain.size() != 0) return false;
-  if (s->concrete.size() != 0) {
-    return s->concrete.size() == 1;
+IntSet IntSet::range(Range r) {
+  // must make sure it can be matched back by MatchRange.
+  if (is_one(r->extent)) {
+    return IntSet::single_point(r->min);
+  }
+  if (is_positive_const(r->extent) && is_const(r->min)) {
+    return IntervalSet::make(
+        r->min, ComputeExpr<Sub>(ComputeExpr<Add>(r->extent, r->min), 1));
   }
-  return Match(s->base->extent, 1);
+  return IntervalSet::make(r->min, (r->extent + r->min) - 1);
 }
 
-inline Expr AsNumber(const IntSet& s) {
-  return s->base->min;
+// Check if a is created from b.
+inline bool MatchRange(const IntSet& a,
+                       const Range& b) {
+  const IntervalSet* a_int = a.as<IntervalSet>();
+  if (!a_int) return false;
+  const Interval& i = a_int->i;
+  if (!i.min.same_as(b)) return false;
+  if (is_one(b->extent)) return i.is_single_point();
+  if (is_positive_const(b->extent) && is_const(b->min)) {
+    // deep equality
+    return Equal(
+        ComputeExpr<Sub>(ComputeExpr<Add>(b->extent, b->min), 1),
+        a_int->i.max);
+  }
+  const Sub* sub = i.max.as<Sub>();
+  if (!sub) return false;
+  if (is_one(sub->b)) return false;
+  const Add* add = sub->a.as<Add>();
+  return add &&
+      add->a.same_as(b->min) &&
+      add->b.same_as(b->extent);
 }
 
-// set combination rule by operators
-template<typename T>
-inline IntSet BinaryCombine(IntSet a, IntSet b) {
-  LOG(WARNING) << "cannot evaluate binary op " << T::_type_key;
-  return IntSet::make_all_set();
+inline bool MatchPoint(const IntSet& a,
+                       const Expr& b) {
+  const IntervalSet* a_int = a.as<IntervalSet>();
+  if (!a_int) return false;
+  const Interval& i = a_int->i;
+  return i.is_single_point() && i.min.same_as(b);
 }
 
-template<>
-inline IntSet BinaryCombine<Add>(IntSet a, IntSet b) {
-  auto n = std::make_shared<IntSetNode>(*(a.operator->()));
-  for (size_t i = 0; i < b->domain.size(); ++i) {
-    n->domain.push_back(b->domain[i]);
-    n->stride.push_back(b->stride[i]);
-  }
-
-  if (IsNumber(a)) {
-    n->base = Range::make_with_min_extent(
-        a->base->min + b->base->min,
-        b->base->extent);
-  } else if (IsNumber(b)) {
-    n->base = Range::make_with_min_extent(
-        a->base->min + b->base->min,
-        a->base->extent);
-  } else {
-    n->base = Range::make_with_min_extent(
-        a->base->min + b->base->min,
-        a->base->extent + b->base->extent - 1);
+IntSet Union(const Array<IntSet>& set) {
+  if (set.size() == 1) return set[0];
+  Interval x = set[0].cover_interval().as<IntervalSet>()->i;
+  for (size_t i = 1; i < set.size(); ++i) {
+    x.include(set[i].cover_interval().as<IntervalSet>()->i);
   }
-  return IntSet(n);
+  return IntervalSet::make(x);
 }
 
-inline Range Negation(Range a) {
-  if (Match(a->extent, 1)) {
-    return Range::make_with_min_extent(-a->min, a->extent);
-  } else {
-    return Range::make_with_min_extent(-(a->min + a->extent - 1), a->extent);
+// type traits
+template<typename OP>
+struct is_logical_op {
+  static const bool value = false;
+};
+
+#define TVM_DECLARE_LOGICAL_OP(OP)              \
+  template<>                                    \
+  struct is_logical_op<ir::OP> {                \
+    static const bool value = true;             \
+  };
+
+// interval related.
+template<typename OP>
+inline IntSet CombineInterval(Interval a, Interval b) {
+  if (a.is_single_point() && b.is_single_point()) {
+    return IntSet::single_point(ComputeExpr<OP>(a.min, b.min));
   }
+  LOG(WARNING) << "Return Everything in CombineInterval " << OP::_type_key;
+  return IntSet::everything();
 }
 
-inline IntSet Negation(IntSet a) {
-  CHECK_EQ(a->concrete.size(), 0U);
-  auto n = std::make_shared<IntSetNode>();
-  n->base = Negation(a->base);
-  for (size_t i = 0; i < a->domain.size(); ++i) {
-    n->domain.push_back(Negation(a->domain[i]));
-    n->stride.push_back(a->stride[i]);
+template<>
+inline IntSet CombineInterval<Add>(Interval a, Interval b) {
+  if (a.is_single_point() && b.is_single_point()) {
+    return IntSet::single_point(ComputeExpr<Add>(a.min, b.min));
+  }
+  Interval r = Interval::everything();
+  if (a.has_lower_bound() && b.has_lower_bound()) {
+    r.min = ComputeExpr<Add>(a.min, b.min);
   }
-  return IntSet(a);
+  if (a.has_upper_bound() && b.has_upper_bound()) {
+    r.max = ComputeExpr<Add>(a.max, b.max);
+  }
+  return IntervalSet::make(r);
 }
 
 template<>
-inline IntSet BinaryCombine<Sub>(IntSet a, IntSet b) {
-  return BinaryCombine<Add>(a, Negation(b));
+inline IntSet CombineInterval<Sub>(Interval a, Interval b) {
+  if (a.is_single_point() && b.is_single_point()) {
+    return IntSet::single_point(ComputeExpr<Sub>(a.min, b.min));
+  }
+  Interval r = Interval::everything();
+  if (a.has_lower_bound() && b.has_upper_bound()) {
+    r.min = ComputeExpr<Sub>(a.min, b.max);
+  }
+  if (a.has_upper_bound() && b.has_lower_bound()) {
+    r.max = ComputeExpr<Sub>(a.max, b.min);
+  }
+  return IntervalSet::make(r);
 }
 
-inline IntSet BinaryMul(IntSet a, Expr b) {
-  // copy construct
-  if (Match(b, 1)) return a;
-  if (Match(b, -1)) return Negation(a);
-  auto n = std::make_shared<IntSetNode>();
-  n->base = Range::make_with_min_extent(0, 1);
-  n->domain.push_back(a->base);
-  n->stride.push_back(b);
-  for (size_t i = 0; i < a->domain.size(); ++i) {
-    n->domain.push_back(a->domain[i]);
-    n->stride.push_back(a->stride[i] * b);
-  }
-  return IntSet(a);
+template<>
+inline IntSet CombineInterval<Mul>(Interval a, Interval b) {
+  if (a.is_single_point() && b.is_single_point()) {
+    return IntSet::single_point(ComputeExpr<Mul>(a.min, b.min));
+  }
+  if (a.is_single_point() && !b.is_single_point()) {
+    std::swap(a, b);
+  }
+  if (b.is_single_point()) {
+    if (is_zero(b.min)) return IntSet::single_point(0);
+    if (is_one(b.min)) return IntervalSet::make(a);
+    Expr e1 = a.has_lower_bound() ? ComputeExpr<Mul>(a.min, b.min) : a.min;
+    Expr e2 = a.has_upper_bound() ? ComputeExpr<Mul>(a.max, b.min) : a.max;
+    // This is relaxiation
+    // TODO(tqchen): consider convert to StrideSet.
+    if (is_positive_const(b.min)) {
+      return IntervalSet::make(e1, e2);
+    } else if (is_negative_const(b.min)) {
+      return IntervalSet::make(e2, e1);
+    } else if (a.is_bounded()) {
+      Expr cmp = b.min >= make_zero(b.min.type().element_of());
+      return IntervalSet::make(select(cmp, e1, e2), select(cmp, e2, e1));
+    }
+  }
+  LOG(WARNING) << "Return Everything in CombineInterval Mul";
+  return IntSet::everything();
 }
 
 template<>
-inline IntSet BinaryCombine<Mul>(IntSet a, IntSet b) {
-  if (IsNumber(a)) {
-    return BinaryMul(a, AsNumber(b));
-  } else if (IsNumber(b)) {
-    return BinaryMul(b, AsNumber(a));
-  } else {
-    return IntSet::make_all_set();
+inline IntSet CombineInterval<Max>(Interval a, Interval b) {
+  if (a.is_single_point() && b.is_single_point()) {
+    return IntSet::single_point(ComputeExpr<Max>(a.min, b.min));
   }
+  return IntervalSet::make(Interval::make_max(a.min, b.min),
+                           Interval::make_max(a.max, b.max));
 }
 
-}  // namespace
-
-inline const IntSetNode* IntSet::operator->() const {
-  return static_cast<const IntSetNode*>(node_.get());
+template<>
+inline IntSet CombineInterval<Min>(Interval a, Interval b) {
+  if (a.is_single_point() && b.is_single_point()) {
+    return IntSet::single_point(ComputeExpr<Min>(a.min, b.min));
+  }
+  return IntervalSet::make(Interval::make_min(a.min, b.min),
+                           Interval::make_min(a.max, b.max));
 }
 
-TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable)
-.set_dispatch<IntSetNode>([](const IntSetNode *op, IRPrinter *p) {
-    p->stream << "int-set(base=";
-    p->print(op->base);
-    p->stream << ')';
-  });
+template<typename OP>
+inline IntSet CombineInterval_(IntSet a, IntSet b) {
+  return CombineInterval<OP>(
+      a.as<IntervalSet>()->i, b.as<IntervalSet>()->i);
+}
 
-IntSet IntSet::make_range(Range dom) {
-  auto n = std::make_shared<IntSetNode>();
-  n->base = dom;
+// stride related
+inline IntSet AsStrideSet(IntSet a) {
+  if (a.as<StrideSet>()) return a;
+  const IntervalSet* s = a.as<IntervalSet>();
+  CHECK(s->i.is_bounded());
+  std::shared_ptr<StrideSet> n = std::make_shared<StrideSet>();
+  n->base = s->i;
   return IntSet(n);
 }
+template<typename OP>
+inline IntSet CombineSets(IntSet a, IntSet b) {
+  return CombineInterval_<OP>(a.cover_interval(), b.cover_interval());
+}
 
-Range IntSet::GetCoverRange() const {
-  const IntSetNode* s = operator->();
-  CHECK(s != nullptr) << "empty set";
-  if (s->domain.size() == 0 && s->concrete.size() == 0) {
-    return s->base;
+template<>
+inline IntSet CombineSets<Add>(IntSet a, IntSet b) {
+  const IntervalSet* a_int = a.as<IntervalSet>();
+  const IntervalSet* b_int = b.as<IntervalSet>();
+  if (a_int && is_zero(a_int->i.min)) return b;
+  if (b_int && is_zero(b_int->i.min)) return a;
+  a = AsStrideSet(a);
+  b = AsStrideSet(b);
+  const StrideSet* a_stride = a.as<StrideSet>();
+  const StrideSet* b_stride = b.as<StrideSet>();
+  auto n = std::make_shared<StrideSet>(*a_stride);
+  for (size_t i = 0; i < b_stride->extents.size(); ++i) {
+    n->extents.push_back(b_stride->extents[i]);
+    n->strides.push_back(b_stride->strides[i]);
   }
-  LOG(FATAL) << "not yet implemented";
-  return Range();
+  n->base = CombineInterval<Add>(
+      a_stride->base, b_stride->base).as<IntervalSet>()->i;
+  return IntSet(n);
 }
 
-IntSet IntSet::make_point(Expr point) {
-  return IntSet::make_range(Range::make_with_min_extent(point, 1));
+inline IntSet NegateSet(IntSet a) {
+  const IntervalSet* a_int = a.as<IntervalSet>();
+  if (a_int) {
+    if (a_int->i.is_single_point()) {
+      return IntSet::single_point(-a_int->i.min);
+    } else {
+      Interval r = Interval::everything();
+      if (a_int->i.has_upper_bound()) {
+        r.min = -(a_int->i.max);
+      }
+      if (a_int->i.has_lower_bound()) {
+        r.max = -(a_int->i.min);
+      }
+      return IntervalSet::make(r);
+    }
+  } else {
+    return NegateSet(a.cover_interval());
+  }
 }
 
-IntSet IntSet::make_all_set() {
-  LOG(FATAL) << "TODO";
-  return IntSet();
+template<>
+inline IntSet CombineSets<Sub>(IntSet a, IntSet b) {
+  return CombineSets<Add>(a, NegateSet(b));
 }
 
-IntSet Union(const Array<IntSet>& set) {
-  if (set.size() == 1) return set[0];
-  LOG(FATAL) << "TODO";
-  return IntSet();
+TVM_DECLARE_LOGICAL_OP(And);
+TVM_DECLARE_LOGICAL_OP(Or);
+TVM_DECLARE_LOGICAL_OP(EQ);
+TVM_DECLARE_LOGICAL_OP(NE);
+TVM_DECLARE_LOGICAL_OP(GE);
+TVM_DECLARE_LOGICAL_OP(GT);
+TVM_DECLARE_LOGICAL_OP(LE);
+TVM_DECLARE_LOGICAL_OP(LT);
+TVM_DECLARE_LOGICAL_OP(Not);
+
+// generic combine operations of two sets
+template<typename OP>
+inline IntSet Combine(const IntSet& a, const IntSet &b) {
+  if (is_logical_op<OP>::value) {
+    return IntervalSet::make(0, 1);
+  }
+  const IntervalSet* a_int = a.as<IntervalSet>();
+  const IntervalSet* b_int = b.as<IntervalSet>();
+  if (a_int && a_int->i.is_everything()) return a;
+  if (b_int && b_int->i.is_everything()) return b;
+  if (a_int && b_int) {
+    return CombineInterval<OP>(a_int->i, b_int->i);
+  }
+  if (a_int && !(a_int->i.is_bounded())) {
+    return CombineInterval_<OP>(a, b.cover_interval());
+  }
+  if (b_int && !(b_int->i.is_bounded())) {
+    return CombineInterval_<OP>(a.cover_interval(), b);
+  }
+  return CombineSets<OP>(a, b);
 }
 
+// Implementation of Evaluations and passing.
 void PassUp(const SplitNode* s,
             const std::unordered_map<IterVar, Range>& dom_map,
             const IntSet& outer,
@@ -215,33 +358,21 @@ void PassUp(const SplitNode* s,
   if (dom_map.count(s->outer) &&
       dom_map.count(s->inner) &&
       dom_map.count(s->parent) &&
-      Match(outer, dom_map.at(s->outer)) &&
-      Match(inner, dom_map.at(s->inner))) {
-    *parent = IntSet::make_range(dom_map.at(s->parent));
+      MatchRange(outer, dom_map.at(s->outer)) &&
+      MatchRange(inner, dom_map.at(s->inner))) {
+    *parent = IntSet::range(dom_map.at(s->parent));
     return;
   }
   Expr factor = dom_map.at(s->inner)->extent;
+  Expr parent_min = dom_map.at(s->parent)->min;
   CHECK(outer.defined());
   CHECK(inner.defined());
   CHECK(factor.defined());
-  // copy construct
-  auto n = std::make_shared<IntSetNode>(*(inner.operator->()));
-
-  if (IsNumber(outer)) {
-    // shift the base offset
-    n->base = Range::make_with_min_extent(
-        AsNumber(outer) * factor + inner->base->min,
-        inner->base->extent);
-  } else {
-    // default use all domains in the data.
-    n->domain.push_back(outer->base);
-    n->stride.push_back(factor);
-    for (size_t i = 0; i < outer->domain.size(); ++i) {
-      n->domain.push_back(outer->domain[i]);
-      n->stride.push_back(outer->stride[i] * factor);
-    }
-  }
-  *parent = IntSet(n);
+
+  *parent = Combine<Add>(
+      Combine<Add>(
+          Combine<Mul>(outer, IntSet::single_point(factor)), inner),
+      IntSet::single_point(parent_min));
 }
 
 void PassUp(const FuseNode* s,
@@ -253,29 +384,51 @@ void PassUp(const FuseNode* s,
   CHECK(dom_map.count(s->inner));
   CHECK(dom_map.count(s->fused));
 
-  if (Match(fused, dom_map.at(s->fused))) {
-    *outer = IntSet::make_range(dom_map.at(s->outer));
-    *inner = IntSet::make_range(dom_map.at(s->inner));
+  if (MatchRange(fused, dom_map.at(s->fused))) {
+    *outer = IntSet::range(dom_map.at(s->outer));
+    *inner = IntSet::range(dom_map.at(s->inner));
     return;
   }
 
-  if (IsNumber(fused)) {
-    Expr value = AsNumber(fused);
+  Expr outer_min = dom_map.at(s->outer)->min;
+  Expr inner_min = dom_map.at(s->inner)->min;
+
+  const IntervalSet* fused_int = fused.as<IntervalSet>();
+
+  if (fused_int && fused_int->i.is_single_point()) {
+    Expr value = fused_int->i.min;
     Expr factor = dom_map.at(s->inner)->extent;
-    *outer = IntSet::make_point(value / factor);
-    *inner = IntSet::make_point(value % factor);
+    Expr v_outer  = value / factor;
+    Expr v_inner  = value % factor;
+    if (!is_zero(outer_min)) v_outer = v_outer + outer_min;
+    if (!is_zero(inner_min)) v_inner = v_inner + inner_min;
+    *outer = IntSet::single_point(v_outer);
+    *inner = IntSet::single_point(v_inner);
   } else {
     LOG(WARNING) << "use fallback inference rule in fuse";
     // simply use the entire set, this rule can be enhanced.
-    *outer = IntSet::make_range(dom_map.at(s->outer));
-    *inner = IntSet::make_range(dom_map.at(s->inner));
+    *outer = IntSet::range(dom_map.at(s->outer));
+    *inner = IntSet::range(dom_map.at(s->inner));
+    return;
+  }
+}
+
+
+void PassUp(const RebaseNode* s,
+            const std::unordered_map<IterVar, Range>& dom_map,
+            const IntSet& rebased,
+            IntSet* parent) {
+  CHECK(dom_map.count(s->parent));
+  if (MatchRange(rebased, dom_map.at(s->rebased))) {
+    *parent = IntSet::range(dom_map.at(s->parent));
     return;
   }
+  Expr parent_min = dom_map.at(s->parent)->min;
+  *parent = Combine<Add>(rebased, IntSet::single_point(parent_min));
 }
 
-namespace {
-// evaluator to evaluate the int set
-class IRSetEvaluator {
+// Evaluator to evalute the epxression.
+class IntSetEvaluator {
  public:
   inline IntSet Eval(Expr expr) {
     static const FType& f = vtable();
@@ -283,11 +436,11 @@ class IRSetEvaluator {
       return f(expr, expr, this);
     } else {
       LOG(WARNING) << "cannot evaluate set type " << expr->type_key();
-      return IntSet::make_all_set();
+      return IntSet::everything();
     }
   }
 
-  using FType = tvm::IRFunctor<IntSet (const NodeRef&, const Expr&, IRSetEvaluator *)>;
+  using FType = tvm::IRFunctor<IntSet (const NodeRef&, const Expr&, IntSetEvaluator *)>;
   static FType& vtable() {  // NOLINT(*)
     static FType inst; return inst;
   }
@@ -295,76 +448,84 @@ class IRSetEvaluator {
   std::unordered_map<const Variable*, IntSet> dom_map;
 };
 
-inline IntSet ConstOp(const NodeRef&, const Expr& e, IRSetEvaluator*) {
-  return IntSet::make_point(e);
+inline IntSet ConstOp(const NodeRef&, const Expr& e, IntSetEvaluator*) {
+  return IntSet::single_point(e);
 }
 
-TVM_STATIC_IR_FUNCTOR(IRSetEvaluator, vtable)
+TVM_STATIC_IR_FUNCTOR(IntSetEvaluator, vtable)
 .set_dispatch<IntImm>(ConstOp)
 .set_dispatch<UIntImm>(ConstOp)
 .set_dispatch<FloatImm>(ConstOp);
 
-TVM_STATIC_IR_FUNCTOR(IRSetEvaluator, vtable)
-.set_dispatch<Variable>([](const Variable* op, const Expr& e, IRSetEvaluator* m) {
+TVM_STATIC_IR_FUNCTOR(IntSetEvaluator, vtable)
+.set_dispatch<Variable>([](const Variable* op, const Expr& e, IntSetEvaluator* m) {
     auto it = m->dom_map.find(op);
     if (it != m->dom_map.end()) {
       return it->second;
     } else {
-      return IntSet::make_point(e);
+      return IntSet::single_point(e);
     }
   });
 
 // binary operator
 template<typename T>
-inline IntSet Binary(const T* op, const Expr& e, IRSetEvaluator* m) {
+inline IntSet Binary(const T* op, const Expr& e, IntSetEvaluator* m) {
   IntSet a = m->Eval(op->a);
   IntSet b = m->Eval(op->b);
-  if (IsNumber(a) && IsNumber(b)) {
-    if (Match(a, op->a) &&
-        Match(b, op->b)) {
-      return IntSet::make_point(e);
-    } else {
-      return IntSet::make_point(T::make(AsNumber(a), AsNumber(b)));
-    }
-  } else {
-    return BinaryCombine<T>(a, b);
+  if (MatchPoint(a, op->a) && MatchPoint(b, op->b)) {
+    return IntSet::single_point(e);
   }
+  IntSet r = Combine<T>(a, b);
+  return r;
 }
 
-TVM_STATIC_IR_FUNCTOR(IRSetEvaluator, vtable)
+TVM_STATIC_IR_FUNCTOR(IntSetEvaluator, vtable)
 .set_dispatch<Add>(Binary<Add>)
 .set_dispatch<Sub>(Binary<Sub>)
 .set_dispatch<Mul>(Binary<Mul>)
 .set_dispatch<Div>(Binary<Div>)
 .set_dispatch<Mod>(Binary<Mod>)
 .set_dispatch<Min>(Binary<Min>)
-.set_dispatch<Max>(Binary<Max>);
-
-// use simply bound for logical expressions for now.
-inline IntSet Logical(const NodeRef&, const Expr& e, IRSetEvaluator*) {
-  return IntSet::make_range(Range::make_with_min_extent(0, 2));
-}
-
-TVM_STATIC_IR_FUNCTOR(IRSetEvaluator, vtable)
-.set_dispatch<EQ>(Logical)
-.set_dispatch<NE>(Logical)
-.set_dispatch<LT>(Logical)
-.set_dispatch<LE>(Logical)
-.set_dispatch<GT>(Logical)
-.set_dispatch<GE>(Logical)
-.set_dispatch<And>(Logical)
-.set_dispatch<Or>(Logical);
-
-}  // namespace
+.set_dispatch<Max>(Binary<Max>)
+.set_dispatch<EQ>(Binary<EQ>)
+.set_dispatch<NE>(Binary<NE>)
+.set_dispatch<LT>(Binary<LT>)
+.set_dispatch<LE>(Binary<LE>)
+.set_dispatch<GT>(Binary<GT>)
+.set_dispatch<GE>(Binary<GE>)
+.set_dispatch<And>(Binary<And>)
+.set_dispatch<Or>(Binary<Or>);
 
 IntSet EvalSet(Expr e,
                const Map<IterVar, IntSet>& dom_map) {
-  IRSetEvaluator m;
+  IntSetEvaluator m;
   for (auto kv : dom_map) {
     m.dom_map[kv.first->var.as<Variable>()] = kv.second;
   }
   return m.Eval(e);
 }
 
+IntSet EvalSet(Range r,
+               const Map<IterVar, IntSet>& dom_map) {
+  IntSetEvaluator m;
+  for (auto kv : dom_map) {
+    m.dom_map[kv.first->var.as<Variable>()] = kv.second;
+  }
+  IntSet min_set = m.Eval(r->min);
+  IntSet ext_set = m.Eval(r->extent).cover_interval();
+  const Interval& ei = ext_set.as<IntervalSet>()->i;
+  if (!ei.has_upper_bound()) return IntSet::everything();
+  ext_set = IntervalSet::make(0, ComputeExpr<Sub>(ei.max, 1));
+  return Combine<Add>(min_set, ext_set);
+}
+
+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 schedule
 }  // namespace tvm

src/schedule/int_set.h

@@ -22,35 +22,48 @@ class IntSet : public NodeRef {
  public:
   /*! \brief constructor */
   IntSet() {}
-  // constructor from not deontainer.
+  // constructor from not container.
   explicit IntSet(std::shared_ptr<Node> n) : NodeRef(n) {}
-  /*! \return whether the set is empty */
-  inline bool is_empty() const {
-    return !defined();
-  }
-  /*!
-   * \return a range that covers the IntSet
-   */
-  Range GetCoverRange() const;
   /*!
    * \brief access the internal node container
    * \return the pointer to the internal node container
    */
   inline const IntSetNode* operator->() const;
   /*!
-   * \param dom The domain to be created.
-   * \return create integer set from existing domain
+   * \brief Find a range that covers the region.
+   * \param max_range The range to be covered.
+   * \return The covering range.
+   */
+  Range cover_range(Range max_range) const;
+  /*!
+   * \brief find an interval that covers the set.
+   * \return The covering interval set.
    */
-  static IntSet make_range(Range dom);
+  IntSet cover_interval() const;
+  /*! \return Whether the set represent everything  */
+  bool is_everything() const;
+  /*! \return Whether the set is a single point */
+  bool is_single_point() const;
+  /*! \return Whether the set contains everything */
+  static IntSet everything();
   /*!
-   * \param point
-   * \return create integer set that only contains one point
+   * \brief construct a point set.
+   * \param point The point in the set.
+   * \return construct a single point set
    */
-  static IntSet make_point(Expr point);
+  static IntSet single_point(Expr point);
   /*!
-   * \return create integer set that represents everything
+   * \brief Construct a set representing a range.
+   * \param r The range
+   * \return constructed set.
    */
-  static IntSet make_all_set();
+  static IntSet range(Range r);
+};
+
+/*!
+ * \brief Base class of all IntSet containers.
+ */
+struct IntSetNode : public Node {
 };
 
 /*!
@@ -63,6 +76,18 @@ class IntSet : public NodeRef {
  */
 IntSet EvalSet(Expr e,
                const Map<IterVar, IntSet>& dom_map);
+
+/*!
+ * \brief Find an symbolic integer set that contains is union over
+ *  all the possible conditional values in dom_map.
+ *
+ * \param r The initial range.
+ * \param dom_map The domain of each variable.
+ * \return An integer set that can cover all the possible values.
+ */
+IntSet EvalSet(Range r,
+               const Map<IterVar, IntSet>& dom_map);
+
 /*!
  * \brief Conditional upward message passing.
  *
@@ -99,6 +124,23 @@ void PassUp(const FuseNode* s,
             const IntSet& fused,
             IntSet* outer,
             IntSet* inner);
+
+/*!
+ * \brief Conditional upward message passing.
+ *
+ * Get domain of parent, condition on domain of children.
+ * Domain is represented as IntSet.
+ *
+ * \param s The Fuse relation node.
+ * \param dom_map The old domain result from downward message passing.
+ *    Contains the domain set if all the children are full set.
+ * \param rebased domain of rebased iteration.
+ * \param parent The result domain of parent iteration.
+ */
+void PassUp(const RebaseNode* s,
+            const std::unordered_map<IterVar, Range>& dom_map,
+            const IntSet& fused,
+            IntSet* parent);
 /*!
  * \brief Create an union set of all sets
  * \param sets The sets to be unioned
@@ -106,6 +148,11 @@ void PassUp(const FuseNode* s,
  */
 IntSet Union(const Array<IntSet>& sets);
 
+// implementation
+inline const IntSetNode* IntSet::operator->() const {
+  return static_cast<const IntSetNode*>(node_.get());
+}
+
 }  // namespace schedule
 }  // namespace tvm
 

src/schedule/schedule_lang.cc

@@ -81,7 +81,7 @@ Stage& Stage::compute_at(Stage parent, IterVar scope) {   // NOLINT(*)
     }
   }
   CHECK(found)
-      << "Cannot compute at a iteration variable that is not part of parent leaf vars";
+      << "Cannot find the specified axis in parent stage's leaf_iter_vars";
   return *this;
 }
 
@@ -165,7 +165,6 @@ Stage& Stage::tile(IterVar x_parent, IterVar y_parent,
   return *this;
 }
 
-
 Schedule::Schedule(Array<Operation> ops) {
   auto n = std::make_shared<ScheduleNode>();
   n->roots = ops;
@@ -203,9 +202,53 @@ IterVarRelation FuseNode::make(
   return IterVarRelation(n);
 }
 
+IterVarRelation RebaseNode::make(IterVar parent, IterVar rebased) {
+  auto n = std::make_shared<RebaseNode>();
+  n->parent = parent;
+  n->rebased = rebased;
+  return IterVarRelation(n);
+}
+
+void Schedule::normalize() {
+  std::unordered_map<IterVar, IterVar> rebase_map;
+  std::unordered_map<const Node*, int> attach_mark;
+
+
+  for (Stage s : (*this)->stages) {
+    if (s->attach_type == kScope) {
+      attach_mark[s->attach_stage.get()] = 1;
+    }
+  }
+
+  for (Stage s : (*this)->stages) {
+    if (!attach_mark.count(s.get())) continue;
+    auto root_iter_vars = s->op->root_iter_vars();
+    ArrayNode* leaf_vars = s->leaf_iter_vars.CopyOnWrite();
+
+    for (IterVar iv : root_iter_vars) {
+      size_t idx = FindIterVar(leaf_vars, iv);
+      if (idx < leaf_vars->data.size()) {
+        // insert rebase
+        IterVar rebased(Range(), iv->var->name_hint + ".rb");
+        s->relations.push_back(RebaseNode::make(iv, rebased));
+        leaf_vars->data[idx] = rebased.node_;
+        rebase_map[iv] = rebased;
+      }
+    }
+  }
+  // remap the parent relation
+  for (Stage s : (*this)->stages) {
+    if (s->attach_type != kScope) continue;
+    if (rebase_map.count(s->attach_ivar)) {
+      s->attach_ivar = rebase_map.at(s->attach_ivar);
+    }
+  }
+}
+
 TVM_REGISTER_NODE_TYPE(StageNode);
 TVM_REGISTER_NODE_TYPE(SplitNode);
 TVM_REGISTER_NODE_TYPE(FuseNode);
+TVM_REGISTER_NODE_TYPE(RebaseNode);
 TVM_REGISTER_NODE_TYPE(ScheduleNode);
 
 }  // namespace tvm

src/pass/schedule_ops.cc → src/schedule/schedule_ops.cc

@@ -8,12 +8,44 @@
 #include <tvm/ir_visitor.h>
 #include <tvm/schedule_pass.h>
 
-#include "./scope.h"
-#include "./ir_util.h"
-#include "../schedule/graph.h"
+#include "../pass/ir_util.h"
+#include "./int_set.h"
+#include "./graph.h"
 
 namespace tvm {
-namespace ir {
+namespace schedule {
+
+using namespace ir;
+
+/*!
+ * \brief message passing to find if IterVar is related to reduction.
+ * \param s The stage to be used.
+ * \param p_state The message passing state
+ *     IterVar->flag
+ */
+void PassDownFlag(const Stage& s,
+                   std::unordered_map<IterVar, int>* p_state) {
+  auto& state = *p_state;
+  for (IterVarRelation rel : s->relations) {
+    if (rel.as<SplitNode>()) {
+      const SplitNode* s = rel.as<SplitNode>();
+      int flag = state.at(s->parent);
+      state[s->outer] = flag;
+      state[s->inner] = flag;
+    } else if (rel.as<FuseNode>()) {
+      const FuseNode* s = rel.as<FuseNode>();
+      int flag_outer = state.at(s->outer);
+      int flag_inner = state.at(s->inner);
+      state[s->fused] = flag_outer | flag_inner;
+    } else if (rel.as<RebaseNode>()) {
+      const RebaseNode* s = rel.as<RebaseNode>();
+      int flag = state.at(s->parent);
+      state[s->rebased] = flag;
+    } else {
+      LOG(FATAL) << "unknown relation type";
+    }
+  }
+}
 
 /*!
  * \brief use message passing to calculate the assignment of each Var inside the loop body.
@@ -37,7 +69,7 @@ void PassUpOffset(const Stage& s,
       state[s->parent] = inner + outer * factor;
       // add min if they exist
       if (!is_zero(parent_min)) {
-        state[s->parent] = parent_min + state[s->parent];
+        state[s->parent] = state[s->parent] + parent_min;
       }
     } else if (rel.as<FuseNode>()) {
       const FuseNode* s = rel.as<FuseNode>();
@@ -49,10 +81,20 @@ void PassUpOffset(const Stage& s,
       state[s->inner] = value % factor;
       // add min if they exist
       if (!is_zero(outer_min)) {
-        state[s->outer] = outer_min + state[s->outer];
+        state[s->outer] = state[s->outer] + outer_min;
       }
       if (!is_zero(inner_min)) {
-        state[s->inner] = outer_min + state[s->inner];
+        state[s->inner] = state[s->inner] + inner_min;
+      }
+    } else if (rel.as<RebaseNode>()) {
+      const RebaseNode* s = rel.as<RebaseNode>();
+      Expr value = state.at(s->rebased);
+      Expr parent_min = dom_map.at(s->parent)->min;
+      // add min if they exist
+      if (!is_zero(parent_min)) {
+        state[s->parent] = value + parent_min;
+      } else {
+        state[s->parent] = value;
       }
     } else {
       LOG(FATAL) << "unknown relation type";
@@ -60,76 +102,54 @@ void PassUpOffset(const Stage& s,
   }
 }
 
-/*!
- * \brief split the expr by addition.
- * \param expr The expression to be splitted.
- * \param loop_level The loop level of each Variable
- * \param result vector of (level, expr)
- *   The level gives the mimimum loop level this expression need to be computed.
- *   The Expr gives the expression content.
- */
-void SplitByAdd(Expr expr,
-                const std::unordered_map<const Variable*, size_t>& loop_level,
-                std::vector<std::pair<size_t, Expr> > *result) {
-  const Add* op = expr.as<Add>();
-  if (op != nullptr) {
-    SplitByAdd(op->a, loop_level, result);
-    SplitByAdd(op->b, loop_level, result);
-  } else {
-    size_t max_level = 0;
-    auto fvisit = [&max_level, &loop_level](const NodeRef& n) {
-      const Variable* op = n.as<Variable>();
-      if (op != nullptr) {
-        auto it = loop_level.find(op);
-        if (it != loop_level.end()) {
-          max_level = std::max(max_level, it->second);
-        }
-      }
-    };
-    PostOrderVisit(expr, fvisit);
-    result->push_back(std::make_pair(max_level, expr));
-  }
-}
-
-/*!
- * \brief Make the loop nest of the correspondings schedule.
- * \param sch The schedule.
- * \param dom_map The domain map.
- *
- * \return a nested representation of loop statements.
- *  The flattened Stmt are ordered from outmost to inner most order.
- */
-std::vector<std::vector<Stmt> > MakeLoopNest(
-    const Stage& sch,
-    const Map<IterVar, Range>& dom_map) {
-  // optional, use let to define some CSE in dom_map.
+std::vector<std::vector<Stmt> >
+MakeLoopNest(const Stage& sch,
+             const Map<IterVar, Range>& dom_map,
+             size_t begin_loop,
+             bool reduce_init_loop,
+             std::unordered_map<IterVar, Expr>* p_value_map,
+             const std::unordered_map<IterVar, bool>& skip_iter) {
   auto leaf_iter_vars = sch->leaf_iter_vars;
-  std::unordered_map<IterVar, Expr> offset;
-  std::unordered_map<const Variable*, size_t> loop_level;
   Stmt no_op = Evaluate::make(0);
   // create the loop nest
   std::vector<std::vector<Stmt> > nest;
   nest.resize(leaf_iter_vars.size() + 1);
+  std::unordered_map<IterVar, Expr>& value_map = *p_value_map;
 
-  for (size_t i = 0; i < leaf_iter_vars.size(); ++i) {
+  for (size_t i = begin_loop; i < leaf_iter_vars.size(); ++i) {
     auto iv = leaf_iter_vars[i];
+    if (skip_iter.count(iv) && skip_iter.at(iv)) {
+      // skip this iteration.
+      value_map[iv] = iv->var;
+      continue;
+    }
+
     Range dom = dom_map.at(iv);
     // initialize the offset and loop_level
-    offset[iv] = iv->var;
-    loop_level[iv->var.as<Variable>()] = i + 1;
+    Var var = iv->var;
+    if (reduce_init_loop) {
+      var = Var(iv->var->name_hint + ".init", iv->var.type());
+    }
     // Mark the iter var in the IR, to remember the point
     if (iv->thread_tag.length() == 0) {
-      if (is_zero(dom->min)) {
+      if (is_one(dom->extent)) {
+        nest[i + 1].emplace_back(
+            LetStmt::make(var, dom->min, no_op));
+        value_map[iv] = dom->min;
+      } else if (is_zero(dom->min)) {
         nest[i + 1].emplace_back(
-            For::make(iv->var, 0, dom->extent,
+            For::make(var, 0, dom->extent,
                       ForType::Serial, DeviceAPI::None, no_op));
+        value_map[iv] = var;
       } else {
         Var idx(iv->var->name_hint + ".idx", iv->var.type());
         nest[i + 1].emplace_back(
             For::make(idx, 0, dom->extent,
                       ForType::Serial, DeviceAPI::None, no_op));
+        Expr new_value = dom->min + idx;
+        value_map[iv] = new_value;
         nest[i + 1].emplace_back(
-            LetStmt::make(iv->var, dom->min + idx, no_op));
+            LetStmt::make(var, new_value, no_op));
       }
     } else {
       // Always restrict threaded IterVar to starts from 0.
@@ -137,69 +157,73 @@ std::vector<std::vector<Stmt> > MakeLoopNest(
       // annotate the extent of the IterVar
       nest[i + 1].emplace_back(
           AttrStmt::make(iv, "thread_extent", dom->extent, no_op));
+      value_map[iv] = var;
+    }
+    if (!reduce_init_loop) {
+      // annotate the extent of the IterVar
+      nest[i + 1].emplace_back(
+          AttrStmt::make(iv, "scope", iv->var, no_op));
     }
-    // annotate the extent of the IterVar
-    nest[i + 1].emplace_back(
-        AttrStmt::make(iv, "scope", iv->var, no_op));
   }
   // message passing to get offset of root iter vars.
-  PassUpOffset(sch, dom_map, &offset);
-
-  for (IterVar iv : sch->op->root_iter_vars()) {
-    Expr value = offset.at(iv);
-    if (!value.same_as(iv->var)) {
-      using Entry = std::pair<size_t, Expr>;
-      std::vector<Entry> splits;
-      SplitByAdd(value, loop_level, &splits);
+  PassUpOffset(sch, dom_map, &value_map);
+  return nest;
+}
 
-      Expr offset = 0;
-      size_t nsplit_left = splits.size() - 1;
-      for (size_t i = 0; i <= leaf_iter_vars.size(); ++i) {
-        size_t hit = 0;
-        for (const auto& kv : splits) {
-          if (kv.first == i) {
-            if (is_zero(offset)) {
-              offset = kv.second;
-            } else {
-              offset = offset + kv.second;
-              ++hit;
-            }
-          }
-        }
-        nsplit_left -= hit;
-        if (hit != 0) {
-          std::ostringstream os;
-          os << iv->var->name_hint << ".at.l" << i;
-          Var base_offset(os.str());
-          if (nsplit_left == 0) {
-            base_offset = iv->var;
-          }
-          nest[i].emplace_back(
-              LetStmt::make(base_offset, offset, no_op));
-          offset = base_offset;
-        }
-      }
-      Range dom = dom_map.at(iv);
-      if (!offset.same_as(iv->var)) {
-        // define the iv->var
-        nest.back().emplace_back(
-            LetStmt::make(iv->var, offset, no_op));
+Stmt MakeLoop(const Stage& s,
+              const Map<IterVar, Range>& dom_map,
+              Stmt provide,
+              Stmt init) {
+  std::unordered_map<IterVar, Expr> value_map;
+  auto nest = MakeLoopNest(s, dom_map, 0, false, &value_map, {});
+  provide = Substitute(provide, value_map);
+  if (init.defined()) {
+    // try to find the location to insert the initialization.
+    // Fuse the initialization and provide loop when possible.
+    std::unordered_map<IterVar, int> reduce_state;
+    const ComputeOpNode* compute = s->op.as<ComputeOpNode>();
+    for (IterVar iv : compute->reduce_axis) {
+      reduce_state[iv] = 2;
+    }
+    for (IterVar iv : compute->axis) {
+      reduce_state[iv] = 1;
+    }
+    // find which iter var is related to reduction and which is related to axis.
+    PassDownFlag(s, &reduce_state);
+    auto leaf_iter_vars = s->leaf_iter_vars;
+    std::unordered_map<IterVar, Expr> init_value_map;
+    // first first loop that is related to reduction.
+    size_t begin_loop = leaf_iter_vars.size();
+    for (size_t i = 0; i < leaf_iter_vars.size(); ++i) {
+      auto iv = leaf_iter_vars[i];
+      int flag = reduce_state.at(iv);
+      if ((flag & 2) != 0) {
+        begin_loop = i; break;
       }
-      Expr condition = (iv->var - dom->min) < dom->extent;
-      // Boundary condition checking
-      // Need better boundary condition here.
-      nest.back().emplace_back(IfThenElse::make(condition, no_op));
+      init_value_map[iv] = value_map.at(iv);
     }
+    // skip loops that does not relates to axis.
+    std::unordered_map<IterVar, bool> skip_iter;
+    for (size_t i = begin_loop; i < leaf_iter_vars.size(); ++i) {
+      auto iv = leaf_iter_vars[i];
+      int flag = reduce_state.at(iv);
+      if ((flag & 1) == 0) skip_iter[iv] = true;
+    }
+    auto init_nest = MakeLoopNest(
+        s, dom_map, begin_loop, true, &init_value_map, skip_iter);
+    init = Substitute(init, init_value_map);
+    init  = MergeNest(init_nest, init);
+    // common nest
+    std::vector<std::vector<Stmt> > common(nest.begin(), nest.begin() + begin_loop);
+    std::vector<std::vector<Stmt> > reduce(nest.begin() + begin_loop, nest.end());
+    provide = MergeNest(reduce, provide);
+    return MergeNest(
+        common, Block::make(init, provide));
+  } else {
+    return MergeNest(nest, provide);
   }
-  return nest;
 }
 
-
-/*!
- * \brief Make pipeline specifically for compute op node.
- * \param op The compute node
- * \param tensors The tensors generated by provide.
- */
 Stmt MakeProvide(const ComputeOpNode* op,
                  const std::vector<Tensor>& tensors) {
   Tensor t = tensors[0];
@@ -210,13 +234,6 @@ Stmt MakeProvide(const ComputeOpNode* op,
   return Provide::make(t->op, t->value_index, op->body, args);
 }
 
-/*!
- * \brief Make pipeline specifically for compute op node.
- * \param op The compute node
- * \param dom_map The domain map
- * \param tensors The tensors generated by provide.
- * \param body The content of the pipeline.
- */
 Stmt MakeRealize(const ComputeOpNode* op,
                  const Map<IterVar, Range>& dom_map,
                  const std::vector<Tensor>& tensors,
@@ -230,6 +247,38 @@ Stmt MakeRealize(const ComputeOpNode* op,
                        bounds, make_const(Bool(1), true), body);
 }
 
+
+void MakeReduction(const ComputeOpNode* op,
+                   const std::vector<Tensor>& tensors,
+                   const Map<IterVar, Range>& dom_map,
+                   Stmt* init,
+                   Stmt* provide) {
+  Stmt no_op = Evaluate::make(0);
+  Tensor t = tensors[0];
+  std::vector<Stmt> nest;
+  Array<Expr>  args;
+  for (IterVar iv : op->axis) {
+    args.push_back(iv->var);
+  }
+  const Reduce* reduce = op->body.as<Reduce>();
+  CHECK(reduce);
+  Expr init_value, update_value;
+  if (reduce->op == "Add") {
+    init_value = make_zero(reduce->type);
+    update_value = Add::make(t(args), reduce->source);
+  } else if (reduce->op == "Max") {
+    init_value = reduce->type.min();
+    update_value = Max::make(t(args), reduce->source);
+  } else if (reduce->op == "Min") {
+    init_value = reduce->type.max();
+    update_value = Min::make(t(args), reduce->source);
+  } else {
+    LOG(FATAL) << "Unsupported reduction " << reduce->op;
+  }
+  *init = Provide::make(t->op, t->value_index, init_value, args);
+  *provide = Provide::make(t->op, t->value_index, update_value, args);
+}
+
 Stmt MakePipeline(const Stage& sch,
                   const Map<IterVar, Range>& dom_map,
                   Stmt consumer) {
@@ -238,14 +287,20 @@ Stmt MakePipeline(const Stage& sch,
     tensors.emplace_back(sch->op.output(i));
   }
 
-  Stmt provide;
-  if (sch->op.as<ComputeOpNode>()) {
-    provide = MakeProvide(sch->op.as<ComputeOpNode>(), tensors);
+  Stmt init, provide;
+
+  const ComputeOpNode* compute = sch->op.as<ComputeOpNode>();
+  if (compute) {
+    if (compute->reduce_axis.size() == 0) {
+      provide = MakeProvide(compute, tensors);
+    } else {
+      MakeReduction(compute, tensors, dom_map, &init, &provide);
+    }
   } else {
     LOG(FATAL) << "not supported op " << sch->op->type_key();
   }
-  std::vector<std::vector<Stmt> > nest = MakeLoopNest(sch, dom_map);
-  Stmt producer = MergeNest(nest, provide);
+
+  Stmt producer = MakeLoop(sch, dom_map, provide, init);
   producer = ProducerConsumer::make(sch->op, true, producer);
 
   Stmt pipeline = producer;
@@ -306,7 +361,6 @@ Stmt InjectInline(const Operation op, Stmt body) {
   return Inline(body, op, args, compute->body);
 }
 
-
 Stmt ScheduleOps(
     Schedule sch, Map<IterVar, Range> dom_map) {
   Stmt body = Stmt();
@@ -330,5 +384,5 @@ Stmt ScheduleOps(
   return body;
 }
 
-}  // namespace ir
+}  // namespace schedule
 }  // namespace tvm

tests/python/integration/test_ewise.py

@@ -18,7 +18,8 @@ def test_add():
 
     # one line to build the function.
     codes = []
-    fadd = tvm.build(s, args=[A, B, C],
+    fadd = tvm.build(s,
+                     args=[A, B, C],
                      target="cuda", name="myadd",
                      record_codes=codes)
     for c in codes:

tests/python/integration/test_reduce.py

+import tvm
+import numpy as np
+
+def test_sum():
+    # graph
+    n = tvm.Var('n')
+    m = tvm.Var('m')
+    A = tvm.placeholder((n, m), name='A')
+    k = tvm.IterVar((0, m))
+    B = tvm.compute((n,), lambda i: tvm.sum(A[i, k], axis=k), name='B')
+    # schedule
+    s = tvm.Schedule(B.op)
+    # create iter var and assign them tags.
+    num_thread = 1
+    block_x = tvm.IterVar(thread_tag="blockIdx.x")
+    thread_x = tvm.IterVar((0, num_thread), thread_tag="threadIdx.x")
+    _, x = s[B].split(B.op.axis[0], factor=num_thread, outer=block_x)
+    _, x = s[B].split(x, outer=thread_x)
+
+    tvm.init_opencl()
+    codes = []
+    fsum = tvm.build(s,
+                     args=[A, B],
+                     target="opencl", name="myadd",
+                     record_codes=codes)
+    for c in codes:
+        print(c)
+    num_device = 1
+    for i in range(num_device):
+        ctx = tvm.opencl(i)
+        if not ctx.enabled:
+            continue
+        # launch the kernel.
+        n = 1028
+        m = 129
+        #a = tvm.nd.array(np.zeros((n, m)).astype(A.dtype), ctx)
+        a = tvm.nd.array(np.random.uniform(size=(n, m)).astype(A.dtype), ctx)
+        b  = tvm.nd.array(np.zeros(n, dtype=B.dtype), ctx)
+        fsum(a, b)
+        np.testing.assert_allclose(
+            b.asnumpy(), np.sum(a.asnumpy(), axis=1), rtol=1e-4)
+
+
+if __name__ == "__main__":
+    test_sum()

tests/python/unittest/test_codegen_device.py

@@ -18,8 +18,7 @@ def test_add_pipeline():
 
     # compile to IR
     bounds = tvm.schedule.InferBound(s)
-    stmt = tvm.ir_pass.ScheduleOps(s, bounds)
-
+    stmt = tvm.schedule.ScheduleOps(s, bounds)
     Ab = tvm.Buffer(A.shape, A.dtype, name='A')
     Bb = tvm.Buffer(B.shape, B.dtype, name='B')
     Cb = tvm.Buffer(C.shape, C.dtype, name='C')

tests/python/unittest/test_codegen_makeapi.py

@@ -10,12 +10,13 @@ def test_makeapi():
     s = tvm.Schedule(C.op)
 
     bounds = tvm.schedule.InferBound(s)
-    stmt = tvm.ir_pass.ScheduleOps(s, bounds)
+    stmt = tvm.schedule.ScheduleOps(s, bounds)
 
     Ab = tvm.Buffer(A.shape, A.dtype, name='A')
     Bb = tvm.Buffer(B.shape, B.dtype, name='B')
     Cb = tvm.Buffer(C.shape, C.dtype, name='C')
     stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, B:Bb, C:Cb})
+
     num_packed_args = 2
     f = tvm.codegen.MakeAPI(stmt, "myadd", [n, Ab, Bb, Cb], num_packed_args)
     assert(f.handle_data_type[Ab.data].dtype == Ab.dtype)

tests/python/unittest/test_lang_tensor.py

@@ -26,7 +26,7 @@ def test_tensor_reduce():
     B = tvm.placeholder((n, l), name='B')
     T = tvm.compute((m, n, l), lambda i, j, k: A[i, k] * B[j, k])
     rv = tvm.IterVar((0, A.shape[1]), name="k")
-    C = tvm.compute((m, n), lambda i, j: tvm.sum(T(i, j, rv+1), rdom=rv))
+    C = tvm.compute((m, n), lambda i, j: tvm.sum(T(i, j, rv+1), axis=rv))
     # json load save
     C_json = tvm.save_json(C)
     C_loaded = tvm.load_json(C_json)

tests/python/unittest/test_pass_storage_flatten.py

@@ -12,7 +12,7 @@ def test_flatten2():
     s[A1].compute_at(s[A2], xo)
     bounds = tvm.schedule.InferBound(s)
     assert isinstance(bounds, tvm.collections.Map)
-    stmt = tvm.ir_pass.ScheduleOps(s, bounds)
+    stmt = tvm.schedule.ScheduleOps(s, bounds)
 
     print(stmt)
     Ab = tvm.Buffer(A.shape, A.dtype, name='A')

tests/python/unittest/test_pass_schedule_ops.py → tests/python/unittest/test_schedule_schedule_ops.py

@@ -11,7 +11,7 @@ def test_schedule0():
 
     bounds = tvm.schedule.InferBound(s)
     assert isinstance(bounds, tvm.collections.Map)
-    stmt = tvm.ir_pass.ScheduleOps(s, bounds)
+    stmt = tvm.schedule.ScheduleOps(s, bounds)
     print(stmt)
 
 def test_schedule1():
@@ -24,7 +24,7 @@ def test_schedule1():
     xo, xi = s[A1].split(A1.op.axis[0], 8)
     bounds = tvm.schedule.InferBound(s)
     assert isinstance(bounds, tvm.collections.Map)
-    stmt = tvm.ir_pass.ScheduleOps(s, bounds)
+    stmt = tvm.schedule.ScheduleOps(s, bounds)
     print(stmt)
 
 def test_schedule2():
@@ -39,7 +39,7 @@ def test_schedule2():
     s[A1].compute_at(s[A2], xo)
     bounds = tvm.schedule.InferBound(s)
     assert isinstance(bounds, tvm.collections.Map)
-    stmt = tvm.ir_pass.ScheduleOps(s, bounds)
+    stmt = tvm.schedule.ScheduleOps(s, bounds)
     print(stmt)