[Autodiff] Optimize and eliminate the Jacobian tensor for te.autodiff

include/tvm/arith/analyzer.h

@@ -52,6 +52,13 @@ class Analyzer;
 
 using tir::Var;
 
+enum DivMode {
+  /*! \brief Truncated division. */
+  kTruncDiv,
+  /*! \brief Floor division. */
+  kFloorDiv
+};
+
 /*!
  * \brief Constant integer up and lower bound(inclusive).
  *  Useful for value bound analysis.

include/tvm/arith/int_solver.h

@@ -251,6 +251,15 @@ class IntConstraintsTransform : public ObjectRef {
   TVM_DLL IntConstraintsTransform(IntConstraints src, IntConstraints dst,
                                   Map<Var, PrimExpr> src_to_dst, Map<Var, PrimExpr> dst_to_src);
 
+  /*!
+   * \brief Chain-compose two IntConstraintsTransform together.
+   *        this->dst must be the same as other->src.
+   * @param other another IntConstraintsTransform whose src is same as this->dst.
+   * @return composed IntConstraintsTransform(this->src, other->dst)
+   *         with its variables and ranges are properly modified.
+   */
+  IntConstraintsTransform operator+(const IntConstraintsTransform& other) const;
+
   TVM_DEFINE_OBJECT_REF_METHODS(IntConstraintsTransform, ObjectRef, IntConstraintsTransformNode);
 };
 
