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-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new e46e8d409 [REFACTOR] Fold known comparison queries into
RewriteSimplifier (#79)
e46e8d409 is described below
commit e46e8d40999b45a76a15bb1e57ed18f7e4c89c10
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 26 09:56:52 2026 -0400
[REFACTOR] Fold known comparison queries into RewriteSimplifier (#79)
Move direct comparison reasoning into a private KnownComparisons helper
owned by RewriteSimplifier. Retain normalized comparison facts,
expression keys, and constant-offset matching, with ingestion, reset,
temporary-query state, and copying managed by the rewriter. Original
binding and range relations remain available alongside simplified
substitutions.
Retire the separate public KnownComparisonAnalyzer and route the
existing comparison API and canonical-simplifier queries through the
owning rewriter. Ordered fact ingestion replaces per-fact context
snapshots and raw-context query scans.
---
include/tvm/sym/analyzer.h | 61 +---
src/sym/analyzer.cc | 11 +-
src/sym/known_comparison_analyzer.cc | 588 -----------------------------------
src/sym/known_comparisons.h | 365 ++++++++++++++++++++++
src/sym/rewrite_simplify.cc | 42 ++-
src/sym/rewrite_simplify.h | 11 +-
6 files changed, 420 insertions(+), 658 deletions(-)
diff --git a/include/tvm/sym/analyzer.h b/include/tvm/sym/analyzer.h
index 29f526c15..7a1923efa 100644
--- a/include/tvm/sym/analyzer.h
+++ b/include/tvm/sym/analyzer.h
@@ -57,6 +57,7 @@ namespace sym {
// Forward declare the analyzer object and its reference handle.
class AnalyzerObj;
class Analyzer;
+enum class CompareResult : int;
namespace detail {
class TempConstraintScope;
} // namespace detail
@@ -270,6 +271,14 @@ class RewriteSimplifier {
*/
TVM_TIRX_DLL PrimExpr operator()(const PrimExpr& expr);
+ /*!
+ * \brief Compare signed integer expressions using active premises and
constant offsets.
+ * \param lhs The left-hand side of the comparison.
+ * \param rhs The right-hand side of the comparison.
+ * \return The directly known ordering, or kUnknown when it cannot be
established.
+ */
+ TVM_TIRX_DLL CompareResult TryCompare(const PrimExpr& lhs, const PrimExpr&
rhs);
+
/*!
* \brief Read cumulative rewrite statistics.
* \return A snapshot of the current counters.
@@ -302,8 +311,10 @@ class RewriteSimplifier {
friend class AnalyzerObj;
// Clear context-derived facts before replay; retain query settings and
diagnostics.
void Reset();
- // Record a value binding after AnalyzerObj validates and orders the
premises.
- void IngestBinding(const Var& var, const PrimExpr& expr);
+ // Keep the original relation alongside the simplified substitution.
+ void IngestBinding(const Var& var, const PrimExpr& value, const PrimExpr&
original);
+ // Record a range relation after AnalyzerObj validates and orders the
premises.
+ void IngestBinding(const Var& var, const Range& range);
// Consume a predicate prepared against the preceding premise prefix.
void IngestPredicate(const PrimExpr& constraint);
friend class CanonicalSimplifier;
@@ -391,48 +402,6 @@ inline constexpr CompareResult operator|(CompareResult
lhs, CompareResult rhs) {
return CompareResult(static_cast<int>(lhs) | static_cast<int>(rhs));
}
-/*!
- * \brief Using previously specified knowns, compare the expressions provided
- *
- * Compare directly related expressions, allowing constant offsets.
- */
-class KnownComparisonAnalyzer {
- public:
- /*!
- * \brief Compare expressions using recorded relations and constant offsets.
- *
- * \param lhs The left-hand side of the comparison
- *
- * \param rhs The right-hand side of the comparison
- *
- * Directly known comparisons also allow constant offsets on either side.
- *
- * \return The most specific result that can be proven about the
- * comparison. If nothing can be proven, returns kUnknown.
- */
- TVM_TIRX_DLL CompareResult TryCompare(const PrimExpr& lhs, const PrimExpr&
rhs);
-
- private:
- friend class AnalyzerObj;
- // Clear context-derived facts before replay; retain query settings and
diagnostics.
- void Reset();
- // Record a value binding after AnalyzerObj validates and orders the
premises.
- void IngestBinding(const Var& var, const PrimExpr& expr);
- // Record a half-open variable range from the accepted premise prefix.
- void IngestBinding(const Var& var, const Range& range);
- // Consume a predicate prepared against the preceding premise prefix.
- void IngestPredicate(const PrimExpr& constraint);
- // Borrow the owning analyzer for cross-backend queries.
- explicit KnownComparisonAnalyzer(AnalyzerObj* parent);
- // Exported for destruction through the inline Analyzer object allocator.
- TVM_TIRX_DLL ~KnownComparisonAnalyzer();
- // Copy durable facts/settings into this backend without sharing mutable
query state.
- void CopyFrom(const KnownComparisonAnalyzer& other);
- class Impl;
- // Owned backend state; its definition stays inside the compiler library.
- std::unique_ptr<Impl> impl_;
-};
-
/*!
* \brief Integer set analyzer.
*/
@@ -596,8 +565,6 @@ class TVM_TIRX_DLL AnalyzerObj : public ffi::Object {
CanonicalSimplifier canonical_simplify;
/*! \brief sub-analyzer: int set */
IntSetAnalyzer int_set;
- /*! \brief sub-analyzer known comparisons */
- KnownComparisonAnalyzer known_comparisons;
/*! \brief sub-analyzer using Z3 */
Z3Prover z3_prover;
/*! \brief Construct an analyzer with no accepted premises and default query
settings. */
@@ -743,7 +710,7 @@ class TVM_TIRX_DLL AnalyzerObj : public ffi::Object {
void IngestConstraint(const Constraint& constraint);
// Normalize and derive predicate facts using only the already valid prefix.
std::vector<PrimExpr> PreparePredicate(const PrimExpr& condition);
- // Fixed position of the last successfully accepted premise environment.
+ // Accepted premise prefix, advanced after each fact during synchronization.
ConstraintContext last_snapshot_;
// Query-local premises, managed in LIFO order by TempConstraintScope.
std::vector<Constraint> temp_constraints_;
diff --git a/src/sym/analyzer.cc b/src/sym/analyzer.cc
index b9a684338..b3da3cfde 100644
--- a/src/sym/analyzer.cc
+++ b/src/sym/analyzer.cc
@@ -55,7 +55,6 @@ AnalyzerObj::AnalyzerObj()
rewrite_simplify(this),
canonical_simplify(this),
int_set(this),
- known_comparisons(this),
z3_prover(this) {}
ConstraintContext AnalyzerObj::GetConstraintContext() const {
@@ -122,7 +121,6 @@ void AnalyzerObj::SyncConstraintContext(const
ConstraintContext& context) {
rewrite_simplify.Reset();
canonical_simplify.Reset();
int_set.Reset();
- known_comparisons.Reset();
z3_prover.Reset();
next->ForEach([&](const Constraint& fact) { IngestConstraint(fact); });
}
@@ -136,10 +134,9 @@ void AnalyzerObj::IngestBinding(const Var& var, const
PrimExpr& expr) {
value = rewrite_simplify(value);
const_int_bound.IngestBinding(var, value);
modular_set.IngestBinding(var, value);
- rewrite_simplify.IngestBinding(var, value);
+ rewrite_simplify.IngestBinding(var, value, expr);
canonical_simplify.IngestBinding(var, value);
int_set.IngestBinding(var, value);
- known_comparisons.IngestBinding(var, expr);
z3_prover.IngestBinding(var, expr);
}
@@ -149,7 +146,7 @@ void AnalyzerObj::IngestBinding(const Var& var, const
Range& range) {
} else {
const_int_bound.IngestBinding(var, range);
int_set.IngestBinding(var, range);
- known_comparisons.IngestBinding(var, range);
+ rewrite_simplify.IngestBinding(var, range);
z3_prover.IngestBinding(var, range);
}
}
@@ -190,7 +187,6 @@ void AnalyzerObj::IngestConstraint(const Constraint& fact) {
modular_set.IngestPredicate(condition);
rewrite_simplify.IngestPredicate(condition);
int_set.IngestPredicate(condition);
- known_comparisons.IngestPredicate(condition);
z3_prover.IngestPredicate(condition);
}
}
@@ -354,7 +350,6 @@ Analyzer AnalyzerObj::Clone() const {
cloned->rewrite_simplify.CopyFrom(this->rewrite_simplify);
cloned->canonical_simplify.CopyFrom(this->canonical_simplify);
cloned->int_set.CopyFrom(this->int_set);
- cloned->known_comparisons.CopyFrom(this->known_comparisons);
cloned->z3_prover.CopyFrom(this->z3_prover);
cloned->last_snapshot_ = last_snapshot_;
return cloned;
@@ -416,7 +411,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def_method("sym.AnalyzerCanProveEqual", &AnalyzerObj::CanProveEqual)
.def("sym.AnalyzerTryCompare",
[](Analyzer analyzer, const PrimExpr& lhs, const PrimExpr& rhs) {
- return
static_cast<int64_t>(analyzer->known_comparisons.TryCompare(lhs, rhs));
+ return
static_cast<int64_t>(analyzer->rewrite_simplify.TryCompare(lhs, rhs));
})
.def("sym.AnalyzerIsZ3Enabled",
[](Analyzer analyzer) { return analyzer->z3_prover.IsEnabled(); })
diff --git a/src/sym/known_comparison_analyzer.cc
b/src/sym/known_comparison_analyzer.cc
deleted file mode 100644
index 3d3d0087d..000000000
--- a/src/sym/known_comparison_analyzer.cc
+++ /dev/null
@@ -1,588 +0,0 @@
-/*
- * 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/sym/known_comparison_analyzer.cc
- */
-
-#include <tvm/ir/prim/expr.h>
-#include <tvm/ir/prim/pattern_match.h>
-#include <tvm/sym/analyzer.h>
-
-#include <optional>
-#include <vector>
-
-#include "constraint_extract.h"
-#include "constraint_helpers.h"
-
-namespace tvm {
-namespace sym {
-using namespace tvm::prim;
-
-using prim::IsConstInt;
-
-class KnownComparisonAnalyzer::Impl {
- public:
- explicit Impl(AnalyzerObj* owner = nullptr) : owner_(owner) {}
- AnalyzerObj* owner_;
- /* \brief Using previously specified knowns, compare the expressions provided
- *
- * \param lhs The left-hand side of the comparison
- *
- * \param rhs The right-hand side of the comparison
- *
- * Directly known comparisons also allow constant offsets on either side.
- *
- * \return The most specific result that can be proven about the
- * comparison. If nothing can be proven, returns kUnknown.
- */
- CompareResult TryCompare(const PrimExpr& lhs, const PrimExpr& rhs) const;
-
- /*! \brief Bind a variable as being equal to a known expression
- *
- * \param var The variable of interest.
- * \param expr The bound expression
- * \param allow_override Whether to allow override of existing information.
- */
- void IngestBinding(const Var& var, const PrimExpr& expr, bool allow_override
= false);
-
- /*! \brief Bind a variable as being within a specified range
- *
- * \param var The variable of interest.
- * \param range The known range
- * \param allow_override Whether to allow override of existing information.
- */
- void IngestBinding(const Var& var, const Range& expr, bool allow_override =
false);
-
- /*!
- * \brief Add a predicate to the active derived state.
- * \param constraint A constraint expression.
- */
- void AddPredicate(const PrimExpr& expr);
-
- void CopyFrom(const Impl& other) {
- expr_to_key = other.expr_to_key;
- prev_bindings_ = other.prev_bindings_;
- knowns_ = other.knowns_;
- scoped_knowns_ = other.scoped_knowns_;
- }
-
- private:
- /* \brief Internal representation of a PrimExpr
- *
- * The Key enum serves two purposes.
- *
- * 1. Providing efficiency, as compared to a PrimExpr. Two keys are
- * equal if and only if the corresponding PrimExprs would satisfy
- * prim::ExprDeepEqual. This allows two expressions to be checked for
- * equivalency, without requiring a call to prim::ExprDeepEqual for
- * each comparison.
- *
- * 2. Providing type-safety, as compared to using `size_t` directly.
- * Requiring an explicit conversion from an integer to a Key
- * prevents accidental comparisons, especially if both loop
- * iterators and Keys are used in the same scope.
- *
- * A Key should only be obtained using the methods `ExprToKey` and
- * `ExprToPreviousKey`.
- */
- enum class Key : size_t {};
-
- /*! \brief Convert an expression to internal representation
- *
- * If the expression has previously been converted to the internal
- * representation, returns the same Key as has been used previously.
- * Otherwise, generate and return a new Key.
- *
- * \param expr The PrimExpr to be converted
- *
- * \returns The Key representing the expression
- *
- * \see ExprToPreviousKey
- */
- Key ExprToKey(const PrimExpr& expr);
-
- /*! \brief Convert an expression to internal representation
- *
- * If the expression has previously been converted to the internal
- * representation, returns the same Key as has been used previously.
- * Otherwise, return `std::nullopt`.
- *
- * \param expr The PrimExpr to be converted
- *
- * \returns The Key representing the expression, if one exists.
- *
- * \see ExprToKey
- */
- std::optional<Key> ExprToPreviousKey(const PrimExpr& expr) const;
-
- /*! \brief The mapping from expression to Key
- *
- * Should not be used directly. Instead, use the helper functions
- * `ExprToKey` and `ExprToPreviousKey`.
- *
- * \see ExprToKey
- * \see ExprToPreviousKey
- */
- std::unordered_map<PrimExpr, Key, ffi::StructuralHash, ffi::StructuralEqual>
expr_to_key;
-
- /*! \brief Internal representation of a comparison operator */
- struct Comparison {
- /*! \brief Construct a comparison that represents `lhs OP rhs +
- * offset`, where the operation is specified by the CompareResult.
- */
- Comparison(Key lhs, Key rhs, int64_t offset, CompareResult result);
-
- /*! \brief Move the specified expression to the LHS.
- *
- * \param new_lhs The argument that should be moved to the LHS of the
- * comparison.
- *
- * \return If possible, returns a comparison that is equivalent to
- * the current comparison, but with the specified LHS. If not
- * possible, returns nullopt.
- */
- std::optional<Comparison> WithLHS(Key new_lhs) const;
-
- // The LHS of the comparison
- Key lhs_;
-
- // The RHS of the comparison, not including any constant offset.
- Key rhs_;
-
- // Additive offset on rhs
- int64_t offset_{0};
-
- // The comparison operator.
- CompareResult result_{CompareResult::kInconsistent};
- };
-
- /*! \brief Generate a Comparison representing the given expression */
- std::optional<Comparison> FromExpr(const PrimExpr& expr);
-
- /*! \brief Utility function used by IngestBinding and AddPredicate
- *
- * \param expr The comparison expression, to be converted into
- * internal Comparison objects.
- *
- * \param vec The vector to which the Comparison objects should be
- * appended.
- */
- void AddKnown(const PrimExpr& expr, std::vector<Comparison>* vec);
-
- /*! Collect known comparisons between LHS and RHS, without propagation
- *
- * Allows the internal representation to handle any constant
- * offsets, without searching for a sequence of inequalities.
- *
- * \param lhs_key The left-hand side of the comparison
- *
- * \param rhs_key The right-hand side of the comparison
- *
- * \returns A subset of `knowns_` and `scoped_knowns_`, filtered to
- * only include comparisons between `lhs_key` and `rhs_key`,
- * normalized such that `lhs_key` is on the left-hand side.
- */
- std::vector<Comparison> CollectDirectComparisons(Key lhs_key, Key rhs_key)
const;
-
- /*! \brief Combine a set of comparisons that share a LHS and RHS
- *
- * \param lhs_to_rhs The comparisons to merge. These should all
- * have the same LHS and RHS. This parameter will typically be the
- * result from `CollectDirectComparisons`.
- *
- * \param offset The constant offset in the comparison being proven.
- * This is extracted from any additive/subtractive constants in the
- * `PrimExpr` arguments to `TryCompare`.
- *
- * \returns The possible comparisons between LHS and RHS provided
- * inequalities.
- */
- CompareResult MergeComparisons(const std::vector<Comparison>& lhs_to_rhs,
int64_t offset) const;
-
- /*! \brief Previous Range bindings
- *
- * Tracked separatedly to handle the `allow_override` option used by
- * all sub-analyzers when binding variables.
- */
- ffi::Map<Var, Range> prev_bindings_;
-
- /*! \brief Known comparisons based on definitionally-true statements
- *
- * For example, a Let binding, or the range of an iterator. These
- * known statements are always true, based on the definition site of
- * the variable. e.g. A loop iterator may never exceed the bounds
- * of its loop.
- */
- std::vector<Comparison> knowns_;
-
- /*! \brief Known comparisons based on scoped conditions
- *
- * For example, the condition of an If. These known
- * statements may only be used within the scope of the statement
- * from which they were derived. e.g. After exiting an If,
- * the condition may no longer be true.
- */
- std::vector<Comparison> scoped_knowns_;
-};
-
-namespace {
-
-// Internal utility, return the CompareResult resulting from swapping
-// the left-hand side with the right-hand side.
-CompareResult Reverse(CompareResult res) {
- switch (res) {
- case CompareResult::kInconsistent:
- return CompareResult::kInconsistent;
- case CompareResult::kEQ:
- return CompareResult::kEQ;
- case CompareResult::kLT:
- return CompareResult::kGT;
- case CompareResult::kLE:
- return CompareResult::kGE;
- case CompareResult::kGT:
- return CompareResult::kLT;
- case CompareResult::kGE:
- return CompareResult::kLE;
- case CompareResult::kNE:
- return CompareResult::kNE;
- case CompareResult::kUnknown:
- return CompareResult::kUnknown;
- default:
- TVM_FFI_THROW(InternalError) << "Invalid CompareResult: " <<
static_cast<int>(res);
- }
-}
-
-// Internal utility, extract constant offsets out of the two sides of
-// a comparison. Given lhs and rhs, return a tuple of three elements
-// (lhs_inner, rhs_inner, offset), such that (lhs OP rhs) and
-// (lhs_inner OP rhs_inner + offset) are equivalent.
-std::tuple<PrimExpr, PrimExpr, int64_t> ExtractOffsets(const PrimExpr& lhs,
const PrimExpr& rhs) {
- auto extract_offset = [](const PrimExpr& expr) -> std::pair<PrimExpr,
ffi::BigInt> {
- PVar<PrimExpr> x;
- PVar<IntImm> c;
- if ((x + c).Match(expr)) {
- return {x.Eval(), c.Eval()->value};
- } else if ((x - c).Match(expr)) {
- return {x.Eval(), -c.Eval()->value};
- } else if (c.Match(expr)) {
- return {0, c.Eval()->value};
- } else {
- return {expr, 0};
- }
- };
-
- auto lhs_split = extract_offset(lhs);
- auto rhs_split = extract_offset(rhs);
- if (auto offset = (rhs_split.second - lhs_split.second).as<int64_t>();
offset.has_value()) {
- return {lhs_split.first, rhs_split.first, *offset};
- }
- return {lhs, rhs, 0};
-}
-
-} // namespace
-
-std::optional<KnownComparisonAnalyzer::Impl::Comparison>
KnownComparisonAnalyzer::Impl::FromExpr(
- const PrimExpr& expr) {
- CompareResult res;
- PVar<PrimExpr> x, y;
- if ((x <= y).Match(expr)) {
- res = CompareResult::kLE;
- } else if ((x >= y).Match(expr)) {
- res = CompareResult::kGE;
- } else if ((x < y).Match(expr)) {
- res = CompareResult::kLT;
- } else if ((x > y).Match(expr)) {
- res = CompareResult::kGT;
- } else if ((x == y).Match(expr)) {
- res = CompareResult::kEQ;
- } else if ((x != y).Match(expr)) {
- res = CompareResult::kNE;
- } else {
- return std::nullopt;
- }
-
- PrimExpr lhs_expr = x.Eval();
- PrimExpr rhs_expr = y.Eval();
-
- if (lhs_expr.as<IntImmObj>() && rhs_expr.as<IntImmObj>()) {
- return std::nullopt;
- }
-
- auto [lhs, rhs, offset] = ExtractOffsets(lhs_expr, rhs_expr);
- ffi::BigInt normalized_offset = offset;
- if (res == CompareResult::kLT) --normalized_offset;
- if (res == CompareResult::kGT) ++normalized_offset;
- if (!normalized_offset.as<int64_t>().has_value()) return std::nullopt;
- Key lhs_key = ExprToKey(lhs);
- Key rhs_key = ExprToKey(rhs);
-
- return Comparison(lhs_key, rhs_key, offset, res);
-}
-
-KnownComparisonAnalyzer::Impl::Comparison::Comparison(Key lhs, Key rhs,
int64_t offset,
- CompareResult result)
- : lhs_(lhs), rhs_(rhs), offset_(offset), result_(result) {
- // Integer strict comparisons become non-strict comparisons with a unit
offset.
- if (result_ == CompareResult::kLT) {
- result_ = CompareResult::kLE;
- offset_ -= 1;
- }
- if (result_ == CompareResult::kGT) {
- result_ = CompareResult::kGE;
- offset_ += 1;
- }
-}
-
-std::optional<KnownComparisonAnalyzer::Impl::Key>
KnownComparisonAnalyzer::Impl::ExprToPreviousKey(
- const PrimExpr& expr) const {
- auto it = expr_to_key.find(expr);
- if (it != expr_to_key.end()) {
- return it->second;
- } else {
- return std::nullopt;
- }
-}
-
-KnownComparisonAnalyzer::Impl::Key
KnownComparisonAnalyzer::Impl::ExprToKey(const PrimExpr& expr) {
- if (auto prev = ExprToPreviousKey(expr)) {
- return prev.value();
- } else {
- Key new_key = Key(expr_to_key.size());
- expr_to_key[expr] = new_key;
- return new_key;
- }
-}
-
-std::optional<KnownComparisonAnalyzer::Impl::Comparison>
-KnownComparisonAnalyzer::Impl::Comparison::WithLHS(Key new_lhs) const {
- if (new_lhs == lhs_) {
- return *this;
- } else if (new_lhs == rhs_) {
- auto reversed = (-ffi::BigInt(offset_)).as<int64_t>();
- if (!reversed.has_value()) return std::nullopt;
- return Comparison(rhs_, lhs_, *reversed, Reverse(result_));
- } else {
- return std::nullopt;
- }
-}
-
-KnownComparisonAnalyzer::KnownComparisonAnalyzer(AnalyzerObj* parent)
- : impl_(std::make_unique<Impl>(parent)) {}
-KnownComparisonAnalyzer::~KnownComparisonAnalyzer() {}
-
-void KnownComparisonAnalyzer::Reset() {
- // Drop comparisons, binding metadata and their expression keys together.
- *impl_ = Impl(impl_->owner_);
-}
-
-void KnownComparisonAnalyzer::CopyFrom(const KnownComparisonAnalyzer& other) {
- impl_->CopyFrom(*other.impl_);
-}
-
-CompareResult KnownComparisonAnalyzer::TryCompare(const PrimExpr& lhs, const
PrimExpr& rhs) {
- detail::TempConstraintScope query(impl_->owner_);
- const auto& facts = detail::TempConstraintScope::Query(impl_->owner_);
- if (facts.empty()) return impl_->TryCompare(lhs, rhs);
- // Temporary binding metadata and expression keys belong to query scratch.
- Impl temporary;
- temporary.CopyFrom(*impl_);
- for (const Constraint& fact : facts) {
- if (const auto* binding = fact.as<BindingConstraintObj>())
- temporary.IngestBinding(binding->var, binding->value, true);
- else if (const auto* predicate = fact.as<PredicateConstraintObj>())
- temporary.AddPredicate(predicate->condition);
- }
- return temporary.TryCompare(lhs, rhs);
-}
-
-void KnownComparisonAnalyzer::IngestBinding(const Var& var, const PrimExpr&
expr) {
- impl_->IngestBinding(var, expr, true);
-}
-void KnownComparisonAnalyzer::IngestBinding(const Var& var, const Range&
range) {
- impl_->IngestBinding(var, range, true);
-}
-
-void KnownComparisonAnalyzer::IngestPredicate(const PrimExpr& constraint) {
- impl_->AddPredicate(constraint);
-}
-
-void KnownComparisonAnalyzer::Impl::AddKnown(const PrimExpr& expr,
std::vector<Comparison>* vec) {
- for (const auto& subexpr : ExtractConstraints(expr, false)) {
- if (SideEffect(expr) <= CallEffectKind::kPure) {
- if (auto cmp = FromExpr(subexpr)) {
- vec->push_back(cmp.value());
- }
- }
- }
-}
-
-void KnownComparisonAnalyzer::Impl::IngestBinding(const Var& var, const Range&
range,
- bool allow_override) {
- auto it = prev_bindings_.find(var);
- if (it != prev_bindings_.end()) {
- prim::ExprDeepEqual expr_equal;
- bool differs_from_previous = !expr_equal(range->begin,
(*it).second->begin) ||
- !expr_equal(range->extent,
(*it).second->extent);
- if (differs_from_previous) {
- TVM_FFI_ICHECK(allow_override) << "Binding of variable " << var << " as
" << range
- << " conflicts with previous binding as "
<< (*it).second;
- if (auto key = ExprToPreviousKey(var.as_or_throw<PrimExpr>())) {
- knowns_.erase(std::remove_if(knowns_.begin(), knowns_.end(),
- [&](const auto& known) { return
known.lhs_ == key.value(); }),
- knowns_.end());
- }
- }
- }
-
- prev_bindings_.Set(var, range);
-
- PrimExpr prim_var = var.as_or_throw<PrimExpr>();
- if (IsConstInt(range->extent, 1)) {
- AddKnown(prim_var == range->begin, &knowns_);
- } else {
- AddKnown(prim_var >= range->begin, &knowns_);
- AddKnown(prim_var < range->begin + range->extent, &knowns_);
- }
-}
-
-void KnownComparisonAnalyzer::Impl::IngestBinding(const Var& var, const
PrimExpr& expr,
- bool allow_override) {
- IngestBinding(var, Range::FromMinExtent(expr, 1), allow_override);
-}
-
-void KnownComparisonAnalyzer::Impl::AddPredicate(const PrimExpr& expr) {
- AddKnown(expr, &scoped_knowns_);
-}
-
-CompareResult KnownComparisonAnalyzer::Impl::TryCompare(const PrimExpr&
lhs_expr,
- const PrimExpr&
rhs_expr) const {
- // Currently only supports integer checks
- if (!lhs_expr.ty().MatchesCode(DLDataTypeCode::kDLInt) ||
- !rhs_expr.ty().MatchesCode(DLDataTypeCode::kDLInt)) {
- return CompareResult::kUnknown;
- }
-
- // Bail out early if possible. This int check should have been
- // constant-folded earlier, so this check shouldn't occur.
- auto* x_int = lhs_expr.as<IntImmObj>();
- auto* y_int = rhs_expr.as<IntImmObj>();
- if (x_int && y_int) {
- if (x_int->value < y_int->value) {
- return CompareResult::kLT;
- } else if (x_int->value > y_int->value) {
- return CompareResult::kGT;
- } else {
- return CompareResult::kEQ;
- }
- }
-
- auto [lhs, rhs, offset] = ExtractOffsets(lhs_expr, rhs_expr);
- auto lhs_key = ExprToPreviousKey(lhs);
- auto rhs_key = ExprToPreviousKey(rhs);
-
- if (!lhs_key.has_value() || !rhs_key.has_value()) {
- return CompareResult::kUnknown;
- }
-
- auto lhs_to_rhs = CollectDirectComparisons(lhs_key.value(), rhs_key.value());
- return MergeComparisons(lhs_to_rhs, offset);
-}
-
-std::vector<KnownComparisonAnalyzer::Impl::Comparison>
-KnownComparisonAnalyzer::Impl::CollectDirectComparisons(Key lhs_key, Key
rhs_key) const {
- std::vector<Comparison> output;
-
- auto append_known = [&](Comparison cmp) {
- if (auto normalized = cmp.WithLHS(lhs_key)) {
- if (normalized.value().rhs_ == rhs_key) {
- output.push_back(normalized.value());
- }
- }
- };
-
- for (const auto& known : knowns_) {
- append_known(known);
- }
- for (const auto& known : scoped_knowns_) {
- append_known(known);
- }
-
- return output;
-}
-
-CompareResult KnownComparisonAnalyzer::Impl::MergeComparisons(
- const std::vector<Comparison>& lhs_to_rhs, int64_t offset) const {
- // Just because we found a comparison involving LHS and RHS doesn't
- // mean that it's useful. e.g. Knowing that `x < y` doesn't let us
- // prove whether `x + 5 < y`.
- CompareResult result = CompareResult::kUnknown;
- for (const auto& cmp : lhs_to_rhs) {
- switch (cmp.result_) {
- case CompareResult::kInconsistent:
- result = CompareResult::kInconsistent;
- break;
-
- case CompareResult::kEQ:
- if (offset == cmp.offset_) {
- result = result & CompareResult::kEQ;
- } else {
- result = result & CompareResult::kNE;
- }
- break;
-
- case CompareResult::kLE:
- if (cmp.offset_ < offset) {
- result = result & CompareResult::kLT;
- } else if (cmp.offset_ <= offset) {
- result = result & CompareResult::kLE;
- }
- break;
-
- case CompareResult::kGE:
- if (cmp.offset_ > offset) {
- result = result & CompareResult::kGT;
- } else if (cmp.offset_ >= offset) {
- result = result & CompareResult::kGE;
- }
- break;
-
- case CompareResult::kNE:
- if (offset == cmp.offset_) {
- result = result & CompareResult::kNE;
- }
- break;
-
- case CompareResult::kUnknown:
- break;
-
- case CompareResult::kGT:
- case CompareResult::kLT:
- TVM_FFI_THROW(InternalError)
- << "Internal error, normalized comparisons should only include <=
and >=";
-
- default:
- TVM_FFI_THROW(InternalError) << "Invalid CompareResult: " <<
static_cast<int>(cmp.result_);
- }
- }
-
- return result;
-}
-
-} // namespace sym
-} // namespace tvm
diff --git a/src/sym/known_comparisons.h b/src/sym/known_comparisons.h
new file mode 100644
index 000000000..c4b586849
--- /dev/null
+++ b/src/sym/known_comparisons.h
@@ -0,0 +1,365 @@
+/*
+ * 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 sym/known_comparisons.h
+ * \brief Private directly known comparison state for the rewrite simplifier.
+ */
+#ifndef TVM_SYM_KNOWN_COMPARISONS_H_
+#define TVM_SYM_KNOWN_COMPARISONS_H_
+
+#include <tvm/ir/prim/expr.h>
+#include <tvm/ir/prim/op.h>
+#include <tvm/ir/prim/pattern_match.h>
+#include <tvm/sym/analyzer.h>
+
+#include <optional>
+#include <tuple>
+#include <unordered_map>
+#include <vector>
+
+#include "constraint_extract.h"
+
+namespace tvm {
+namespace sym {
+
+// Retain normalized direct comparisons, without transitive inference or calls
+// back into the analyzer. The rewriter owns ingestion, reset and query scopes.
+class KnownComparisons {
+ public:
+ CompareResult TryCompare(const PrimExpr& lhs, const PrimExpr& rhs) const;
+ void IngestBinding(const Var& var, const PrimExpr& expr);
+ void IngestBinding(const Var& var, const Range& range);
+ void IngestPredicate(const PrimExpr& expr) { AddKnown(expr,
&scoped_knowns_); }
+
+ private:
+ // Structural keys avoid comparing expression trees for every stored fact.
+ enum class Key : size_t {};
+ Key ExprToKey(const PrimExpr& expr);
+ std::optional<Key> ExprToPreviousKey(const PrimExpr& expr) const;
+ std::unordered_map<PrimExpr, Key, ffi::StructuralHash, ffi::StructuralEqual>
expr_to_key_;
+
+ // lhs OP rhs + offset. Strict integer comparisons normalize to <= or >=.
+ struct Comparison {
+ Comparison(Key lhs, Key rhs, int64_t offset, CompareResult result);
+ std::optional<Comparison> WithLHS(Key new_lhs) const;
+ Key lhs_;
+ Key rhs_;
+ int64_t offset_{0};
+ CompareResult result_{CompareResult::kInconsistent};
+ };
+
+ static CompareResult Reverse(CompareResult result);
+ static std::tuple<PrimExpr, PrimExpr, int64_t> ExtractOffsets(const
PrimExpr& lhs,
+ const
PrimExpr& rhs);
+ std::optional<Comparison> FromExpr(const PrimExpr& expr);
+ void AddKnown(const PrimExpr& expr, std::vector<Comparison>* vec);
+ std::vector<Comparison> CollectDirectComparisons(Key lhs_key, Key rhs_key)
const;
+ CompareResult MergeComparisons(const std::vector<Comparison>& lhs_to_rhs,
int64_t offset) const;
+
+ // Keep the defining Var separate from expression keys: type promotion may
+ // wrap it in a Cast, but a temporary binding must still replace its facts.
+ std::unordered_map<Var, std::vector<Comparison>> bindings_;
+ std::vector<Comparison> scoped_knowns_;
+};
+
+// Internal utility, return the CompareResult resulting from swapping
+// the left-hand side with the right-hand side.
+inline CompareResult KnownComparisons::Reverse(CompareResult res) {
+ switch (res) {
+ case CompareResult::kInconsistent:
+ return CompareResult::kInconsistent;
+ case CompareResult::kEQ:
+ return CompareResult::kEQ;
+ case CompareResult::kLT:
+ return CompareResult::kGT;
+ case CompareResult::kLE:
+ return CompareResult::kGE;
+ case CompareResult::kGT:
+ return CompareResult::kLT;
+ case CompareResult::kGE:
+ return CompareResult::kLE;
+ case CompareResult::kNE:
+ return CompareResult::kNE;
+ case CompareResult::kUnknown:
+ return CompareResult::kUnknown;
+ default:
+ TVM_FFI_THROW(InternalError) << "Invalid CompareResult: " <<
static_cast<int>(res);
+ }
+}
+
+// Internal utility, extract constant offsets out of the two sides of
+// a comparison. Given lhs and rhs, return a tuple of three elements
+// (lhs_inner, rhs_inner, offset), such that (lhs OP rhs) and
+// (lhs_inner OP rhs_inner + offset) are equivalent.
+inline std::tuple<PrimExpr, PrimExpr, int64_t>
KnownComparisons::ExtractOffsets(
+ const PrimExpr& lhs, const PrimExpr& rhs) {
+ auto extract_offset = [](const PrimExpr& expr) -> std::pair<PrimExpr,
ffi::BigInt> {
+ PVar<PrimExpr> x;
+ PVar<IntImm> c;
+ if ((x + c).Match(expr)) {
+ return {x.Eval(), c.Eval()->value};
+ } else if ((x - c).Match(expr)) {
+ return {x.Eval(), -c.Eval()->value};
+ } else if (c.Match(expr)) {
+ return {0, c.Eval()->value};
+ } else {
+ return {expr, 0};
+ }
+ };
+
+ auto lhs_split = extract_offset(lhs);
+ auto rhs_split = extract_offset(rhs);
+ if (auto offset = (rhs_split.second - lhs_split.second).as<int64_t>();
offset.has_value()) {
+ return {lhs_split.first, rhs_split.first, *offset};
+ }
+ return {lhs, rhs, 0};
+}
+
+inline std::optional<KnownComparisons::Comparison> KnownComparisons::FromExpr(
+ const PrimExpr& expr) {
+ CompareResult res;
+ PVar<PrimExpr> x, y;
+ if ((x <= y).Match(expr)) {
+ res = CompareResult::kLE;
+ } else if ((x >= y).Match(expr)) {
+ res = CompareResult::kGE;
+ } else if ((x < y).Match(expr)) {
+ res = CompareResult::kLT;
+ } else if ((x > y).Match(expr)) {
+ res = CompareResult::kGT;
+ } else if ((x == y).Match(expr)) {
+ res = CompareResult::kEQ;
+ } else if ((x != y).Match(expr)) {
+ res = CompareResult::kNE;
+ } else {
+ return std::nullopt;
+ }
+
+ PrimExpr lhs_expr = x.Eval();
+ PrimExpr rhs_expr = y.Eval();
+
+ if (!lhs_expr.ty().MatchesCode(DLDataTypeCode::kDLInt) ||
+ !rhs_expr.ty().MatchesCode(DLDataTypeCode::kDLInt)) {
+ return std::nullopt;
+ }
+ if (lhs_expr.as<IntImmObj>() && rhs_expr.as<IntImmObj>()) {
+ return std::nullopt;
+ }
+
+ auto [lhs, rhs, offset] = ExtractOffsets(lhs_expr, rhs_expr);
+ ffi::BigInt normalized_offset = offset;
+ if (res == CompareResult::kLT) --normalized_offset;
+ if (res == CompareResult::kGT) ++normalized_offset;
+ if (!normalized_offset.as<int64_t>().has_value()) return std::nullopt;
+ Key lhs_key = ExprToKey(lhs);
+ Key rhs_key = ExprToKey(rhs);
+
+ return Comparison(lhs_key, rhs_key, offset, res);
+}
+
+inline KnownComparisons::Comparison::Comparison(Key lhs, Key rhs, int64_t
offset,
+ CompareResult result)
+ : lhs_(lhs), rhs_(rhs), offset_(offset), result_(result) {
+ // Integer strict comparisons become non-strict comparisons with a unit
offset.
+ if (result_ == CompareResult::kLT) {
+ result_ = CompareResult::kLE;
+ offset_ -= 1;
+ }
+ if (result_ == CompareResult::kGT) {
+ result_ = CompareResult::kGE;
+ offset_ += 1;
+ }
+}
+
+inline std::optional<KnownComparisons::Key>
KnownComparisons::ExprToPreviousKey(
+ const PrimExpr& expr) const {
+ auto it = expr_to_key_.find(expr);
+ if (it != expr_to_key_.end()) {
+ return it->second;
+ } else {
+ return std::nullopt;
+ }
+}
+
+inline KnownComparisons::Key KnownComparisons::ExprToKey(const PrimExpr& expr)
{
+ if (auto prev = ExprToPreviousKey(expr)) {
+ return prev.value();
+ } else {
+ Key new_key = Key(expr_to_key_.size());
+ expr_to_key_[expr] = new_key;
+ return new_key;
+ }
+}
+
+inline std::optional<KnownComparisons::Comparison>
KnownComparisons::Comparison::WithLHS(
+ Key new_lhs) const {
+ if (new_lhs == lhs_) {
+ return *this;
+ } else if (new_lhs == rhs_) {
+ auto reversed = (-ffi::BigInt(offset_)).as<int64_t>();
+ if (!reversed.has_value()) return std::nullopt;
+ return Comparison(rhs_, lhs_, *reversed, Reverse(result_));
+ } else {
+ return std::nullopt;
+ }
+}
+
+inline void KnownComparisons::AddKnown(const PrimExpr& expr,
std::vector<Comparison>* vec) {
+ // Only stable facts may enter the expression-key table. Queries then need no
+ // side-effect traversal before the integer fast path or key lookup.
+ if (SideEffect(expr) > CallEffectKind::kPure) return;
+ for (const auto& subexpr : ExtractConstraints(expr, false)) {
+ if (auto cmp = FromExpr(subexpr)) vec->push_back(*cmp);
+ }
+}
+
+inline void KnownComparisons::IngestBinding(const Var& var, const Range&
range) {
+ auto& knowns = bindings_[var];
+ knowns.clear();
+ PrimExpr prim_var = var.as_or_throw<PrimExpr>();
+ if (prim::IsConstInt(range->extent, 1)) {
+ AddKnown(prim_var == range->begin, &knowns);
+ } else {
+ AddKnown(prim_var >= range->begin, &knowns);
+ AddKnown(prim_var < range->begin + range->extent, &knowns);
+ }
+}
+
+inline void KnownComparisons::IngestBinding(const Var& var, const PrimExpr&
expr) {
+ IngestBinding(var, Range::FromMinExtent(expr, 1));
+}
+
+inline CompareResult KnownComparisons::TryCompare(const PrimExpr& lhs_expr,
+ const PrimExpr& rhs_expr)
const {
+ // Currently only supports integer checks
+ if (!lhs_expr.ty().MatchesCode(DLDataTypeCode::kDLInt) ||
+ !rhs_expr.ty().MatchesCode(DLDataTypeCode::kDLInt)) {
+ return CompareResult::kUnknown;
+ }
+
+ // Bail out early if possible. This int check should have been
+ // constant-folded earlier, so this check shouldn't occur.
+ auto* x_int = lhs_expr.as<IntImmObj>();
+ auto* y_int = rhs_expr.as<IntImmObj>();
+ if (x_int && y_int) {
+ if (x_int->value < y_int->value) {
+ return CompareResult::kLT;
+ } else if (x_int->value > y_int->value) {
+ return CompareResult::kGT;
+ } else {
+ return CompareResult::kEQ;
+ }
+ }
+
+ auto [lhs, rhs, offset] = ExtractOffsets(lhs_expr, rhs_expr);
+ auto lhs_key = ExprToPreviousKey(lhs);
+ auto rhs_key = ExprToPreviousKey(rhs);
+
+ if (!lhs_key.has_value() || !rhs_key.has_value()) {
+ return CompareResult::kUnknown;
+ }
+
+ auto lhs_to_rhs = CollectDirectComparisons(lhs_key.value(), rhs_key.value());
+ return MergeComparisons(lhs_to_rhs, offset);
+}
+
+inline std::vector<KnownComparisons::Comparison>
KnownComparisons::CollectDirectComparisons(
+ Key lhs_key, Key rhs_key) const {
+ std::vector<Comparison> output;
+
+ auto append_known = [&](Comparison cmp) {
+ if (auto normalized = cmp.WithLHS(lhs_key)) {
+ if (normalized.value().rhs_ == rhs_key) {
+ output.push_back(normalized.value());
+ }
+ }
+ };
+
+ for (const auto& [var, knowns] : bindings_) {
+ for (const auto& known : knowns) append_known(known);
+ }
+ for (const auto& known : scoped_knowns_) {
+ append_known(known);
+ }
+
+ return output;
+}
+
+inline CompareResult KnownComparisons::MergeComparisons(const
std::vector<Comparison>& lhs_to_rhs,
+ int64_t offset) const {
+ // Just because we found a comparison involving LHS and RHS doesn't
+ // mean that it's useful. e.g. Knowing that `x < y` doesn't let us
+ // prove whether `x + 5 < y`.
+ CompareResult result = CompareResult::kUnknown;
+ for (const auto& cmp : lhs_to_rhs) {
+ switch (cmp.result_) {
+ case CompareResult::kInconsistent:
+ result = CompareResult::kInconsistent;
+ break;
+
+ case CompareResult::kEQ:
+ if (offset == cmp.offset_) {
+ result = result & CompareResult::kEQ;
+ } else {
+ result = result & CompareResult::kNE;
+ }
+ break;
+
+ case CompareResult::kLE:
+ if (cmp.offset_ < offset) {
+ result = result & CompareResult::kLT;
+ } else if (cmp.offset_ <= offset) {
+ result = result & CompareResult::kLE;
+ }
+ break;
+
+ case CompareResult::kGE:
+ if (cmp.offset_ > offset) {
+ result = result & CompareResult::kGT;
+ } else if (cmp.offset_ >= offset) {
+ result = result & CompareResult::kGE;
+ }
+ break;
+
+ case CompareResult::kNE:
+ if (offset == cmp.offset_) {
+ result = result & CompareResult::kNE;
+ }
+ break;
+
+ case CompareResult::kUnknown:
+ break;
+
+ case CompareResult::kGT:
+ case CompareResult::kLT:
+ TVM_FFI_THROW(InternalError)
+ << "Internal error, normalized comparisons should only include <=
and >=";
+
+ default:
+ TVM_FFI_THROW(InternalError) << "Invalid CompareResult: " <<
static_cast<int>(cmp.result_);
+ }
+ }
+
+ return result;
+}
+
+} // namespace sym
+} // namespace tvm
+
+#endif // TVM_SYM_KNOWN_COMPARISONS_H_
diff --git a/src/sym/rewrite_simplify.cc b/src/sym/rewrite_simplify.cc
index 4b0584ad3..887fee4e5 100644
--- a/src/sym/rewrite_simplify.cc
+++ b/src/sym/rewrite_simplify.cc
@@ -49,6 +49,7 @@ TVM_FFI_INLINE bool IsVectorExpr(const ExprObj* expr) {
PrimType ty = expr->ty.as_or_throw<PrimType>();
return ty.IsScalableVector() || ty.IsFixedLengthVector();
}
+
} // namespace
TVM_FFI_STATIC_INIT_BLOCK() { RewriteSimplifierStatsObj::RegisterReflection();
}
@@ -164,22 +165,27 @@ CompareResult RewriteSimplifier::Impl::TryCompare(const
PrimExpr& x, const PrimE
output == CompareResult::kGT;
};
- output = CompareResult(output & TryCompareUsingConstIntBounds(x, y));
+ output = output & TryCompare(x - y, 0);
if (is_finished()) return output;
- output = CompareResult(output & TryCompareUsingKnownInequalities(x, y));
+ output = output & analyzer_->rewrite_simplify.TryCompare(x, y);
return output;
}
-CompareResult RewriteSimplifier::Impl::TryCompareUsingConstIntBounds(const
PrimExpr& x,
- const
PrimExpr y) {
- return TryCompare(x - y, 0);
-}
-
-CompareResult RewriteSimplifier::Impl::TryCompareUsingKnownInequalities(const
PrimExpr& x,
- const
PrimExpr& y) {
- return analyzer_->known_comparisons.TryCompare(x, y);
+CompareResult RewriteSimplifier::Impl::TryCompareKnown(const PrimExpr& lhs,
+ const PrimExpr& rhs)
const {
+ const auto& facts = detail::TempConstraintScope::Query(analyzer_);
+ if (facts.empty()) return known_comparisons_.TryCompare(lhs, rhs);
+ // Temporary bindings, predicates and expression keys belong to query
scratch.
+ KnownComparisons temporary = known_comparisons_;
+ for (const Constraint& fact : facts) {
+ if (const auto* binding = fact.as<BindingConstraintObj>())
+ temporary.IngestBinding(binding->var, binding->value);
+ else if (const auto* predicate = fact.as<PredicateConstraintObj>())
+ temporary.IngestPredicate(predicate->condition);
+ }
+ return temporary.TryCompare(lhs, rhs);
}
// try to prove x equals val
@@ -2387,12 +2393,24 @@ PrimExpr RewriteSimplifier::operator()(const PrimExpr&
expr) {
return res;
}
-void RewriteSimplifier::IngestBinding(const Var& var, const PrimExpr& info) {
- impl_->SetBinding(var, info, true);
+CompareResult RewriteSimplifier::TryCompare(const PrimExpr& lhs, const
PrimExpr& rhs) {
+ detail::TempConstraintScope query(impl_->owner());
+ return impl_->TryCompareKnown(lhs, rhs);
+}
+
+void RewriteSimplifier::IngestBinding(const Var& var, const PrimExpr& value,
+ const PrimExpr& original) {
+ impl_->SetBinding(var, value, true);
+ impl_->known_comparisons_.IngestBinding(var, original);
+}
+
+void RewriteSimplifier::IngestBinding(const Var& var, const Range& range) {
+ impl_->known_comparisons_.IngestBinding(var, range);
}
void RewriteSimplifier::IngestPredicate(const PrimExpr& constraint) {
impl_->AddPredicate(constraint, false);
+ impl_->known_comparisons_.IngestPredicate(constraint);
}
ffi::ObjectRef RewriteSimplifier::GetStatsCounters() const { return
impl_->GetStatsCounters(); }
diff --git a/src/sym/rewrite_simplify.h b/src/sym/rewrite_simplify.h
index 2be51fd64..a708575a2 100644
--- a/src/sym/rewrite_simplify.h
+++ b/src/sym/rewrite_simplify.h
@@ -35,6 +35,7 @@
#include <vector>
#include "const_fold.h"
+#include "known_comparisons.h"
#include "simplify_base.h"
namespace tvm {
@@ -124,6 +125,11 @@ class RewriteSimplifier::Impl : public SimplifierBase {
void AddPredicate(const PrimExpr& constraint, bool is_assume);
+ KnownComparisons known_comparisons_;
+
+ // Direct premise matching does not invoke simplification or other analyzers.
+ CompareResult TryCompareKnown(const PrimExpr& lhs, const PrimExpr& rhs)
const;
+
RewriteSimplifierStats GetStatsCounters() const { return
RewriteSimplifierStats(stats_); }
void ResetStatsCounters() { stats_ = {}; }
@@ -135,12 +141,14 @@ class RewriteSimplifier::Impl : public SimplifierBase {
void ResetContext() {
var_map_.clear();
literal_constraints_.clear();
+ known_comparisons_ = {};
// Limits and cumulative diagnostics survive environment changes. Recursion
// belongs to the active query and cannot be reset underneath that query.
}
void CopyFrom(const Impl& other) {
var_map_ = other.var_map_;
literal_constraints_ = other.literal_constraints_;
+ known_comparisons_ = other.known_comparisons_;
maximum_rewrite_steps_ = other.maximum_rewrite_steps_;
}
@@ -221,9 +229,6 @@ class RewriteSimplifier::Impl : public SimplifierBase {
PrimExpr ApplyRewriteRules(prim::Not node, InplaceMode inplace_mode);
private:
- CompareResult TryCompareUsingKnownInequalities(const PrimExpr& x, const
PrimExpr& y);
- CompareResult TryCompareUsingConstIntBounds(const PrimExpr& x, const
PrimExpr y);
-
// Whether x >= val
bool CanProveGreaterEqual(const PrimExpr& x, int64_t val) {
return analyzer_->CanProveGreaterEqual(x, val);
