[Arith] Inequalities solver

.gitignore

@@ -196,6 +196,7 @@ tvm_t.*
 .python_history
 .pytest_cache
 .local
+cmake-build-debug
 
 # Visual Studio Code
 .vscode

include/tvm/arith/analyzer.h

@@ -464,11 +464,16 @@ class TVM_DLL Analyzer {
    * \brief Simplify expr.
    *
    * \param expr The expression to be simplified.
+   * \param steps The simplification runs in the order of
+   *        rewrite_simplify (step 1) -> canonical_simplify (step 2) ->
+   *        rewrite_simplify (step 3) -> canonical_simplify (step 4) -> ...
+   *        param steps controls how many steps to run.
+   *        Default is 2, i.e., rewrite_simplify + canonical_simplify.
    * \return The result.
    *
    * \note Analyzer will call into sub-analyzers to get the result.
    */
-  PrimExpr Simplify(const PrimExpr& expr);
+  PrimExpr Simplify(const PrimExpr& expr, size_t steps = 2);
 };
 
 }  // namespace arith

include/tvm/arith/int_solver.h

@@ -26,17 +26,99 @@
 
 #include <tvm/ir/expr.h>
 #include <tvm/tir/expr.h>
+#include <tvm/tir/op.h>
 
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
+#include "analyzer.h"
+
 namespace tvm {
 namespace arith {
 
 using tir::IterVar;
 using tir::Var;
 using tir::VarNode;
 
+/*!
+ * \brief Represent integer grouped bounds which are classified into
+ *        lower bounds (include), upper bounds (include) and equalities.
+ *        It also contains coefficient as a multiplier for the bounds, i.e.,
+ *        coef * var >= lower
+ *        coef * var == equal
+ *        coef * var <= upper
+ * \sa IntGrpBounds
+ */
+class IntGrpBoundsNode : public Object {
+ public:
+  PrimExpr coef;
+  Array<PrimExpr> lower;
+  Array<PrimExpr> equal;
+  Array<PrimExpr> upper;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("coef", &coef);
+    v->Visit("lower", &lower);
+    v->Visit("equal", &equal);
+    v->Visit("upper", &upper);
+  }
+
+  bool SEqualReduce(const IntGrpBoundsNode* other, SEqualReducer eq) const {
+    return eq(coef, other->coef) && eq(lower, other->lower) && eq(equal, other->equal) &&
+           eq(upper, other->upper);
+  }
+
+  void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce(coef);
+    hash_reduce(lower);
+    hash_reduce(equal);
+    hash_reduce(upper);
+  }
+
+  static constexpr const bool _type_has_method_sequal_reduce = true;
+  static constexpr const char* _type_key = "arith.IntGrpBounds";
+  TVM_DECLARE_FINAL_OBJECT_INFO(IntGrpBoundsNode, Object);
+};
+
+/*!
+ * \brief Managed reference to IntGrpBoundsNode.
+ * \sa IntGrpBoundsNode
+ */
+class IntGrpBounds : public ObjectRef {
+ public:
+  /*!
+   * \brief Constructor by fields
+   * \param coef The coefficient. Must be integer.
+   *        coef * var >= lower
+   *        coef * var == equal
+   *        coef * var >= upper
+   * \param lower the lower bounds (include)
+   * \param equal equalities
+   * \param upper the upper bounds (include)
+   */
+  TVM_DLL IntGrpBounds(PrimExpr coef, Array<PrimExpr> lower, Array<PrimExpr> equal,
+                       Array<PrimExpr> upper);
+
+  /*!
+   * \brief Construct bounds from a range.
+   * \param r The range
+   * \return constructed bounds.
+   */
+  static IntGrpBounds range(const Range& r);
+
+  /*!
+   * \brief Perform substitution on all components of the struct.
+   */
+  IntGrpBounds Substitute(const Map<Var, PrimExpr>& subst) const;
+
+  Range FindBestRange(const Map<Var, Range>& vranges_addl = {}) const;
+
+  IntGrpBounds operator+(const Range& r);
+
+  TVM_DEFINE_OBJECT_REF_METHODS(IntGrpBounds, ObjectRef, IntGrpBoundsNode);
+};
+
 /*!
  * \brief Represent integer constrains including (integer) variables, their ranges and
  *        the relations between them (either equations or inequalities).
@@ -161,6 +243,8 @@ class IntConstraintsTransform : public ObjectRef {
   TVM_DEFINE_OBJECT_REF_METHODS(IntConstraintsTransform, ObjectRef, IntConstraintsTransformNode);
 };
 
+typedef std::pair<Map<Var, IntGrpBounds>, Array<PrimExpr>> PartialSolvedInequalities;
+
 /*!
  * \brief Obtain Smith Normal Form of linear equation A x = y.
  *        Smith Normal Form of matrix A_{mxn} is S_{mxn} = U_{mxm} A_{mxn} V_{nxn},
@@ -191,6 +275,30 @@ void SmithNormalFormDiag(std::vector<std::vector<int64_t>>* S, std::vector<std::
  */
 IntConstraintsTransform SolveLinearEquations(const IntConstraints& system_to_solve);
 