@@ -306,6 +315,16 @@ IntConstraintsTransform SolveLinearEquations(const IntConstraints& system_to_sol
  */
 PartialSolvedInequalities SolveLinearInequalities(const IntConstraints& system_to_solve);
 
+/*!
+ * \brief Combine the information into an array of (in)equalities.
+ * \param variables The variables in \p bounds.
+ *        It is used to determine the iteration order to avoid indeterministic results.
+ * \param bounds grouped boundary of the variables.
+ * \param relations other relations.
+ */
+Array<PrimExpr> AsConditions(const Array<Var>& variables, const Map<Var, IntGroupBounds>& bounds,
+                             const Array<PrimExpr>& relations);
+
 /*!
  * \brief Solve linear inequalities and infer the range of each variable.
  * \param system_to_solve the variables to solve, their ranges, and a list of inequalities.

src/arith/canonical_simplify.cc

@@ -59,13 +59,6 @@ class CanonicalExprNode : public PrimExprNode {
   TVM_DECLARE_BASE_OBJECT_INFO(CanonicalExprNode, PrimExprNode);
 };
 
-enum DivMode {
-  /*! \brief Truncated division. */
-  kTruncDiv,
-  /*! \brief Floor division. */
-  kFloorDiv
-};
-
 inline PrimExpr ModImpl(PrimExpr a, PrimExpr b, DivMode mode) {
   if (mode == kTruncDiv) {
     return truncmod(a, b);

src/arith/int_constraints.cc

@@ -38,6 +38,32 @@
 namespace tvm {
 namespace arith {
 
+Array<PrimExpr> AsConditions(const Array<Var>& variables, const Map<Var, IntGroupBounds>& bounds,
+                             const Array<PrimExpr>& relations) {
+  Array<PrimExpr> res;
+  // use variables to keep the order of iteration
+  // so as to get rid of any non-determinism.
+  CHECK_EQ(variables.size(), bounds.size());
+  for (const auto v : variables) {
+    CHECK(bounds.count(v));
+    const auto& bnds = bounds[v];
+    PrimExpr lhs = bnds->coef * v;
+    for (const PrimExpr& rhs : bnds->equal) {
+      res.push_back(tir::EQ(lhs, rhs));
+    }
+    for (const PrimExpr& rhs : bnds->lower) {
+      res.push_back(tir::GE(lhs, rhs));
+    }
+    for (const PrimExpr& rhs : bnds->upper) {
+      res.push_back(tir::LE(lhs, rhs));
+    }
+  }
+  for (const PrimExpr& e : relations) {
+    res.push_back(e);
+  }
+  return res;
+}
+
 IntGroupBounds::IntGroupBounds(PrimExpr coef, Array<PrimExpr> lower, Array<PrimExpr> equal,
                                Array<PrimExpr> upper) {
   CHECK(coef.dtype().is_int() || coef.dtype().is_uint())
@@ -231,6 +257,26 @@ IntConstraintsTransform::IntConstraintsTransform(IntConstraints src, IntConstrai
   data_ = std::move(node);
 }
 
+IntConstraintsTransform IntConstraintsTransform::operator+(
+    const IntConstraintsTransform& other) const {
+  CHECK(other->src.same_as(operator->()->dst));
+  Map<Var, PrimExpr> dst_to_src;
+  Map<Var, PrimExpr> src_to_dst;
+
+  Analyzer ana_first;
+  ana_first.Bind(operator->()->src->ranges);
+  for (auto p : other->dst_to_src) {
+    dst_to_src.Set(p.first, ana_first.Simplify(Substitute(p.second, operator->()->dst_to_src)));
+  }
+
+  Analyzer ana_second;
+  ana_second.Bind(other->dst->ranges);
+  for (auto p : operator->()->src_to_dst) {
+    src_to_dst.Set(p.first, ana_second.Simplify(Substitute(p.second, other->src_to_dst)));
+  }
+  return IntConstraintsTransform(operator->()->src, other->dst, src_to_dst, dst_to_src);
+}
+
 TVM_REGISTER_NODE_TYPE(IntConstraintsTransformNode);
 
 TVM_REGISTER_GLOBAL("arith.IntConstraintsTransform")

src/arith/solve_linear_inequality.cc

@@ -94,35 +94,6 @@ struct ExprLess {
   }
 };
 
-/*!
- * \brief Combine the information into an array of (in)equalities.
- */
-Array<PrimExpr> as_conditions(const Array<Var>& variables, const Map<Var, IntGroupBounds>& bounds,
-                              const Array<PrimExpr>& relations) {
-  Array<PrimExpr> res;
-  // use variables to keep the order of iteration
-  // so as to get rid of any non-determinism.
-  CHECK_EQ(variables.size(), bounds.size());
-  for (const auto v : variables) {
-    CHECK(bounds.count(v));
-    const auto& bnds = bounds[v];
-    PrimExpr lhs = bnds->coef * v;
-    for (const PrimExpr& rhs : bnds->equal) {
-      res.push_back(tir::EQ(lhs, rhs));
-    }
-    for (const PrimExpr& rhs : bnds->lower) {
-      res.push_back(tir::GE(lhs, rhs));
-    }
-    for (const PrimExpr& rhs : bnds->upper) {
-      res.push_back(tir::LE(lhs, rhs));
-    }
-  }
-  for (const PrimExpr& e : relations) {
-    res.push_back(e);
-  }
-  return res;
-}
-
 void DebugPrint(
     const std::unordered_set<PrimExpr, StructuralHash, StructuralEqual>& current_ineq_set,
     const std::unordered_set<PrimExpr, StructuralHash, StructuralEqual>& next_ineq_set,
@@ -491,7 +462,7 @@ IntConstraints SolveInequalitiesToRange(const IntConstraints& inequalities) {
   arith::Analyzer analyzer;
   analyzer.Bind(vranges);
   for (const PrimExpr& old_cond :
-       as_conditions(inequalities->variables, solved_bounds, solved_other_relations)) {
+       AsConditions(inequalities->variables, solved_bounds, solved_other_relations)) {
     if (!analyzer.CanProve(old_cond)) {
       // those not represented in vranges (res_ranges)
       res_relations.push_back(old_cond);
@@ -584,7 +555,7 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& inequ
 
   // Add the original conditions (with variables substituted) to the resulting conditions
   for (const PrimExpr& old_cond :
-       as_conditions(inequalities->variables, solved_bounds, solved_other_relations)) {
+       AsConditions(inequalities->variables, solved_bounds, solved_other_relations)) {
     PrimExpr new_cond = analyzer.Simplify(Substitute(old_cond, res_src_to_dst));
     if (!is_const_int(new_cond, 1)) {
       // those not represented in vranges (res_ranges)
@@ -615,7 +586,7 @@ TVM_REGISTER_GLOBAL("arith.SolveInequalitiesAsCondition")
         LOG(FATAL) << "arith.SolveInequalitiesAsCondition expects 1 or 3 arguments, gets "
                    << args.size();
       }
-      *ret = as_conditions(problem->variables, ret_ineq.first, ret_ineq.second);
+      *ret = AsConditions(problem->variables, ret_ineq.first, ret_ineq.second);
     });
 
 TVM_REGISTER_GLOBAL("arith.SolveInequalitiesToRange").set_body([](TVMArgs args, TVMRetValue* ret) {