This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 000f558c72 [REFACTOR][Arith] Delete IntGroupBounds::Substitute (#20305)
000f558c72 is described below
commit 000f558c72d124cc056906294a779dc42c2b4e75
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Sep 10 08:11:56 2026 -0400
[REFACTOR][Arith] Delete IntGroupBounds::Substitute (#20305)
Delete the redundant `IntGroupBounds::Substitute` helper and use
reflected pre-order `StructuralMap` at its only call site.
Also use checked `as_or_throw<T>()` conversion for arith StructuralMap
results and move values out of successful map lookups. `IterMarkNode`
DAG semantics are unchanged.
Validation: LLVM 17 CPU build; focused arith tests (11 passed, 1
skipped); C++ tests (133/133 passed); clang-format.
---
include/tvm/arith/int_solver.h | 5 -----
src/arith/int_constraints.cc | 21 ++++-----------------
src/arith/solve_linear_equation.cc | 5 +++--
src/arith/solve_linear_inequality.cc | 19 ++++++++++---------
4 files changed, 17 insertions(+), 33 deletions(-)
diff --git a/include/tvm/arith/int_solver.h b/include/tvm/arith/int_solver.h
index d7a968bf89..c2b2574109 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<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.
diff --git a/src/arith/int_constraints.cc b/src/arith/int_constraints.cc
index 0d42fb0d50..a16553e8b2 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<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);
}
diff --git a/src/arith/solve_linear_equation.cc
b/src/arith/solve_linear_equation.cc
index 90c3404e78..c00af8a58b 100644
--- a/src/arith/solve_linear_equation.cc
+++ b/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);
diff --git a/src/arith/solve_linear_inequality.cc
b/src/arith/solve_linear_inequality.cc
index 44b97a30ab..54f1fbc80b 100644
--- a/src/arith/solve_linear_inequality.cc
+++ b/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);