+/*!
+ * \brief Solve linear inequalities.
+ * \param system_to_solve the variables to solve, their ranges, and a list of inequalities.
+ * \return A map of variables and their solved bounds,
+ *         and constrains that cannot be solved to bounds.
+ */
+PartialSolvedInequalities SolveLinearInequalities(const IntConstraints& system_to_solve);
+
+/*!
+ * \brief Solve linear inequalities.
+ * \param system_to_solve the variables to solve, their ranges, and a list of inequalities.
+ * \return The result ranges for each variables.
+ *         Constrains that cannot be transformed to Range will be stored in relations.
+ */
+IntConstraints SolveInequalitiesToRange(const IntConstraints& system_to_solve);
+
+/*!
+ * \brief Solve linear inequalities.
+ * \param system_to_solve the variables to solve, their ranges, and a list of inequalities.
+ * \return Solved ranges are deskewed to be started from zero.
+ *         New variables and the mapping are created accordingly.
+ */
+IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& system_to_solve);
+
 }  // namespace arith
 }  // namespace tvm
 #endif  // TVM_ARITH_INT_SOLVER_H_

python/tvm/arith/__init__.py

@@ -20,4 +20,4 @@
 from .analyzer import ModularSet, ConstIntBound, Analyzer
 from .bound import deduce_bound
 from .pattern import detect_linear_equation, detect_clip_bound
-from .int_solver import solve_linear_equations
+from .int_solver import solve_linear_equations, solve_linear_inequalities

python/tvm/arith/int_solver.py

@@ -20,6 +20,52 @@
 from . import _ffi_api
 
 
