[REFACTOR][Arith] Delete IntGroupBounds::Substitute

include/tvm/arith/int_solver.h

@@ -102,11 +102,6 @@ class IntGroupBounds : public ffi::ObjectRef {
    */
   static IntGroupBounds FromRange(const Range& r);
 
-  /*!
-   * \brief Perform substitution on all components of the struct.
-   */
-  IntGroupBounds Substitute(const ffi::Map<Var, PrimExpr>& subst) const;
-
   /*!
    * \brief Find the best range from the grouped bounds.
    * \param vranges_addl additional variable ranges that help infer the best range.

src/arith/int_constraints.cc

@@ -116,19 +116,6 @@ IntGroupBounds IntGroupBounds::operator+(const Range& r) {
   return IntGroupBounds(coef, lower, equal, upper);
 }
 
-IntGroupBounds IntGroupBounds::Substitute(const ffi::Map<Var, PrimExpr>& subst) const {
-  auto f_subst = [&subst](const Var& var) -> ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
-    if (auto repl = subst.Get(var)) return ffi::Any(repl.value());
-    return ffi::Unchanged();
-  };
-  auto apply_fun = [&f_subst](const PrimExpr& e) {
-    return ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(e, f_subst).cast<PrimExpr>();
-  };
-  return IntGroupBounds(apply_fun(operator->()->coef), operator->()->lower.Map(apply_fun),
-                                                       operator->()->equal.Map(apply_fun),
-                                                       operator->()->upper.Map(apply_fun));
-}
-
 Range IntGroupBounds::FindBestRange(const ffi::Map<Var, Range>& vranges_addl) const {
   Analyzer analyzer;
   analyzer->Bind(vranges_addl);
@@ -276,25 +263,25 @@ IntConstraintsTransform IntConstraintsTransform::operator+(
   Analyzer ana_first;
   ana_first->Bind(operator->()->src->ranges);
   auto f_dst_to_src = [this](const Var& var) -> ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
-    if (auto repl = operator->()->dst_to_src.Get(var)) return ffi::Any(repl.value());
+    if (auto repl = operator->()->dst_to_src.Get(var)) return ffi::Any(*std::move(repl));
     return ffi::Unchanged();
   };
   for (auto p : other->dst_to_src) {
     dst_to_src.Set(p.first, ana_first->Simplify(ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(
                                                     p.second, f_dst_to_src)
-                                                    .cast<PrimExpr>()));
+                                                    .as_or_throw<PrimExpr>()));
   }
 
   Analyzer ana_second;
   ana_second->Bind(other->dst->ranges);
   auto f_src_to_dst = [&other](const Var& var) -> ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
-    if (auto repl = other->src_to_dst.Get(var)) return ffi::Any(repl.value());
+    if (auto repl = other->src_to_dst.Get(var)) return ffi::Any(*std::move(repl));
     return ffi::Unchanged();
   };
   for (auto p : operator->()->src_to_dst) {
     src_to_dst.Set(p.first, ana_second->Simplify(ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(
                                                      p.second, f_src_to_dst)
-                                                     .cast<PrimExpr>()));
+                                                     .as_or_throw<PrimExpr>()));
   }
   return IntConstraintsTransform(operator->()->src, other->dst, src_to_dst, dst_to_src);
 }

src/arith/solve_linear_equation.cc

@@ -33,6 +33,7 @@
 #include <tvm/tirx/op.h>
 
 #include <unordered_set>
+#include <utility>
 
 #include "int_operator.h"
 
@@ -450,12 +451,12 @@ IntConstraintsTransform SolveLinearEquations(const IntConstraints& system_to_sol
 
   // Add the rest conditions
   auto f_subst = [&old_to_new_map](const Var& var) -> ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
-    if (auto repl = old_to_new_map.Get(var)) return ffi::Any(repl.value());
+    if (auto repl = old_to_new_map.Get(var)) return ffi::Any(*std::move(repl));
     return ffi::Unchanged();
   };
   for (const PrimExpr& cond : rest) {
     new_relations.push_back(
-        ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(cond, f_subst).cast<PrimExpr>());
+        ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(cond, f_subst).as_or_throw<PrimExpr>());
   }
 
   IntConstraints solution(new_vars, new_ranges, new_relations);

src/arith/solve_linear_inequality.cc

@@ -33,6 +33,8 @@
 #include <tvm/tirx/expr_functor.h>
 #include <tvm/tirx/op.h>
 
+#include <utility>
+
 #include "int_operator.h"
 
 namespace tvm {
@@ -468,9 +470,13 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& inequ
   }
   analyzer->Bind(vranges);
 
+  auto subst = [&res_src_to_dst](const Var& var) -> ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
+    if (auto repl = res_src_to_dst.Get(var)) return ffi::Any(*std::move(repl));
+    return ffi::Unchanged();
+  };
   auto f_dst_to_src =
       [&res_dst_to_src](const Var& var) -> ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
-    if (auto repl = res_dst_to_src.Get(var)) return ffi::Any(repl.value());
+    if (auto repl = res_dst_to_src.Get(var)) return ffi::Any(*std::move(repl));
     return ffi::Unchanged();
   };
 
@@ -480,7 +486,7 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& inequ
     const PrimVar& var = *it;
     auto bnd = solved_bounds[var];
     // Note that we replace old vars with new ones
-    bnd = bnd.Substitute(res_src_to_dst);
+    bnd = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(bnd, subst).as_or_throw<IntGroupBounds>();
 
     if (is_one(bnd->coef) && !bnd->equal.empty()) {
       // There is an equation of the form `v == expr`,
@@ -520,7 +526,7 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& inequ
                            analyzer->Simplify(var.as_or_throw<PrimExpr>() -
                                               ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(
                                                   best_range->min, f_dst_to_src)
-                                                  .cast<PrimExpr>()));
+                                                  .as_or_throw<PrimExpr>()));
 
         // Add the new var to the resulting axis
         auto range = Range(IntImm(new_var->ty.as_or_throw<PrimType>(), 0), best_range->extent);
@@ -534,15 +540,10 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& inequ
   }
 
   // Add the original conditions (with variables substituted) to the resulting conditions
-  auto f_src_to_dst =
-      [&res_src_to_dst](const Var& var) -> ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
-    if (auto repl = res_src_to_dst.Get(var)) return ffi::Any(repl.value());
-    return ffi::Unchanged();
-  };
   for (const PrimExpr& old_cond :
        AsConditions(inequalities->variables, solved_bounds, solved_other_relations)) {
     PrimExpr new_cond = analyzer->Simplify(
-        ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(old_cond, f_src_to_dst).cast<PrimExpr>());
+        ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(old_cond, subst).as_or_throw<PrimExpr>());
     if (!is_const_int(new_cond, 1)) {
       // those not represented in vranges (res_ranges)
       res_relations.push_back(new_cond);