diff --git a/include/tvm/arith/int_solver.h b/include/tvm/arith/int_solver.h index d7a968bf89cb..c2b257410938 100644 --- a/include/tvm/arith/int_solver.h +++ b/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& subst) const; - /*! * \brief Find the best range from the grouped bounds. * \param vranges_addl additional variable ranges that help infer the best range. diff --git a/src/arith/int_constraints.cc b/src/arith/int_constraints.cc index 0d42fb0d503c..a16553e8b24b 100644 --- a/src/arith/int_constraints.cc +++ b/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& subst) const { - auto f_subst = [&subst](const Var& var) -> ffi::Expected> { - 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(e, f_subst).cast(); - }; - 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& 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> { - 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( p.second, f_dst_to_src) - .cast())); + .as_or_throw())); } Analyzer ana_second; ana_second->Bind(other->dst->ranges); auto f_src_to_dst = [&other](const Var& var) -> ffi::Expected> { - 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( p.second, f_src_to_dst) - .cast())); + .as_or_throw())); } return IntConstraintsTransform(operator->()->src, other->dst, src_to_dst, dst_to_src); } diff --git a/src/arith/solve_linear_equation.cc b/src/arith/solve_linear_equation.cc index 90c3404e7823..c00af8a58b16 100644 --- a/src/arith/solve_linear_equation.cc +++ b/src/arith/solve_linear_equation.cc @@ -33,6 +33,7 @@ #include #include +#include #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> { - 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(cond, f_subst).cast()); + ffi::StructuralMap(cond, f_subst).as_or_throw()); } IntConstraints solution(new_vars, new_ranges, new_relations); diff --git a/src/arith/solve_linear_inequality.cc b/src/arith/solve_linear_inequality.cc index 44b97a30abb6..54f1fbc80b3c 100644 --- a/src/arith/solve_linear_inequality.cc +++ b/src/arith/solve_linear_inequality.cc @@ -33,6 +33,8 @@ #include #include +#include + #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> { + 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> { - 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(bnd, subst).as_or_throw(); 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() - ffi::StructuralMap( best_range->min, f_dst_to_src) - .cast())); + .as_or_throw())); // Add the new var to the resulting axis auto range = Range(IntImm(new_var->ty.as_or_throw(), 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> { - 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(old_cond, f_src_to_dst).cast()); + ffi::StructuralMap(old_cond, subst).as_or_throw()); if (!is_const_int(new_cond, 1)) { // those not represented in vranges (res_ranges) res_relations.push_back(new_cond);