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 08db12de5f [REFACTOR][SYM] Retire optional simplifier extensions
(#20712)
08db12de5f is described below
commit 08db12de5f5dc34ae03a50b9e8a580b236459579
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 14 07:07:14 2026 -0400
[REFACTOR][SYM] Retire optional simplifier extensions (#20712)
---
include/tvm/sym/analyzer.h | 112 +--
include/tvm/tirx/type.h | 6 +-
python/tvm/relax/transform/__init__.py | 1 -
python/tvm/relax/transform/transform.py | 17 -
python/tvm/sym/__init__.py | 1 -
python/tvm/sym/analyzer.py | 44 +-
src/ir/expr.cc | 3 +-
src/relax/ir/dataflow_matcher.cc | 65 +-
src/relax/ir/dataflow_matcher.h | 6 +-
src/relax/transform/adjust_matmul_order.cc | 356 -------
src/s_tir/transform/stmt_simplify.cc | 1 -
src/sym/analyzer.cc | 24 +-
src/sym/conjunctive_normal_form.cc | 445 ---------
src/sym/conjunctive_normal_form.h | 50 -
...on_analyzer.cc => known_comparison_analyzer.cc} | 392 +-------
src/sym/rewrite_simplify.cc | 264 +----
src/sym/rewrite_simplify.h | 19 -
src/tirx/transform/stmt_simplify.cc | 1 -
.../relax/test_transform_adjust_matmul_order.py | 1034 --------------------
tests/python/sym/test_sym_analyzer_object.py | 26 +-
tests/python/sym/test_sym_rewrite_simplify.py | 36 -
.../tirx-transform/test_tir_transform_simplify.py | 25 -
22 files changed, 132 insertions(+), 2796 deletions(-)
diff --git a/include/tvm/sym/analyzer.h b/include/tvm/sym/analyzer.h
index c0d4988f9a..ce90fcbb20 100644
--- a/include/tvm/sym/analyzer.h
+++ b/include/tvm/sym/analyzer.h
@@ -306,90 +306,6 @@ class RewriteSimplifier {
TVM_TIRX_DLL std::function<void()> EnterConstraint(const PrimExpr&
constraint,
bool is_assume = false);
- /*! \brief Flags to enable more computationally-intensive simplifications
- *
- * These simplifications may be required for specific schedules, but
- * would impose too high a compile-time cost to enable by default.
- * They can be enabled on an as-needed basis by calling
- * `RewriteSimplifier::SetEnabledExtensions` prior to using
- * `RewriteSimplifier::operator()`.
- *
- * Flags are defined as powers of two to allow future expansion. To
- * enable multiple extensions, a user should pass a bitwise OR of the
- * flags for each desired extension.
- */
- enum Extension {
- // No extensions enabled
- kNone = 0,
-
- /* When simplifying an inequality, attempt to use scope-based knowns.
- *
- * Example:
- * if_then_else(i<j && j<k, i<k, false) => if_then_else(i<j && j<k, true,
false)
- */
- kTransitivelyProveInequalities = (1 << 0),
-
- /* When simplifying a boolean expression, convert to an AND of ORs
- * (conjunctive normal form).
- *
- * Example:
- * (a && b) || c => (a || c) && (b || c)
- */
- kConvertBooleanToAndOfOrs = (1 << 1),
-
- /* When simplifying a boolean AND or a boolean OR, simplify each
- * branch under the assumption that the other branch does not
- * already dominate the result. That is, simplify each branch of
- * (A && B) under the assumption that the other branch is true,
- * and simplify each branch of (A || B) under the assumption that
- * the other branch is false.
- *
- * Example:
- * (n < 10) && (n < 5) => (n < 10)
- * (n < 10) || (n < 5) => (n < 5)
- */
- kApplyConstraintsToBooleanBranches = (1 << 2),
-
- /* Special handling for expressions `(A+B)*C < (A*B)*D`
- *
- * Expressions of the form `(A+B)*C < (A*B)*D` can occur occur
- * when comparing the number of operations required for two
- * different orderings in which matrix multiplications can be
- * performed. Proving or disproving this conditional allows an
- * optimal order of execution to be selected, even for dynamic
- * argument shapes.
- *
- * The default behavior of `ConstIntBounds` assumes that each term
- * in an expression is independent, and is insufficient to prove
- * these inequalities. For example, the maximum value of `(A+B)*C
- * - (A*B)*D` is determined by taking the maximum value of
- * `(A+B)*C` and subtracting the minimum value of `(A*B)*D`.
- * While this algorithm can be applied in all cases, the bound it
- * provides is looser than strictly required.
- *
- * This extension adds a check for this case. When `A`, `B`, `C`,
- * and `D` are all positive values, as is the case for tensor
- * shapes, the inequality can be written as `1/A + 1/B < D/C`. If
- * this inequality holds for the minimum values of `A`, `B`, and
- * `D`, along with the maximum value of `C`, then the inequality
- * holds for all values.
- *
- * This extension requires little to no performance overhead, and
- * may be enabled by default in future releases.
- */
- kComparisonOfProductAndSum = (1 << 3),
- };
-
- /*! \brief Enable an optional extension or extensions
- *
- * \param flags A bitwise OR of all optional extensions that should
- * be enabled.
- */
- TVM_TIRX_DLL void SetEnabledExtensions(Extension flags);
-
- /*! \brief Return the currently enabled extensions */
- TVM_TIRX_DLL Extension GetEnabledExtensions() const;
-
/*! \brief Return the statistics counters */
TVM_TIRX_DLL ffi::ObjectRef GetStatsCounters() const;
@@ -482,10 +398,9 @@ inline constexpr CompareResult operator|(CompareResult
lhs, CompareResult rhs) {
/*!
* \brief Using previously specified knowns, compare the expressions provided
*
- * Given known expressions [(a OP b), (b OP c), ..., (y OP z)], search
- * for a known result for `(a OP z)`.
+ * Compare directly related expressions, allowing constant offsets.
*/
-class TransitiveComparisonAnalyzer {
+class KnownComparisonAnalyzer {
public:
/* \brief Using previously specified knowns, compare the expressions provided
*
@@ -493,19 +408,12 @@ class TransitiveComparisonAnalyzer {
*
* \param rhs The right-hand side of the comparison
*
- * \param propagate_inequalities If true, attempt to find a sequence
- * of transitive inequalities that allow the lhs and rhs to be
- * compared. If false, only use the known comparison that have been
- * directly provided. Using `propagate_inequalities = false` is
- * roughly equivalent to comparing against all known inequality
- * expressions using `prim::ExprDeepEqual`, but also allows for constant
- * offsets on either side of the inequality.
+ * 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,
- bool propagate_inequalities = true);
+ TVM_TIRX_DLL CompareResult TryCompare(const PrimExpr& lhs, const PrimExpr&
rhs);
/*! \brief Bind a variable as being equal to a known expression
*
@@ -534,9 +442,9 @@ class TransitiveComparisonAnalyzer {
private:
friend class AnalyzerObj;
friend class ConstraintContext;
- TransitiveComparisonAnalyzer();
- TVM_TIRX_DLL ~TransitiveComparisonAnalyzer();
- void CopyFrom(const TransitiveComparisonAnalyzer& other);
+ KnownComparisonAnalyzer();
+ TVM_TIRX_DLL ~KnownComparisonAnalyzer();
+ void CopyFrom(const KnownComparisonAnalyzer& other);
class Impl;
/*! \brief Internal impl */
std::unique_ptr<Impl> impl_;
@@ -725,8 +633,8 @@ class TVM_TIRX_DLL AnalyzerObj : public ffi::Object {
CanonicalSimplifier canonical_simplify;
/*! \brief sub-analyzer: int set */
IntSetAnalyzer int_set;
- /*! \brief sub-analyzer transitive comparisons */
- TransitiveComparisonAnalyzer transitive_comparisons;
+ /*! \brief sub-analyzer known comparisons */
+ KnownComparisonAnalyzer known_comparisons;
/*! \brief sub-analyzer using Z3 */
Z3Prover z3_prover;
/*! \brief constructor */
@@ -873,7 +781,7 @@ class TVM_TIRX_DLL AnalyzerObj : public ffi::Object {
*
* The returned analyzer carries the same accumulated facts (variable
* bounds, modular sets, rewrite/canonical bindings, integer-set domains,
- * literal constraints and transitive comparisons) as this one, but owns
+ * literal constraints and known comparisons) as this one, but owns
* its own state: binding or simplifying on either analyzer afterwards does
* not affect the other. This is the deep copy that handle-copying an
* Analyzer does not provide.
diff --git a/include/tvm/tirx/type.h b/include/tvm/tirx/type.h
index 05f3844821..7fc164e82f 100644
--- a/include/tvm/tirx/type.h
+++ b/include/tvm/tirx/type.h
@@ -93,8 +93,7 @@ class TensorTypeNode : public TypeNode {
refl::ObjectDef<TensorTypeNode>()
.def_ro("shape", &TensorTypeNode::shape,
refl::AttachFieldFlag::SEqHashDefSimple())
.def_ro("dtype", &TensorTypeNode::dtype)
- .def_ro("scope", &TensorTypeNode::scope,
- refl::DefaultValue(runtime::MemoryScope::Global()))
+ .def_ro("scope", &TensorTypeNode::scope,
refl::DefaultValue(runtime::MemoryScope::Global()))
.def_ro("strides", &TensorTypeNode::strides,
refl::AttachFieldFlag::SEqHashDefSimple(),
refl::DefaultValue(ffi::Array<PrimExpr>()))
.def_ro("elem_offset", &TensorTypeNode::elem_offset,
@@ -140,8 +139,7 @@ class TensorTypeNode : public TypeNode {
class TensorType : public Type {
public:
TVM_TIRX_DLL TensorType(ffi::Array<PrimExpr> shape, PrimType dtype =
PrimType::Float(32),
- runtime::MemoryScope scope = {},
- ffi::Array<PrimExpr> strides = {},
+ runtime::MemoryScope scope = {},
ffi::Array<PrimExpr> strides = {},
ffi::Optional<PrimExpr> elem_offset = std::nullopt,
int data_alignment = 0, int offset_factor = 0,
ffi::Optional<Layout> layout = std::nullopt,
Location loc = UnknownLoc());
diff --git a/python/tvm/relax/transform/__init__.py
b/python/tvm/relax/transform/__init__.py
index e12604b20f..88ce17d982 100644
--- a/python/tvm/relax/transform/__init__.py
+++ b/python/tvm/relax/transform/__init__.py
@@ -18,7 +18,6 @@
"""Relax transformations."""
from .transform import (
- AdjustMatmulOrder,
AllocateWorkspace,
AlterOpImpl,
AnnotateTIROpPattern,
diff --git a/python/tvm/relax/transform/transform.py
b/python/tvm/relax/transform/transform.py
index abc6af1876..07e66d0342 100644
--- a/python/tvm/relax/transform/transform.py
+++ b/python/tvm/relax/transform/transform.py
@@ -1450,23 +1450,6 @@ def UpdateParamType(ty_func: Callable[[Var], Type |
None]):
return _ffi_api.UpdateParamType(ty_func) # type: ignore
-def AdjustMatmulOrder():
- """Reorder `x*(A*B)` to `(x*A)*B`
-
- Useful for optimizing LoRA computations, where `matmul(x,
- LoraA*LoraB)` may be computed as `matmul(matmul(x, LoraA),
- LoraB)`, reducing the total memory usage.
-
-
- Returns
- -------
- ret : tvm.ir.transform.Pass
- The corresponding pass.
- """
-
- return _ffi_api.AdjustMatmulOrder() # type: ignore
-
-
def ExpandMatmulOfSum():
"""Expand `matmul(x, A+B)` to `matmul(x,A) + matmul(x,B)`
diff --git a/python/tvm/sym/__init__.py b/python/tvm/sym/__init__.py
index 3a3e091eff..59274c093e 100644
--- a/python/tvm/sym/__init__.py
+++ b/python/tvm/sym/__init__.py
@@ -29,7 +29,6 @@ from .analyzer import (
ConstIntBound,
Analyzer,
ProofStrength,
- Extension,
CompareResult,
allow_uint_as_index,
)
diff --git a/python/tvm/sym/analyzer.py b/python/tvm/sym/analyzer.py
index 560b45a364..ce34101cf0 100644
--- a/python/tvm/sym/analyzer.py
+++ b/python/tvm/sym/analyzer.py
@@ -52,19 +52,6 @@ class CompareResult(enum.IntEnum):
UNKNOWN = 7
-class Extension(enum.Flag):
- """Extensions enabled for RewriteSimplifier
-
- Values should match `RewriteSimplifier::Extensions`
- """
-
- NoExtensions = 0
- TransitivelyProveInequalities = 1 << 0
- ConvertBooleanToAndOfOrs = 1 << 1
- ApplyConstraintsToBooleanBranches = 1 << 2
- ComparisonOfProductAndSum = 1 << 3
-
-
@tvm_ffi.register_object("sym.ModularSet")
class ModularSet(Object):
"""Represent range of (coeff * x + base) for x in Z"""
@@ -485,10 +472,8 @@ class Analyzer(Object):
"""
return _ffi_api.AnalyzerCanProveEqual(self, lhs, rhs)
- def try_compare(
- self, lhs: ir.Expr, rhs: ir.Expr, propagate_inequalities: bool = True
- ) -> CompareResult:
- """Compare lhs and rhs using previously provided known comparisons.
+ def try_compare(self, lhs: ir.Expr, rhs: ir.Expr) -> CompareResult:
+ """Compare lhs and rhs using directly known comparisons and constant
offsets.
Parameters
----------
@@ -498,36 +483,13 @@ class Analyzer(Object):
rhs : Expr
The right-hand side of the comparison.
- propagate_inequalities : bool
- If true, attempt to find a sequence of transitive inequalities that
- allow lhs and rhs to be compared.
-
Returns
-------
result : CompareResult
The most specific result that can be proven about the comparison.
Returns ``CompareResult.UNKNOWN`` when nothing can be proven.
"""
- return CompareResult(_ffi_api.AnalyzerTryCompare(self, lhs, rhs,
propagate_inequalities))
-
- @property
- def enabled_extensions(self) -> Extension:
- """Return the currently enabled extensions"""
- value = _ffi_api.AnalyzerGetEnabledExtensions(self)
- return Extension(value)
-
- @enabled_extensions.setter
- def enabled_extensions(self, flags: int | Extension):
- """Enable extensions for the analyzer
-
- Parameters
- ----------
- flags: Union[int,Extension]
-
- The extensions to enable.
- """
- flags = Extension(flags).value
- _ffi_api.AnalyzerSetEnabledExtensions(self, flags)
+ return CompareResult(_ffi_api.AnalyzerTryCompare(self, lhs, rhs))
@contextmanager
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index 59ec239ee2..81b93398ff 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -718,7 +718,8 @@ TensorLoad::TensorLoad(Expr source, ffi::Array<PrimExpr>
indices, Location loc)
TensorLoad TensorLoad::Unchecked(PrimType ty, Expr source,
ffi::Array<PrimExpr> indices,
Location loc) {
- TVM_FFI_CHECK(ty.defined() && ty.as<PrimTypeNode>(), TypeError) <<
"TensorLoad requires a concrete PrimType";
+ TVM_FFI_CHECK(ty.defined() && ty.as<PrimTypeNode>(), TypeError)
+ << "TensorLoad requires a concrete PrimType";
TVM_FFI_CHECK(source.defined(), TypeError) << "TensorLoad requires a defined
source";
auto node = ffi::make_object<TensorLoadNode>(std::move(source));
node->ty = std::move(ty);
diff --git a/src/relax/ir/dataflow_matcher.cc b/src/relax/ir/dataflow_matcher.cc
index 93d5d09450..ce5f59ab4c 100644
--- a/src/relax/ir/dataflow_matcher.cc
+++ b/src/relax/ir/dataflow_matcher.cc
@@ -26,6 +26,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/ir/analysis.h>
#include <tvm/ir/prim/op.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/dataflow_matcher.h>
@@ -36,6 +37,7 @@
#include <tvm/runtime/logging.h>
#include <tvm/sym/analyzer.h>
+#include <algorithm>
#include <array>
#include <cstddef>
#include <limits>
@@ -428,6 +430,67 @@ bool DFPatternMatcher::VisitDFPattern_(const
UnorderedTuplePatternNode* op, cons
return false;
}
+// Check accumulated shape equalities without changing the constraints
retained for backtracking.
+// Bindings live only in this check, and use variable identity rather than
printed names.
+static bool ShapeConstraintsConsistent(const PrimExpr& condition) {
+ std::vector<PrimExpr> constraints{condition};
+ for (size_t i = 0; i < constraints.size();) {
+ if (const auto* conjunction = constraints[i].as<prim::AndNode>()) {
+ PrimExpr lhs = conjunction->a;
+ PrimExpr rhs = conjunction->b;
+ constraints[i] = std::move(lhs);
+ constraints.push_back(std::move(rhs));
+ } else {
+ ++i;
+ }
+ }
+
+ Analyzer analyzer;
+ // Follow every possible local alias before binding, including chains longer
than
+ // the analyzer's default two simplification passes.
+ int simplify_steps = 2 * static_cast<int>(constraints.size());
+ // Each successful binding eliminates a variable from an equality. Revisit
earlier
+ // equalities that may become directly bindable after a later constraint is
processed.
+ for (size_t iteration = 0; iteration < constraints.size(); ++iteration) {
+ bool changed = false;
+ for (const PrimExpr& constraint : constraints) {
+ PrimExpr simplified = analyzer->Simplify(constraint, simplify_steps);
+ if (prim::IsZero(simplified)) return false;
+ if (const auto* equality = simplified.as<prim::EQNode>()) {
+ PrimExpr lhs = equality->a;
+ PrimExpr rhs = equality->b;
+ if (!lhs.as<tvm::VarNode>()) std::swap(lhs, rhs);
+ if (const auto* var = lhs.as<tvm::VarNode>();
+ var && SideEffect(rhs) <= CallEffectKind::kPure) {
+ auto variables = tvm::UndefinedVars(rhs);
+ if (std::none_of(variables.begin(), variables.end(),
+ [var](const tvm::Var& used) { return used.get() ==
var; })) {
+ analyzer->Bind(ffi::GetRef<tvm::Var>(var), rhs);
+ changed = true;
+ }
+ }
+ }
+ }
+ if (!changed) break;
+ }
+ // Composite dimensions may not admit a variable binding. Check each
remaining
+ // equality against its peers, so repeated expressions also detect
contradictions.
+ for (PrimExpr& constraint : constraints) {
+ constraint = analyzer->Simplify(constraint, simplify_steps);
+ if (prim::IsZero(constraint)) return false;
+ }
+ for (size_t i = 0; i < constraints.size(); ++i) {
+ if (prim::IsOne(constraints[i])) continue;
+ PrimExpr others = IntImm::Bool(true);
+ for (size_t j = 0; j < constraints.size(); ++j) {
+ if (i != j) others = others && constraints[j];
+ }
+ With<sym::ConstraintContext> scope(analyzer, others);
+ if (prim::IsZero(analyzer->Simplify(constraints[i], simplify_steps)))
return false;
+ }
+ return true;
+}
+
bool DFPatternMatcher::VisitDFPattern_(const TypePatternNode* op, const Expr&
expr0) {
if (!VisitDFPattern(op->pattern, expr0)) {
return false;
@@ -446,7 +509,7 @@ bool DFPatternMatcher::VisitDFPattern_(const
TypePatternNode* op, const Expr& ex
if (auto* as_int = symbolic_expr_condition_.as<IntImmNode>()) {
return static_cast<bool>(as_int->value);
} else {
- return true;
+ return ShapeConstraintsConsistent(symbolic_expr_condition_);
}
}
diff --git a/src/relax/ir/dataflow_matcher.h b/src/relax/ir/dataflow_matcher.h
index 677b5526b8..98a4adab74 100644
--- a/src/relax/ir/dataflow_matcher.h
+++ b/src/relax/ir/dataflow_matcher.h
@@ -42,11 +42,7 @@ class DFPatternMatcher : public DFPatternFunctor<bool(const
DFPattern&, const Ex
using var2val_t = ffi::Map<Var, Expr>;
explicit DFPatternMatcher() : DFPatternMatcher(var2val_t{}) {}
- explicit DFPatternMatcher(var2val_t var2val) : var2val_(std::move(var2val)) {
- // Propagate shape constraints across conjunctions regardless of operand
order.
- analyzer_->rewrite_simplify.SetEnabledExtensions(
- sym::RewriteSimplifier::kApplyConstraintsToBooleanBranches);
- }
+ explicit DFPatternMatcher(var2val_t var2val) : var2val_(std::move(var2val))
{}
bool Match(const DFPattern& pattern, const Expr& expr);
ffi::Map<DFPattern, Expr> GetMemo() { return memo_; }
diff --git a/src/relax/transform/adjust_matmul_order.cc
b/src/relax/transform/adjust_matmul_order.cc
deleted file mode 100644
index 56d2346833..0000000000
--- a/src/relax/transform/adjust_matmul_order.cc
+++ /dev/null
@@ -1,356 +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/relax/transform/adjust_matmul_order.cc
- * \brief Re-order `matmul(matmul(A,B), x)` to `matmul(A, matmul(B,x))`
- */
-
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/relax/analysis.h>
-#include <tvm/relax/dataflow_matcher.h>
-#include <tvm/relax/expr.h>
-#include <tvm/relax/expr_functor.h>
-#include <tvm/relax/op/linear_algebra.h>
-#include <tvm/relax/op/manipulate.h>
-#include <tvm/relax/transform.h>
-
-#include <optional>
-#include <unordered_set>
-#include <vector>
-
-#include "../op/op_common.h"
-#include "../op/tensor/linear_algebra.h"
-#include "../op/tensor/manipulate.h"
-
-namespace tvm {
-namespace relax {
-using namespace tvm::prim;
-
-namespace {
-
-ffi::Array<PrimExpr> GetBatchPrefix(const ffi::Array<PrimExpr>& shape) {
- if (shape.size() <= 2) return {};
- return {shape.begin(), shape.end() - 2};
-}
-
-PrimExpr ProductDims(const ffi::Array<PrimExpr>& dims) {
- PrimExpr product = IntImm::Int64(1);
- for (const auto& dim : dims) product = product * dim;
- return product;
-}
-
-bool IsLastTwoDimsSwap(const Expr& expr) {
- const auto* call = expr.as<CallNode>();
- if (call == nullptr) return false;
-
- const auto* attrs = call->attrs.as<PermuteDimsAttrs>();
- const auto* input_type = GetTypeAs<TensorTypeNode>(call->args[0]);
- if (attrs == nullptr || input_type == nullptr || input_type->ndim < 2)
return false;
-
- size_t ndim = input_type->ndim;
- if (!attrs->axes.has_value()) return ndim == 2;
-
- const auto& axes = attrs->axes.value();
- if (axes.size() != ndim) return false;
- for (size_t i = 0; i < axes.size(); ++i) {
- int64_t axis = axes[i];
- if (axis < 0) axis += ndim;
- size_t expected = i;
- if (i == ndim - 2) {
- expected = ndim - 1;
- } else if (i == ndim - 1) {
- expected = ndim - 2;
- }
- if (axis != static_cast<int64_t>(expected)) return false;
- }
- return true;
-}
-
-ffi::Optional<ffi::Array<PrimExpr>> InferBatchedMatmulBroadcastPrefix(
- sym::AnalyzerObj* analyzer, const ffi::Array<PrimExpr>& x1, const
ffi::Array<PrimExpr>& x2) {
- auto infer_result = InferBinaryBroadcastShape(analyzer, x1, x2);
- if (infer_result.status ==
BinaryBroadcastShapeInferResult::Status::kSuccess) {
- return infer_result.shape;
- }
- return std::nullopt;
-}
-
-std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr, ffi::Map<DFPattern,
Expr>)>> CreatePatterns(
- const Function& func) {
- auto compile_time_arr = ComputableAtCompileTime(func);
- std::unordered_set<Var> compile_time_lookup(compile_time_arr.begin(),
compile_time_arr.end());
-
- ffi::TypedFunction<bool(Expr)> is_compile_time = [compile_time_lookup](Expr
arg) -> bool {
- if (auto as_var = arg.as<Var>()) {
- return compile_time_lookup.count(as_var.value());
- } else {
- return false;
- }
- };
- ffi::TypedFunction<bool(Expr)> is_runtime = [is_compile_time](Expr arg) ->
bool {
- return !is_compile_time(arg);
- };
-
- DFPattern pat_a = WildcardPattern();
- DFPattern pat_b = WildcardPattern();
- DFPattern pat_c = WildcardPattern();
-
- auto pat_matmul = IsOp("relax.matmul");
- auto pat_permute_dims = IsOp("relax.permute_dims");
-
- auto pat_matmul_on_lhs = pat_matmul(pat_matmul(pat_a, pat_b), pat_c);
- auto pat_matmul_on_rhs = pat_matmul(pat_a, pat_matmul(pat_b, pat_c));
-
- auto pat_permuted_inner_matmul_on_lhs = pat_permute_dims(pat_matmul(pat_b,
pat_a));
- auto pat_permuted_inner_matmul_on_rhs = pat_permute_dims(pat_matmul(pat_c,
pat_b));
- auto pat_permuted_matmul_on_lhs =
pat_matmul(pat_permuted_inner_matmul_on_lhs, pat_c);
- auto pat_permuted_matmul_on_rhs = pat_matmul(pat_a,
pat_permuted_inner_matmul_on_rhs);
-
- auto pat = pat_matmul_on_lhs | pat_matmul_on_rhs |
pat_permuted_matmul_on_lhs |
- pat_permuted_matmul_on_rhs;
-
- PrimExpr symbolic_var_constraints = IntImm::Bool(true);
- auto upper_bounds =
- func->GetAttr<ffi::Map<ffi::String,
Any>>(tvm::relax::attr::kTirVarUpperBound);
- auto lower_bounds =
- func->GetAttr<ffi::Map<ffi::String,
Any>>(tvm::relax::attr::kTirVarLowerBound);
-
- if (upper_bounds || lower_bounds) {
- ffi::Map<ffi::String, tvm::Var> name_lookup;
- for (const auto& tir_var : TIRVarsInType(GetType(func))) {
- name_lookup.Set(tir_var->name, tir_var);
- symbolic_var_constraints = symbolic_var_constraints && (0 <=
tir_var.as_or_throw<PrimExpr>());
- }
-
- // Add lower bound constraints
- if (lower_bounds) {
- for (const auto& [key, obj_bound] : lower_bounds.value()) {
- auto tir_var_name = key;
- if (auto opt_var = name_lookup.Get(tir_var_name)) {
- auto var = opt_var.value();
- auto expr_bound = obj_bound.cast<PrimExpr>();
- symbolic_var_constraints =
- symbolic_var_constraints && (expr_bound <=
var.as_or_throw<PrimExpr>());
- }
- }
- }
-
- // Add upper bound constraints
- if (upper_bounds) {
- for (const auto& [key, obj_bound] : upper_bounds.value()) {
- auto tir_var_name = key;
- if (auto opt_var = name_lookup.Get(tir_var_name)) {
- auto var = opt_var.value();
- auto expr_bound = obj_bound.cast<PrimExpr>();
- symbolic_var_constraints =
- symbolic_var_constraints && (var.as_or_throw<PrimExpr>() <
expr_bound);
- }
- }
- }
- }
-
- auto rewriter = [=](Expr expr, ffi::Map<DFPattern, Expr> matches) -> Expr {
- auto expr_a = matches[pat_a];
- auto expr_b = matches[pat_b];
- auto expr_c = matches[pat_c];
- auto out_dtype = expr.as<CallNode>()->attrs.as<MatmulAttrs>()->out_dtype;
-
- // If all three components are compile-time, the order doesn't
- // matter as the entire expression can be lifted out and
- // pre-computed.
- if (is_compile_time(expr_a) && is_compile_time(expr_b) &&
is_compile_time(expr_c)) {
- return expr;
- }
-
- auto get_shape = [](Expr expr) -> ffi::Optional<ffi::Array<PrimExpr>> {
- auto ty = expr->ty.as<TensorTypeNode>();
- if (ty) {
- return ty->GetShape();
- } else {
- return std::nullopt;
- }
- };
-
- auto opt_shape_a = get_shape(expr_a);
- if (!opt_shape_a) return expr;
- auto opt_shape_b = get_shape(expr_b);
- if (!opt_shape_b) return expr;
- auto opt_shape_c = get_shape(expr_c);
- if (!opt_shape_c) return expr;
-
- auto shape_a = opt_shape_a.value();
- auto shape_b = opt_shape_b.value();
- auto shape_c = opt_shape_c.value();
-
- auto permute_last_two_dims = [&](Expr expr) -> Expr {
- auto opt_shape = get_shape(expr);
- if (!opt_shape) return expr;
-
- size_t ndim = opt_shape.value().size();
- TVM_FFI_ICHECK_GE(ndim, 2);
-
- ffi::Optional<ffi::Array<int64_t>> axes;
-
- if (ndim == 2) {
- // Pass none axes to permute_dims for simple transpose of 2D tensors.
- axes = std::nullopt;
- } else {
- ffi::Array<int64_t> axes_array;
- for (size_t i = 0; i < ndim; ++i) axes_array.push_back(i);
- axes_array.Set(ndim - 1, ndim - 2);
- axes_array.Set(ndim - 2, ndim - 1);
- axes = ffi::Optional<ffi::Array<int64_t>>(axes_array);
- }
- return permute_dims(std::move(expr), axes);
- };
-
- auto transpose_shape_last_two_dims = [&](ffi::Array<PrimExpr>& shape) {
- PrimExpr last_dim_shape = shape[shape.size() - 1];
- shape.Set(shape.size() - 1, shape[shape.size() - 2]);
- shape.Set(shape.size() - 2, last_dim_shape);
- };
-
- if (matches.count(pat_permuted_matmul_on_lhs)) {
- if (!IsLastTwoDimsSwap(matches[pat_permuted_inner_matmul_on_lhs]))
return expr;
- if (shape_a.size() < 2 || shape_b.size() < 2) return expr;
- expr_a = permute_last_two_dims(expr_a);
- expr_b = permute_last_two_dims(expr_b);
- transpose_shape_last_two_dims(shape_a);
- transpose_shape_last_two_dims(shape_b);
- } else if (matches.count(pat_permuted_matmul_on_rhs)) {
- if (!IsLastTwoDimsSwap(matches[pat_permuted_inner_matmul_on_rhs]))
return expr;
- if (shape_b.size() < 2 || shape_c.size() < 2) return expr;
- expr_b = permute_last_two_dims(expr_b);
- expr_c = permute_last_two_dims(expr_c);
- transpose_shape_last_two_dims(shape_b);
- transpose_shape_last_two_dims(shape_c);
- }
-
- // If two of the three are compile-time, group those two values
- // together, to allow them to be lifted out and pre-computed.
- if (is_compile_time(expr_a) && is_compile_time(expr_b)) {
- return matmul(matmul(expr_a, expr_b, std::nullopt), expr_c, out_dtype);
- } else if (is_compile_time(expr_b) && is_compile_time(expr_c)) {
- return matmul(expr_a, matmul(expr_b, expr_c, std::nullopt), out_dtype);
- }
-
- // Otherwise, select the order that reduces the total number of
- // operations required, assuming a naive matmul (see below).
-
- if (shape_a.size() == 1) {
- shape_a = {IntImm(shape_a[0].ty(), 1), shape_a[0]};
- }
- if (shape_b.size() == 1) {
- if (matches.count(pat_matmul_on_lhs)) {
- shape_b = {shape_b[0], IntImm(shape_b[0].ty(), 1)};
- } else if (matches.count(pat_matmul_on_rhs)) {
- shape_b = {IntImm(shape_b[0].ty(), 1), shape_b[0]};
- } else {
- TVM_FFI_THROW(InternalError)
- << "OrPattern " << pat << " matched, but neither " <<
pat_matmul_on_lhs << " nor "
- << pat_matmul_on_rhs << " matched";
- }
- }
- if (shape_c.size() == 1) {
- shape_c = {shape_c[0], IntImm(shape_c[0].ty(), 1)};
- }
-
- PrimExpr size_N = shape_a[shape_a.size() - 2]; // row of A
- PrimExpr size_R = shape_a[shape_a.size() - 1]; // col of A and row of B
- PrimExpr size_M = shape_c[shape_c.size() - 2]; // row of C and col of B
- PrimExpr size_B = shape_c[shape_c.size() - 1]; // col of C
-
- sym::Analyzer analyzer;
- auto prefix_a = GetBatchPrefix(shape_a);
- auto prefix_b = GetBatchPrefix(shape_b);
- auto prefix_c = GetBatchPrefix(shape_c);
-
- auto opt_prefix_ab = InferBatchedMatmulBroadcastPrefix(analyzer.get(),
prefix_a, prefix_b);
- if (!opt_prefix_ab) return expr;
- auto opt_prefix_bc = InferBatchedMatmulBroadcastPrefix(analyzer.get(),
prefix_b, prefix_c);
- if (!opt_prefix_bc) return expr;
- auto opt_prefix_outer_lhs =
- InferBatchedMatmulBroadcastPrefix(analyzer.get(),
opt_prefix_ab.value(), prefix_c);
- if (!opt_prefix_outer_lhs) return expr;
- auto opt_prefix_outer_rhs =
- InferBatchedMatmulBroadcastPrefix(analyzer.get(), prefix_a,
opt_prefix_bc.value());
- if (!opt_prefix_outer_rhs) return expr;
-
- PrimExpr batch_ab = ProductDims(opt_prefix_ab.value());
- PrimExpr batch_bc = ProductDims(opt_prefix_bc.value());
- PrimExpr batch_outer_lhs = ProductDims(opt_prefix_outer_lhs.value());
- PrimExpr batch_outer_rhs = ProductDims(opt_prefix_outer_rhs.value());
-
- // Compare naive matmul FLOPs for two evaluation orders of
- // matmul(A, matmul(B, C)) vs matmul(matmul(A, B), C)
- //
- // Matrix dims (last two axes): A [N, R], B [R, M], C [M, B_last]
- // Each matmul uses the broadcasted batch prefix of its operands.
- //
- // LHS first — matmul(matmul(A, B), C):
- // batch_ab * N * R * M + batch_outer_lhs * N * M * B_last
- PrimExpr ops_with_lhs_first =
- batch_ab * size_N * size_R * size_M + batch_outer_lhs * size_N *
size_M * size_B;
- // RHS first — matmul(A, matmul(B, C)):
- // batch_bc * R * M * B_last + batch_outer_rhs * N * R * B_last
- PrimExpr ops_with_rhs_first =
- batch_bc * size_R * size_M * size_B + batch_outer_rhs * size_N *
size_R * size_B;
-
-
analyzer->rewrite_simplify.SetEnabledExtensions(static_cast<sym::RewriteSimplifier::Extension>(
- analyzer->rewrite_simplify.GetEnabledExtensions() |
- sym::RewriteSimplifier::Extension::kComparisonOfProductAndSum));
- With<sym::ConstraintContext> func_attr_constraint(analyzer,
symbolic_var_constraints);
- With<sym::ConstraintContext> analyzer_constraint(
- analyzer, batch_ab > 0 && batch_bc > 0 && batch_outer_lhs > 0 &&
batch_outer_rhs > 0 &&
- size_N > 0 && size_R > 0 && size_M > 0 && size_B > 0);
-
- if (analyzer->CanProve(ops_with_lhs_first < ops_with_rhs_first)) {
- return matmul(matmul(expr_a, expr_b, std::nullopt), expr_c, out_dtype);
- } else if (analyzer->CanProve(ops_with_rhs_first < ops_with_lhs_first)) {
- return matmul(expr_a, matmul(expr_b, expr_c, std::nullopt), out_dtype);
- }
-
- // If we cannot determine which order is best, keep the existing order.
- return expr;
- };
-
- return {pat, rewriter};
-}
-
-} // namespace
-
-namespace transform {
-Pass AdjustMatmulOrder() {
- auto pass_func = [=](Function func, IRModule mod, PassContext pc) {
- auto [pattern, rewriter] = CreatePatterns(func);
- return RewriteCall(pattern, rewriter, func);
- };
- return CreateFunctionPass(pass_func, 1, "AdjustMatmulOrder");
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("relax.transform.AdjustMatmulOrder",
AdjustMatmulOrder);
-}
-
-} // namespace transform
-} // namespace relax
-} // namespace tvm
diff --git a/src/s_tir/transform/stmt_simplify.cc
b/src/s_tir/transform/stmt_simplify.cc
index 2066cc6051..0f1c2cd4dc 100644
--- a/src/s_tir/transform/stmt_simplify.cc
+++ b/src/s_tir/transform/stmt_simplify.cc
@@ -94,7 +94,6 @@ class StmtSimplifier : public StmtExprMutatorWithAnalyzer {
};
Function StmtSimplifier::Run(Function func) {
-
analyzer_->rewrite_simplify.SetEnabledExtensions(sym::RewriteSimplifier::kNone);
MarkTensorParamShapes(func);
auto* n = func.CopyOnWrite();
// Shared string literals no longer enter the primitive analyzer. Inline
their
diff --git a/src/sym/analyzer.cc b/src/sym/analyzer.cc
index 129d0f755e..89a4f8b2ce 100644
--- a/src/sym/analyzer.cc
+++ b/src/sym/analyzer.cc
@@ -63,7 +63,7 @@ void AnalyzerObj::Bind(const Var& var, const PrimExpr& expr,
bool allow_override
this->rewrite_simplify.Update(var, new_expr, allow_override);
this->canonical_simplify.Update(var, new_expr, allow_override);
this->int_set.Update(var, this->int_set(new_expr), allow_override);
- this->transitive_comparisons.Bind(var, expr, allow_override);
+ this->known_comparisons.Bind(var, expr, allow_override);
this->z3_prover.Bind(var, expr, allow_override);
}
@@ -74,7 +74,7 @@ void AnalyzerObj::Bind(const Var& var, const Range& range,
bool allow_override)
} else {
this->const_int_bound.Bind(var, range, allow_override);
this->int_set.Bind(var, range, allow_override);
- this->transitive_comparisons.Bind(var, range, allow_override);
+ this->known_comparisons.Bind(var, range, allow_override);
this->z3_prover.Bind(var, range, allow_override);
}
// skip modular_set
@@ -147,7 +147,7 @@ void ConstraintContext::EnterWithScope() {
recovery_functions_.push_back(
analyzer_->rewrite_simplify.EnterConstraint(constraint_, is_assume_));
recovery_functions_.push_back(analyzer_->int_set.EnterConstraint(constraint_));
-
recovery_functions_.push_back(analyzer_->transitive_comparisons.EnterConstraint(constraint_));
+
recovery_functions_.push_back(analyzer_->known_comparisons.EnterConstraint(constraint_));
recovery_functions_.push_back(analyzer_->z3_prover.EnterConstraint(constraint_,
is_assume_));
}
@@ -289,7 +289,7 @@ 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->transitive_comparisons.CopyFrom(this->transitive_comparisons);
+ cloned->known_comparisons.CopyFrom(this->known_comparisons);
cloned->z3_prover.CopyFrom(this->z3_prover);
return cloned;
}
@@ -377,10 +377,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
})
.def_method("sym.AnalyzerCanProveEqual", &AnalyzerObj::CanProveEqual)
.def("sym.AnalyzerTryCompare",
- [](Analyzer analyzer, const PrimExpr& lhs, const PrimExpr& rhs,
- bool propagate_inequalities) {
- return static_cast<int64_t>(
- analyzer->transitive_comparisons.TryCompare(lhs, rhs,
propagate_inequalities));
+ [](Analyzer analyzer, const PrimExpr& lhs, const PrimExpr& rhs) {
+ return
static_cast<int64_t>(analyzer->known_comparisons.TryCompare(lhs, rhs));
})
.def("sym.AnalyzerIsZ3Enabled",
[](Analyzer analyzer) { return analyzer->z3_prover.IsEnabled(); })
@@ -397,15 +395,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
analyzer->z3_prover.SetRLimit(static_cast<unsigned>(rlimit));
})
.def("sym.AnalyzerGetZ3Stats",
- [](Analyzer analyzer) { return analyzer->z3_prover.GetStats(); })
- .def("sym.AnalyzerGetEnabledExtensions",
- [](Analyzer analyzer) {
- return
static_cast<std::int64_t>(analyzer->rewrite_simplify.GetEnabledExtensions());
- })
- .def("sym.AnalyzerSetEnabledExtensions", [](Analyzer analyzer, int64_t
flags) {
- analyzer->rewrite_simplify.SetEnabledExtensions(
- static_cast<RewriteSimplifier::Extension>(flags));
- });
+ [](Analyzer analyzer) { return analyzer->z3_prover.GetStats(); });
}
} // namespace sym
diff --git a/src/sym/conjunctive_normal_form.cc
b/src/sym/conjunctive_normal_form.cc
deleted file mode 100644
index 56f96d73cb..0000000000
--- a/src/sym/conjunctive_normal_form.cc
+++ /dev/null
@@ -1,445 +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/conjunctive_normal_form.cc
- */
-
-#include "conjunctive_normal_form.h"
-
-#include <tvm/ir/prim/expr.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/sym/analyzer.h>
-
-#include <optional>
-#include <unordered_map>
-#include <unordered_set>
-#include <utility>
-#include <vector>
-
-#include "pattern_match.h"
-#include "rewrite_simplify.h"
-
-namespace tvm {
-namespace sym {
-
-namespace {
-/* \brief A utility for simplifying expressions using conjunctive/disjuctive
normal forms */
-class AndOfOrs {
- public:
- /*! \brief Construct the simplifier
- *
- * Convert a PrimExpr to the internal representation.
- *
- * \param expr The PrimExpr to be simplified.
- */
- explicit AndOfOrs(const PrimExpr& expr);
-
- /*! \brief Convert internal representation to PrimExpr */
- PrimExpr ToPrimExpr() const;
-
- /*! \brief Simplify the internal representation */
- void Simplify(AnalyzerObj* analyzer);
-
- private:
- /*! \brief Internal utility, simplify within each group of expressions
- *
- * For each pair of values within a chunk, attempt to simplify them into
- * a single expression.
- *
- * For example,
- * before = (a == 5) && ((b < 10) || (b > 10))
- * after = (a == 5) && ((b != 10) || false)
- */
- void SimplifyWithinChunks(AnalyzerObj* analyzer);
-
- /*! \brief Internal utility, simplify across groups of expressions
- *
- * For each pair of chunks, if the two chunks differ by only a single
- * term, attempt to simplify those differing terms.
- *
- * For example,
- * before = ((a == 5) || (b <= 10)) && ((a == 5) || (b >= 10))
- * after = ((a == 5) || (b == 10)) && ((a == 5) || true)
- */
- void SimplifyAcrossChunks(AnalyzerObj* analyzer);
-
- /*! \brief Remove instances of true/false from internal representation
- *
- * To avoid invalidating iterators, `SimplifyWithinChunks` and
- * `SimplifyAcrossChunks` may replace keys, but may not remove keys
- * from the internal representation. For example, `(a < 5) && (a <
- * 10)` would be simplified to `(a < 5) && true`. The
- * `RemoveTrueFalse` function removes these leftover instances of
- * true/false.
- */
- void RemoveTrueFalse();
-
- /*! \brief Internal utility function used to convert to internal form */
- static void VisitAndExpressions(const PrimExpr& expr,
- std::function<void(const PrimExpr&)>
callback);
- /*! \brief Internal utility function used to convert to internal form */
- static void VisitOrExpressions(const PrimExpr& expr,
- std::function<void(const PrimExpr&)>
callback);
-
- /* \brief Type-safe wrapper class that represents an PrimExpr
- *
- * Because integer indices are used frequently through this class,
- * maintaining a separation between integer indices used to access
- * specific elements of the internal representation, and unique
- * identifiers used to represent expressions PrimExpr, is useful.
- */
- enum class Key : size_t {};
-
- /*! \brief Convert a PrimExpr to a Key */
- Key GetKey(const PrimExpr& expr);
-
- /*! \brief Convert a Key to a PrimExpr */
- PrimExpr GetExpr(Key key) const;
-
- /*! \brief Attempt to simplify (a && b)
- *
- * If successful, will overwrite the parameters `a` and `b` with the
- * simplified form.
- */
- void TrySimplifyOr(Key* a, Key* b, AnalyzerObj* analyzer);
-
- /*! \brief Attempt to simplify (a || b)
- *
- * If successful, will overwrite the parameters `a` and `b` with the
- * simplified form.
- */
- void TrySimplifyAnd(Key* a, Key* b, AnalyzerObj* analyzer);
-
- /*! \brief The internal representation
- *
- * `chunks[i][j]` is the j-th expression in the i-th OR-group.
- */
- std::vector<std::vector<Key>> chunks_;
-
- /*! \brief Mapping from internal Key to PrimExpr */
- std::unordered_map<Key, PrimExpr, ffi::StructuralHash, ffi::StructuralEqual>
key_to_expr_;
-
- /*! \brief Mapping from PrimExpr to internal Key */
- std::unordered_map<PrimExpr, Key, ffi::StructuralHash, ffi::StructuralEqual>
expr_to_key_;
-
- /*! \brief Cached key representing IntImm::Bool(true) */
- Key key_true_;
-
- /*! \brief Cached key representing IntImm::Bool(false) */
- Key key_false_;
-};
-
-AndOfOrs::AndOfOrs(const PrimExpr& expr)
- : key_true_(GetKey(IntImm::Bool(true))),
key_false_(GetKey(IntImm::Bool(false))) {
- VisitAndExpressions(expr, [&](const PrimExpr& outer_expr) {
- std::vector<Key> or_components;
- VisitOrExpressions(outer_expr, [&](const PrimExpr& inner_expr) {
- Key key = GetKey(inner_expr);
- bool is_duplicate = std::any_of(or_components.begin(),
or_components.end(),
- [&](Key prev) { return prev == key; });
- if (!is_duplicate) {
- or_components.push_back(key);
- }
- });
-
- bool is_permutation =
- std::any_of(chunks_.begin(), chunks_.end(), [&](const
std::vector<Key>& prev_components) {
- return or_components.size() == prev_components.size() &&
- std::is_permutation(prev_components.begin(),
prev_components.end(),
- or_components.begin());
- });
- if (!is_permutation) {
- chunks_.push_back(std::move(or_components));
- }
- });
-}
-
-void AndOfOrs::VisitAndExpressions(const PrimExpr& expr,
- std::function<void(const PrimExpr&)>
callback) {
- PVar<PrimExpr> x, y, z;
- if ((x && y).Match(expr)) {
- // These are separate AND conditions, recurse into them in case
- // they contain AND internally.
- VisitAndExpressions(x.Eval(), callback);
- VisitAndExpressions(y.Eval(), callback);
- } else if ((x || y).Match(expr)) {
- // This may be the bottom-most breakdown, but either x or y may
- // themselves contain AND. (e.g. (A && B) || (C && D) should be
- // split into (A || C), (A || D), (B || C), and (B || D).)
- // Recurse into each, then reconstruct an OR condition.
- VisitAndExpressions(x.Eval(), [&](const PrimExpr& x_part) {
- VisitAndExpressions(y.Eval(), [&](const PrimExpr& y_part) {
callback(x_part || y_part); });
- });
- } else {
- // This is bottom-most breakdown.
- callback(expr);
- }
-}
-
-void AndOfOrs::VisitOrExpressions(const PrimExpr& expr,
- std::function<void(const PrimExpr&)>
callback) {
- PVar<PrimExpr> x, y, z;
- if ((x || y).Match(expr)) {
- // These are separate OR conditions, recurse into them in case
- // they contain OR internally.
- VisitOrExpressions(x.Eval(), callback);
- VisitOrExpressions(y.Eval(), callback);
- } else if ((x && y).Match(expr)) {
- // This may be the bottom-most breakdown, but either x or y may
- // themselves contain OR. (e.g. (A || B) && (C || D) should be
- // split into (A && C), (A && D), (B && C), and (B && D).)
- // Recurse into each, then reconstruct an AND condition.
- VisitOrExpressions(x.Eval(), [&](const PrimExpr& x_part) {
- VisitOrExpressions(y.Eval(), [&](const PrimExpr& y_part) {
callback(x_part && y_part); });
- });
- } else {
- // This is bottom-most breakdown.
- callback(expr);
- }
-}
-
-AndOfOrs::Key AndOfOrs::GetKey(const PrimExpr& expr) {
- auto it = expr_to_key_.find(expr);
- if (it != expr_to_key_.end()) {
- return it->second;
- }
-
- Key key{expr_to_key_.size()};
- expr_to_key_[expr] = key;
- key_to_expr_.insert_or_assign(key, expr);
- return key;
-}
-
-PrimExpr AndOfOrs::GetExpr(AndOfOrs::Key key) const {
- auto it = key_to_expr_.find(key);
- TVM_FFI_ICHECK(it != key_to_expr_.end());
- return it->second;
-}
-
-PrimExpr AndOfOrs::ToPrimExpr() const {
- PrimExpr expr = IntImm::Bool(true);
- for (const auto& chunk : chunks_) {
- PrimExpr chunk_expr = IntImm::Bool(false);
- for (Key j : chunk) {
- chunk_expr = chunk_expr || GetExpr(j);
- }
- expr = expr && chunk_expr;
- }
- return expr;
-}
-
-void AndOfOrs::TrySimplifyOr(Key* a_ptr, Key* b_ptr, AnalyzerObj* analyzer) {
- Key& a = *a_ptr;
- Key& b = *b_ptr;
- PrimExpr joint = GetExpr(a) || GetExpr(b);
- PrimExpr simplified = analyzer->rewrite_simplify(joint);
- if (!prim::ExprDeepEqual()(simplified, joint)) {
- if (auto* simplified_or = simplified.as<prim::OrNode>()) {
- a = GetKey(simplified_or->a);
- b = GetKey(simplified_or->b);
- } else {
- a = key_false_;
- b = GetKey(simplified);
- }
- }
-}
-
-void AndOfOrs::TrySimplifyAnd(Key* a_ptr, Key* b_ptr, AnalyzerObj* analyzer) {
- Key& a = *a_ptr;
- Key& b = *b_ptr;
- PrimExpr joint = GetExpr(a) && GetExpr(b);
- PrimExpr simplified = analyzer->rewrite_simplify(joint);
- if (!prim::ExprDeepEqual()(simplified, joint)) {
- if (auto* simplified_and = simplified.as<prim::AndNode>()) {
- a = GetKey(simplified_and->a);
- b = GetKey(simplified_and->b);
- } else {
- a = key_true_;
- b = GetKey(simplified);
- }
- }
-}
-
-void AndOfOrs::Simplify(AnalyzerObj* analyzer) {
- SimplifyWithinChunks(analyzer);
- RemoveTrueFalse();
- SimplifyAcrossChunks(analyzer);
- RemoveTrueFalse();
-}
-
-void AndOfOrs::SimplifyWithinChunks(AnalyzerObj* analyzer) {
- for (auto& chunk : chunks_) {
- for (size_t expr_i = 0; expr_i < chunk.size(); expr_i++) {
- for (size_t expr_j = expr_i + 1; expr_j < chunk.size(); expr_j++) {
- Key& key_i = chunk[expr_i];
- Key& key_j = chunk[expr_j];
-
- TrySimplifyOr(&key_i, &key_j, analyzer);
- }
- }
- }
-}
-
-void AndOfOrs::SimplifyAcrossChunks(AnalyzerObj* analyzer) {
- for (size_t i_and = 0; i_and < chunks_.size(); i_and++) {
- for (size_t j_and = i_and + 1; j_and < chunks_.size(); j_and++) {
- auto& i_chunk = chunks_[i_and];
- auto& j_chunk = chunks_[j_and];
-
- if (i_chunk.size() == 1 && j_chunk.size() == 1) {
- auto& key_i = i_chunk[0];
- auto& key_j = j_chunk[0];
- TrySimplifyAnd(&key_i, &key_j, analyzer);
- continue;
- }
- std::unordered_set<Key> j_set(j_chunk.begin(), j_chunk.end());
-
- std::optional<size_t> i_distinct_index;
- for (size_t i = 0; i < i_chunk.size(); i++) {
- if (!j_set.count(i_chunk[i])) {
- i_distinct_index = i;
- break;
- }
- }
-
- if (!i_distinct_index.has_value()) {
- // I = (i_0 || i_1 || ... || i_N)
- // J = (i_0 || i_1 || ... || i_N || j_0 || ... || j_N)
- // I && J == I == I && true
-
- j_chunk = {key_true_};
- continue;
- }
-
- std::unordered_set<Key> i_set(i_chunk.begin(), i_chunk.end());
-
- std::optional<size_t> j_distinct_index;
- for (size_t j = 0; j < j_chunk.size(); j++) {
- if (!i_set.count(j_chunk[j])) {
- j_distinct_index = j;
- break;
- }
- }
-
- if (!j_distinct_index.has_value()) {
- // I = (i_0 || ... || i_N || j_0 || ... || j_N)
- // J = (j_0 || ... || j_N)
- // I && J == J == true && J
-
- i_chunk = {key_true_};
- continue;
- }
-
- if (i_chunk.size() == j_chunk.size()) {
- size_t num_shared_exprs = 0;
- for (const auto& j_key : j_chunk) {
- if (i_set.count(j_key)) {
- ++num_shared_exprs;
- }
- }
-
- if (num_shared_exprs + 1 == i_chunk.size()) {
- // All but one of the expressions are shared. If the AND
- // of the distinct expressions can be simplified, we can
- // replace.
- //
- // (A or B) and (A or C) => A or (B and C)
- auto& key_i = i_chunk[i_distinct_index.value()];
- auto& key_j = j_chunk[j_distinct_index.value()];
-
- // When attempting to simplify (B and C), the analyzer may
- // assume that A is false.
- PrimExpr known = [&]() {
- PrimExpr known = IntImm::Bool(true);
- for (const auto& key : i_chunk) {
- if (&key != &key_i) {
- known = known && analyzer->Simplify(!GetExpr(key));
- }
- }
- return known;
- }();
-
- With<ConstraintContext> context(analyzer, known);
- TrySimplifyAnd(&key_i, &key_j, analyzer);
- }
- }
- }
- }
-}
-
-void AndOfOrs::RemoveTrueFalse() {
- for (auto& chunk : chunks_) {
- // Any occurrence of True inside an OR makes the entire expression True.
- if (std::any_of(chunk.begin(), chunk.end(), [&](Key key) { return key ==
key_true_; })) {
- chunk = {key_true_};
- } else {
- // Any occurrence of False inside an OR can be removed
- chunk.erase(
- std::remove_if(chunk.begin(), chunk.end(), [&](Key key) { return key
== key_false_; }),
- chunk.end());
- }
- }
-
- // Any occurence of False inside an AND makes the entire expression False.
- if (std::any_of(chunks_.begin(), chunks_.end(),
- [&](const std::vector<Key>& chunk) { return chunk.size() ==
0; })) {
- chunks_ = {{}};
- } else {
- // Any occurrence of True inside an AND can be removed.
- chunks_.erase(std::remove_if(chunks_.begin(), chunks_.end(),
- [&](const std::vector<Key>& chunk) {
- return chunk.size() == 1 && chunk[0] ==
key_true_;
- }),
- chunks_.end());
- }
-}
-
-// Helper utility for temporarily disabling the
-// kConvertBooleanToAndOfOrs flag on an analyzer, to prevent infinite
-// recursion.
-class DisableAndOfOrRecursion {
- public:
- explicit DisableAndOfOrRecursion(AnalyzerObj* analyzer)
- : analyzer_(analyzer),
cached_flags_(analyzer->rewrite_simplify.GetEnabledExtensions()) {
- auto new_flags = static_cast<RewriteSimplifier::Extension>(
- cached_flags_ & (~RewriteSimplifier::kConvertBooleanToAndOfOrs));
- analyzer->rewrite_simplify.SetEnabledExtensions(new_flags);
- }
- ~DisableAndOfOrRecursion() {
analyzer_->rewrite_simplify.SetEnabledExtensions(cached_flags_); }
-
- DisableAndOfOrRecursion(const DisableAndOfOrRecursion&) = delete;
- DisableAndOfOrRecursion& operator=(const DisableAndOfOrRecursion&) = delete;
-
- private:
- AnalyzerObj* analyzer_;
- RewriteSimplifier::Extension cached_flags_;
-};
-
-} // namespace
-
-PrimExpr SimplifyAsAndOfOrs(const PrimExpr& expr, AnalyzerObj* analyzer) {
- DisableAndOfOrRecursion context(analyzer);
- AndOfOrs repr(analyzer->Simplify(expr));
- repr.Simplify(analyzer);
- return repr.ToPrimExpr();
-}
-
-} // namespace sym
-} // namespace tvm
diff --git a/src/sym/conjunctive_normal_form.h
b/src/sym/conjunctive_normal_form.h
deleted file mode 100644
index 489feef7ed..0000000000
--- a/src/sym/conjunctive_normal_form.h
+++ /dev/null
@@ -1,50 +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 conjunctive_normal_form.h
- *
- * \brief Centralized location for simplifying into specific forms
- */
-
-#ifndef TVM_SYM_CONJUNCTIVE_NORMAL_FORM_H_
-#define TVM_SYM_CONJUNCTIVE_NORMAL_FORM_H_
-
-#include <tvm/ir/prim/expr.h>
-
-namespace tvm {
-namespace sym {
-
-class AnalyzerObj;
-class Analyzer;
-
-/*! \brief Convert boolean expression to AND of ORs and simplify
- *
- * \param expr The PrimExpr to be simplified
- *
- * \param analyzer The analyzer with which to simplify
- *
- * \return The simplified expression
- */
-PrimExpr SimplifyAsAndOfOrs(const PrimExpr& expr, AnalyzerObj* analyzer);
-
-} // namespace sym
-} // namespace tvm
-
-#endif // TVM_SYM_CONJUNCTIVE_NORMAL_FORM_H_
diff --git a/src/sym/transitive_comparison_analyzer.cc
b/src/sym/known_comparison_analyzer.cc
similarity index 53%
rename from src/sym/transitive_comparison_analyzer.cc
rename to src/sym/known_comparison_analyzer.cc
index a1b128467b..16ae5f69b8 100644
--- a/src/sym/transitive_comparison_analyzer.cc
+++ b/src/sym/known_comparison_analyzer.cc
@@ -17,14 +17,13 @@
* under the License.
*/
/*!
- * \file tvm/sym/transitive_comparison_analyzer.cc
+ * \file tvm/sym/known_comparison_analyzer.cc
*/
#include <tvm/ir/prim/expr.h>
#include <tvm/sym/analyzer.h>
#include <optional>
-#include <unordered_set>
#include <vector>
#include "constraint_extract.h"
@@ -36,7 +35,7 @@ using namespace tvm::prim;
using prim::IsConstInt;
-class TransitiveComparisonAnalyzer::Impl {
+class KnownComparisonAnalyzer::Impl {
public:
/* \brief Using previously specified knowns, compare the expressions provided
*
@@ -44,19 +43,12 @@ class TransitiveComparisonAnalyzer::Impl {
*
* \param rhs The right-hand side of the comparison
*
- * \param propagate_inequalities If true, attempt to find a sequence
- * of transitive inequalities that allow the lhs and rhs to be
- * compared. If false, only use the known comparison that have been
- * directly provided. Using `propagate_inequalities = false` is
- * roughly equivalent to comparing against all known values with
- * `prim::ExprDeepEqual`, but also allowing for constant offsets on either
- * side of the inequality.
+ * 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,
- bool propagate_inequalities = true) const;
+ CompareResult TryCompare(const PrimExpr& lhs, const PrimExpr& rhs) const;
/*! \brief Bind a variable as being equal to a known expression
*
@@ -156,11 +148,6 @@ class TransitiveComparisonAnalyzer::Impl {
*/
Comparison(Key lhs, Key rhs, int64_t offset, CompareResult result);
- /*! \brief Utility function to validate that all GT and LT results
- * have been normalized out
- */
- bool IsNormalized() const;
-
/*! \brief Move the specified expression to the LHS.
*
* \param new_lhs The argument that should be moved to the LHS of the
@@ -172,17 +159,6 @@ class TransitiveComparisonAnalyzer::Impl {
*/
std::optional<Comparison> WithLHS(Key new_lhs) const;
- /*! \brief Create the negation of the current comparison */
- Comparison Negated() const;
-
- /*! \brief Check the this comparison implies
- *
- * Returns true if this comparison being true implies that the
- * other comparison must also be true. Returns false if the other
- * comparison cannot be shown to be true.
- */
- bool Implies(const Comparison& other) const;
-
// The LHS of the comparison
Key lhs_;
@@ -224,51 +200,11 @@ class TransitiveComparisonAnalyzer::Impl {
*/
std::vector<Comparison> CollectDirectComparisons(Key lhs_key, Key rhs_key)
const;
- /*! Collect known comparisons between LHS and RHS, with propagation
- *
- * \param lhs_key The left-hand side of the comparison
- *
- * \param rhs_key The right-hand side of the comparison
- *
- * \returns All comparisons between `lhs_key` and `rhs_key`,
- * including the explicitly-provided comparisons in `knowns_` and
- * `scoped_knowns_`, and comparisons provable through a series of
- * comparisons through other values. All comparisons returned are
- * between `lhs_key` and `rhs_key`, and are normalized such that
- * `lhs_key` is on the left-hand side.
- */
- std::vector<Comparison> CollectIndirectComparisons(Key lhs_key, Key rhs_key)
const;
-
- /*! \brief Internal function used by CollectIndirectComparisons
- *
- * Perform a depth-first search through the space of known
- * expressions, starting at the LHS of a comparison. In this
- * search, each expression is a node of a graph, and each known
- * comparison is an edge of the graph.
- *
- * For example, suppose we have previous knowns of (A<=B), (B<=C+1)
- * and (C<=D-5). The expressions [A,B,C,D] are the nodes of the
- * search space. Each comparison is an edge connecting two
- * expressions, such as (B<=C+1) connecting the expressions B and D.
- * If we are attempting to compare expressions A and D, a search
- * starting at expression A could follow each edge until reaching
- * expression D, then combine the comparisons that compose the path
- * into the expression A<=D-4.
- *
- * \param lhs_key The left-hand side of the comparison
- *
- * \param rhs_key The right-hand side of the comparison
- *
- * \returns A vector of comparisons between the two expressions.
- */
- std::vector<Comparison> DFSFromLHS(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` or
- * `CollectIndirectComparisons`.
+ * result from `CollectDirectComparisons`.
*
* \param offset The constant offset in the comparison being proven.
* This is extracted from any additive/subtractive constants in the
@@ -332,19 +268,6 @@ CompareResult Reverse(CompareResult res) {
}
}
-// Internal utility, return the CompareResult resulting from negating
-// the comparison.
-CompareResult Negate(CompareResult res) {
- switch (res) {
- case CompareResult::kInconsistent:
- return CompareResult::kInconsistent;
- case CompareResult::kUnknown:
- return CompareResult::kUnknown;
- default:
- return CompareResult(~static_cast<int>(res) &
static_cast<int>(CompareResult::kUnknown));
- }
-}
-
// 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
@@ -374,8 +297,8 @@ std::tuple<PrimExpr, PrimExpr, int64_t>
ExtractOffsets(const PrimExpr& lhs, cons
} // namespace
-std::optional<TransitiveComparisonAnalyzer::Impl::Comparison>
-TransitiveComparisonAnalyzer::Impl::FromExpr(const PrimExpr& expr) {
+std::optional<KnownComparisonAnalyzer::Impl::Comparison>
KnownComparisonAnalyzer::Impl::FromExpr(
+ const PrimExpr& expr) {
CompareResult res;
PVar<PrimExpr> x, y;
if ((x <= y).Match(expr)) {
@@ -412,32 +335,10 @@ TransitiveComparisonAnalyzer::Impl::FromExpr(const
PrimExpr& expr) {
return Comparison(lhs_key, rhs_key, offset, res);
}
-TransitiveComparisonAnalyzer::Impl::Comparison::Comparison(Key lhs, Key rhs,
int64_t offset,
- CompareResult
result)
+KnownComparisonAnalyzer::Impl::Comparison::Comparison(Key lhs, Key rhs,
int64_t offset,
+ CompareResult result)
: lhs_(lhs), rhs_(rhs), offset_(offset), result_(result) {
- // Normalize the comparison to remove LT and GT expressions,
- // reducing the number of operators that must be handled later. By
- // eliminating LT and GT, instead of eliminating LE or GE, a
- // potential off-by-one error is avoided.
- //
- // For floating-point numbers, (x < y + c1) and (y < z + c2) implies
- // that (x < z + (c1 + c2)). For integer types, which the
- // TransitiveComparisonAnalyzer is intended for use with integers,
- // LT or GT can give a tighter constraint, though with a less
- // convenient symmetry.
- //
- // i < j + c1, j < k + c2
- // i <= j + c1 - 1, j <= k + c2 - 1
- // i + 1 - c1 <= j, j <= k + c2 - 1
- // i + 1 - c1 <= k + c2 - 1
- // i <= k + c1 + c2 - 2
- // i < k + (c1 + c2 - 1)
- //
- // By always working with LE and GE comparisons, we avoid needing to
- // handle the offset of one that would be introduced by LT and GT at
- // all points of use. The only point of use for LT and GT is when
- // normalizing comparisons (i.e. this constructor).
-
+ // Integer strict comparisons become non-strict comparisons with a unit
offset.
if (result_ == CompareResult::kLT) {
result_ = CompareResult::kLE;
offset_ -= 1;
@@ -448,8 +349,8 @@
TransitiveComparisonAnalyzer::Impl::Comparison::Comparison(Key lhs, Key rhs, int
}
}
-std::optional<TransitiveComparisonAnalyzer::Impl::Key>
-TransitiveComparisonAnalyzer::Impl::ExprToPreviousKey(const PrimExpr& expr)
const {
+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;
@@ -458,8 +359,7 @@ TransitiveComparisonAnalyzer::Impl::ExprToPreviousKey(const
PrimExpr& expr) cons
}
}
-TransitiveComparisonAnalyzer::Impl::Key
TransitiveComparisonAnalyzer::Impl::ExprToKey(
- const PrimExpr& expr) {
+KnownComparisonAnalyzer::Impl::Key
KnownComparisonAnalyzer::Impl::ExprToKey(const PrimExpr& expr) {
if (auto prev = ExprToPreviousKey(expr)) {
return prev.value();
} else {
@@ -469,14 +369,8 @@ TransitiveComparisonAnalyzer::Impl::Key
TransitiveComparisonAnalyzer::Impl::Expr
}
}
-bool TransitiveComparisonAnalyzer::Impl::Comparison::IsNormalized() const {
- // These < and > should be removed during normalization. See the
- // `Comparison::Comparison` constructor for further details.
- return result_ != CompareResult::kLT && result_ != CompareResult::kGT;
-}
-
-std::optional<TransitiveComparisonAnalyzer::Impl::Comparison>
-TransitiveComparisonAnalyzer::Impl::Comparison::WithLHS(Key new_lhs) const {
+std::optional<KnownComparisonAnalyzer::Impl::Comparison>
+KnownComparisonAnalyzer::Impl::Comparison::WithLHS(Key new_lhs) const {
if (new_lhs == lhs_) {
return *this;
} else if (new_lhs == rhs_) {
@@ -488,85 +382,29 @@
TransitiveComparisonAnalyzer::Impl::Comparison::WithLHS(Key new_lhs) const {
}
}
-TransitiveComparisonAnalyzer::Impl::Comparison
-TransitiveComparisonAnalyzer::Impl::Comparison::Negated() const {
- return Comparison(lhs_, rhs_, offset_, Negate(result_));
-}
-
-bool TransitiveComparisonAnalyzer::Impl::Comparison::Implies(
- const TransitiveComparisonAnalyzer::Impl::Comparison& other) const {
- TVM_FFI_ICHECK(lhs_ == other.lhs_);
- TVM_FFI_ICHECK(rhs_ == other.rhs_);
- TVM_FFI_ICHECK(IsNormalized());
- TVM_FFI_ICHECK(other.IsNormalized());
-
- if (result_ == other.result_ && offset_ == other.offset_) {
- // if c1 == c2, x != y + c1 => x != y + c2
- // if c1 == c2, x == y + c1 => x == y + c2
- return true;
- }
-
- if (other.result_ == CompareResult::kLE && offset_ <= other.offset_) {
- if (result_ == CompareResult::kEQ || result_ == CompareResult::kLE) {
- // if c1 <= c2, x <= y + c1 => x <= y + c2
- // if c1 <= c2, x == y + c1 => x <= y + c2
- return true;
- }
- }
-
- if (other.result_ == CompareResult::kGE && offset_ >= other.offset_) {
- if (result_ == CompareResult::kEQ || result_ == CompareResult::kGE) {
- // if c1 >= c2, x == y + c1 => x >= y + c2
- // if c1 >= c2, x >= y + c1 => x >= y + c2
- return true;
- }
- }
-
- if (other.result_ == CompareResult::kNE) {
- if (result_ == CompareResult::kEQ && offset_ != other.offset_) {
- // if c1 != c2, x == y + c1 => x != y + c2
- return true;
- }
-
- if (result_ == CompareResult::kLE && offset_ < other.offset_) {
- // if c1 < c2, x <= y + c1 => x < y + c2 => x != y + c2
- return true;
- }
-
- if (result_ == CompareResult::kGE && offset_ > other.offset_) {
- // if c1 != c2, x >= y + c1 => x > y + c2 => x != y + c2
- return true;
- }
- }
+KnownComparisonAnalyzer::KnownComparisonAnalyzer() :
impl_(std::make_unique<Impl>()) {}
+KnownComparisonAnalyzer::~KnownComparisonAnalyzer() {}
- return false;
-}
-
-TransitiveComparisonAnalyzer::TransitiveComparisonAnalyzer() :
impl_(std::make_unique<Impl>()) {}
-TransitiveComparisonAnalyzer::~TransitiveComparisonAnalyzer() {}
-
-void TransitiveComparisonAnalyzer::CopyFrom(const
TransitiveComparisonAnalyzer& other) {
+void KnownComparisonAnalyzer::CopyFrom(const KnownComparisonAnalyzer& other) {
impl_->CopyFrom(*other.impl_);
}
-CompareResult TransitiveComparisonAnalyzer::TryCompare(const PrimExpr& lhs,
const PrimExpr& rhs,
- bool
propagate_inequalities) {
- return impl_->TryCompare(lhs, rhs, propagate_inequalities);
+CompareResult KnownComparisonAnalyzer::TryCompare(const PrimExpr& lhs, const
PrimExpr& rhs) {
+ return impl_->TryCompare(lhs, rhs);
}
-void TransitiveComparisonAnalyzer::Bind(const Var& var, const PrimExpr& expr,
bool allow_override) {
+void KnownComparisonAnalyzer::Bind(const Var& var, const PrimExpr& expr, bool
allow_override) {
impl_->Bind(var, expr, allow_override);
}
-void TransitiveComparisonAnalyzer::Bind(const Var& var, const Range& range,
bool allow_override) {
+void KnownComparisonAnalyzer::Bind(const Var& var, const Range& range, bool
allow_override) {
impl_->Bind(var, range, allow_override);
}
-std::function<void()> TransitiveComparisonAnalyzer::EnterConstraint(const
PrimExpr& constraint) {
+std::function<void()> KnownComparisonAnalyzer::EnterConstraint(const PrimExpr&
constraint) {
return impl_->EnterConstraint(constraint);
}
-void TransitiveComparisonAnalyzer::Impl::AddKnown(const PrimExpr& expr,
- std::vector<Comparison>*
vec) {
+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)) {
@@ -576,8 +414,7 @@ void TransitiveComparisonAnalyzer::Impl::AddKnown(const
PrimExpr& expr,
}
}
-void TransitiveComparisonAnalyzer::Impl::Bind(const Var& var, const Range&
range,
- bool allow_override) {
+void KnownComparisonAnalyzer::Impl::Bind(const Var& var, const Range& range,
bool allow_override) {
auto it = prev_bindings_.find(var);
if (it != prev_bindings_.end()) {
prim::ExprDeepEqual expr_equal;
@@ -605,12 +442,12 @@ void TransitiveComparisonAnalyzer::Impl::Bind(const Var&
var, const Range& range
}
}
-void TransitiveComparisonAnalyzer::Impl::Bind(const Var& var, const PrimExpr&
expr,
- bool allow_override) {
+void KnownComparisonAnalyzer::Impl::Bind(const Var& var, const PrimExpr& expr,
+ bool allow_override) {
Bind(var, Range::FromMinExtent(expr, 1), allow_override);
}
-std::function<void()>
TransitiveComparisonAnalyzer::Impl::EnterConstraint(const PrimExpr& expr) {
+std::function<void()> KnownComparisonAnalyzer::Impl::EnterConstraint(const
PrimExpr& expr) {
size_t old_literal_size = scoped_knowns_.size();
AddKnown(expr, &scoped_knowns_);
size_t new_literal_size = scoped_knowns_.size();
@@ -622,9 +459,8 @@ std::function<void()>
TransitiveComparisonAnalyzer::Impl::EnterConstraint(const
return frecover;
}
-CompareResult TransitiveComparisonAnalyzer::Impl::TryCompare(const PrimExpr&
lhs_expr,
- const PrimExpr&
rhs_expr,
- bool
propagate_inequalities) const {
+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)) {
@@ -653,18 +489,12 @@ CompareResult
TransitiveComparisonAnalyzer::Impl::TryCompare(const PrimExpr& lhs
return CompareResult::kUnknown;
}
- auto lhs_to_rhs = [&]() {
- if (propagate_inequalities) {
- return CollectIndirectComparisons(lhs_key.value(), rhs_key.value());
- } else {
- return CollectDirectComparisons(lhs_key.value(), rhs_key.value());
- }
- }();
+ auto lhs_to_rhs = CollectDirectComparisons(lhs_key.value(), rhs_key.value());
return MergeComparisons(lhs_to_rhs, offset);
}
-std::vector<TransitiveComparisonAnalyzer::Impl::Comparison>
-TransitiveComparisonAnalyzer::Impl::CollectDirectComparisons(Key lhs_key, Key
rhs_key) const {
+std::vector<KnownComparisonAnalyzer::Impl::Comparison>
+KnownComparisonAnalyzer::Impl::CollectDirectComparisons(Key lhs_key, Key
rhs_key) const {
std::vector<Comparison> output;
auto append_known = [&](Comparison cmp) {
@@ -685,161 +515,7 @@
TransitiveComparisonAnalyzer::Impl::CollectDirectComparisons(Key lhs_key, Key rh
return output;
}
-std::vector<TransitiveComparisonAnalyzer::Impl::Comparison>
-TransitiveComparisonAnalyzer::Impl::CollectIndirectComparisons(Key lhs_key,
Key rhs_key) const {
- auto output = DFSFromLHS(lhs_key, rhs_key);
- for (Comparison cmp : DFSFromLHS(rhs_key, lhs_key)) {
- auto opt_normalized = cmp.WithLHS(lhs_key);
- TVM_FFI_ICHECK(opt_normalized.has_value());
- output.push_back(opt_normalized.value());
- }
- return output;
-}
-
-std::vector<TransitiveComparisonAnalyzer::Impl::Comparison>
-TransitiveComparisonAnalyzer::Impl::DFSFromLHS(Key lhs_key, Key rhs_key) const
{
- // Everything in `to_visit` has lhs as its lhs.
- std::unordered_set<Key> seen;
- std::unordered_set<Key> to_visit;
- std::unordered_map<Key, std::vector<Comparison>> compared_to_lhs;
-
- // Utility function to add a new known statement
- auto declare_known = [&](Comparison cmp) {
- std::vector<Comparison>& knowns = compared_to_lhs[cmp.rhs_];
-
- // The comparison adds no new information, no modification
- // required.
- for (auto& prev_known : knowns) {
- if (prev_known.Implies(cmp)) {
- return;
- }
- }
-
- // New information may require visiting a new expression.
- if (cmp.rhs_ != rhs_key && !seen.count(cmp.rhs_)) {
- to_visit.insert(cmp.rhs_);
- seen.insert(cmp.rhs_);
- }
-
- // This comparison is a stronger version of a previous constraint.
- // Therefore, replace the old version entirely.
- for (auto& prev_known : knowns) {
- if (cmp.Implies(prev_known)) {
- prev_known = cmp;
- return;
- }
- }
-
- // Neither a superset nor a subset of previously known
- // constraints, must be tracked separately.
- knowns.push_back(cmp);
- };
-
- // Initialize the search based on any known (in)equalities that use
- // the LHS of the comparison.
- for (const auto& known : knowns_) {
- if (auto normalized = known.WithLHS(lhs_key)) {
- declare_known(normalized.value());
- }
- }
- for (const auto& known : scoped_knowns_) {
- if (auto normalized = known.WithLHS(lhs_key)) {
- declare_known(normalized.value());
- }
- }
-
- // Walk through the space of all comparisons that can be made with
- // LHS.
- while (to_visit.size()) {
- Key middle_key = *to_visit.begin();
- to_visit.erase(to_visit.begin());
-
- std::vector<Comparison>& prev_knowns_using_middle =
compared_to_lhs.at(middle_key);
- TVM_FFI_ICHECK(compared_to_lhs.count(middle_key));
-
- std::vector<Comparison> new_knowns_using_lhs;
-
- auto attempt_transitive = [&](Comparison cmp) {
- TVM_FFI_ICHECK(cmp.IsNormalized());
-
- Key right_key = cmp.rhs_;
-
- if (right_key == lhs_key) {
- return;
- }
-
- for (const auto& prev : prev_knowns_using_middle) {
- CompareResult new_result = CompareResult::kUnknown;
- auto new_offset = (ffi::BigInt(prev.offset_) +
cmp.offset_).as<int64_t>();
- if (!new_offset.has_value()) continue;
-
- if (prev.result_ == CompareResult::kEQ) {
- // x == y + c1 && y OP z + c2, x OP z + (c1 + c2)
- new_result = cmp.result_;
- } else if (cmp.result_ == CompareResult::kEQ) {
- // x OP y + c1 && y == z + c2, x OP z + (c1 + c2)
- new_result = prev.result_;
- } else if (prev.result_ == cmp.result_ &&
- (prev.result_ == CompareResult::kLE || prev.result_ ==
CompareResult::kGE)) {
- // x <= y + c1 && y <= z + c2, x <= z + (c1 + c2)
- // x >= y + c1 && y >= z + c2, x >= z + (c1 + c2)
- //
- // This condition is much simpler to write than the
- // equivalent handling of < or of >, which is why the
- // inequalities are normalized to <= and to >=. See
- // `TransitiveComparisonAnalyzer::Impl::Comparison::Comparison`
- // for further details.
- new_result = prev.result_;
- }
-
- if (new_result != CompareResult::kUnknown) {
- Comparison new_known(lhs_key, right_key, *new_offset, new_result);
- new_knowns_using_lhs.push_back(new_known);
- }
- }
- };
-
- // Attempt to prove a new comparison using one of the original
- // known comparisons. We want to find a known such that
- // `(LHS OP1 middle) && (middle OP2 right)` can be simplified
- // into `(LHS OP3 right)`.
- //
- // Note: The right side is this step is not necessarily the RHS of
- // the comparison we're trying to prove, as we may need to find
- // intermediate comparisons first. For example, if we know that
- // `a<=b`, `b<=c`, and `c<=d`, and we wish to prove that `a<=d`,
- // we must first combine `a<=b` and `b<=c` into `a<=c`. During
- // this first step, `b` is the "middle" and `c` is the "right".
- // The next step can then combind `a<=c` and `c<=d` into `a<=d`.
- for (const auto& known : knowns_) {
- if (auto cmp = known.WithLHS(middle_key)) {
- attempt_transitive(cmp.value());
- }
- }
-
- for (const auto& known : scoped_knowns_) {
- if (auto cmp = known.WithLHS(middle_key)) {
- attempt_transitive(cmp.value());
- }
- }
-
- // Collect together all new knowns, marking new nodes for visiting
- // as needed.
- for (const auto& new_known : new_knowns_using_lhs) {
- declare_known(new_known);
- }
- }
-
- if (auto it = compared_to_lhs.find(rhs_key); it != compared_to_lhs.end()) {
- return it->second;
- } else {
- // There are known comparisons involving the LHS and the RHS, but
- // no path that connects the two expressions.
- return {};
- }
-}
-
-CompareResult TransitiveComparisonAnalyzer::Impl::MergeComparisons(
+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
diff --git a/src/sym/rewrite_simplify.cc b/src/sym/rewrite_simplify.cc
index fa59082a4f..23e225a501 100644
--- a/src/sym/rewrite_simplify.cc
+++ b/src/sym/rewrite_simplify.cc
@@ -36,7 +36,6 @@
#include <tuple>
#include <utility>
-#include "conjunctive_normal_form.h"
#include "const_fold.h"
#include "constraint_extract.h"
#include "pattern_match.h"
@@ -168,9 +167,6 @@ CompareResult RewriteSimplifier::Impl::TryCompare(const
PrimExpr& x, const PrimE
if (is_finished()) return output;
output = CompareResult(output & TryCompareUsingKnownInequalities(x, y));
- if (is_finished()) return output;
-
- output = CompareResult(output & TryComparisonOfProductAndSum(x, y));
return output;
}
@@ -182,152 +178,7 @@ CompareResult
RewriteSimplifier::Impl::TryCompareUsingConstIntBounds(const PrimE
CompareResult RewriteSimplifier::Impl::TryCompareUsingKnownInequalities(const
PrimExpr& x,
const
PrimExpr& y) {
- bool propagate_inequalities = enabled_extensions_ &
kTransitivelyProveInequalities;
- return analyzer_->transitive_comparisons.TryCompare(x, y,
propagate_inequalities);
-}
-
-CompareResult RewriteSimplifier::Impl::TryComparisonOfProductAndSum(const
PrimExpr& x,
- const
PrimExpr& y) {
- bool check_comparison_of_product_and_sum = enabled_extensions_ &
kComparisonOfProductAndSum;
- if (!check_comparison_of_product_and_sum) {
- return CompareResult::kUnknown;
- }
-
- auto opt_special_case =
- [&]() -> std::optional<std::tuple<PrimExpr, PrimExpr, PrimExpr,
PrimExpr>> {
- // Match expressions of the form `(A+B)*C - (A*B)*D`. Depending on
- // previous simplifications, the exact form of the expression may vary.
- PVar<PrimExpr> A, B, C, D;
-
- // diff is `(A+B)*C - (A*B)*D`.
- PrimExpr difference = x - y;
- PrimExpr diff = Mutate(difference).ValueOrUnchanged(difference);
-
- if (PMatchesOneOf{
- (A + B) * C + (A * B) * D,
- (A + B) * C + (B * A) * D,
- (A * B) * D + (A + B) * C,
- (B * A) * D + (A + B) * C,
- }
- .Match(diff)) {
- return std::tuple{A.Eval(), B.Eval(), C.Eval(), -D.Eval()};
- } else if (PMatchesOneOf{
- (A + B) * C + (A * B),
- (A + B) * C + (B * A),
- (A * B) + (A + B) * C,
- (B * A) + (A + B) * C,
- }
- .Match(diff)) {
- return std::tuple{A.Eval(), B.Eval(), C.Eval(), IntImm::Int32(-1)};
- } else {
- return std::nullopt;
- }
- }();
-
- if (!opt_special_case.has_value()) {
- return CompareResult::kUnknown;
- }
- auto [A, B, C, D] = *opt_special_case;
-
- auto A_bound = analyzer_->const_int_bound(A);
- auto B_bound = analyzer_->const_int_bound(B);
- auto C_bound = analyzer_->const_int_bound(C);
- auto D_bound = analyzer_->const_int_bound(D);
-
- auto negate = [](ConstIntBound bound) {
- return ConstIntBound(-bound->max_value, -bound->min_value);
- };
- auto is_negative = [](const ConstIntBound& bound) { return bound->max_value
< 0; };
- auto is_positive = [](const ConstIntBound& bound) { return bound->min_value
> 0; };
-
- // If D is negative, then we'll be providing an upper bound for
- // `(A*B)*D`, rather than a lower bound. To avoid code duplication,
- // flip all the signs here, find a lower bound, then flip the sign
- // to produce the upper bound of the original expression.
- //
- // Before: (A+B)*C < (A*B)*D
- // After: (A*B)*(-D) < (A + B)*(-C)
- bool is_upper_bound = is_negative(D_bound);
- if (is_upper_bound) {
- C_bound = negate(C_bound);
- D_bound = negate(D_bound);
- }
-
- // Before: (A+B)*C < (A*B)*D
- // After: ((-A) + (-B))*(-C) < ((-A)*(-B))*D
- if (is_negative(C_bound)) {
- A_bound = negate(A_bound);
- B_bound = negate(B_bound);
- C_bound = negate(C_bound);
- }
-
- bool all_terms_positive = (is_positive(A_bound) && is_positive(B_bound) &&
is_positive(C_bound) &&
- is_positive(D_bound));
- if (!all_terms_positive) {
- return CompareResult::kUnknown;
- }
-
- // (A + B) * C < (A * B) * D
- // (A + B) * C / (A*B*C*D) < (A * B) * D / (A*B*C*D)
- // 1/(A*D) + 1/(B*D) < 1/C
- // (A*B*C*D) * ( (A+B)/(A*B*D) - 1/C )
- // (A*B*C*D) * ( (1/A + 1/B)/D - 1/C )
- // (A*B*C*D) * (1/(A*D) + 1/(B*D) - 1/C)
- //
- // The constant (A*B*C*D) is positive, and its minimum value is the
- // product of the minimum values of A, B, C, and D. If the reciprocal
- // term (1/(A*D) + 1/(B*D) - 1/C) is positive, then this constant can
- // be used to provide a lower bound on the expression.
-
- bool reciprocal_term_is_positive = [&]() {
- if (D_bound->max_value == ConstIntBound::kPosInf) {
- // If D can grow without bound, the `1/(A*D)` and `1/(B*D)`
- // terms will approach zero, at which point the `-1/C` term
- // will determine the sign the sign.
- return false;
- }
-
- if (std::min(A_bound->max_value, B_bound->max_value) * D_bound->max_value
<=
- C_bound->min_value) {
- // 1/(A*D) + 1/(B*D) - 1/C is positive if 1/C < 1/(A*D) + 1/(B*D).
- // Since each term is positive, this condition can hold if either
- // A*D <= C or B*D <= C.
- return true;
- }
- if (A_bound->max_value != ConstIntBound::kPosInf &&
- B_bound->max_value != ConstIntBound::kPosInf) {
- // Even if neither term is sufficient on its own, if both A and B
- // have known upper bounds, the inequality 1/C < 1/(A*D) + 1/(B*D)
- // may still be provable.
- //
- // The maximum value of the LHS is found when C is minimized. The
- // minimum value of the RHS is found when A, B, and D are
- // maximized. If the condition holds in this case, then it holds
- // in all cases.
- //
- // 1/C_min < 1/(A_max * D_max) + 1/(B_max*D_max)
- // A_max*B_max*D_max < C_min*B_max + C_min*A_max
- // A_max*B_max*D_max < C_min*(A_max + B_max)
- //
- if (A_bound->max_value * B_bound->max_value * D_bound->max_value <
- C_bound->min_value * (A_bound->max_value + B_bound->max_value)) {
- return true;
- }
- }
- return false;
- }();
-
- if (!reciprocal_term_is_positive) {
- return CompareResult::kUnknown;
- }
-
- if (is_upper_bound) {
- // If we flipped the sign of the original expression, flip the sign of
- // the resulting set of possible values.
- return CompareResult::kLT;
- } else {
- return CompareResult::kGT;
- }
+ return analyzer_->known_comparisons.TryCompare(x, y);
}
// try to prove x equals val
@@ -543,12 +394,6 @@ std::function<void()>
RewriteSimplifier::Impl::EnterConstraint(const PrimExpr& c
return frecover;
}
-void RewriteSimplifier::Impl::SetEnabledExtensions(Extension flags) {
enabled_extensions_ = flags; }
-
-RewriteSimplifier::Extension RewriteSimplifier::Impl::GetEnabledExtensions()
const {
- return enabled_extensions_;
-}
-
UnchangedOr<PrimExpr> RewriteSimplifier::Impl::Mutate_(const prim::SubNode* op,
InplaceMode
inplace_mode) {
PrimExpr ret =
@@ -2152,59 +1997,13 @@ PrimExpr
RewriteSimplifier::Impl::ApplyRewriteRules(prim::Not ret, InplaceMode i
UnchangedOr<PrimExpr> RewriteSimplifier::Impl::Mutate_(const prim::AndNode* op,
InplaceMode
inplace_mode) {
- PrimExpr ret{ffi::UnsafeInit{}};
- // If this extension isn't enabled, just delegate out.
- if (!(enabled_extensions_ & kApplyConstraintsToBooleanBranches)) {
- ret = SimplifierBase::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<PrimExpr>(op));
- } else {
- PrimExpr a = op->a;
- PrimExpr b = op->b;
-
- // Alternate which branch is used as the constraint, and which is
- // being simplified. Because some sub-analyzers expect their
- // constraints to already be simplified, each branch may require
- // more than one update. The loop condition allows each branch to
- // be visited up to twice, but only performs the second visit if
- // necessary.
- size_t iterations_since_update = 0;
- for (size_t i = 0; i < 4; i++) {
- PrimExpr& to_update = (i % 2 == 0) ? a : b;
- const PrimExpr& constraint = (i % 2 == 0) ? b : a;
-
- With<ConstraintContext> context(analyzer_, constraint);
- // Keep this iteration out of place: convergence uses same_as() to
detect updates.
- // An in-place change would preserve identity and could stop the loop
early.
- PrimExpr updated = Mutate(to_update,
InplaceMode::kDisallow).ValueOrUnchanged(to_update);
-
- if (!to_update.same_as(updated)) {
- to_update = updated;
- iterations_since_update = 0;
- } else {
- iterations_since_update++;
- if (iterations_since_update >= 2) {
- break;
- }
- }
- }
-
- // Only construct a new object if a change has been made.
- // Otherwise, follow ExprMutator's convention of returning the
- // original object.
- if (a.same_as(op->a) && b.same_as(op->b)) {
- ret = ffi::GetRef<PrimExpr>(op);
- } else {
- ret = prim::And(a, b);
- }
- }
+ PrimExpr ret =
+ SimplifierBase::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<PrimExpr>(op));
op = ret.as<prim::AndNode>();
if (auto const_res = TryConstFold<prim::And>(op->a, op->b)) return
*std::move(const_res);
if (auto match = TryMatchLiteralConstraint(ret)) return match.value();
- if ((enabled_extensions_ & RewriteSimplifier::kConvertBooleanToAndOfOrs) &&
- !recursively_visiting_boolean_) {
- return SimplifyAsAndOfOrs(ret, analyzer_);
- }
// Pattern var to match any expression
PVar<PrimExpr> x, y, z;
@@ -2303,58 +2102,12 @@ UnchangedOr<PrimExpr>
RewriteSimplifier::Impl::Mutate_(const prim::OrNode* op,
InplaceMode
inplace_mode) {
PrimExpr orig = ffi::GetRef<PrimExpr>(op);
- PrimExpr ret{ffi::UnsafeInit{}};
- // If this extension isn't enabled, just delegate out.
- if (!(enabled_extensions_ & kApplyConstraintsToBooleanBranches)) {
- ret = SimplifierBase::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<PrimExpr>(op));
- } else {
- PrimExpr a = op->a;
- PrimExpr b = op->b;
-
- // Alternate which branch is used as the constraint, and which
- // is being simplified. Because some sub-analyzers expect their
- // constraints to already be simplified, each branch may require
- // more than update. The loop condition allows each branch to be
- // visited up to twice, but only if performs the second visit if
- // necessary.
- size_t iterations_since_update = 0;
- for (size_t i = 0; i < 4; i++) {
- PrimExpr& to_update = (i % 2 == 0) ? a : b;
- const PrimExpr& constraint = (i % 2 == 0) ? b : a;
-
- With<ConstraintContext> context(analyzer_,
NormalizeBooleanOperators(prim::Not(constraint)));
- // Keep this iteration out of place: convergence uses same_as() to
detect updates.
- // An in-place change would preserve identity and could stop the loop
early.
- PrimExpr updated = Mutate(to_update,
InplaceMode::kDisallow).ValueOrUnchanged(to_update);
-
- if (!to_update.same_as(updated)) {
- to_update = updated;
- iterations_since_update = 0;
- } else {
- iterations_since_update++;
- if (iterations_since_update >= 2) {
- break;
- }
- }
- }
-
- // Only construct a new object if a change has been made.
- // Otherwise, follow ExprMutator's convention of returning the
- // original object.
- if (a.same_as(op->a) && b.same_as(op->b)) {
- ret = ffi::GetRef<PrimExpr>(op);
- } else {
- ret = prim::Or(a, b);
- }
- }
+ PrimExpr ret =
+ SimplifierBase::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<PrimExpr>(op));
op = ret.as<prim::OrNode>();
if (auto const_res = TryConstFold<prim::Or>(op->a, op->b)) return
*std::move(const_res);
if (auto match = TryMatchLiteralConstraint(ret)) return match.value();
- if ((enabled_extensions_ & RewriteSimplifier::kConvertBooleanToAndOfOrs) &&
- !recursively_visiting_boolean_) {
- return SimplifyAsAndOfOrs(ret, analyzer_);
- }
// Pattern var to match any expression
PVar<PrimExpr> x, y, z;
@@ -2632,13 +2385,6 @@ std::function<void()>
RewriteSimplifier::EnterConstraint(const PrimExpr& constra
return impl_->EnterConstraint(constraint, is_assume);
}
-void RewriteSimplifier::SetEnabledExtensions(Extension flags) {
- impl_->SetEnabledExtensions(flags);
-}
-RewriteSimplifier::Extension RewriteSimplifier::GetEnabledExtensions() const {
- return impl_->GetEnabledExtensions();
-}
-
ffi::ObjectRef RewriteSimplifier::GetStatsCounters() const { return
impl_->GetStatsCounters(); }
void RewriteSimplifier::ResetStatsCounters() { impl_->ResetStatsCounters(); }
diff --git a/src/sym/rewrite_simplify.h b/src/sym/rewrite_simplify.h
index f67ca79a7d..a374f9189e 100644
--- a/src/sym/rewrite_simplify.h
+++ b/src/sym/rewrite_simplify.h
@@ -124,16 +124,6 @@ class RewriteSimplifier::Impl : public SimplifierBase {
std::function<void()> EnterConstraint(const PrimExpr& constraint, bool
is_assume);
- /*! \brief Enable an optional extension or extensions
- *
- * \param flags A bitwise OR of all optional extensions that should
- * be enabled.
- */
- void SetEnabledExtensions(Extension flags);
-
- /*! \brief Return the currently enabled extensions */
- Extension GetEnabledExtensions() const;
-
RewriteSimplifierStats GetStatsCounters() const { return
RewriteSimplifierStats(stats_); }
void ResetStatsCounters() { stats_ = {}; }
@@ -143,7 +133,6 @@ class RewriteSimplifier::Impl : public SimplifierBase {
void CopyFrom(const Impl& other) {
var_map_ = other.var_map_;
literal_constraints_ = other.literal_constraints_;
- enabled_extensions_ = other.enabled_extensions_;
maximum_rewrite_steps_ = other.maximum_rewrite_steps_;
}
@@ -168,13 +157,6 @@ class RewriteSimplifier::Impl : public SimplifierBase {
std::vector<PrimExpr> literal_constraints_;
- // Optionally enabled extensions
- Extension enabled_extensions_{kNone};
-
- /*! Whether the simplifier is current
- */
- bool recursively_visiting_boolean_{false};
-
// maximum number of recursion allowed during a single pass.
static const constexpr int64_t kMaxRecurDepth = 5;
/*!
@@ -233,7 +215,6 @@ class RewriteSimplifier::Impl : public SimplifierBase {
private:
CompareResult TryCompareUsingKnownInequalities(const PrimExpr& x, const
PrimExpr& y);
CompareResult TryCompareUsingConstIntBounds(const PrimExpr& x, const
PrimExpr y);
- CompareResult TryComparisonOfProductAndSum(const PrimExpr& x, const
PrimExpr& y);
// Whether x >= val
bool CanProveGreaterEqual(const PrimExpr& x, int64_t val) {
diff --git a/src/tirx/transform/stmt_simplify.cc
b/src/tirx/transform/stmt_simplify.cc
index f845881e08..f80354f537 100644
--- a/src/tirx/transform/stmt_simplify.cc
+++ b/src/tirx/transform/stmt_simplify.cc
@@ -77,7 +77,6 @@ class StmtSimplifier : public StmtExprMutatorWithAnalyzer {
};
Function StmtSimplifier::Run(Function func) {
-
analyzer_->rewrite_simplify.SetEnabledExtensions(sym::RewriteSimplifier::kNone);
MarkTensorParamShapes(func);
auto* n = func.CopyOnWrite();
// Shared string literals no longer enter the primitive analyzer. Inline
their
diff --git a/tests/python/relax/test_transform_adjust_matmul_order.py
b/tests/python/relax/test_transform_adjust_matmul_order.py
deleted file mode 100644
index 2fce58f328..0000000000
--- a/tests/python/relax/test_transform_adjust_matmul_order.py
+++ /dev/null
@@ -1,1034 +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.
-
-import inspect
-
-import numpy as np
-import pytest
-import torch
-
-import tvm
-import tvm.testing
-from tvm import relax
-from tvm.script import ir as I
-from tvm.script import relax as R
-from tvm.script import tirx as T
-
-
-class Base:
- def test_compare(self):
- transform = relax.transform.AdjustMatmulOrder()
-
- if inspect.isclass(self.Expected) and issubclass(self.Expected,
Exception):
- with pytest.raises(self.Expected):
- transform(self.Before)
- else:
- after = transform(self.Before)
- tvm.ir.assert_structural_equal(self.Expected, after)
-
-
-class TestLHS(Base):
- """Prefer (x*A)*B instead of x*(A*B)
-
- LHS first - (x*A)*B:
- ops = 1*16*2 + 1*2*32 = 96
- RHS first - x*(A*B):
- ops = 16*2*32 + 1*16*32 = 1536
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([16, 2]),
- B: R.Tensor([2, 32]),
- ) -> R.Tensor([32]):
- weight: R.Tensor([16, 32]) = R.matmul(A, B)
- out: R.Tensor([32]) = R.matmul(x, weight)
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([16, 2]),
- B: R.Tensor([2, 32]),
- ) -> R.Tensor([32]):
- x: R.Tensor([2]) = R.matmul(x, A)
- x: R.Tensor([32]) = R.matmul(x, B)
- return x
-
-
-class TestRHS(Base):
- """Prefer A*(B*x) instead of (A*B)*x
-
- LHS first - (A*B)*x:
- ops = 32*2*16 + 32*16*1 = 1536
- RHS first - A*(B*x):
- ops = 2*16*1 + 32*2*1 = 96
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([32, 2]),
- B: R.Tensor([2, 16]),
- ) -> R.Tensor([32]):
- weight: R.Tensor([32, 16]) = R.matmul(A, B)
- out: R.Tensor([32]) = R.matmul(weight, x)
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([32, 2]),
- B: R.Tensor([2, 16]),
- ) -> R.Tensor([32]):
- x: R.Tensor([2]) = R.matmul(B, x)
- x: R.Tensor([32]) = R.matmul(A, x)
- return x
-
-
-class TestIdempotentLHS(Base):
- """The transform shouldn't undo itself if re-applied"""
-
- Before = TestLHS.Expected
- Expected = TestLHS.Expected
-
-
-class TestIdempotentRHS(Base):
- """The transform shouldn't undo itself if re-applied"""
-
- Before = TestRHS.Expected
- Expected = TestRHS.Expected
-
-
-class TestPreserveCompileTimeMatmulOnLHS(Base):
- """Prefer x*(A*B) if (A*B) can be pre-computed
-
- If both `A` and `B` are known at compile-time, they can be lifted
- out and computed at compile-time. Therefore, optimization should
- avoid breaking apart the `(A*B)` expression.
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([32, 2]),
- B: R.Tensor([2, 16]),
- ) -> R.Tensor([32]):
- R.func_attr({"num_input": 1})
-
- weight = R.matmul(A, B)
- out = R.matmul(weight, x)
- return out
-
- Expected = Before
-
-
-class TestPreserveCompileTimeMatmulOnRHS(Base):
- """Prefer (A*B)*x if (A*B) can be pre-computed
-
- If both `A` and `B` are known at compile-time, they can be lifted
- out and computed at compile-time. Therefore, optimization should
- avoid breaking apart the `(A*B)` expression.
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([16, 2]),
- B: R.Tensor([2, 32]),
- ) -> R.Tensor([32]):
- R.func_attr({"num_input": 1})
-
- weight = R.matmul(A, B)
- out = R.matmul(x, weight)
- return out
-
- Expected = Before
-
-
-lora_r_before = T.dynamic("lora_r", "int64")
-lora_r_expected = T.dynamic("lora_r", "int64")
-
-
-class TestLHSDynamic(Base):
- """Prefer (x*A)*B instead of x*(A*B)
-
- This case appears when evaluating LoRA-tuned models with a dynamic
- rank.
-
- LHS first - (x*A)*B:
- ops = 1*16*lora_r + 1*lora_r*32 = 48*lora_r
- RHS first - x*(A*B):
- ops = 16*lora_r*32 + 1*16*32 = 512*lora_r + 512
-
- 48*lora_r can be proved to be less than 512*lora_r + 512, so the LHS first
is preferred.
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([16, lora_r_before]),
- B: R.Tensor([lora_r_before, 32]),
- ) -> R.Tensor([32]):
- weight: R.Tensor([16, 32]) = R.matmul(A, B)
- out: R.Tensor([32]) = R.matmul(x, weight)
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([16, lora_r_expected]),
- B: R.Tensor([lora_r_expected, 32]),
- ) -> R.Tensor([32]):
- x: R.Tensor([lora_r_expected]) = R.matmul(x, A)
- x: R.Tensor([32]) = R.matmul(x, B)
- return x
-
-
-lora_r_before = T.dynamic("lora_r", "int64")
-lora_r_expected = T.dynamic("lora_r", "int64")
-
-
-class TestRHSDynamic(Base):
- """Prefer A*(B*x) instead of (A*B)*x
-
- LHS first - (A*B)*x:
- ops = 32*lora_r*16 + 32*16*1 = 512*lora_r + 512
- RHS first - A*(B*x):
- ops = lora_r*16*1 + 32*lora_r*1 = 48*lora_r
-
- 48*lora_r can be proved to be less than 512*lora_r + 512, so the RHS first
is preferred.
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([32, lora_r_before]),
- B: R.Tensor([lora_r_before, 16]),
- ) -> R.Tensor([32]):
- weight: R.Tensor([32, 16]) = R.matmul(A, B)
- out: R.Tensor([32]) = R.matmul(weight, x)
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([32, lora_r_expected]),
- B: R.Tensor([lora_r_expected, 16]),
- ) -> R.Tensor([32]):
- x: R.Tensor([lora_r_expected]) = R.matmul(B, x)
- x: R.Tensor([32]) = R.matmul(A, x)
- return x
-
-
-class TestIdempotentLHSDynamic(Base):
- """The transform shouldn't undo itself if re-applied"""
-
- Before = TestLHSDynamic.Expected
- Expected = TestLHSDynamic.Expected
-
-
-class TestIdempotentRHSDynamic(Base):
- """The transform shouldn't undo itself if re-applied"""
-
- Before = TestRHSDynamic.Expected
- Expected = TestRHSDynamic.Expected
-
-
-batch_size_before = T.dynamic("batch_size", "int64")
-lora_r_before = T.dynamic("lora_r", "int64")
-
-
-class TestDynamicWithBatchSymbolic1(Base):
- """When both batch_size and lora_r are symbolic and it cannot be proven
which
- is cheaper, LHS or RHS, maintain the existing order.
-
- `Before` computes `x * (A * B)` with
- `x: [batch_size, 1, 16]`, `A: [16, lora_r]`, `B: [lora_r, 32]`.
-
- RHS first - x * (A * B):
- 16*lora_r*32 + batch_size*1*16*32 = 512*(lora_r + batch_size)
-
- LHS first - (x * A) * B:
- batch_size*1*16*lora_r + batch_size*1*lora_r*32 = 48*batch_size*lora_r
-
- When `batch_size` and `lora_r` are known at compile-time:
- - satisfy the inequality 48*batch_size*lora_r < 512*(lora_r +
batch_size),
- the LHS first is preferred.
- - satisfy the inequality 512*(lora_r + batch_size) <
48*batch_size*lora_r,
- the RHS first is preferred.
-
- Without bounds on `batch_size` and `lora_r`, neither side is provably
cheaper.
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([batch_size_before, 1, 16]),
- A: R.Tensor([16, lora_r_before]),
- B: R.Tensor([lora_r_before, 32]),
- ) -> R.Tensor([batch_size_before, 1, 32]):
- weight: R.Tensor([16, 32]) = R.matmul(A, B)
- out: R.Tensor([batch_size_before, 1, 32]) = R.matmul(x, weight)
- return out
-
- Expected = Before
-
-
-class TestDynamicWithBatchConcrete1LHSFirst(Base):
- """With concrete shapes, LHS first is provably cheaper.
-
- batch_size=4, lora_r=16:
- LHS first: 48*4*16 = 3072
- RHS first: 512*(16 + 4) = 10240
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([4, 1, 16]),
- A: R.Tensor([16, 16]),
- B: R.Tensor([16, 32]),
- ) -> R.Tensor([4, 1, 32]):
- weight: R.Tensor([16, 32]) = R.matmul(A, B)
- out: R.Tensor([4, 1, 32]) = R.matmul(x, weight)
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- x: R.Tensor([4, 1, 16]),
- A: R.Tensor([16, 16]),
- B: R.Tensor([16, 32]),
- ) -> R.Tensor([4, 1, 32]):
- weight: R.Tensor([4, 1, 16]) = R.matmul(x, A)
- out: R.Tensor([4, 1, 32]) = R.matmul(weight, B)
- return out
-
-
-class TestDynamicWithBatchConcrete1RHSFirst(Base):
- """With concrete shapes, RHS first is provably cheaper.
-
- batch_size=64, lora_r=16:
- LHS first: 48*64*16 = 49152
- RHS first: 512*(16 + 64) = 40960
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([64, 1, 16]),
- A: R.Tensor([16, 16]),
- B: R.Tensor([16, 32]),
- ) -> R.Tensor([64, 1, 32]):
- weight: R.Tensor([64, 1, 16]) = R.matmul(x, A)
- out: R.Tensor([64, 1, 32]) = R.matmul(weight, B)
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- x: R.Tensor([64, 1, 16]),
- A: R.Tensor([16, 16]),
- B: R.Tensor([16, 32]),
- ) -> R.Tensor([64, 1, 32]):
- weight: R.Tensor([16, 32]) = R.matmul(A, B)
- out: R.Tensor([64, 1, 32]) = R.matmul(x, weight)
- return out
-
-
-batch_size_before = T.dynamic("batch_size", "int64")
-lora_r_before = T.dynamic("lora_r", "int64")
-
-
-class TestDynamicWithBatchSymbolic2(Base):
- """When both batch_size and lora_r are symbolic and it cannot be proven
which
- is cheaper, LHS or RHS, maintain the existing order.
-
- `Before` computes `(A * B) * x` with
- `A: [32, lora_r]`, `B: [lora_r, 16]`, `x: [batch_size, 16, 1]`.
-
- LHS first - (A * B) * x:
- 32*lora_r*16 + batch_size*32*16*1 = 512*(lora_r + batch_size)
-
- RHS first - A * (B * x):
- batch_size*lora_r*16*1 + batch_size*32*lora_r*1 = 48*batch_size*lora_r
-
- When `batch_size` and `lora_r` are known at compile-time:
- - satisfy the inequality 48*batch_size*lora_r < 512*(lora_r +
batch_size),
- the RHS first is preferred.
- - satisfy the inequality 512*(lora_r + batch_size) <
48*batch_size*lora_r,
- the LHS first is preferred.
-
- Without bounds on `batch_size` and `lora_r`, neither side is provably
cheaper.
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([batch_size_before, 16, 1]),
- A: R.Tensor([32, lora_r_before]),
- B: R.Tensor([lora_r_before, 16]),
- ) -> R.Tensor([batch_size_before, 32, 1]):
- weight: R.Tensor([32, 16]) = R.matmul(A, B)
- out: R.Tensor([batch_size_before, 32, 1]) = R.matmul(weight, x)
- return out
-
- Expected = Before
-
-
-class TestDynamicWithBatchConcrete2RHSFirst(Base):
- """With concrete shapes, RHS first is provably cheaper.
-
- batch_size=4, lora_r=16:
- RHS first: 48*4*16 = 3072
- LHS first: 512*(16 + 4) = 10240
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([4, 16, 1]),
- A: R.Tensor([32, 16]),
- B: R.Tensor([16, 16]),
- ) -> R.Tensor([4, 32, 1]):
- weight: R.Tensor([32, 16]) = R.matmul(A, B)
- out: R.Tensor([4, 32, 1]) = R.matmul(weight, x)
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- x: R.Tensor([4, 16, 1]),
- A: R.Tensor([32, 16]),
- B: R.Tensor([16, 16]),
- ) -> R.Tensor([4, 32, 1]):
- weight: R.Tensor([4, 16, 1]) = R.matmul(B, x)
- out: R.Tensor([4, 32, 1]) = R.matmul(A, weight)
- return out
-
-
-class TestDynamicWithBatchConcrete2LHSFirst(Base):
- """With concrete shapes, LHS first is provably cheaper.
-
- batch_size=64, lora_r=16:
- RHS first: 48*64*16 = 49152
- LHS first: 512*(16 + 64) = 40960
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([64, 16, 1]),
- A: R.Tensor([32, 16]),
- B: R.Tensor([16, 16]),
- ) -> R.Tensor([64, 32, 1]):
- weight: R.Tensor([64, 16, 1]) = R.matmul(B, x)
- out: R.Tensor([64, 32, 1]) = R.matmul(A, weight)
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- x: R.Tensor([64, 16, 1]),
- A: R.Tensor([32, 16]),
- B: R.Tensor([16, 16]),
- ) -> R.Tensor([64, 32, 1]):
- weight: R.Tensor([32, 16]) = R.matmul(A, B)
- out: R.Tensor([64, 32, 1]) = R.matmul(weight, x)
- return out
-
-
-M_before = T.dynamic("M", "int64")
-N_before = T.dynamic("N", "int64")
-P_before = T.dynamic("P", "int64")
-Q_before = T.dynamic("Q", "int64")
-
-
-class TestNoOpForFullyDynamicOnLHS(Base):
- """Keep existing order if no benefit can be proven
-
- Like `TestNoOpForFullyDynamicOnRHS`, except the input has the LHS
- computed first.
-
- Here, it is uncertain whether a reordering would improve the
- number of operations.
-
- LHS first: (M+P)*N*Q
- RHS first: (N+Q)*M*P
-
- After simplifying, the LHS should be performed first if the
- following inequality holds:
-
- 1/M + 1/P < 1/N + 1/Q
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- A: R.Tensor([M_before, N_before]),
- B: R.Tensor([N_before, P_before]),
- C: R.Tensor([P_before, Q_before]),
- ):
- out = R.matmul(R.matmul(A, B), C)
- return out
-
- Expected = Before
-
-
-M_before = T.dynamic("M", "int64")
-N_before = T.dynamic("N", "int64")
-P_before = T.dynamic("P", "int64")
-Q_before = T.dynamic("Q", "int64")
-
-
-class TestNoOpForFullyDynamicOnRHS(Base):
- """Keep existing order if no benefit can be proven
-
- Like `TestNoOpForFullyDynamicOnLHS`, except the input has the RHS
- computed first.
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- A: R.Tensor([M_before, N_before]),
- B: R.Tensor([N_before, P_before]),
- C: R.Tensor([P_before, Q_before]),
- ):
- out = R.matmul(A, R.matmul(B, C))
- return out
-
- Expected = Before
-
-
-class TestRHSPermuteDims(Base):
- """Prefer (x*A)*B instead of x*(A*B)
-
- Like `TestRHS`, but the weights on the RHS are transposed.
-
- Before: x * (BT * AT)
- ops = 16*2*32 + 1*16*32 = 1536
- After: (x * BT) * AT
- ops = 1*16*2 + 1*2*32 = 96
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([32, 2]),
- B: R.Tensor([2, 16]),
- ) -> R.Tensor([32]):
- linear_weight: R.Tensor([32, 16]) = R.matmul(A, B)
- matmul_weight: R.Tensor([16, 32]) = R.permute_dims(linear_weight)
- out: R.Tensor([32]) = R.matmul(x, matmul_weight)
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([32, 2]),
- B: R.Tensor([2, 16]),
- ) -> R.Tensor([32]):
- B_transpose = R.permute_dims(B)
- x: R.Tensor([2]) = R.matmul(x, B_transpose)
- A_transpose = R.permute_dims(A)
- x: R.Tensor([32]) = R.matmul(x, A_transpose)
- return x
-
-
-class TestRHSPermuteDimsIdentity(Base):
- """Do not treat an explicit identity permutation as a transpose.
-
- `TestRHSPermuteDims` above covers the real transpose case. Here, the
- explicit axes preserve the inner matmul's order, so reassociation must not
- insert transposes for its operands.
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([2]),
- A: R.Tensor([2, 1]),
- B: R.Tensor([1, 2]),
- ) -> R.Tensor([2]):
- linear_weight: R.Tensor([2, 2]) = R.matmul(A, B)
- matmul_weight: R.Tensor([2, 2]) = R.permute_dims(linear_weight,
axes=[0, 1])
- out: R.Tensor([2]) = R.matmul(x, matmul_weight)
- return out
-
- Expected = Before
-
-
-class TestRHSPermuteDimsNonMatrixAxes(Base):
- """Do not rewrite permutations that move a batch axis."""
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([4, 1, 4]),
- A: R.Tensor([4, 4, 1]),
- B: R.Tensor([4, 1, 4]),
- ) -> R.Tensor([4, 1, 4]):
- weight: R.Tensor([4, 4, 4]) = R.matmul(A, B)
- permuted: R.Tensor([4, 4, 4]) = R.permute_dims(weight, axes=[1, 0,
2])
- out: R.Tensor([4, 1, 4]) = R.matmul(x, permuted)
- return out
-
- Expected = Before
-
-
-class TestLHSPermuteDimsNonMatrixAxes(Base):
- """Apply the same batch-axis guard to the left-hand pattern."""
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- A: R.Tensor([4, 4, 1]),
- B: R.Tensor([4, 1, 4]),
- x: R.Tensor([4, 4, 1]),
- ) -> R.Tensor([4, 4, 1]):
- weight: R.Tensor([4, 4, 4]) = R.matmul(A, B)
- permuted: R.Tensor([4, 4, 4]) = R.permute_dims(weight, axes=[1, 0,
2])
- out: R.Tensor([4, 4, 1]) = R.matmul(permuted, x)
- return out
-
- Expected = Before
-
-
-lora_r_before = T.dynamic("lora_r", "int64")
-lora_r_expected = T.dynamic("lora_r", "int64")
-
-
-class TestRHSPermuteDimsDynamic(Base):
- """Prefer (x*A)*B instead of x*(A*B)
-
- Like `TestRHSPermuteDims`, but the weights on the RHS have a
- dynamic shape.
-
- Before: x * (BT * AT)
- ops = 16*lora_r*32 + 1*16*32 = 512*lora_r + 512
- After: (x * BT) * AT
- ops = 1*16*lora_r + 1*lora_r*32 = 48*lora_r
-
- 48*lora_r can be proved to be less than 512*lora_r + 512, so the After is
preferred.
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([32, lora_r_before]),
- B: R.Tensor([lora_r_before, 16]),
- ) -> R.Tensor([32]):
- linear_weight: R.Tensor([32, 16]) = R.matmul(A, B)
- matmul_weight: R.Tensor([16, 32]) = R.permute_dims(linear_weight)
- out: R.Tensor([32]) = R.matmul(x, matmul_weight)
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- x: R.Tensor([16]),
- A: R.Tensor([32, lora_r_expected]),
- B: R.Tensor([lora_r_expected, 16]),
- ) -> R.Tensor([32]):
- B_transpose = R.permute_dims(B)
- x: R.Tensor([lora_r_expected]) = R.matmul(x, B_transpose)
- A_transpose = R.permute_dims(A)
- x: R.Tensor([32]) = R.matmul(x, A_transpose)
- return x
-
-
-lora_r_before = T.dynamic("lora_r", "int64")
-batch_size_before = T.dynamic("batch_size", "int64")
-lora_r_expected = T.dynamic("lora_r", "int64")
-batch_size_expected = T.dynamic("batch_size", "int64")
-
-
-class TestRHSPermuteDimsWithDynamicBatch(Base):
- """Prefer (x*A)*B instead of x*(A*B)
-
- Like `TestRHSPermuteDims`, but both the weights on the RHS and the
- activations on the LHS have a dynamic dimension.
-
- Unlike most of the tests for this transform, the
- `tir_vars_upper_bound` attribute is required. In order to make a
- change, `AdjustMatmulOrder` must first prove that the modified
- execution order reduces the number of computations.
-
- ops_left_to_right = (batch_size + lora_r)*4096*4096
- ops_right_to_left = (4096 + 4096)*batch_size*lora_r
-
- Without an upper bound on batch_size and`lora_r`, we cannot prove which
- of these is the preferred execution order.
-
- With the upper bound, TVM can determine the preferred order using
- the following arithmetic reasoning.
-
- (batch_size + lora_r)*4096*4096 > (4096 + 4096)*batch_size*lora_r
- (batch_size + lora_r)*2048 > batch_size*lora_r
- 1/batch_size + 1/lora_r > 1/2048
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([batch_size_before, 4096]),
- A: R.Tensor([4096, lora_r_before]),
- B: R.Tensor([lora_r_before, 4096]),
- ) -> R.Tensor([batch_size_before, 4096]):
- R.func_attr(
- {
- "tir_var_upper_bound": {"lora_r": 2048, "batch_size":
2048},
- }
- )
- linear_weight: R.Tensor([4096, 4096]) = R.matmul(A, B)
- matmul_weight: R.Tensor([4096, 4096]) =
R.permute_dims(linear_weight)
- out: R.Tensor([batch_size_before, 4096]) = R.matmul(x,
matmul_weight)
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- x: R.Tensor([batch_size_expected, 4096]),
- A: R.Tensor([4096, lora_r_expected]),
- B: R.Tensor([lora_r_expected, 4096]),
- ) -> R.Tensor([batch_size_expected, 4096]):
- R.func_attr(
- {
- "tir_var_upper_bound": {"lora_r": 2048, "batch_size":
2048},
- }
- )
- B_transpose = R.permute_dims(B)
- x: R.Tensor([batch_size_expected, lora_r_expected]) = R.matmul(x,
B_transpose)
- A_transpose = R.permute_dims(A)
- x: R.Tensor([batch_size_expected, 4096]) = R.matmul(x, A_transpose)
- return x
-
-
-lora_r_before = T.dynamic("lora_r", "int64")
-lora_r_expected = T.dynamic("lora_r", "int64")
-
-
-class TestRHSPermuteDimsDynamicWithSquareMatrix(Base):
- """Prefer (x*A)*B instead of x*(A*B)
-
- Like `TestRHSPermuteDims`, but the weights on the RHS have a
- dynamic shape.
-
- Before: x * (BT * AT)
- ops = 32*lora_r*32 + 1*32*32 = 1024*lora_r + 1024
- After: (x * BT) * AT
- ops = 1*32*lora_r + 1*lora_r*32 = 64*lora_r
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- x: R.Tensor([32]),
- A: R.Tensor([32, lora_r_before]),
- B: R.Tensor([lora_r_before, 32]),
- ) -> R.Tensor([32]):
- linear_weight: R.Tensor([32, 32]) = R.matmul(A, B)
- matmul_weight: R.Tensor([32, 32]) = R.permute_dims(linear_weight)
- out: R.Tensor([32]) = R.matmul(x, matmul_weight)
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- x: R.Tensor([32]),
- A: R.Tensor([32, lora_r_expected]),
- B: R.Tensor([lora_r_expected, 32]),
- ) -> R.Tensor([32]):
- B_transpose = R.permute_dims(B)
- x: R.Tensor([lora_r_expected]) = R.matmul(x, B_transpose)
- A_transpose = R.permute_dims(A)
- x: R.Tensor([32]) = R.matmul(x, A_transpose)
- return x
-
-
-class TestBatchedBroadcastPreferLHSFirst(Base):
- """Use broadcasted batch prefix per matmul, not independent prefix
products.
-
- Example with broadcast batch axes: A:[2,1,1], B:[2,1,2], C:[2,2,3].
-
- LHS first: (A * B) * C
- ops = 2*1*1*2 + 2*1*2*3 = 16
- RHS first: A * (B * C)
- ops = 2*1*2*3 + 2*1*1*3 = 18
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- A: R.Tensor([2, 1, 1]),
- B: R.Tensor([2, 1, 2]),
- C: R.Tensor([2, 2, 3]),
- ) -> R.Tensor([2, 1, 3]):
- out: R.Tensor([2, 1, 3]) = R.matmul(A, R.matmul(B, C))
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- A: R.Tensor([2, 1, 1]),
- B: R.Tensor([2, 1, 2]),
- C: R.Tensor([2, 2, 3]),
- ) -> R.Tensor([2, 1, 3]):
- temp: R.Tensor([2, 1, 2]) = R.matmul(A, B)
- out: R.Tensor([2, 1, 3]) = R.matmul(temp, C)
- return out
-
-
-class TestBatchedSharedPrefixPreferLHSFirst(Base):
- """All operands share a nontrivial batch prefix [2, 3].
-
- Shapes: A:[2,3,4,5], B:[2,3,5,6], C:[2,3,6,7]
-
- LHS first:
- ops = 6*4*5*6 + 6*4*6*7 = 1728
- RHS first:
- ops = 6*5*6*7 + 6*4*5*7 = 2100
- """
-
- @I.ir_module
- class Before:
- @R.function
- def main(
- A: R.Tensor([2, 3, 4, 5]),
- B: R.Tensor([2, 3, 5, 6]),
- C: R.Tensor([2, 3, 6, 7]),
- ) -> R.Tensor([2, 3, 4, 7]):
- out: R.Tensor([2, 3, 4, 7]) = R.matmul(A, R.matmul(B, C))
- return out
-
- @I.ir_module
- class Expected:
- @R.function
- def main(
- A: R.Tensor([2, 3, 4, 5]),
- B: R.Tensor([2, 3, 5, 6]),
- C: R.Tensor([2, 3, 6, 7]),
- ) -> R.Tensor([2, 3, 4, 7]):
- temp: R.Tensor([2, 3, 4, 6]) = R.matmul(A, B)
- out: R.Tensor([2, 3, 4, 7]) = R.matmul(temp, C)
- return out
-
-
[email protected]("lhs_first", [False, True])
[email protected]("compile_time,transpose", [(False, False), (False,
True), (True, False)])
[email protected]("out_dtype", [None, "float32"])
-def test_preserve_outer_out_dtype(lhs_first, compile_time, transpose,
out_dtype):
- # Compile-time grouping should override the cheaper evaluation order.
- shapes = [(2, 20), (20, 2), (2, 10)]
- if lhs_first == compile_time:
- shapes = [(10, 2), (2, 20), (20, 2)]
- inner_indices = (1, 2) if lhs_first else (0, 1)
- if transpose:
- for i in inner_indices:
- shapes[i] = shapes[i][::-1]
-
- def build_module(expected):
- bb = relax.BlockBuilder()
- a, b, c = [
- relax.Var(name, relax.TensorType(shape, "float16"))
- for name, shape in zip("abc", shapes)
- ]
- params = [c, a, b] if lhs_first else [a, b, c]
- attrs = {"num_input": 1} if compile_time else None
- with bb.function("main", params, attrs=attrs):
- with bb.dataflow():
- operands = [a, b, c]
- if expected and transpose:
- for i in inner_indices:
- operands[i] = relax.op.permute_dims(operands[i])
- a, b, c = operands
- if expected:
- if lhs_first:
- out = relax.op.matmul(relax.op.matmul(a, b), c,
out_dtype=out_dtype)
- else:
- out = relax.op.matmul(a, relax.op.matmul(b, c),
out_dtype=out_dtype)
- else:
- x, y = (b, c) if lhs_first else (a, b)
- inner = bb.emit(relax.op.matmul(y, x) if transpose else
relax.op.matmul(x, y))
- if transpose:
- inner = bb.emit(relax.op.permute_dims(inner))
- out = (
- relax.op.matmul(a, inner, out_dtype=out_dtype)
- if lhs_first
- else relax.op.matmul(inner, c, out_dtype=out_dtype)
- )
- output = bb.emit_output(out)
- bb.emit_func_output(output)
- return bb.finalize()
-
- before = build_module(expected=False)
- expected = build_module(expected=True)
- transform = relax.transform.AdjustMatmulOrder()
- after = transform(before)
- assert after["main"].ret_ty.dtype == (out_dtype or "float16")
- tvm.ir.assert_structural_equal(after, expected)
- tvm.ir.assert_structural_equal(transform(after), after)
-
-
-class TestAdjustMatmulOrderAttentionBlock:
- """AdjustMatmulOrder preserves numerics on a batched attention block.
-
- Covers ND `permute_dims` (swap last two axes) inside `matmul(q, kt)`,
- regression for issue #19576.
- """
-
- def _build_attention_module(self, batch, seq, dim):
- """Minimal batched attention block exercising ND permute_dims +
matmul."""
- bb = relax.BlockBuilder()
- x = relax.Var("x", relax.TensorType((batch, seq, dim), "float32"))
- wq = relax.Var("wq", relax.TensorType((dim, dim), "float32"))
- wk = relax.Var("wk", relax.TensorType((dim, dim), "float32"))
- wv = relax.Var("wv", relax.TensorType((dim, dim), "float32"))
- wo = relax.Var("wo", relax.TensorType((dim, dim), "float32"))
- with bb.function("main", [x, wq, wk, wv, wo]):
- with bb.dataflow():
- q = bb.emit(relax.op.matmul(x, wq))
- k = bb.emit(relax.op.matmul(x, wk))
- v = bb.emit(relax.op.matmul(x, wv))
- kt = bb.emit(relax.op.permute_dims(k, axes=[0, 2, 1]))
- scores = bb.emit(relax.op.matmul(q, kt))
- scale = bb.emit(relax.const(1.0 / np.sqrt(dim), "float32"))
- scores = bb.emit(relax.op.multiply(scores, scale))
- attn = bb.emit(relax.op.nn.softmax(scores, axis=-1))
- out = bb.emit(relax.op.matmul(attn, v))
- proj = bb.emit_output(relax.op.matmul(out, wo))
- bb.emit_func_output(proj)
- return bb.finalize()
-
- def _run_relax_main(self, mod, inputs):
- exe = relax.build(mod, target="llvm")
- vm = relax.VirtualMachine(exe, device=tvm.cpu())
- args = [tvm.runtime.tensor(arr, device=tvm.cpu()) for arr in inputs]
- return vm["main"](*args).numpy()
-
- def _torch_attention_ref(self, x_np, w_np, dim):
- x = torch.from_numpy(x_np)
- w = torch.from_numpy(w_np)
- with torch.no_grad():
- q = torch.matmul(x, w)
- k = torch.matmul(x, w)
- v = torch.matmul(x, w)
- scores = torch.matmul(q, k.transpose(-2, -1))
- scores = scores * (1.0 / np.sqrt(dim))
- attn = torch.nn.functional.softmax(scores, dim=-1)
- out = torch.matmul(attn, v)
- out = torch.matmul(out, w)
- return out.detach().numpy()
-
- @pytest.mark.parametrize("batch,seq,dim", [(2, 16, 64)])
- def test_attention_block_numerics(self, batch, seq, dim):
- mod = self._build_attention_module(batch, seq, dim)
- mod_opt = relax.transform.AdjustMatmulOrder()(mod)
-
- x_np = np.random.randn(batch, seq, dim).astype("float32")
- w_np = np.random.randn(dim, dim).astype("float32")
- inputs = [x_np, w_np, w_np, w_np, w_np]
-
- ref = self._torch_attention_ref(x_np, w_np, dim)
- out_before = self._run_relax_main(mod, inputs)
- out_after = self._run_relax_main(mod_opt, inputs)
-
- tvm.testing.assert_allclose(out_before, ref, rtol=1e-3, atol=1e-3)
- tvm.testing.assert_allclose(out_after, ref, rtol=1e-3, atol=1e-3)
- tvm.testing.assert_allclose(out_before, out_after, rtol=1e-5,
atol=1e-5)
-
- def test_identity_permute_dims_numerics(self):
- bb = relax.BlockBuilder()
- x = relax.Var("x", relax.TensorType((2,), "int32"))
- A = relax.Var("A", relax.TensorType((2, 1), "int32"))
- B = relax.Var("B", relax.TensorType((1, 2), "int32"))
- with bb.function("main", [x, A, B]):
- with bb.dataflow():
- linear_weight = bb.emit(relax.op.matmul(A, B))
- identity_weight = bb.emit(relax.op.permute_dims(linear_weight,
axes=[0, 1]))
- out = bb.emit_output(relax.op.matmul(x, identity_weight))
- bb.emit_func_output(out)
- mod = bb.finalize()
- mod_opt = relax.transform.AdjustMatmulOrder()(mod)
-
- inputs = [
- np.array([-1, 3], dtype="int32"),
- np.array([[2], [-4]], dtype="int32"),
- np.array([[1, -3]], dtype="int32"),
- ]
- expected = np.array([-14, 42], dtype="int32")
-
- out_before = self._run_relax_main(mod, inputs)
- out_after = self._run_relax_main(mod_opt, inputs)
-
- np.testing.assert_array_equal(out_before, expected)
- np.testing.assert_array_equal(out_after, expected)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/sym/test_sym_analyzer_object.py
b/tests/python/sym/test_sym_analyzer_object.py
index d74fd81ecd..dd37a3ebd1 100644
--- a/tests/python/sym/test_sym_analyzer_object.py
+++ b/tests/python/sym/test_sym_analyzer_object.py
@@ -21,7 +21,7 @@ import tvm
import tvm.testing
from tvm import tirx
from tvm.runtime import Object
-from tvm.sym.analyzer import CompareResult, Extension
+from tvm.sym.analyzer import CompareResult
def test_analyzer_is_ffi_object_with_persistent_state():
@@ -146,7 +146,7 @@ def test_analyzer_object_set_maximum_rewrite_steps():
relaxed.rewrite_simplify(expr)
-def test_analyzer_object_try_compare_transitive():
+def test_analyzer_object_try_compare_direct():
analyzer = tvm.sym.Analyzer()
x = tirx.Var("x", "int32")
y = tirx.Var("y", "int32")
@@ -158,21 +158,8 @@ def test_analyzer_object_try_compare_transitive():
with analyzer.constraint_scope(y < z):
# Direct known comparison.
assert analyzer.try_compare(x, y) == CompareResult.LT
- # Transitive chain x < y < z is found only when propagation is
enabled.
- assert analyzer.try_compare(x, z) == CompareResult.LT
- assert analyzer.try_compare(x, z, propagate_inequalities=False) ==
CompareResult.UNKNOWN
-
-
-def test_analyzer_object_enabled_extensions_round_trip():
- analyzer = tvm.sym.Analyzer()
-
- assert analyzer.enabled_extensions == Extension.NoExtensions
-
- analyzer.enabled_extensions = Extension.ComparisonOfProductAndSum
- assert analyzer.enabled_extensions == Extension.ComparisonOfProductAndSum
-
- analyzer.enabled_extensions = Extension.NoExtensions
- assert analyzer.enabled_extensions == Extension.NoExtensions
+ # No direct comparison connects x and z.
+ assert analyzer.try_compare(x, z) == CompareResult.UNKNOWN
def test_analyzer_object_rewrite_simplify_stats():
@@ -238,7 +225,6 @@ def test_analyzer_object_clone_copies_every_sub_analyzer():
analyzer.update(x, tvm.sym.ModularSet(4, 0))
analyzer.bind(w, tirx.const(4, "int64"))
analyzer.update(v, tvm.sym.IntervalSet(2, 9))
- analyzer.enabled_extensions = Extension.ComparisonOfProductAndSum
clone = analyzer.clone()
@@ -246,7 +232,6 @@ def test_analyzer_object_clone_copies_every_sub_analyzer():
assert clone.modular_set(x).coeff == 4
tvm.ir.assert_structural_equal(clone.simplify(w + 1), tirx.const(5,
"int64"))
assert clone.int_set(v).max_value.value == 9
- assert clone.enabled_extensions == Extension.ComparisonOfProductAndSum
assert clone.try_compare(x, tirx.const(0, "int64")) == CompareResult.GE
t = tirx.Var("t", "int64")
@@ -254,7 +239,6 @@ def test_analyzer_object_clone_copies_every_sub_analyzer():
clone.update(v, tvm.sym.IntervalSet(0, 3), override=True)
clone.bind(w, tirx.const(8, "int64"), allow_override=True)
clone.bind(t, tvm.ir.Range(0, 4))
- clone.enabled_extensions = Extension.NoExtensions
assert analyzer.modular_set(x).coeff == 4
assert clone.modular_set(x).coeff == 8
@@ -262,8 +246,6 @@ def test_analyzer_object_clone_copies_every_sub_analyzer():
assert clone.int_set(v).max_value.value == 3
tvm.ir.assert_structural_equal(analyzer.simplify(w + 1), tirx.const(5,
"int64"))
tvm.ir.assert_structural_equal(clone.simplify(w + 1), tirx.const(9,
"int64"))
- assert analyzer.enabled_extensions == Extension.ComparisonOfProductAndSum
- assert clone.enabled_extensions == Extension.NoExtensions
assert clone.try_compare(t, tirx.const(0, "int64")) == CompareResult.GE
assert analyzer.try_compare(t, tirx.const(0, "int64")) ==
CompareResult.UNKNOWN
diff --git a/tests/python/sym/test_sym_rewrite_simplify.py
b/tests/python/sym/test_sym_rewrite_simplify.py
index b6efb06e02..ef4c2f5bf4 100644
--- a/tests/python/sym/test_sym_rewrite_simplify.py
+++ b/tests/python/sym/test_sym_rewrite_simplify.py
@@ -68,11 +68,8 @@ class TestCase:
class BaseCompare:
- extensions = tvm.sym.Extension.NoExtensions
-
def test_simplify(self, test_case):
analyzer = tvm.sym.Analyzer()
- analyzer.enabled_extensions = self.extensions
if inspect.isclass(test_case.expected) and
issubclass(test_case.expected, Exception):
with pytest.raises(test_case.expected):
@@ -1172,39 +1169,6 @@ class TestComparisons(BaseCompare):
)
-class TestComparisonOfProductAndSum(BaseCompare):
- extensions = tvm.sym.Extension.ComparisonOfProductAndSum
-
- x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"),
tvm.tirx.Var("z", "int32")
-
- test_case = tvm.testing.parameter(
- # Special inequality cases
- TestCase(
- x * y < (x + y) * 2048,
- tvm.tirx.const(1, "bool"),
- [x > 0, y > 0, x < 2048],
- ),
- TestCase(
- x * y < (x + y) * 2048,
- tvm.tirx.const(1, "bool"),
- [x > 0, y > 0, x < 4096, y < 4096],
- ),
- TestCase(
- # Both sides are divisible by 8192
- x * y * 8192 < (y + x) * 16777216,
- tvm.tirx.const(1, "bool"),
- [x > 0, y > 0, x < 4096, y < 4096],
- ),
- TestCase(
- # The two sides have co-prime factors, but the bounds are
- # still sufficient to prove the inequality.
- x * y * 59 < (y + x) * 176128,
- tvm.tirx.const(1, "bool"),
- [x > 0, y > 0, x < 4096, y < 4096],
- ),
- )
-
-
class TestLogical(BaseCompare):
x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"),
tvm.tirx.Var("z", "int32")
diff --git a/tests/python/tirx-transform/test_tir_transform_simplify.py
b/tests/python/tirx-transform/test_tir_transform_simplify.py
index d94384377b..83737ad91e 100644
--- a/tests/python/tirx-transform/test_tir_transform_simplify.py
+++ b/tests/python/tirx-transform/test_tir_transform_simplify.py
@@ -614,31 +614,6 @@ def test_left_shift_by_negative_value():
tvm.ir.assert_structural_equal(after, expected)
-def test_suppress_transitively_provable_condition():
- @T.function(private=True)
- def before(A: T.Tensor(1, "bool"), i: T.int32, j: T.int32, k: T.int32):
- if i < j and j < k:
- A[0] = i < k
-
- expected = before
-
- after = _apply_simplify(before)
- tvm.ir.assert_structural_equal(after, expected)
-
-
-def test_suppress_rewrite_as_and_of_ors():
- """Only rewrite into AND of OR when allowed"""
-
- @T.function(private=True)
- def before(A: T.Tensor(3, "bool")):
- T.evaluate(A[0] or (A[1] and A[2]))
-
- expected = before
-
- after = _apply_simplify(before)
- tvm.ir.assert_structural_equal(after, expected)
-
-
def test_conditional_floor_mod():
"""A regression test for negative floormod denominator
