[Arith] linear system and equation solver (#5171)

* [arith] linear system and equation solver

Co-authored-by: Sergei Grechanik <sergei.grechanik+h@gmail.com>

* avoid constructing analyzer every time

* generate random test cases and address comments

Co-authored-by: Sergei Grechanik <sergei.grechanik@gmail.com>

* rename linear_system to int_constraints

* add comments and use random seed

* message for reporting failure with seed

* add SEqualReduce to IntConstraints; allow variables & ranges to be None

Co-authored-by: Sergei Grechanik <sergei.grechanik+h@gmail.com>
Co-authored-by: Sergei Grechanik <sergei.grechanik@gmail.com>

include/tvm/arith/analyzer.h

@@ -424,6 +424,12 @@ class Analyzer {
    */
   void Bind(const Var& var, const Range& range);
   /*!
+   * \brief Bind all the vars in the Map
+   *
+   * \param variables The {variable -> range} map.
+   */
+  void Bind(const Map<Var, Range>& variables);
+  /*!
    * \brief Whether can we prove expr >= val.
 
    *  Non-negative proof is very useful in integer analysis

include/tvm/arith/int_solver.h

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/arith/int_solver.h
+ * \brief integer constraints data structures and solvers
+ */
+#ifndef TVM_ARITH_INT_SOLVER_H_
+#define TVM_ARITH_INT_SOLVER_H_
+
+#include <tvm/ir/expr.h>
+#include <tvm/tir/expr.h>
+#include <unordered_map>
+#include <vector>
+
+namespace tvm {
+namespace arith {
+
+using tir::Var;
+using tir::VarNode;
+using tir::IterVar;
+
+/*!
+ * \brief Represent integer constrains including (integer) variables, their ranges and
+ *        the relations between them (either equations or inequalities).
+ * \sa LinearSystem
+ */
+class IntConstraintsNode : public Object {
+ public:
+  // e.g., \alpha, \beta, must be integers
+  Array<Var> variables;
+  // e.g., 1 <= \alpha <= N, etc.
+  // it is absolutely ok to include ranges for parameters
+  // (variables that are not in this->variables) in this map
+  Map<Var, Range> ranges;
+  // linear equalities or inequalities
+  // e.g., A \alpha = \beta or A \alpha <= \beta
+  Array<PrimExpr> relations;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("variables", &variables);
+    v->Visit("ranges", &ranges);
+    v->Visit("relations", &relations);
+  }
+
+  bool SEqualReduce(const IntConstraintsNode* other, SEqualReducer equal) const {
+    return
+        equal(variables, other->variables) &&
+        equal(ranges, other->ranges) &&
+        equal(relations, other->relations);
+  }
+
+  void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce(variables);
+    hash_reduce(ranges);
+    hash_reduce(relations);
+  }
+
+  static constexpr const bool _type_has_method_sequal_reduce = true;
+  static constexpr const char* _type_key = "arith.IntConstraints";
+  TVM_DECLARE_FINAL_OBJECT_INFO(IntConstraintsNode, Object);
+};
+
+/*!
+ * \brief Managed reference to IntConstraintsNode.
+ * \sa IntConstraintsNode
+ */
+class IntConstraints : public ObjectRef {
+ public:
+  /*!
+   * \brief Constructor by fields
+   * \param variables The variables in the constraints, must be integers.
+   * \param ranges    The ranges of the variables.
+   * \param relations The linear relations between the variables
+   *                  (either equations or inequalities)
+   */
+  TVM_DLL IntConstraints(Array<Var> variables,
+                         Map<Var, Range> ranges,
+                         Array<PrimExpr> relations);
+
+  TVM_DEFINE_OBJECT_REF_METHODS(IntConstraints, ObjectRef, IntConstraintsNode);
+};
+
+/*!
+ * \brief We can have different set of variables to represent the same constraints.
+ *        For example, the following two systems are equivalent,
+ *        {a + b = 0 | a >= 0, b >= 0} and
+ *        {m - n = 0 | m >= 0, n <= 0}
+ *        This data structure represents the transformation
+ *        between two equivalent linear systems.
+ *        In the above example,
+ *        src        : {a + b = 0 | a >= 0, b >= 0}
+ *        dst        : {m - n = 0 | m >= 0, n <= 0}
+ *        src_to_dst : {a -> m, b -> -n}
+ *        dst_to_src : {m -> a, n -> -b}
+ * \sa IntConstraintsTransform
+ */
+class IntConstraintsTransformNode : public Object {
+ public:
+  IntConstraints src;
+  IntConstraints dst;
+  Map<Var, PrimExpr> src_to_dst;
+  Map<Var, PrimExpr> dst_to_src;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("src", &src);
+    v->Visit("dst", &dst);
+    v->Visit("src_to_dst", &src_to_dst);
+    v->Visit("dst_to_src", &dst_to_src);
+  }
+
+  bool SEqualReduce(const IntConstraintsTransformNode* other, SEqualReducer equal) const {
+    return
+        equal(src, other->src) &&
+        equal(dst, other->dst) &&
+        equal(src_to_dst, other->src_to_dst) &&
+        equal(dst_to_src, other->dst_to_src);
+  }
+
+  void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce(src);
+    hash_reduce(dst);
+    hash_reduce(src_to_dst);
+    hash_reduce(dst_to_src);
+  }
+
+  static constexpr const bool _type_has_method_sequal_reduce = true;
+  static constexpr const char* _type_key = "arith.IntConstraintsTransform";
+  TVM_DECLARE_FINAL_OBJECT_INFO(IntConstraintsTransformNode, Object);
+};
+
+/*!
+ * \brief Managed reference to IntConstraintsTransformNode.
+ * \sa IntConstraintsTransformNode
+ */
+class IntConstraintsTransform : public ObjectRef {
+ public:
+  /*!
+   * \brief Constructor by fields
+   * \param src        source integer constraints, e.g., {a + b = 0 | a >= 0, b >= 0}
+   * \param dst        integer constraints equivalent to the source,
+   *                   e.g., {m - n = 0 | m >= 0, n <= 0}
+   * \param src_to_dst mapping from variables in the \p src to the variables in the \p dst,
+   *                   e.g., {a -> m, b -> -n}
+   * \param dst_to_src mapping from variables in the \p dst to the variables in the \p src,
+   *                   e.g., {m -> a, n -> -b}
+   */
+  TVM_DLL IntConstraintsTransform(IntConstraints src,
+                                  IntConstraints dst,
+                                  Map<Var, PrimExpr> src_to_dst,
+                                  Map<Var, PrimExpr> dst_to_src);
+
+  TVM_DEFINE_OBJECT_REF_METHODS(IntConstraintsTransform, ObjectRef, IntConstraintsTransformNode);
+};
+
+/*!
+ * \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},
+ *        in which S_{mxn} is diag(s1, s2, ..., sr, 0, ..., 0) and r is the rank of A.
+ *        NOTE: Although in standard Smith Normal Form the diagonal elements satisfy
+ *              s_i | s_{i+1} (| means divides), the implement here does not guarantee it.
+ *        TODO(yzhliu): From sergei-grechanik:
+ *          computing the proper Smith normal form may improve stability of automatic differentiation
+ *          (generating the same gradient code for slightly different but equivalent input code
+ *        U_{mxm} and V_{nxn} are invertible matrices.
+ *        This function modifies \p S to be S_{mxn}, \p V to be V_{nxn},
+ *        \p y to be U_{mxm} y_{mx1} and \p x to be V^{-1} x.
+ * \param S  the original A_{mxn}, it will be modified to S_{mxn}
+ * \param V  an identity matrix, it will be modified to V_{nxn}
+ * \param x  the x in A x = y. it will be modified to V^{-1}_{nxn} x_{nx1}
+ * \param y  the y in A x = y. it will be modified to U_{mxm} y_{mx1}
+ */
+void SmithNormalFormDiag(std::vector<std::vector<int64_t>> *S,
+                         std::vector<std::vector<int64_t>> *V,
+                         std::vector<PrimExpr>* x,
+                         std::vector<PrimExpr> *y);
+
+/*!
+ * \brief Solve linear equations.
+ * \param system_to_solve the variables to solve, their ranges, and a list of equations.
+ * \return  A new linear system, with less variables (if \p system_to_solve is NOT of full rank),
+ *          or no variable (if \p system_to_solve is of full rank),
+ *          or an empty linear system (if \p system_to_solve is unsolvable).
+ *          It also provides the ranges of the variables in the new system,
+ *          as well as inequalities inferred from the \p system_to_solve.
+ *          You can get the mapping from the original variables to the solution via ret->src_to_dst.
+ */
+IntConstraintsTransform SolveLinearEquations(const IntConstraints &system_to_solve);
+
+}  // namespace arith
+}  // namespace tvm
+#endif  // TVM_ARITH_INT_SOLVER_H_

