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 c7ccd26f7c [Refactor][IR] Move ExprDeepEqual into shared primitive
expressions (#20356)
c7ccd26f7c is described below
commit c7ccd26f7c022eb82dc335a6d820d23732e27f60
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Sep 16 13:48:25 2026 -0400
[Refactor][IR] Move ExprDeepEqual into shared primitive expressions (#20356)
Move `ExprDeepEqual` from TIRX analysis to `tvm::prim`, using the shared
IR expression functor. TensorLoad sources are compared directly by
identity, allowing shared expression sources without a dialect-specific
cast.
Update C++ consumers and expose the Python/FFI entry point as
`tvm.ir.prim.expr_deep_equal` / `ir.prim.expr_deep_equal`. Existing
variable identity, operand order, types, Call attributes, and nested
expression comparisons are preserved.
---
include/tvm/arith/analyzer.h | 2 +-
include/tvm/ir/prim/expr.h | 21 +++++++++++
include/tvm/tirx/analysis.h | 20 -----------
include/tvm/topi/detail/constant_utils.h | 3 +-
python/tvm/ir/prim/__init__.py | 41 ++++++++++++++++++++++
.../frontend/torch/exported_program_translator.py | 2 +-
python/tvm/tirx/analysis/analysis.py | 37 -------------------
src/arith/canonical_simplify.cc | 4 +--
src/arith/conjunctive_normal_form.cc | 4 +--
src/arith/const_int_bound.cc | 3 +-
src/arith/int_set.cc | 4 +--
src/arith/iter_affine_map.cc | 6 ++--
src/arith/pattern_match.h | 3 +-
src/arith/rewrite_simplify.cc | 6 ++--
src/arith/transitive_comparison_analyzer.cc | 8 ++---
src/arith/z3_prover.cc | 2 +-
src/backend/vulkan/codegen/codegen_spirv.cc | 2 +-
src/backend/vulkan/codegen/codegen_spirv.h | 2 +-
src/{tirx/analysis => ir/prim}/deep_equal.cc | 33 ++++++-----------
src/relax/backend/vm/vm_shape_lower.cc | 3 +-
src/s_tir/schedule/analysis/reducer.cc | 3 +-
.../schedule/primitive/cache_index_helpers.cc | 7 ++--
src/s_tir/schedule/primitive/cache_index_helpers.h | 9 ++---
src/s_tir/schedule/primitive/cache_read_write.cc | 3 +-
src/s_tir/schedule/primitive/compute_inline.cc | 3 +-
.../schedule/primitive/layout_transformation.cc | 5 +--
src/s_tir/transform/thread_storage_sync.cc | 2 +-
src/target/llvm/codegen_llvm.h | 2 +-
src/target/source/codegen_c.h | 2 +-
src/tirx/analysis/var_use_def_analysis.h | 3 +-
src/tirx/analysis/verify_ssa.cc | 2 +-
src/tirx/ir/buffer.cc | 2 +-
src/tirx/ir/specialize.cc | 3 +-
src/tirx/script/printer/block.cc | 2 +-
src/tirx/script/printer/stmt.cc | 2 +-
src/tirx/transform/common_subexpr_elim.cc | 15 ++++----
src/tirx/transform/stmt_simplify.cc | 4 +--
src/tirx/transform/storage_rewrite.cc | 4 +--
src/tirx/transform/vectorize_loop.cc | 2 +-
tests/cpp/expr_test.cc | 18 ++++++++++
tests/cpp/pattern_match_test.cc | 16 ++++-----
.../test_ir_prim_expr_deep_equal.py} | 6 ++--
.../s_tir/schedule/test_tir_schedule_analysis.py | 2 +-
tests/python/tirx-base/test_tir_constructor.py | 2 +-
44 files changed, 174 insertions(+), 151 deletions(-)
diff --git a/include/tvm/arith/analyzer.h b/include/tvm/arith/analyzer.h
index 47be76e898..9b2de84d3a 100644
--- a/include/tvm/arith/analyzer.h
+++ b/include/tvm/arith/analyzer.h
@@ -496,7 +496,7 @@ class TransitiveComparisonAnalyzer {
* 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 `ExprDeepEqual`, but also allows for constant
+ * expressions using `prim::ExprDeepEqual`, but also allows for constant
* offsets on either side of the inequality.
*
* \return The most specific result that can be proven about the
diff --git a/include/tvm/ir/prim/expr.h b/include/tvm/ir/prim/expr.h
index b5569d31da..99740c8288 100644
--- a/include/tvm/ir/prim/expr.h
+++ b/include/tvm/ir/prim/expr.h
@@ -580,6 +580,27 @@ inline std::unordered_map<K, V> as_unordered_map(const
ffi::Map<K, V>& dmap) {
}
return ret;
}
+
+/*!
+ * \brief Compare two expressions recursively and check if they are equal
+ * to each other without var remapping.
+ *
+ * This function does not remap variable bindings, it will not
+ * return true for (let x = 1 in x + 1) vs (let y = 1 in y + 1), unless
x.same_as(y).
+ *
+ * Use StructuralEqual for such cases.
+ *
+ * Due to the restriction of not remapping variables, this function can run
+ * faster than StructuralEqual and can be used as a utility function during
arithmetic
+ * simplifications.
+ *
+ * \sa StructuralEqual
+ */
+struct ExprDeepEqual {
+ public:
+ TVM_DLL bool operator()(const PrimExpr& lhs, const PrimExpr& rhs) const;
+};
+
} // namespace prim
namespace ffi {
diff --git a/include/tvm/tirx/analysis.h b/include/tvm/tirx/analysis.h
index b0e32b65e3..1379046a21 100644
--- a/include/tvm/tirx/analysis.h
+++ b/include/tvm/tirx/analysis.h
@@ -39,26 +39,6 @@ namespace tvm {
namespace tirx {
-/*!
- * \brief Compare two expressions recursively and check if they are equal
- * to each other without var remapping.
- *
- * This function does not remap variable bindings, it will not
- * return true for (let x = 1 in x + 1) vs (let y = 1 in y + 1), unless
x.same_as(y).
- *
- * Use StructuralEqual for such cases.
- *
- * Due to the restriction of not remapping variables, this function can run
- * faster than StructuralEqual and can be used as a utility function during
arithmetic
- * simplifications.
- *
- * \sa StructuralEqual
- */
-struct ExprDeepEqual {
- public:
- TVM_DLL bool operator()(const PrimExpr& lhs, const PrimExpr& rhs) const;
-};
-
/*!
* \brief Visit the PrimFuncs in the IRModule
* \tparam FLambda The type of the PrimFunc visitor
diff --git a/include/tvm/topi/detail/constant_utils.h
b/include/tvm/topi/detail/constant_utils.h
index 1a69c918b3..54e3a0c21e 100644
--- a/include/tvm/topi/detail/constant_utils.h
+++ b/include/tvm/topi/detail/constant_utils.h
@@ -28,7 +28,6 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/logging.h>
#include <tvm/te/operation.h>
-#include <tvm/tirx/analysis.h>
#include <string>
#include <vector>
@@ -130,7 +129,7 @@ inline std::vector<int64_t>
GetConstInt64Values(ffi::Array<PrimExpr> exprs,
* \return result True if both expressions are equal, else false
*/
inline bool EqualCheck(PrimExpr lhs, PrimExpr rhs) {
- tvm::tirx::ExprDeepEqual expr_equal;
+ tvm::prim::ExprDeepEqual expr_equal;
bool result = expr_equal(lhs, rhs);
if (!result) {
PrimExpr t = tvm::arith::Analyzer()->Simplify(lhs - rhs);
diff --git a/python/tvm/ir/prim/__init__.py b/python/tvm/ir/prim/__init__.py
index 72923c0b48..f554ec655d 100644
--- a/python/tvm/ir/prim/__init__.py
+++ b/python/tvm/ir/prim/__init__.py
@@ -16,6 +16,47 @@
# under the License.
"""Primitive expression nodes shared by TVM IR dialects."""
+from ..expr import Expr
+from . import _ffi_api
+
+
+def expr_deep_equal(lhs: Expr, rhs: Expr) -> bool:
+ """Deeply compare two nested expressions.
+
+ Parameters
+ ----------
+ lhs : Expr
+ The left operand.
+
+ rhs : Expr
+ The right operand.
+
+ Returns
+ -------
+ result : bool
+ The comparison result
+
+ Note
+ ----
+
+ This function does not remap variable bindings, it will not
+ return true for (let x = 1 in x + 1) vs (let y = 1 in y + 1), unless
x.same_as(y).
+ Use py:func:`tvm_ffi.structural_equal` to handle structural variable
remapping.
+
+ Due to the restriction of not remapping variables, this function can run
+ faster than StructuralEqual and can be used as a utility function during
arithmetic
+ simplifications.
+
+ Always consider py:func:`tvm_ffi.structural_equal` first, which handles
+ the structural remapping.
+
+ See Also
+ --------
+ tvm_ffi.structural_equal
+ """
+ return _ffi_api.expr_deep_equal(lhs, rhs) # type: ignore
+
+
_EXPR_NAMES = {
"StringImm",
"Cast",
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py
b/python/tvm/relax/frontend/torch/exported_program_translator.py
index dbf4993251..86c936723b 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -1171,7 +1171,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
actual_dim = dim if dim >= 0 else len(in_shape) + dim
dim_expr = in_shape[actual_dim]
if tvm.ir.is_prim_expr(dim_expr):
- if tvm.tirx.analysis.expr_deep_equal(end_val, dim_expr):
+ if tvm.ir.prim.expr_deep_equal(end_val, dim_expr):
return x
axes = [dim]
diff --git a/python/tvm/tirx/analysis/analysis.py
b/python/tvm/tirx/analysis/analysis.py
index c1da4f57df..f48710a683 100644
--- a/python/tvm/tirx/analysis/analysis.py
+++ b/python/tvm/tirx/analysis/analysis.py
@@ -27,43 +27,6 @@ from ..function import PrimFunc
from . import _ffi_api
-def expr_deep_equal(lhs: Expr, rhs: Expr) -> bool:
- """Deeply compare two nested expressions.
-
- Parameters
- ----------
- lhs : Expr
- The left operand.
-
- rhs : Expr
- The right operand.
-
- Returns
- -------
- result : bool
- The comparison result
-
- Note
- ----
-
- This function does not remap variable bindings, it will not
- return true for (let x = 1 in x + 1) vs (let y = 1 in y + 1), unless
x.same_as(y).
- Use py:func:`tvm_ffi.structural_equal` to handle structural variable
remapping.
-
- Due to the restriction of not remapping variables, this function can run
- faster than StructuralEqual and can be used as a utility function during
arithmetic
- simplifications.
-
- Always consider py:func:`tvm_ffi.structural_equal` first, which handles
- the structural remapping.
-
- See Also
- --------
- tvm_ffi.structural_equal
- """
- return _ffi_api.expr_deep_equal(lhs, rhs) # type: ignore
-
-
def verify_ssa(func: PrimFunc) -> bool:
"""Verify if the func is in SSA form.
diff --git a/src/arith/canonical_simplify.cc b/src/arith/canonical_simplify.cc
index ba6d5a7a0d..f0bea1a5a9 100644
--- a/src/arith/canonical_simplify.cc
+++ b/src/arith/canonical_simplify.cc
@@ -25,7 +25,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/expected.h>
#include <tvm/ir/cow.h>
-#include <tvm/tirx/analysis.h>
+#include <tvm/ir/prim/expr.h>
#include <tvm/tirx/op.h>
#include "const_fold.h"
@@ -239,7 +239,7 @@ class SplitExpr : public PrimExpr {
inline bool SplitExprNode::IndexEqual(const SplitExpr& other) const {
if (index.same_as(other->index)) return true;
- return tirx::ExprDeepEqual()(index, other->index);
+ return prim::ExprDeepEqual()(index, other->index);
}
inline bool SplitExprNode::DivModeCompatibleTo(DivMode mode) const {
diff --git a/src/arith/conjunctive_normal_form.cc
b/src/arith/conjunctive_normal_form.cc
index e52446ea4e..37c9a4aa7a 100644
--- a/src/arith/conjunctive_normal_form.cc
+++ b/src/arith/conjunctive_normal_form.cc
@@ -250,7 +250,7 @@ void AndOfOrs::TrySimplifyOr(Key* a_ptr, Key* b_ptr,
AnalyzerObj* analyzer) {
Key& b = *b_ptr;
PrimExpr joint = GetExpr(a) || GetExpr(b);
PrimExpr simplified = analyzer->rewrite_simplify(joint);
- if (!ExprDeepEqual()(simplified, joint)) {
+ if (!prim::ExprDeepEqual()(simplified, joint)) {
if (auto* simplified_or = simplified.as<prim::OrNode>()) {
a = GetKey(simplified_or->a);
b = GetKey(simplified_or->b);
@@ -266,7 +266,7 @@ void AndOfOrs::TrySimplifyAnd(Key* a_ptr, Key* b_ptr,
AnalyzerObj* analyzer) {
Key& b = *b_ptr;
PrimExpr joint = GetExpr(a) && GetExpr(b);
PrimExpr simplified = analyzer->rewrite_simplify(joint);
- if (!ExprDeepEqual()(simplified, joint)) {
+ if (!prim::ExprDeepEqual()(simplified, joint)) {
if (auto* simplified_and = simplified.as<prim::AndNode>()) {
a = GetKey(simplified_and->a);
b = GetKey(simplified_and->b);
diff --git a/src/arith/const_int_bound.cc b/src/arith/const_int_bound.cc
index 78c6d63c37..678603189d 100644
--- a/src/arith/const_int_bound.cc
+++ b/src/arith/const_int_bound.cc
@@ -27,6 +27,7 @@
#include <tvm/ir/expr_functor.h>
#include <tvm/ir/op.h>
#include <tvm/ir/prim/builtin.h>
+#include <tvm/ir/prim/expr.h>
#include <tvm/tirx/builtin.h>
#include <algorithm>
@@ -158,7 +159,7 @@ class ConstIntBoundAnalyzer::Impl
Entry Dispatch(const Expr& expr) final {
PrimExpr prim_expr = expr.as_or_throw<PrimExpr>();
Entry res = ExprFunctor::Dispatch(expr);
- tirx::ExprDeepEqual equal;
+ prim::ExprDeepEqual equal;
// a linear search over additional info
// assume we won't have a lot of conditions
for (const BoundInfo& info : additional_info_) {
diff --git a/src/arith/int_set.cc b/src/arith/int_set.cc
index cfda30a96c..3fe11eb61d 100644
--- a/src/arith/int_set.cc
+++ b/src/arith/int_set.cc
@@ -709,12 +709,12 @@ void IntSetAnalyzer::Impl::Update(const Var& var, const
IntSet& info, bool can_o
if (it != dom_map_.end()) {
const IntSet& old_info = (*it).second;
- TVM_FFI_ICHECK(ExprDeepEqual()(old_info.min(), info.min()))
+ TVM_FFI_ICHECK(prim::ExprDeepEqual()(old_info.min(), info.min()))
<< "Trying to update var \'" << var << "\'"
<< " with a different minimum value: "
<< "original=" << old_info.min() << ", new=" << info.min();
- TVM_FFI_ICHECK(ExprDeepEqual()(old_info.max(), info.max()))
+ TVM_FFI_ICHECK(prim::ExprDeepEqual()(old_info.max(), info.max()))
<< "Trying to update var \'" << var << "\'"
<< " with a different maximum value: "
<< "original=" << old_info.max() << ", new=" << info.max();
diff --git a/src/arith/iter_affine_map.cc b/src/arith/iter_affine_map.cc
index aac16ec945..16a14412df 100644
--- a/src/arith/iter_affine_map.cc
+++ b/src/arith/iter_affine_map.cc
@@ -433,7 +433,7 @@ class IterMapRewriter : public tvm::ExprMutator {
static bool IterSplitEqual(const IterSplitExpr& lhs, const IterSplitExpr&
rhs,
bool check_scale = true) {
- tirx::ExprDeepEqual equal;
+ prim::ExprDeepEqual equal;
if (!lhs->source.same_as(rhs->source)) return false;
if (!equal(lhs->lower_factor, rhs->lower_factor)) return false;
if (check_scale && !equal(lhs->scale, rhs->scale)) return false;
@@ -443,7 +443,7 @@ class IterMapRewriter : public tvm::ExprMutator {
struct IterSumEqual {
bool operator()(const IterSumExpr& lhs, const IterSumExpr& rhs) const {
- tirx::ExprDeepEqual equal;
+ prim::ExprDeepEqual equal;
if (lhs->args.size() != rhs->args.size()) return false;
if (!equal(lhs->base, rhs->base)) return false;
for (size_t i = 0; i < lhs->args.size(); ++i) {
@@ -1251,7 +1251,7 @@ class IterMapRewriter : public tvm::ExprMutator {
PrimExpr SplitFloorModConst(IterSplitExpr lhs, PrimExpr base, PrimExpr rhs);
static void AddToLhs(IterSumExprNode* lhs, IterSplitExpr rhs, int sign) {
- tirx::ExprDeepEqual equal;
+ prim::ExprDeepEqual equal;
for (size_t i = 0; i < lhs->args.size(); ++i) {
IterSplitExpr lvalue = lhs->args[i];
if (lvalue->source.same_as(rhs->source) && equal(lvalue->lower_factor,
rhs->lower_factor) &&
diff --git a/src/arith/pattern_match.h b/src/arith/pattern_match.h
index 10a93a489e..fa105b158e 100644
--- a/src/arith/pattern_match.h
+++ b/src/arith/pattern_match.h
@@ -68,7 +68,6 @@
#include <tvm/ffi/cast.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/analysis.h>
#include <tvm/tirx/builtin.h>
#include <cmath>
@@ -160,7 +159,7 @@ class PEqualChecker<PrimExpr> {
public:
bool operator()(const PrimExpr& lhs, const PrimExpr& rhs) const {
if (lhs.same_as(rhs)) return true;
- return tirx::ExprDeepEqual()(lhs, rhs);
+ return prim::ExprDeepEqual()(lhs, rhs);
}
};
diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc
index 97b50d2d50..c6054a81ad 100644
--- a/src/arith/rewrite_simplify.cc
+++ b/src/arith/rewrite_simplify.cc
@@ -29,6 +29,8 @@
#include <tvm/ffi/expected.h>
#include <tvm/ir/op.h>
#include <tvm/ir/prim/builtin.h>
+#include <tvm/ir/prim/expr.h>
+#include <tvm/tirx/analysis.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
@@ -387,7 +389,7 @@ void RewriteSimplifier::Impl::Update(const Var& var, const
PrimExpr& info, bool
if (!can_override) {
auto it = var_map_.find(var);
if (it != var_map_.end()) {
- TVM_FFI_ICHECK(ExprDeepEqual()(it->second, info))
+ TVM_FFI_ICHECK(prim::ExprDeepEqual()(it->second, info))
<< "Trying to update var \'" << var << "\'"
<< " with a different value: "
<< "original=" << it->second << ", new=" << info;
@@ -1751,7 +1753,7 @@ ffi::Optional<PrimExpr>
RewriteSimplifier::Impl::TryMatchLiteralConstraint(
const PrimExpr& expr) const {
PrimExpr negation = prim::Not(expr);
- ExprDeepEqual expr_equal;
+ prim::ExprDeepEqual expr_equal;
for (const auto& constraint : literal_constraints_) {
if (expr_equal(constraint, expr)) {
return MakeConst(expr.ty(), true);
diff --git a/src/arith/transitive_comparison_analyzer.cc
b/src/arith/transitive_comparison_analyzer.cc
index a0c4750633..ee96ca427d 100644
--- a/src/arith/transitive_comparison_analyzer.cc
+++ b/src/arith/transitive_comparison_analyzer.cc
@@ -48,7 +48,7 @@ class TransitiveComparisonAnalyzer::Impl {
* 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
- * `ExprDeepEqual`, but also allowing for constant offsets on either
+ * `prim::ExprDeepEqual`, but also allowing for constant offsets on either
* side of the inequality.
*
* \return The most specific result that can be proven about the
@@ -96,8 +96,8 @@ class TransitiveComparisonAnalyzer::Impl {
*
* 1. Providing efficiency, as compared to a PrimExpr. Two keys are
* equal if and only if the corresponding PrimExprs would satisfy
- * ExprDeepEqual. This allows two expressions to be checked for
- * equivalency, without requiring a call to ExprDeepEqual for
+ * prim::ExprDeepEqual. This allows two expressions to be checked for
+ * equivalency, without requiring a call to prim::ExprDeepEqual for
* each comparison.
*
* 2. Providing type-safety, as compared to using `size_t` directly.
@@ -570,7 +570,7 @@ void TransitiveComparisonAnalyzer::Impl::Bind(const Var&
var, const Range& range
bool allow_override) {
auto it = prev_bindings_.find(var);
if (it != prev_bindings_.end()) {
- ExprDeepEqual expr_equal;
+ prim::ExprDeepEqual expr_equal;
bool differs_from_previous = !expr_equal(range->min, (*it).second->min) ||
!expr_equal(range->extent,
(*it).second->extent);
if (differs_from_previous) {
diff --git a/src/arith/z3_prover.cc b/src/arith/z3_prover.cc
index 7121781a56..ba14e09727 100644
--- a/src/arith/z3_prover.cc
+++ b/src/arith/z3_prover.cc
@@ -125,7 +125,7 @@ class Z3Prover::Impl : tvm::ExprFunctor<z3::expr(const
Expr&)> {
/// @brief Memoized PrimExpr -> slot in z3_pool_. Holds no Z3 handles, so
/// its pointer-hashed bucket order cannot affect Z3 object lifetime.
- std::unordered_map<PrimExpr, size_t, StructuralHash, ExprDeepEqual> memo_;
+ std::unordered_map<PrimExpr, size_t, StructuralHash, prim::ExprDeepEqual>
memo_;
/// @brief Slots owning the memoized Z3 handles, plus a free-slot stack.
/// Handles are created and released only at fixed points of the execution
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc
b/src/backend/vulkan/codegen/codegen_spirv.cc
index 345ac548de..5c222ad117 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -501,7 +501,7 @@ spirv::Value CodeGenSPIRV::Dispatch_(const CallNode* op) {
PrimExpr index_d = op->args[1].as_or_throw<PrimExpr>();
PrimExpr index_a = op->args[3].as_or_throw<PrimExpr>();
PrimExpr index_b = op->args[5].as_or_throw<PrimExpr>();
- tvm::tirx::ExprDeepEqual expr_equal;
+ tvm::prim::ExprDeepEqual expr_equal;
PrimExpr index_c = op->args[7].as_or_throw<PrimExpr>();
bool is_equal = ((buffer_d == buffer_c) && expr_equal(index_d, index_c));
spirv::SType& fragment_type_d = fragment_info_[buffer_d].stype;
diff --git a/src/backend/vulkan/codegen/codegen_spirv.h
b/src/backend/vulkan/codegen/codegen_spirv.h
index f252f566c7..b88e5f828f 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.h
+++ b/src/backend/vulkan/codegen/codegen_spirv.h
@@ -226,7 +226,7 @@ class CodeGenSPIRV : public
tirx::ExprFunctor<spirv::Value(const Expr&)>,
arith::Analyzer analyzer_;
// deep comparison of PrimExpr
- ExprDeepEqual deep_equal_;
+ prim::ExprDeepEqual deep_equal_;
// binding of let variables. Enables duplicate var defs that map to same
value
std::unordered_map<Var, const prim::LetNode*> let_binding_;
diff --git a/src/tirx/analysis/deep_equal.cc b/src/ir/prim/deep_equal.cc
similarity index 90%
rename from src/tirx/analysis/deep_equal.cc
rename to src/ir/prim/deep_equal.cc
index 6edc3d4158..cff330402d 100644
--- a/src/tirx/analysis/deep_equal.cc
+++ b/src/ir/prim/deep_equal.cc
@@ -18,17 +18,17 @@
*/
/*!
- * \file tirx/analysis/deep_equal.cc
+ * \file ir/prim/deep_equal.cc
* \brief Deep equality checking.
*/
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/expr_functor.h>
+#include <tvm/ir/expr_functor.h>
+#include <tvm/ir/prim/expr.h>
namespace tvm {
-namespace tirx {
+namespace prim {
#define DEFINE_DEEP_EQUAL_BIN_EXPR(OpNode)
\
bool Dispatch_(const OpNode* plhs, const PrimExpr& rhs) final {
\
@@ -44,7 +44,7 @@ namespace tirx {
plhs->value == prhs->value;
\
}
-class ExprDeepEqualChecker : private ExprFunctor<bool(const Expr&, const
PrimExpr&)> {
+class ExprDeepEqualChecker : private tvm::ExprFunctor<bool(const Expr&, const
PrimExpr&)> {
public:
static bool Check(const PrimExpr& lhs, const PrimExpr& rhs) {
// quick path without constructing the object
@@ -117,15 +117,6 @@ class ExprDeepEqualChecker : private
ExprFunctor<bool(const Expr&, const PrimExp
return true;
}
- bool ArrayDeepEqual(const ffi::Array<IterVar>& lhs, const
ffi::Array<IterVar>& rhs) {
- // for iter var, we require pointer equality
- if (lhs.size() != rhs.size()) return false;
- for (size_t i = 0; i < lhs.size(); i++) {
- if (!lhs[i].same_as(rhs[i])) return true;
- }
- return true;
- }
-
bool OptionalDeepEqual(const ffi::Optional<PrimExpr>& lhs, const
ffi::Optional<PrimExpr>& rhs) {
if (lhs.same_as(rhs)) return true;
if (!lhs.has_value() && rhs.has_value()) return false;
@@ -140,11 +131,9 @@ class ExprDeepEqualChecker : private
ExprFunctor<bool(const Expr&, const PrimExp
bool Dispatch_(const TensorLoadNode* plhs, const PrimExpr& rhs) final {
const auto* prhs = rhs.as<TensorLoadNode>();
- // we run pointer comparison of the buffer
+ // we run pointer comparison of the source
return plhs->ty.as_or_throw<PrimType>() ==
prhs->ty.as_or_throw<PrimType>() &&
- plhs->source.as_or_throw<tvm::tirx::BufferVar>().same_as(
- prhs->source.as_or_throw<tvm::tirx::BufferVar>()) &&
- ArrayDeepEqual(plhs->indices, prhs->indices);
+ plhs->source.same_as(prhs->source) && ArrayDeepEqual(plhs->indices,
prhs->indices);
}
bool Dispatch_(const prim::LetNode* plhs, const PrimExpr& rhs) final {
@@ -229,10 +218,10 @@ bool ExprDeepEqual::operator()(const PrimExpr& lhs, const
PrimExpr& rhs) const {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def(
- "tirx.analysis.expr_deep_equal",
- [](const PrimExpr& lhs, const PrimExpr& rhs) { return
ExprDeepEqual()(lhs, rhs); });
+ refl::GlobalDef().def("ir.prim.expr_deep_equal", [](const PrimExpr& lhs,
const PrimExpr& rhs) {
+ return ExprDeepEqual()(lhs, rhs);
+ });
}
-} // namespace tirx
+} // namespace prim
} // namespace tvm
diff --git a/src/relax/backend/vm/vm_shape_lower.cc
b/src/relax/backend/vm/vm_shape_lower.cc
index b30c83d293..3f22764e1e 100644
--- a/src/relax/backend/vm/vm_shape_lower.cc
+++ b/src/relax/backend/vm/vm_shape_lower.cc
@@ -23,6 +23,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/prim/expr.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/backend.h>
#include <tvm/relax/expr_functor.h>
@@ -74,7 +75,7 @@ struct MatchShapeTodoItem {
/*! \brief Slot map used for shape lowering. */
using PrimExprSlotMap =
- std::unordered_map<PrimExpr, PrimExprSlot*, ffi::StructuralHash,
tirx::ExprDeepEqual>;
+ std::unordered_map<PrimExpr, PrimExprSlot*, ffi::StructuralHash,
prim::ExprDeepEqual>;
using LiveVarSet = std::unordered_set<Var, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>;
diff --git a/src/s_tir/schedule/analysis/reducer.cc
b/src/s_tir/schedule/analysis/reducer.cc
index db32a20909..4a36f541b3 100644
--- a/src/s_tir/schedule/analysis/reducer.cc
+++ b/src/s_tir/schedule/analysis/reducer.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/ir/prim/expr.h>
#include <tvm/te/operation.h>
#include "../utils.h"
@@ -674,7 +675,7 @@ bool MatchReducer(const te::CommReducer& reducer, const
ffi::Array<PrimExpr>& id
const ffi::Array<PrimExpr>& combined_values,
const ffi::Array<TensorLoad>& buf_loads,
ffi::Array<PrimExpr>* lhs,
ffi::Array<PrimExpr>* rhs) {
- ExprDeepEqual equal;
+ prim::ExprDeepEqual equal;
TVM_FFI_ICHECK_EQ(identities.size(), combined_values.size());
int n_buffers = identities.size();
for (int i = 0; i < n_buffers; ++i) {
diff --git a/src/s_tir/schedule/primitive/cache_index_helpers.cc
b/src/s_tir/schedule/primitive/cache_index_helpers.cc
index d1f7fd0ae9..d34b103be6 100644
--- a/src/s_tir/schedule/primitive/cache_index_helpers.cc
+++ b/src/s_tir/schedule/primitive/cache_index_helpers.cc
@@ -28,7 +28,7 @@
#include <tvm/arith/analyzer.h> // For the arith::Analyzer::Simplify() method
simplifying terms
#include <tvm/ffi/cast.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/analysis.h> // For the ExprDeepEqual analysis
+#include <tvm/tirx/analysis.h>
#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>
@@ -399,7 +399,7 @@ ffi::Optional<VisitInterrupt>
DirectSubexpr::Visit(ffi::AnyView expr_value) {
* \brief Decides if two terms are equal syntactically
*/
bool EqualTerms(const PrimExpr& a, const PrimExpr& b) {
- ExprDeepEqual deep_equal_;
+ prim::ExprDeepEqual deep_equal_;
return deep_equal_(a, b);
}
@@ -439,7 +439,8 @@ std::vector<std::pair<PrimExpr, size_t>>
SyntacticToSemanticComputations(
return result;
}
- support::OrderedMap<PrimExpr, std::pair<PrimExpr, size_t>,
ffi::StructuralHash, ExprDeepEqual>
+ support::OrderedMap<PrimExpr, std::pair<PrimExpr, size_t>,
ffi::StructuralHash,
+ prim::ExprDeepEqual>
norm_table;
norm_table.reserve(table.size());
diff --git a/src/s_tir/schedule/primitive/cache_index_helpers.h
b/src/s_tir/schedule/primitive/cache_index_helpers.h
index 156c244b98..2159381816 100644
--- a/src/s_tir/schedule/primitive/cache_index_helpers.h
+++ b/src/s_tir/schedule/primitive/cache_index_helpers.h
@@ -29,7 +29,6 @@
#include <tvm/ffi/extra/structural_hash.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/analysis.h> // For the ExprDeepEqual analysis
#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h> // For the class StmtExprVisitor
@@ -48,10 +47,12 @@ namespace tirx {
a number (which is the number of time that it is computed)
It is important to note that the hash used is a ffi::StructuralHash
(and not an
ffi::ObjectPtrHash) as we need to hash similarly deeply equal terms. The
comparison used is
- ExprDeepEqual, which is stricter than ffi::StructuralEqual (as it does not do
variables remapping),
- so it is compatible with ffi::StructuralHash (intended to be used with
ffi::StructuralEqual).
+ prim::ExprDeepEqual, which is stricter than ffi::StructuralEqual (as it does
not do variables
+ remapping), so it is compatible with ffi::StructuralHash (intended to be used
with
+ ffi::StructuralEqual).
*/
-using ComputationTable = support::OrderedMap<PrimExpr, size_t,
ffi::StructuralHash, ExprDeepEqual>;
+using ComputationTable =
+ support::OrderedMap<PrimExpr, size_t, ffi::StructuralHash,
prim::ExprDeepEqual>;
/*!
* \brief A cache of computations is made of a pair of two hashtables, which
respectively associate
diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc
b/src/s_tir/schedule/primitive/cache_read_write.cc
index 25d4523235..2f7df83719 100644
--- a/src/s_tir/schedule/primitive/cache_read_write.cc
+++ b/src/s_tir/schedule/primitive/cache_read_write.cc
@@ -20,6 +20,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/ir/prim/expr.h>
#include <unordered_set>
@@ -1627,7 +1628,7 @@ class ReIndexCollector : public StmtExprVisitor {
return;
} else if (!std::equal(buffer_access_indices_.value().begin(),
buffer_access_indices_.value().end(),
indices.begin(), indices.end(),
- ExprDeepEqual())) {
+ prim::ExprDeepEqual())) {
throw MakeScheduleError<InvalidBufferAccessError>(
mod_, buffer_, block_,
InvalidBufferAccessError::ErrorKind::kNonUniqueAccess);
}
diff --git a/src/s_tir/schedule/primitive/compute_inline.cc
b/src/s_tir/schedule/primitive/compute_inline.cc
index a6741a63cf..19e0f77b47 100644
--- a/src/s_tir/schedule/primitive/compute_inline.cc
+++ b/src/s_tir/schedule/primitive/compute_inline.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/stmt.h>
#include "../utils.h"
@@ -897,7 +898,7 @@ class ReverseComputeInliner : public BaseInliner {
if (buffer_load_indices_.empty()) {
buffer_load_indices_ = indices;
} else if (!std::equal(buffer_load_indices_.begin(),
buffer_load_indices_.end(),
- indices.begin(), indices.end(), ExprDeepEqual())) {
+ indices.begin(), indices.end(),
prim::ExprDeepEqual())) {
// Failure: indices are not consistent in different BufferLoads
return false;
}
diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc
b/src/s_tir/schedule/primitive/layout_transformation.cc
index 0ca025a7f5..522bf4dc8d 100644
--- a/src/s_tir/schedule/primitive/layout_transformation.cc
+++ b/src/s_tir/schedule/primitive/layout_transformation.cc
@@ -21,6 +21,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/logging.h>
#include <optional>
@@ -202,7 +203,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
PrimExpr index =
ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(op->indices[i], f_substitute)
.as_or_throw<PrimExpr>();
bool is_loop_over_axis = index.same_as(loop->loop_var) &&
is_const_int(loop->min, 0) &&
- ExprDeepEqual()(loop->extent, buffer_dim) &&
+ prim::ExprDeepEqual()(loop->extent,
buffer_dim) &&
loop->kind == ForKind::kSerial;
if (!is_loop_over_axis) {
return false;
@@ -372,7 +373,7 @@ class TransformLayoutPlanner : public StmtExprVisitor {
const ffi::Array<PrimExpr>& old_indices = info.store->indices;
TVM_FFI_ICHECK_EQ(old_indices.size(), op->indices.size());
- ExprDeepEqual expr_equal;
+ prim::ExprDeepEqual expr_equal;
for (size_t i = 0; i < old_indices.size(); i++) {
if (!expr_equal(old_indices[i], op->indices[i])) {
return false;
diff --git a/src/s_tir/transform/thread_storage_sync.cc
b/src/s_tir/transform/thread_storage_sync.cc
index adbb82e296..564d3d8984 100644
--- a/src/s_tir/transform/thread_storage_sync.cc
+++ b/src/s_tir/transform/thread_storage_sync.cc
@@ -234,7 +234,7 @@ class ThreadSyncPlanner : public StorageAccessVisitor {
if (prev_intset.IsSinglePoint() && curr_intset.IsSinglePoint()) {
PrimExpr prev_index = prev_intset.PointValue();
PrimExpr curr_index = curr_intset.PointValue();
- has_same_index = ExprDeepEqual()(prev_index, curr_index);
+ has_same_index = prim::ExprDeepEqual()(prev_index, curr_index);
if (thread_index_var != nullptr) {
auto f_uses_thread_index = [=](const tvm::tirx::VarNode* parameter) {
return parameter == thread_index_var;
diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h
index 013db97297..6839811799 100644
--- a/src/target/llvm/codegen_llvm.h
+++ b/src/target/llvm/codegen_llvm.h
@@ -572,7 +572,7 @@ class CodeGenLLVM : public
tirx::ExprFunctor<llvm::Value*(const Expr&)>,
// set of volatile buffer.
std::unordered_set<const VarNode*> volatile_buf_;
// deep comparison of PrimExpr
- ExprDeepEqual deep_equal_;
+ prim::ExprDeepEqual deep_equal_;
// binding of let variables. Enables duplicate var defs that map to same
value
std::unordered_map<Var, const prim::LetNode*> let_binding_;
// debug info for function being compiled
diff --git a/src/target/source/codegen_c.h b/src/target/source/codegen_c.h
index 290367778e..b5ac52aa67 100644
--- a/src/target/source/codegen_c.h
+++ b/src/target/source/codegen_c.h
@@ -353,7 +353,7 @@ class CodeGenC : public tirx::ExprFunctor<void(const Expr&,
std::ostream&)>,
std::unordered_set<const VarNode*> volatile_buf_;
// deep comparison of PrimExpr
- ExprDeepEqual deep_equal_;
+ prim::ExprDeepEqual deep_equal_;
// binding of let variables. Enables duplicate var defs that map to same
value
std::unordered_map<Var, const prim::LetNode*> let_binding_;
diff --git a/src/tirx/analysis/var_use_def_analysis.h
b/src/tirx/analysis/var_use_def_analysis.h
index c4da205ea9..076acbe161 100644
--- a/src/tirx/analysis/var_use_def_analysis.h
+++ b/src/tirx/analysis/var_use_def_analysis.h
@@ -24,6 +24,7 @@
#ifndef TVM_TIR_ANALYSIS_VAR_USE_DEF_ANALYSIS_H_
#define TVM_TIR_ANALYSIS_VAR_USE_DEF_ANALYSIS_H_
+#include <tvm/ir/prim/expr.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/stmt_functor.h>
@@ -53,7 +54,7 @@ class VarUseDefAnalyzer : public StmtExprVisitor {
std::unordered_map<const VarNode*, int> buffer_def_count_;
private:
- ExprDeepEqual deep_equal_;
+ prim::ExprDeepEqual deep_equal_;
std::unordered_map<const VarNode*, const prim::LetNode*> let_binding_;
ffi::Optional<VisitInterrupt> Visit_(const AttrStmtNode* op) final;
diff --git a/src/tirx/analysis/verify_ssa.cc b/src/tirx/analysis/verify_ssa.cc
index 163d2f7b97..3da1254111 100644
--- a/src/tirx/analysis/verify_ssa.cc
+++ b/src/tirx/analysis/verify_ssa.cc
@@ -130,7 +130,7 @@ class SSAVerifier final : public StmtExprVisitor {
// whether we are in match scope, where a var can occur multiple times.
bool match_scope_{false};
// deep equal
- ExprDeepEqual deep_equal_;
+ prim::ExprDeepEqual deep_equal_;
// def map, for let, maps to the bind value, for others maps to self.
std::unordered_map<Var, Expr> def_map_;
};
diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc
index 623f461e72..a3835314b2 100644
--- a/src/tirx/ir/buffer.cc
+++ b/src/tirx/ir/buffer.cc
@@ -349,7 +349,7 @@ inline std::pair<bool, PrimExpr>
MergeMulModInner(arith::AnalyzerObj* analyzer,
const PrimExpr* search_ptr = inner;
PrimExpr mult_inner; // The inner multiplication factor
PrimExpr no_opt_sum; // Sum of the exprs that cannot be optimized
- tirx::ExprDeepEqual expr_equal;
+ prim::ExprDeepEqual expr_equal;
while (true) {
auto inner_div_ptr = search_ptr->as<IndexDiv>();
diff --git a/src/tirx/ir/specialize.cc b/src/tirx/ir/specialize.cc
index b775a396ba..c79976ab84 100644
--- a/src/tirx/ir/specialize.cc
+++ b/src/tirx/ir/specialize.cc
@@ -24,6 +24,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/prim/expr.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/layout.h>
@@ -355,7 +356,7 @@ class PrimFuncSpecializer : public StmtExprMutator {
void UpdateSpecializeVarMap(const PrimFunc& func, const Var& param, const
BufferVar& specific_buf,
VarMap* var_map) {
// preliminaries
- tirx::ExprDeepEqual equal;
+ prim::ExprDeepEqual equal;
auto opt_buffer = param.as<BufferVar>();
TVM_FFI_CHECK(opt_buffer, ValueError)
diff --git a/src/tirx/script/printer/block.cc b/src/tirx/script/printer/block.cc
index a8627a51e5..3a45eda314 100644
--- a/src/tirx/script/printer/block.cc
+++ b/src/tirx/script/printer/block.cc
@@ -48,7 +48,7 @@ Doc PrintBlock(IRDocsifier d, tirx::SBlock block, AccessPath
block_p, //
std::vector<int> remap_vars_indices;
auto add_remapped_iter_var = [&](int i) -> bool {
if (realize && d->cfg->syntax_sugar) {
- tirx::ExprDeepEqual expr_equal;
+ prim::ExprDeepEqual expr_equal;
tirx::IterVar iter_var = block->iter_vars[i];
PrimExpr value = realize->iter_values[i];
if (iter_var->iter_type == tirx::IterVarType::kDataPar ||
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index 2901c01c85..7e1e26b581 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -315,7 +315,7 @@ ffi::Optional<ExprDoc> TryDeclBufferSugarWithParent(const
tirx::BufferVar& child
if (!parent_doc.has_value()) return std::nullopt;
ExprDoc pdoc = parent_doc.value();
- tirx::ExprDeepEqual expr_equal;
+ prim::ExprDeepEqual expr_equal;
// Check elem_offset equality
bool same_elem_offset = expr_equal(child->elem_offset, parent->elem_offset);
diff --git a/src/tirx/transform/common_subexpr_elim.cc
b/src/tirx/transform/common_subexpr_elim.cc
index 13afd243eb..be4a8bd8b4 100644
--- a/src/tirx/transform/common_subexpr_elim.cc
+++ b/src/tirx/transform/common_subexpr_elim.cc
@@ -101,7 +101,7 @@ namespace tirx {
* Used by CSERewriter to look up whether a visited expression should be
* replaced by a previously-introduced CSE variable.
*/
-using ExprRemapTable = std::unordered_map<PrimExpr, Var, ffi::StructuralHash,
ExprDeepEqual>;
+using ExprRemapTable = std::unordered_map<PrimExpr, Var, ffi::StructuralHash,
prim::ExprDeepEqual>;
/*!
* \brief Map from statement (by pointer identity) to a list of Bind
@@ -188,7 +188,7 @@ class CSEPlanner : public StmtExprVisitor {
* \brief Node in the expression DAG built during the bottom-up scan.
*
* The planner maintains one ExprEntry per structurally-unique eligible
- * expression (keyed by ExprDeepEqual). Since expressions are recorded
+ * expression (keyed by prim::ExprDeepEqual). Since expressions are recorded
* bottom-up (children before parents), the DAG children are naturally
* discovered when a node is first added. Fields like expr_depth are
* computed incrementally from children — no separate traversal needed.
@@ -243,7 +243,7 @@ class CSEPlanner : public StmtExprVisitor {
};
/*!
- * \brief Expression table keyed by structural equality (ExprDeepEqual).
+ * \brief Expression table keyed by structural equality
(prim::ExprDeepEqual).
*
* An insertion-ordered map so that iteration visits entries in discovery
* (program) order. This makes the plan — and hence cse_v numbering —
@@ -251,7 +251,8 @@ class CSEPlanner : public StmtExprVisitor {
* StructuralHash hashes free variables by object identity, which varies
* between processes (ASLR).
*/
- using ExprTable = support::OrderedMap<PrimExpr, ExprEntry,
ffi::StructuralHash, ExprDeepEqual>;
+ using ExprTable =
+ support::OrderedMap<PrimExpr, ExprEntry, ffi::StructuralHash,
prim::ExprDeepEqual>;
// ------------------------------------------------------------------
// Eligibility predicates
@@ -311,7 +312,7 @@ class CSEPlanner : public StmtExprVisitor {
/*!
* \brief Replace all occurrences of `target` in `body` with `replacement`.
*
- * Uses structural equality (ExprDeepEqual) to find matches. Stops recursing
+ * Uses structural equality (prim::ExprDeepEqual) to find matches. Stops
recursing
* into a sub-tree once a match is found (the replacement is a leaf Var).
*
* \param body The expression to transform.
@@ -322,7 +323,7 @@ class CSEPlanner : public StmtExprVisitor {
static PrimExpr SubstituteSubexpr(const PrimExpr& body, const PrimExpr&
target,
const PrimExpr& replacement) {
struct Replacer : public ExprMutator {
- ExprDeepEqual eq;
+ prim::ExprDeepEqual eq;
PrimExpr target, replacement;
Expr Dispatch(const Expr& e) final {
if (auto prim = e.as<PrimExpr>(); prim && eq(prim.value(), target))
return replacement;
@@ -441,7 +442,7 @@ class CSEPlanner : public StmtExprVisitor {
* child `x+y` with multiplicity 2). expr_depth is 1 + max child depth.
*/
void CollectChildren(ExprEntry& entry, std::initializer_list<PrimExpr>
ast_children) {
- ExprDeepEqual eq;
+ prim::ExprDeepEqual eq;
int max_child_depth = 0;
for (const PrimExpr& child : ast_children) {
auto it = table_.find(child);
diff --git a/src/tirx/transform/stmt_simplify.cc
b/src/tirx/transform/stmt_simplify.cc
index e92e3bea39..dc249a3261 100644
--- a/src/tirx/transform/stmt_simplify.cc
+++ b/src/tirx/transform/stmt_simplify.cc
@@ -215,7 +215,7 @@ class StmtSimplifier : public IRMutatorWithAnalyzer {
if (const TensorLoadNode* load = store->value.as<TensorLoadNode>()) {
BufferVar buffer = load->source.as_or_throw<tvm::tirx::BufferVar>();
if (buffer.same_as(store->buffer) && ArrayDeepEqual(load->indices,
store->indices) &&
- tirx::ExprDeepEqual()(buffer->elem_offset,
store->buffer->elem_offset) &&
+ prim::ExprDeepEqual()(buffer->elem_offset,
store->buffer->elem_offset) &&
ArrayDeepEqual(buffer->shape, store->buffer->shape) &&
ArrayDeepEqual(buffer->strides, store->buffer->strides)) {
return Evaluate(0);
@@ -230,7 +230,7 @@ class StmtSimplifier : public IRMutatorWithAnalyzer {
return false;
}
for (size_t i = 0; i < lhs.size(); i++) {
- if (!tirx::ExprDeepEqual()(lhs[i], rhs[i])) {
+ if (!prim::ExprDeepEqual()(lhs[i], rhs[i])) {
return false;
}
}
diff --git a/src/tirx/transform/storage_rewrite.cc
b/src/tirx/transform/storage_rewrite.cc
index 7af54a5973..acb910bb74 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -428,7 +428,7 @@ class InplaceOpVerifier : public StmtExprVisitor {
<< "Store/Load occur to the same buffer " << buf->name
<< " with differing number of indices";
for (size_t i = 0; i < store_->indices.size(); i++) {
- if (!tirx::ExprDeepEqual()(store_->indices[i], op->indices[i])) {
+ if (!prim::ExprDeepEqual()(store_->indices[i], op->indices[i])) {
result_ = false;
return std::nullopt;
}
@@ -818,7 +818,7 @@ class StoragePlanRewriter : public StmtExprMutator {
if (op->buffer->shape.size() != first->buffer->shape.size()) {
return false;
}
- ExprDeepEqual expr_equal;
+ prim::ExprDeepEqual expr_equal;
for (size_t i = 0; i < op->buffer->shape.size(); i++) {
if (!expr_equal(op->buffer->shape[i],
first->buffer->shape[i])) {
return false;
diff --git a/src/tirx/transform/vectorize_loop.cc
b/src/tirx/transform/vectorize_loop.cc
index 9317b067a5..2da6a75d3a 100644
--- a/src/tirx/transform/vectorize_loop.cc
+++ b/src/tirx/transform/vectorize_loop.cc
@@ -1062,7 +1062,7 @@ class Vectorizer : public StmtMutator, public
ExprFunctor<Expr(const Expr&)> {
// analyzer
arith::Analyzer analyzer_;
// deep equal
- ExprDeepEqual deep_equal_;
+ prim::ExprDeepEqual deep_equal_;
// variable to be replaced
Var var_;
// the lanes.
diff --git a/tests/cpp/expr_test.cc b/tests/cpp/expr_test.cc
index 1515d21f71..04bedfce99 100644
--- a/tests/cpp/expr_test.cc
+++ b/tests/cpp/expr_test.cc
@@ -20,6 +20,7 @@
#include <gtest/gtest.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/ir/prim/expr.h>
#include <tvm/ir/source_map.h>
#include <tvm/runtime/logging.h>
#include <tvm/te/operation.h>
@@ -97,3 +98,20 @@ TEST(ExprNodeRef, Basic) {
const prim::MaxNode* op = z.as<prim::MaxNode>();
TVM_FFI_ICHECK(ffi::GetRef<ffi::ObjectRef>(op).same_as(z));
}
+
+TEST(Expr, DeepEqualTensorLoadSourceIdentity) {
+ using namespace tvm;
+ Var source("source", PointerType(PrimType::Float(32)));
+ Var other_source("source", PointerType(PrimType::Float(32)));
+ auto load = [](Expr source, PrimExpr index) {
+ auto node = ffi::make_object<TensorLoadNode>();
+ node->ty = PrimType::Float(32);
+ node->source = source;
+ node->indices = {index};
+ return TensorLoad(node);
+ };
+ prim::ExprDeepEqual equal;
+ EXPECT_TRUE(equal(load(source, 0), load(source, 0)));
+ EXPECT_FALSE(equal(load(source, 0), load(other_source, 0)));
+ EXPECT_FALSE(equal(load(source, 0), load(source, 1)));
+}
diff --git a/tests/cpp/pattern_match_test.cc b/tests/cpp/pattern_match_test.cc
index f53b8acec1..b9280dad6b 100644
--- a/tests/cpp/pattern_match_test.cc
+++ b/tests/cpp/pattern_match_test.cc
@@ -20,7 +20,7 @@
#include "../src/arith/pattern_match.h"
#include <gtest/gtest.h>
-#include <tvm/tirx/analysis.h>
+#include <tvm/ir/prim/expr.h>
TEST(Pattern, Basic) {
using namespace tvm;
@@ -43,12 +43,12 @@ TEST(Pattern, Basic) {
TVM_FFI_ICHECK((px + (py + px)).Match(r));
auto rr = (px + py).Eval();
- TVM_FFI_ICHECK(tirx::ExprDeepEqual()(rr, 1 + y));
- TVM_FFI_ICHECK(tirx::ExprDeepEqual()(px.Eval() + py.Eval(), 1 + y));
+ TVM_FFI_ICHECK(prim::ExprDeepEqual()(rr, 1 + y));
+ TVM_FFI_ICHECK(prim::ExprDeepEqual()(px.Eval() + py.Eval(), 1 + y));
}
{
TVM_FFI_ICHECK((px + max(py, px)).Match((x + 1) + max(y, (x + 1))));
- TVM_FFI_ICHECK(tirx::ExprDeepEqual()(px.Eval(), x + 1));
+ TVM_FFI_ICHECK(prim::ExprDeepEqual()(px.Eval(), x + 1));
}
TVM_FFI_ICHECK(!(px + min(py, px)).Match((x + 1) + max(y, (x + 1))));
@@ -68,7 +68,7 @@ TEST(Pattern, Basic) {
TVM_FFI_ICHECK((!(px > py || px != py)).Match(!(x > y || x != y)));
{
TVM_FFI_ICHECK(select(px >= pz, py, py + pz).Match(prim::Select((x + 1) >=
1, y, y + 1)));
- TVM_FFI_ICHECK(tirx::ExprDeepEqual()(px.Eval(), x + 1));
+ TVM_FFI_ICHECK(prim::ExprDeepEqual()(px.Eval(), x + 1));
}
// bit intrinsics
{
@@ -90,7 +90,7 @@ TEST(Pattern, Basic) {
TVM_FFI_ICHECK(!select(px > pz, py, py).Match(prim::Select(x > 2, y, y +
1)));
{
TVM_FFI_ICHECK(select(px, py, pz).Match(prim::Select(x > 2, y, y + 1)));
- TVM_FFI_ICHECK(tirx::ExprDeepEqual()(pz.Eval(), y + 1));
+ TVM_FFI_ICHECK(prim::ExprDeepEqual()(pz.Eval(), y + 1));
}
// if_then_else
{
@@ -115,7 +115,7 @@ TEST(Pattern, Basic) {
TVM_FFI_ICHECK(ramp(px, PConst<PrimExpr>(1), planes).Match(prim::Ramp(x,
1, 10)));
TVM_FFI_ICHECK(planes.Eval().as<IntImmNode>()->value == 10);
TVM_FFI_ICHECK(ramp(px, PConst<PrimExpr>(1), planes).Match(prim::Ramp(x,
1, scalable_lanes)));
- TVM_FFI_ICHECK(tirx::ExprDeepEqual()(planes.Eval(), scalable_lanes));
+ TVM_FFI_ICHECK(prim::ExprDeepEqual()(planes.Eval(), scalable_lanes));
TVM_FFI_ICHECK(!ramp(px, PConst<PrimExpr>(1), planes).Match(prim::Ramp(x,
2, 10)));
}
// broadcast pattern
@@ -124,7 +124,7 @@ TEST(Pattern, Basic) {
TVM_FFI_ICHECK(planes.Eval().as<IntImmNode>()->value == 10);
TVM_FFI_ICHECK(broadcast(px * py, planes).Match(prim::Broadcast(x * 10,
10)));
TVM_FFI_ICHECK(broadcast(px, planes).Match(prim::Broadcast(x,
scalable_lanes)));
- TVM_FFI_ICHECK(tirx::ExprDeepEqual()(planes.Eval(), scalable_lanes));
+ TVM_FFI_ICHECK(prim::ExprDeepEqual()(planes.Eval(), scalable_lanes));
}
}
diff --git a/tests/python/tirx-analysis/test_tir_analysis_expr_deep_equal.py
b/tests/python/ir/test_ir_prim_expr_deep_equal.py
similarity index 84%
rename from tests/python/tirx-analysis/test_tir_analysis_expr_deep_equal.py
rename to tests/python/ir/test_ir_prim_expr_deep_equal.py
index 0ca88575a4..d7fc2e04d7 100644
--- a/tests/python/tirx-analysis/test_tir_analysis_expr_deep_equal.py
+++ b/tests/python/ir/test_ir_prim_expr_deep_equal.py
@@ -27,9 +27,9 @@ def test_equal_expr():
def func2():
return tvm.tirx.exp(tvm.tirx.truncdiv((x + y + 1) * y, 4))
- assert tvm.tirx.analysis.expr_deep_equal(func1(), func1())
- assert tvm.tirx.analysis.expr_deep_equal(func2(), func2())
- assert not tvm.tirx.analysis.expr_deep_equal(func2(), func1())
+ assert tvm.ir.prim.expr_deep_equal(func1(), func1())
+ assert tvm.ir.prim.expr_deep_equal(func2(), func2())
+ assert not tvm.ir.prim.expr_deep_equal(func2(), func1())
if __name__ == "__main__":
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
index dd87c504d4..c8beaaa90c 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
@@ -22,6 +22,7 @@ from tvm_ffi import structural_walk
import tvm
import tvm.testing
+from tvm.ir.prim import expr_deep_equal
from tvm.s_tir import Schedule
from tvm.s_tir.meta_schedule.testing import te_workload
from tvm.s_tir.schedule.analysis import (
@@ -48,7 +49,6 @@ from tvm.tirx import (
floordiv,
floormod,
)
-from tvm.tirx.analysis import expr_deep_equal
from tvm.tirx.function import TensorIntrin
diff --git a/tests/python/tirx-base/test_tir_constructor.py
b/tests/python/tirx-base/test_tir_constructor.py
index 4df287013c..fd880020d6 100644
--- a/tests/python/tirx-base/test_tir_constructor.py
+++ b/tests/python/tirx-base/test_tir_constructor.py
@@ -20,8 +20,8 @@ import tvm_ffi
import tvm
from tvm import te, topi
+from tvm.ir.prim import expr_deep_equal
from tvm.script import tirx as T
-from tvm.tirx.analysis import expr_deep_equal
def test_expr_constructor():