+@tvm._ffi.register_object("arith.IntGrpBounds")
+class IntGrpBounds(Object):
+    """Represent integer grouped bounds which are classified into
+       lower bounds (include), upper bounds (include) and equalities.
+
+    Parameters
+    ----------
+    coef : tvm.ir.PrimExpr
+        The coefficient. Must be integer.
+        coef * var >= lower
+        coef * var == equal
+        coef * var >= upper
+    lower : List[tvm.ir.PrimExpr]
+        the lower bounds (include)
+    equal : List[tvm.ir.PrimExpr]
+        equalities
+    upper : List[tvm.ir.PrimExpr]
+        the upper bounds (include)
+    """
+    def __init__(self, coef, lower, equal, upper):
+        self.__init_handle_by_constructor__(
+            _ffi_api.IntGrpBounds, coef, lower, equal, upper)
+
+    @staticmethod
+    def make_by_range(rng):
+        """Construct a IntGroupedBounds by Range.
+
+        Parameters
+        ----------
+        rng : tvm.ir.Range
+
+
+        Returns
+        -------
+        ret : Range
+            The constructed range.
+        """
+        return _ffi_api.int_grouped_bounds_by_range(rng)
+
+    def find_best_range(self):
+        """Return the best range from the grouped bounds.
+           None if (-inf, +inf).
+        """
+        return _ffi_api.IntGrpBounds_FindBestRange(self)
+
+
 @tvm._ffi.register_object("arith.IntConstraints")
 class IntConstraints(Object):
     """Represent a set of integer constraints including variables, their ranges and
@@ -97,3 +143,33 @@ def solve_linear_equations(equations, variables=None, ranges=None):
     if isinstance(equations, IntConstraints):
         return _ffi_api.SolveLinearEquations(equations)
     return _ffi_api.SolveLinearEquations(variables, ranges, equations)
+
+
+def solve_linear_inequalities(equations, variables=None, ranges=None, deskew_range=False):
+    """Solve linear inequalities.
+
+    Parameters
+    ----------
+    equations   : List[tvm.ir.PrimExpr] or IntConstraints
+        The inequalities of the variables
+    variables   : Optional[List[tvm.tir.Var]]
+        The variables in the system.
+    ranges      : Optional[Map[tvm.tir.Var, tvm.ir.Range]]
+        The ranges of the variables.
+    deskew_range: Optional[bool]
+        Whether deskew the result ranges to be started from zero.
+        Default false.
+
+    Returns
+    -------
+    ret_ranges: IntConstraints or IntConstraintsTransform
+        The result ranges for each variables.
+        Constrains that cannot be transformed to Range will be stored in IntConstraints.relations.
+        If deskew_range is set (=True), the result ranges will be deskewed to be started from zero.
+        New variables are created accordingly therefore IntConstraintsTransform is returned.
+    """
+    solver = _ffi_api.SolveInequalitiesDeskewRange \
+        if deskew_range else _ffi_api.SolveInequalitiesToRange
+    if isinstance(equations, IntConstraints):
+        return solver(equations)
+    return solver(variables, ranges, equations)

python/tvm/testing.py

@@ -22,6 +22,7 @@
 import tvm
 import tvm.arith
 import tvm.tir
+import tvm.te
 import tvm._ffi
 
 
@@ -188,5 +189,47 @@ def assert_prim_expr_equal(lhs, rhs):
         raise ValueError("{} and {} are not equal".format(lhs, rhs))
 
 
+def check_bool_expr_is_true(bool_expr, vranges, cond=None):
+    """ Check that bool_expr holds given the condition cond
+    for every value of free variables from vranges.
+
+    Parameters
+    ----------
+    bool_expr : tvm.ir.expr.PrimExpr
+        Boolean expression to check
+    vranges: Dict[tvm.tir.expr.Var, tvm.ir.Range]
+        Free variables and their ranges
+    cond: tvm.ir.expr.PrimExpr
+        extra conditions needs to be satisfied.
+    """
+    if cond is not None:
+        bool_expr = tvm.te.any(tvm.tir.Not(cond), bool_expr)
+
+    def _run_expr(expr, vranges):
+        """ Evaluate expr for every value of free variables
+        given by vranges and return the tensor of results.
+        """
+        def _compute_body(*us):
+            vmap = {v: u + r.min for (v, r), u in zip(vranges.items(), us)}
+            return tvm.tir.stmt_functor.substitute(expr, vmap)
+
+        A = tvm.te.compute([r.extent.value for v, r in vranges.items()], _compute_body)
+        args = [tvm.nd.empty(A.shape, A.dtype)]
+        sch = tvm.te.create_schedule(A.op)
+        mod = tvm.build(sch, [A])
+        mod(*args)
+        return args[0].asnumpy()
+
+    res = _run_expr(bool_expr, vranges)
+    if not np.all(res):
+        indices = list(np.argwhere(res == 0)[0])
+        counterex = [(str(v), i + r.min) for (v, r), i in zip(vranges.items(), indices)]
+        counterex = sorted(counterex, key=lambda x: x[0])
+        counterex = ", ".join([v + " = " + str(i) for v, i in counterex])
+        raise AssertionError("Expression {}\nis not true on {}\n"
+                             "Counterexample: {}"
+                             .format(tvm.tir.ir_pass.CanonicalSimplify(bool_expr),
+                                     vranges, counterex))
+
 
 tvm._ffi._init_api("testing", __name__)

src/arith/analyzer.cc

@@ -115,11 +115,15 @@ bool Analyzer::CanProve(const PrimExpr& expr) {
   return false;
 }
 
-PrimExpr Analyzer::Simplify(const PrimExpr& expr) {
+PrimExpr Analyzer::Simplify(const PrimExpr& expr, size_t steps) {
   if (tir::is_const(expr)) return expr;
-  auto res = this->rewrite_simplify(expr);
-  if (tir::is_const(res)) return res;
-  res = this->canonical_simplify(res);
+  PrimExpr res = expr;
+  for (size_t i = 0; i < steps; ++i) {
+    res = this->rewrite_simplify(res);
+    if (tir::is_const(res) || ++i == steps) return res;
+    res = this->canonical_simplify(res);
+    if (tir::is_const(res)) return res;
+  }
   return res;
 }