include/tvm/arith/util.h

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/arith/util.h
+ * \brief Utils for arithmetic analysis.
+ */
+#ifndef TVM_ARITH_UTIL_H_
+#define TVM_ARITH_UTIL_H_
+
+#include <cstdint>
+#include <tuple>
+
+namespace tvm {
+/*! \brief namespace of arithmetic analysis. */
+namespace arith {
+
+/*!
+ * \brief Calculate the extended greatest common divisor for two values.
+ *        See https://en.wikipedia.org/wiki/Extended_Euclidean_algorithm.
+ * \param a an integer number
+ * \param b an integer number
+ * \return 3 integers (div, m, n) where div = gcd(a, b) and a*m + b*n = div
+ */
+std::tuple<int64_t, int64_t, int64_t> xgcd(int64_t a, int64_t b);
+
+}  // namespace arith
+}  // namespace tvm
+#endif  // TVM_ARITH_UTIL_H_

python/tvm/arith/__init__.py

@@ -20,3 +20,4 @@ from .int_set import IntSet, IntervalSet
 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

python/tvm/arith/int_solver.py

+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""integer constraints data structures and solvers"""
+import tvm._ffi
+from tvm.runtime import Object
+from . import _ffi_api
+
+
+@tvm._ffi.register_object("arith.IntConstraints")
+class IntConstraints(Object):
+    """Represent a set of integer constraints including variables, their ranges and
+       the relations between them (either equations or inequalities)
+
+    Parameters
+    ----------
+    variables : List[tvm.tir.Var]
+        The variables in the constraints. Must be integers
+    ranges    : Map[tvm.tir.Var, tvm.ir.Range]
+        The ranges of the variables.
+    relations : List[tvm.ir.PrimExpr]
+        The relations between the variables (either equations or inequalities)
+    """
+    def __init__(self, variables, ranges, relations):
+        self.__init_handle_by_constructor__(
+            _ffi_api.IntConstraints, variables, ranges, relations)
+
+
+@tvm._ffi.register_object("arith.IntConstraintsTransform")
+class IntConstraintsTransform(Object):
+    """We can have different set of variables to represent the same integer constraints.
+       For example, the following two constrains are equivalent,
+       {a + b = 0 | a >= 0, b >= 0} and
+       {m - n = 0 | m >= 0, n <= 0}
+       This data structure represents the transformation
+       between two equivalent integer constraints.
+       In the above example,
+       src        : {a + b = 0 | a >= 0, b >= 0}
+       dst        : {m - n = 0 | m >= 0, n <= 0}
+       src_to_dst : {a -> m, b -> -n}
+       dst_to_src : {m -> a, n -> -b}
+
+    Parameters
+    ----------
+    src : arith.IntConstraints
+        source integer constraints, e.g., {a + b = 0 | a >= 0, b >= 0}
+    dst : arith.IntConstraints
+        integer constraints equivalent to the source, e.g., {m - n = 0 | m >= 0, n <= 0}
+    src_to_dst : Map[tvm.tir.Var, tvm.ir.PrimExpr]
+        mapping from variables in the src to the variables in the dst,
+                e.g., {a -> m, b -> -n}
+    dst_to_src : Map[tvm.tir.Var, tvm.ir.PrimExpr]
+        mapping from variables in the dst to the variables in the src,
+        e.g., {m -> a, n -> -b}
+    """
+    def __init__(self, src, dst, src_to_dst, dst_to_src):
+        self.__init_handle_by_constructor__(
+            _ffi_api.IntConstraintsTransform, src, dst, src_to_dst, dst_to_src)
+
+
+def solve_linear_equations(equations, variables=None, ranges=None):
+    """Solve linear equations.
+
+    Parameters
+    ----------
+    equations: List[tvm.ir.PrimExpr] or IntConstraints
+        The equations 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.
+
+    Returns
+    -------
+    int_constraints_transform : IntConstraintsTransform
+        New integer constraints, with less variables (if the problem is NOT of full rank),
+        or no variable (if the problem is of full rank),
+        or an empty integer constraints (if the problem is unsolvable).
+        It also provides the ranges of the variables in the new system,
+        as well as inequalities inferred from the problem.
+        You can get the mapping from the original variables to the solution via
+        int_constraints_transform.src_to_dst.
+    """
+    if isinstance(equations, IntConstraints):
+        return _ffi_api.SolveLinearEquations(equations)
+    return _ffi_api.SolveLinearEquations(variables, ranges, equations)

src/arith/analyzer.cc

@@ -58,6 +58,11 @@ void Analyzer::Bind(const Var& var, const Range& range) {
   // skip rewrite simplify
 }
 
+void Analyzer::Bind(const Map<Var, Range>& variables) {
+  for (const auto& iter : variables) {
+    this->Bind(iter.first, iter.second);
+  }
+}
 
 void ConstraintContext::EnterWithScope() {
   CHECK(exit_ == nullptr);

src/arith/int_constraints.cc

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file int_constraints.cc
+ * \brief The integer constraints data structures.
+ */
+#include <tvm/arith/int_solver.h>
+#include <tvm/tir/expr.h>
+#include <tvm/tir/expr_functor.h>
+#include <tvm/runtime/registry.h>
+
+#include <utility>
+#include <algorithm>
+#include <unordered_map>
+
+namespace tvm {
+namespace arith {
+
+IntConstraints::IntConstraints(Array<Var> variables,
+                               Map<Var, Range> ranges,
+                               Array<PrimExpr> relations) {
+  ObjectPtr<IntConstraintsNode> node = make_object<IntConstraintsNode>();
+  if (!variables.defined()) {
+    variables = Array<Var>();
+  }
+  if (!ranges.defined()) {
+    ranges = Map<Var, Range>();
+  }
+  CHECK(relations.defined());
+  for (const auto& var : variables) {
+    CHECK(var.dtype().is_int() || var.dtype().is_uint())
+      << "Variables in IntConstraints must be integers";
+  }
+  node->variables = std::move(variables);
+  node->ranges = std::move(ranges);
+  node->relations = std::move(relations);
+  data_ = std::move(node);
+}
+
+TVM_REGISTER_NODE_TYPE(IntConstraintsNode);
+
+TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
+.set_dispatch<IntConstraintsNode>([](const ObjectRef& node, ReprPrinter* p) {
+    auto* op = static_cast<const IntConstraintsNode*>(node.get());
+    p->stream << "IntConstraints("
+              << op->variables
+              << ", " << op->ranges
+              << ", " << op->relations
+              << ")";
+  });
+
+
+IntConstraintsTransform::IntConstraintsTransform(IntConstraints src,
+                                                 IntConstraints dst,
+                                                 Map<Var, PrimExpr> src_to_dst,
+                                                 Map<Var, PrimExpr> dst_to_src) {
+  ObjectPtr<IntConstraintsTransformNode> node = make_object<IntConstraintsTransformNode>();
+  node->src = std::move(src);
+  node->dst = std::move(dst);
+  node->src_to_dst = std::move(src_to_dst);
+  node->dst_to_src = std::move(dst_to_src);
+  data_ = std::move(node);
+}
+
+TVM_REGISTER_NODE_TYPE(IntConstraintsTransformNode);
+
+TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
+.set_dispatch<IntConstraintsTransformNode>([](const ObjectRef& node, ReprPrinter* p) {
+    auto* op = static_cast<const IntConstraintsTransformNode*>(node.get());
+    p->stream << "IntConstraintsTransform("
+              << "\n\t" << op->src
+              << "\n\t" << op->dst
+              << "\n\t" << op->src_to_dst
+              << "\n\t" << op->dst_to_src
+              << "\n)";
+  });
+
+}  // namespace arith
+}  // namespace tvm

src/arith/solve_linear_equation.cc

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/arith/solve_linear_equation.cc
+ * \brief Solve linear equations.
+ */
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/expr.h>
+#include <tvm/arith/analyzer.h>
+#include <tvm/arith/int_solver.h>
+#include <tvm/arith/util.h>
+#include <tvm/tir/op.h>
+#include <tvm/arith/pattern.h>
+#include <tvm/tir/ir_pass.h>
+#include <tvm/runtime/data_type.h>
+
+namespace tvm {
+namespace arith {
+
+using namespace tvm::runtime;
+
+void SmithNormalFormDiag(std::vector<std::vector<int64_t> >* S,
+                         std::vector<std::vector<int64_t> >* V,
+                         std::vector<PrimExpr>* x,
+                         std::vector<PrimExpr>* y) {
+  if (S->empty() || V->empty()) return;
+  size_t m = S->size();
+  size_t n = (*S)[0].size();  // n is # of variables
+  CHECK_EQ(V->size(), n);
+  CHECK_EQ((*V)[0].size(), n);
+
+  for (size_t index = 0; index < std::min(m, n); ++index) {
+    // Here A is partially diagonalized, that is A[i, j] is zero for all i, j
+    // such that (i < index) or (j < index), unless (i == j).
+    // That is, now we are diagonalizing the submatrix with i >= index and j >= index
+
+    // Find a row with a nonzero element in the index-th column
+    // (We also prefer rows where this element has minimal abs value)
+    size_t best_i = index;
+    for (size_t i = best_i; i < m; ++i) {
+      int64_t s_old = (*S)[best_i][index];
+      int64_t s_new = (*S)[i][index];
+      if (s_new != 0) {
+        if (s_old == 0 || std::abs(s_new) < std::abs(s_old)) {
+          best_i = i;
+        }
+      }
+    }
+    // Move the row we found to the index-th position
+    std::swap((*S)[index], (*S)[best_i]);
+    std::swap((*y)[index], (*y)[best_i]);
+
+    // If the index-th diagonal element is still zero, try to find a column with nonzero index-th
+    // element and move it to the index-th position
+    if ((*S)[index][index] == 0) {
+      for (size_t j = index + 1; j < n; ++j) {
+        if ((*S)[index][j] != 0) {
+          for (size_t i = index; i < m; ++i) {
+            std::swap((*S)[i][index], (*S)[i][j]);
+          }
+          // swapping columns corresponds to swapping the corresponding x
+          std::swap((*x)[index], (*x)[j]);
+          for (size_t i = 0; i < n; ++i) {
+            std::swap((*V)[i][index], (*V)[i][j]);
+          }
+          break;
+        }
+      }
+    }
+
+    // If the index-th diagonal element is still zero, then both the index-th row and the index-th
+    // column are completely zero, and we don't need to do anything; just go to the next index
+    if ((*S)[index][index] == 0) {
+      continue;
+    }
+
+    // Now the index-th diagonal element is non-zero and we can zero all the index-th column
+    // below it by subtracting rows from each other
+    for (auto i = index + 1; i < m; ++i) {
+      if ((*S)[i][index] != 0) {
+        int64_t g, a, b;
+        // g = a*matrix[index][index] + b*matrix[i][index]
+        if ((*S)[i][index] % (*S)[index][index] != 0) {
+          std::tie(g, a, b) = xgcd((*S)[index][index], (*S)[i][index]);
+        } else {
+          // Explicitly avoid changing the index-th row. This is important to avoid infinite loop.
+          g = (*S)[index][index];
+          a = 1;
+          b = 0;
+        }
+
+        // Let m = S[index][index], n = S[i][index], then the following is true:
+        //
+        // [ a   n/g ][ m/g  n/g ] = [ 1  0 ]
+        // [ b  -m/g ][ b    -a  ] = [ 0  1 ]
+        //
+        // Note that the two matrices are integer (since g = gcd(m, n)).
+        // We will essentially multiply our matrix on the left by a dilated and transposed version
+        // of the first of these two matrices. The second matrix is not needed here, however we will
+        // use it while zeroing the index-th row.
+
+        int64_t m_g = (*S)[index][index] / g;
+        int64_t n_g = (*S)[i][index] / g;
+
+        // Note that j is the index of the column, not the row
+        for (size_t j = index; j < (*S)[i].size(); ++j) {
+          // Multiply index-th row by a and add the i-th row multiplied by b
+          // This will make the index-th diagonal element equal to the gcd
+          int64_t new_index_j = a*(*S)[index][j] + b*(*S)[i][j];
+          // This transformation performs zeroing of matrix[i][index]
+          int64_t new_i_j = n_g*(*S)[index][j] - m_g*(*S)[i][j];
+          (*S)[index][j] = new_index_j;
+          (*S)[i][j] = new_i_j;
+        }
+        // We have to do the same with rhs
+        PrimExpr ea = te::make_const((*y)[index].dtype(), a);
+        PrimExpr eb = te::make_const((*y)[i].dtype(), b);
+        PrimExpr e_m_g = te::make_const((*y)[i].dtype(), m_g);
+        PrimExpr e_n_g = te::make_const((*y)[index].dtype(), n_g);
+        PrimExpr new_index_rhs = ea*(*y)[index] + eb*(*y)[i];
+        PrimExpr new_i_rhs = e_n_g*(*y)[index] - e_m_g*(*y)[i];
+        (*y)[index] = new_index_rhs;
+        (*y)[i] = new_i_rhs;
+      }
+    }
+
+    bool changed = false;
+
+    // Now we have to zero the elements of the index-th row by manipulating columns.
+    // This is more difficult because column manipulation corresponds to variable manipulation,
+    // but the algorithm is essentially the same as before.
+    for (size_t j = index + 1; j < n; ++j) {
+      if ((*S)[index][j] != 0) {
+        int64_t g, a, b;
+        // g = a*matrix[index][index] + b*matrix[index][j]
+        if ((*S)[index][j] % (*S)[index][index] != 0) {
+          std::tie(g, a, b) = xgcd((*S)[index][index], (*S)[index][j]);
+          // During this phase we may disrupt the zeroness of the index-th column, so we will
+          // have to take some action if this might have happened.
+          changed = true;
+        } else {
+          // Explicitly avoid changing the index-th column. This is important to avoid infinite
+          // loop. Note that here we don't have to set `changed` to true since we don't change the
+          // index-th column.
+          g = (*S)[index][index];
+          a = 1;
+          b = 0;
+        }
+
+        // Let m = S[index][index], n = S[index][j], then the following is true:
+        //
+        // [ a   n/g ][ m/g  n/g ] = [ 1  0 ]
+        // [ b  -m/g ][ b    -a  ] = [ 0  1 ]
+        //
+        // Now we are going to multiply our matrix on the right (to manipulate columns instead of
+        // rows), we will also transform the old_to_new matrix the same way, and we will use the
+        // second matrix to transform new_to_old.
+
+        int64_t m_g = (*S)[index][index] / g;
+        int64_t n_g = (*S)[index][j] / g;
+
+        for (size_t i = index; i < m; ++i) {
+          int64_t new_i_index = a*(*S)[i][index] + b*(*S)[i][j];
+          int64_t new_i_j = n_g*(*S)[i][index] - m_g*(*S)[i][j];
+          (*S)[i][index] = new_i_index;
+          (*S)[i][j] = new_i_j;
+        }
+        // We do exactly the same transformations with V
+        for (size_t i = 0; i < n; ++i) {
+          int64_t new_i_index = a*(*V)[i][index] + b*(*V)[i][j];
+          int64_t new_i_j = n_g*(*V)[i][index] - m_g*(*V)[i][j];
+          (*V)[i][index] = new_i_index;
+          (*V)[i][j] = new_i_j;
+        }
+        // And apply reverse transformations to new_to_old.
+        PrimExpr ea = te::make_const((*x)[j].dtype(), a);
+        PrimExpr eb = te::make_const((*x)[index].dtype(), b);
+        PrimExpr e_m_g = te::make_const((*x)[index].dtype(), m_g);
+        PrimExpr e_n_g = te::make_const((*x)[j].dtype(), n_g);
+        PrimExpr new_index = e_m_g*(*x)[index] + e_n_g*(*x)[j];
+        PrimExpr new_j = eb*(*x)[index] - ea*(*x)[j];
+        (*x)[index] = new_index;
+        (*x)[j] = new_j;
+      }
+    }
+
+    if (changed) {
+      // We might have changed the first column, so we have to zero it once more
+      // (or at least check if it's zero), so just perform this iteration once more.
+      index -= 1;
+    }
+  }
+}
+
+Map<Var, Range> InferRange(const Map<Var, PrimExpr>& vars_to_infer,
+                           const Array<Var>& ori_vars,
+                           const Map<Var, Range>& ori_ranges) {
+  // The resulting ranges
+  Map<Var, Range> new_ranges;
+
+  std::unordered_set<const VarNode*> ori_vset;
+  for (const Var& v : ori_vars) {
+    ori_vset.insert(v.get());
+  }
+
+  std::unordered_map<const VarNode*, IntSet> var_intsets;
+  for (const auto& p : ori_ranges) {
+    if (!ori_vset.count(p.first.get())) {
+      // First of all, fill the new ranges with outer variable ranges
+      new_ranges.Set(p.first, p.second);
+    }
+    // Convert original ranges to IntSets
+    var_intsets[p.first.get()] = IntSet::range(p.second);
+  }
+
+  // Infer ranges for the new variables and add them to the resulting ranges
+  for (const auto& p : vars_to_infer) {
+    const auto& var = p.first;
+    const auto& expr = p.second;
+    Range range = EvalSet(expr, var_intsets).cover_range(Range());
+    if (range.defined()) {
+      new_ranges.Set(var, range);
+    }
+  }
+  return new_ranges;
+}
+
+// pretty print matrix equation
+void DebugPrint(const std::vector<std::vector<int64_t>>& S,
+                const std::vector<std::vector<int64_t>>& V,
+                const std::vector<PrimExpr>& V_inv_x,
+                const std::vector<PrimExpr>& rhs) {
+  std::cout << "S:\n";
+  for (size_t i = 0; i < S.size(); ++i) {
+    for (auto e : S[i]) {
+      std::cout << e << "\t";
+    }
+    std::cout << "\t->\t" << rhs[i];
+    std::cout << "\n";
+  }
+  std::cout << "V:\n";
+  for (const auto& r : V) {
+    for (auto e : r) {
+      std::cout << e << "\t";
+    }
+    std::cout << "\n";
+  }
+  std::cout << "V_inv x:\n" << Array<PrimExpr>(V_inv_x);
+  std::cout << "\n" << std::endl;
+}
+
+IntConstraintsTransform SolveLinearEquations(const IntConstraints &system_to_solve) {
+  // m: # of equations
+  // n: # of variables
+  // we first construct A_{mxn} x_{nx1} = y_{mx1}
+  // then get Smith normal form of matrix A,
+  // S_{mxn} = U_{mxm} A_{mxn} V_{nxn}
+  // => U^{-1} S V^{-1} x = y
+  // S V^{-1} x = U y
+  std::vector<PrimExpr> Uy;  // mx1
+  std::vector<std::vector<int64_t>> S;  // mxn
+  std::vector<std::vector<int64_t>> V;  // nxn
+  std::vector<PrimExpr> V_inv_x;  // V^{-1} x, nx1
+  // Conditions we don't know what to do with
+  std::vector<PrimExpr> rest;
+
+  Analyzer analyzer_problem;
+  analyzer_problem.Bind(system_to_solve->ranges);
+
+  size_t num_vars = system_to_solve->variables.size();
+
+  // initialize V_{nxn} with identity matrix,
+  // initialize V^{-1} x as x
+  for (size_t i = 0; i < num_vars; ++i) {
+    V.emplace_back(num_vars);
+    V.back()[i] = 1;
+    V_inv_x.push_back(system_to_solve->variables[i]);
+  }
+
+  // Transform formulas into rows of the matrix
+  // S_{mxn} V^{-1}_{nxn} x_{nx1} = U y, in which n is # of variables
+  // here we initialize S_{mxn} to be A, U to be identity matrix.
+  for (const PrimExpr& equation : system_to_solve->relations) {
+    if (const tir::EQNode* eq = equation.as<tir::EQNode>()) {
+      // a-b = sum_{i=0}^{n-1} variables[i] * coeff[i] + coeff[n]
+      Array<PrimExpr> coeffs = arith::DetectLinearEquation(
+          analyzer_problem.Simplify(eq->a - eq->b),
+          system_to_solve->variables);
+      if (!coeffs.empty()) {
+        std::vector<int64_t> row;
+        for (size_t j = 0; j < coeffs.size() - 1; ++j) {
+          PrimExpr c = coeffs[j];
+          if (const IntImmNode* ic = c.as<IntImmNode>()) {
+            row.push_back(ic->value);
+          } else {
+            // elements in matrix S V must be integers
+            // ignore equations that we cannot deal with.
+            LOG(WARNING) << "Cannot deal with non-integer coefficients, ignore equation "
+                         << equation;
+            row.clear();
+            break;
+          }
+        }
+
+        if (!row.empty()) {
+          // S V^{-1} (a-b) = Uy
+          // V is identity for now
+          S.push_back(row);
+          Uy.push_back(-coeffs[coeffs.size() - 1]);
+          continue;
+        }
+      }
+    }
+
+    // otherwise
+    rest.push_back(equation);
+  }
+
+  // After diagonalizing, we have
+  // S_{mxn} is the Smith normal form (diagonal matrix)
+  // V_{nxn} is invertible
+  // V_inv_x is V^{-1} \times x
+  // Uy is U \times y
+  SmithNormalFormDiag(&S, &V, &V_inv_x, &Uy);
+
+  Array<Var> new_vars;
+  Array<PrimExpr> new_relations;
+  Map<Var, PrimExpr> new_to_old_map;
+  Map<Var, PrimExpr> old_to_new_map;
+
+  // Simplify right hand sides
+  for (PrimExpr r : Uy) {
+    r = analyzer_problem.Simplify(r);
+  }
+
+  // Create the relations of the existence of a solution
+  for (size_t j = 0; j < S.size(); ++j) {
+    PrimExpr new_relation;
+    if (j >= num_vars || S[j][j] == 0) {
+      // The row of matrix is zero. A solution exists only if the Ub[j] is also zero
+      new_relation = (Uy[j] == 0);
+    } else {
+      // The diagonal element is non-zero. A solution exists only if the diagonal element
+      // is a divisor of the Ub[j]
+      new_relation = (floormod(Uy[j], std::abs(S[j][j])) == 0);
+    }
+    new_relation = analyzer_problem.Simplify(new_relation);
+    if (tir::is_const_int(new_relation, 0)) {
+      // unable to solve the system.
+      return IntConstraintsTransform(
+          system_to_solve,
+          IntConstraints(
+              /*variables=*/{},
+              /*ranges=*/{},
+              /*relations=*/{te::make_zero(DataType::Bool())}),
+          {}, {});
+    } else if (!tir::is_const_int(new_relation, 1)) {
+      new_relations.push_back(new_relation);
+    }
+  }
+
+  Array<PrimExpr> solution_for_V_inv_x;
+  // Now create new variables or directly solve the equations
+  // suppose the rank of A is r, aka r = # of non-zeros in S
+  // the solution of S_{mxn} V^{-1}_{nxn} x_{nx1} = U b
+  // is
+  // x = (pseudo-inverse of A) b + K_{(n)x(n-r)} z_{n-r}
+  //   = V_{nxn} S^{-1}_{nxm} (Ub)_{mxn} + K_{(n)x(n-r)} z_{n-r}
+  // in which K is the right n-r columns of V, z is variable vector
+  // thus,
+  // V^{-1} x = S^{-1}_{nxm} (Ub)_{mxn} +
+  //            [[0, ... 0]_{n-r}, ... [0, ..., 0], diag(1, ..., 1)_{(n-r)x(n-r)}] z_{n-r}
+  for (size_t j = 0; j < num_vars; ++j) {
+    if (j >= S.size() || S[j][j] == 0) {
+      // The j-th variable can take any integer value, create a tvm variable for it
+      PrimExpr to_old = analyzer_problem.Simplify(V_inv_x[j]);
+      std::string name_hint = "n" + std::to_string(new_vars.size());
+      if (const VarNode* v_old = to_old.as<VarNode>()) {
+        name_hint += "_" + v_old->name_hint;
+      }
+      Var v = Var(name_hint, V_inv_x[j].dtype());
+      solution_for_V_inv_x.push_back(v);
+      new_vars.push_back(v);
+      new_to_old_map.Set(v, to_old);
+    } else {
+      // The j-th variable is just a single value, don't create a tvm variable
+      // S^{-1}_{nxm} Uy_{mxn}
+      if (S[j][j] >= 0) {
+        PrimExpr a = te::make_const(Uy[j].dtype(), S[j][j]);
+        solution_for_V_inv_x.push_back(
+            analyzer_problem.Simplify(floordiv(Uy[j], a)));
+      } else {
+        // This is required because some simplifiers
+        // have problems with dividing by negative numbers
+        PrimExpr a = te::make_const(Uy[j].dtype(), -S[j][j]);
+        solution_for_V_inv_x.push_back(
+            analyzer_problem.Simplify(floordiv(-Uy[j], a)));
+      }
+    }
+  }
+
+  // V V^{-1} x = x
+  for (size_t i = 0; i < num_vars; ++i) {
+    PrimExpr e = te::make_zero(system_to_solve->variables[i].dtype());
+    for (size_t j = 0; j < num_vars; ++j) {
+      e = e + te::make_const(e.dtype(), V[i][j])*solution_for_V_inv_x[j];
+    }
+    e = analyzer_problem.Simplify(e);
+    old_to_new_map.Set(system_to_solve->variables[i], e);
+  }
+
+  // The resulting ranges
+  Map<Var, Range> new_ranges = InferRange(
+      new_to_old_map, system_to_solve->variables, system_to_solve->ranges);
+  Analyzer analyzer_solution;
+  analyzer_solution.Bind(new_ranges);
+
+  // We have to transform ranges of the old variables into relations over new variables because
+  // new ranges are not enough usually.
+  for (const auto& p : system_to_solve->ranges) {
+    const Var& old_var = p.first;
+    const Range& old_range = p.second;
+    if (old_to_new_map.count(old_var)) {
+      PrimExpr express_by_new_vars = old_to_new_map[old_var];
+      PrimExpr lower_cond = analyzer_solution.Simplify(
+          old_range->min <= express_by_new_vars);
+      PrimExpr upper_cond = analyzer_solution.Simplify(
+          express_by_new_vars < old_range->min + old_range->extent);
+      if (!tir::is_const_int(lower_cond, 1)) {
+        new_relations.push_back(lower_cond);
+      }
+      if (!tir::is_const_int(upper_cond, 1)) {
+        new_relations.push_back(upper_cond);
+      }
+    }
+  }
+
+  // Add the rest conditions
+  for (const PrimExpr& cond : rest) {
+    new_relations.push_back(Substitute(cond, old_to_new_map));
+  }
+
+  IntConstraints solution(new_vars, new_ranges, new_relations);
+  IntConstraintsTransform transform(
+      system_to_solve, solution, old_to_new_map, new_to_old_map);
+
+  return transform;
+}
+
+TVM_REGISTER_GLOBAL("arith.SolveLinearEquations")
+.set_body([](TVMArgs args, TVMRetValue *ret) {
+    if (args.size() == 1) {
+      *ret = SolveLinearEquations(args[0]);
+    } else if (args.size() == 3) {
+      IntConstraints problem(args[0], args[1], args[2]);
+      *ret = SolveLinearEquations(problem);
+    } else {
+      LOG(FATAL) << "arith.SolveLinearEquations expects 1 or 3 arguments, gets " << args.size();
+    }
+  });
+
+}  // namespace arith
+}  // namespace tvm

src/arith/util.cc

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file util.cc
+ * \brief The utils for arithmetic analysis.
+ */
+#include <tvm/arith/util.h>
+#include <dmlc/logging.h>
+
+namespace tvm {
+namespace arith {
+
+std::tuple<int64_t, int64_t, int64_t> xgcd(int64_t a, int64_t b) {
+  int64_t s = 0, old_s = 1;
+  int64_t t = 1, old_t = 0;
+  int64_t r = b, old_r = a;
+
+  while (r != 0) {
+    int64_t q = old_r / r;
+    std::swap(r, old_r);
+    r -= q * old_r;
+    std::swap(s, old_s);
+    s -= q * old_s;
+    std::swap(t, old_t);
+    t -= q * old_t;
+  }
+
+  CHECK_EQ(a % old_r, 0);
+  CHECK_EQ(b % old_r, 0);
+  CHECK(old_r == old_s*a + old_t*b);
+
+  return std::make_tuple(old_r, old_s, old_t);
+}
+
+}  // namespace arith
+}  // namespace tvm

tests/python/unittest/test_arith_solve_linear_system.py

+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import random
+import numpy as np
+import sys
+import pytest
+import tvm
+from tvm import te, arith, ir, tir
+
+
+def run_expr(expr, vranges):
+    """ Evaluate expr for every value of free variables
+    given by vranges and return the tensor of results.
+    TODO(yzhliu): move to utils
+    """
+    def _compute_body(*us):
+        vmap = {v: u + r.min for (v, r), u in zip(vranges.items(), us)}
+        return tir.ir_pass.Substitute(expr, vmap)
+
+    A = te.compute([r.extent.value for v, r in vranges.items()], _compute_body)
+    args = [tvm.nd.empty(A.shape, A.dtype)]
+    sch = te.create_schedule(A.op)
+    mod = tvm.build(sch, [A])
+    mod(*args)
+    return args[0].asnumpy()
+
+
+def check_bruteforce(bool_expr, vranges, cond=None):
+    """ Check that bool_expr holds given the condition cond
+    for every value of free variables from vranges.
+    TODO(yzhliu): move to utils
+    """
+    if cond is not None:
+        bool_expr = te.any(tir.Not(cond), bool_expr)
+
+    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(tir.ir_pass.CanonicalSimplify(bool_expr), vranges, counterex))
+
+
+def check_solution(solution, vranges={}):
+    """Check that solution is a bijective transformation"""
+    def _check_forward(constraints1, constraints2, varmap, backvarmap):
+        all_vranges = vranges.copy()
+        all_vranges.update({v: r for v, r in constraints1.ranges.items()})
+
+        # Check that the transformation is injective
+        cond_on_vars = tir.const(1, 'bool')
+        for v in constraints1.variables:
+            # variable mapping is consistent
+            v_back = tir.ir_pass.Simplify(tir.ir_pass.Substitute(varmap[v], backvarmap))
+            cond_on_vars = te.all(cond_on_vars, v == v_back)
+        # Also we have to check that the new relations are true when old relations are true
+        cond_subst = tir.ir_pass.Substitute(
+            te.all(tir.const(1, 'bool'), *constraints2.relations), backvarmap)
+        # We have to include relations from vranges too
+        for v in constraints2.variables:
+            if v in constraints2.ranges:
+                r = constraints2.ranges[v]
+                range_cond = te.all(v >= r.min, v < r.min + r.extent)
+                range_cond = tir.ir_pass.Substitute(range_cond, backvarmap)
+                cond_subst = te.all(cond_subst, range_cond)
+        cond_subst = tir.ir_pass.Simplify(cond_subst)
+        check_bruteforce(te.all(cond_subst, cond_on_vars), all_vranges,
+                         cond=te.all(tir.const(1, 'bool'), *constraints1.relations))
+
+    rels = solution.dst.relations
+    if len(rels) == 1 and ir.structural_equal(rels[0], False):
+        # not solvable, skip
+        return
+    _check_forward(solution.src, solution.dst,
+                   solution.src_to_dst, solution.dst_to_src)
+    _check_forward(solution.dst, solution.src,
+                   solution.dst_to_src, solution.src_to_dst)
+
+
+def test_solution_consistency():
+    seed = random.randrange(sys.maxsize)
+    print("\nThis test is intentionally non-deterministic, "
+          "if it fails please report it in github issue together with this seed {}\n".format(seed))
+    random.seed(seed)
+
+    def _check(num_vars, num_formulas, coef=(-5, 5), bounds=(-20, 20)):
+        variables = [te.var("x" + str(i)) for i in range(num_vars)]
+
+        relations = []
+        for i in range(num_formulas):
+            s1 = sum([v*random.randint(coef[0], coef[1]) for v in variables])
+            s1 += random.randint(coef[0], coef[1])
+            s2 = sum([v*random.randint(coef[0], coef[1]) for v in variables])
+            s2 += random.randint(coef[0], coef[1])
+            if random.random() < 0.7:
+                op = tvm.tir.EQ
+            else:
+                # we also make sure it can correctly handle inequalities
+                op = random.choice([tvm.tir.LE, tvm.tir.LT, tvm.tir.GE, tvm.tir.GT])
+            relations.append(op(s1, s2))
+
+        vranges = {v: tvm.ir.expr.Range(bounds[0], bounds[1] + 1) for v in variables}
+        solution = arith.solve_linear_equations(relations, variables, vranges)
+
+        check_solution(solution)
+
+        # leaving some variables as parameters should also be ok
+        for k in [1, 2]:
+            if len(variables) > k:
+                solution = arith.solve_linear_equations(relations, variables[:-k], vranges)
+                param_ranges = {v: vranges[v] for v in variables[-k:]}
+                check_solution(solution, param_ranges)
+
+    for i in range(2):
+        _check(num_vars=1, num_formulas=1)
+    for i in range(2):
+        _check(num_vars=1, num_formulas=2)
+
+    for i in range(2):
+        _check(num_vars=2, num_formulas=1)
+    for i in range(2):
+        _check(num_vars=2, num_formulas=2)
+    for i in range(2):
+        _check(num_vars=2, num_formulas=3)
+
+    for i in range(3):
+        _check(num_vars=3, num_formulas=3, coef=(-2, 2))
+    for i in range(3):
+        _check(num_vars=3, num_formulas=4, coef=(-2, 2))
+
+    for i in range(3):
+        _check(num_vars=4, num_formulas=3, coef=(-1, 1))
+
+    for i in range(3):
+        _check(num_vars=10, num_formulas=2, coef=(-1, 1), bounds=(0, 4))
+    for i in range(3):
+        _check(num_vars=10, num_formulas=3, coef=(0, 1), bounds=(0, 4))
+
+
+def test_empty_var_to_solve():
+    x, y = te.var("x"), te.var("y")
+    equations = [
+        tvm.tir.EQ(x + y, 20),
+        tvm.tir.EQ(x - y, 10),
+    ]
+    solution = arith.solve_linear_equations(equations)
+    assert len(solution.src_to_dst) == 0
+    assert len(solution.dst_to_src) == 0
+    assert len(solution.src.variables) == 0
+    assert len(solution.src.ranges) == 0
+    assert ir.structural_equal(solution.src.relations, equations)
+    assert ir.structural_equal(solution.src, solution.dst)
+
+
+def test_unique_solution():
+    x, y = te.var("x"), te.var("y")
+
+    solution = arith.solve_linear_equations([
+        tvm.tir.EQ(x + y, 20),
+        tvm.tir.EQ(x - y, 10),
+    ], [x, y])
+    assert list(solution.dst.variables) == []
+    assert ir.structural_equal(solution.src_to_dst[x], 15)
+    assert ir.structural_equal(solution.src_to_dst[y], 5)
+
+
+def test_low_rank():
+    x, y, z = te.var("x"), te.var("y"), te.var("z")
+    ranges = {}
+
+    solution = arith.solve_linear_equations([
+        tvm.tir.EQ(x + y + z, 15),
+        tvm.tir.EQ(x + y, 10),
+    ], [x, y, z], ranges)
+    [n0] = solution.dst.variables
+    assert ir.structural_equal(solution.src_to_dst[x], n0 + 10)
+    assert ir.structural_equal(solution.src_to_dst[y], -n0)
+    assert ir.structural_equal(solution.src_to_dst[z], 5)
+
+
+def test_infer_range():
+    x, y = te.var("x"), te.var("y")
+    ranges = {
+        x: tvm.ir.Range.make_by_min_extent(-5, 10),
+        y: tvm.ir.Range.make_by_min_extent(0, 10),
+    }
+
+    solution = arith.solve_linear_equations([
+        tvm.tir.EQ(x + y, 0),
+    ], [x, y], ranges)
+    [n0] = solution.dst.variables
+    assert ir.structural_equal(solution.src_to_dst[x], n0)
+    assert ir.structural_equal(solution.src_to_dst[y], -n0)
+    # inferred from y's range
+    assert ir.structural_equal(solution.dst.ranges[n0].min, -9)
+    assert ir.structural_equal(solution.dst.ranges[n0].extent, 10)
+    # additional inequality is added into the system for x
+    [ineq] = solution.dst.relations
+    assert isinstance(ineq, tvm.tir.LE)
+    assert ir.structural_equal(ineq.a, -5)
+    assert ir.structural_equal(ineq.b, n0)
+
+
+def test_ill_formed():
+    x, y = te.var("x"), te.var("y")
+
+    solution = arith.solve_linear_equations([
+        tvm.tir.EQ(x + y, 0),
+        tvm.tir.EQ(x - y, 0),
+        tvm.tir.EQ(x, 5),
+    ], [x, y], {})
+    assert list(solution.dst.variables) == []
+    [rel] = solution.dst.relations
+    assert ir.structural_equal(rel, False)
+    assert len(solution.src_to_dst) == 0
+    assert len(solution.dst_to_src) == 0
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])