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 1fb1c38665 [IR][Relax] Include expression types in structural identity
(#19933)
1fb1c38665 is described below
commit 1fb1c3866501eee8ddb5f9e01b53603c2fc805dc
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Jul 4 22:39:52 2026 +0800
[IR][Relax] Include expression types in structural identity (#19933)
## Rationale
After `PrimExpr` and `Expr` share one typed expression hierarchy,
expression types are part of semantic identity. Structurally identical
syntax with different types compare and hash differently, while source
spans remain diagnostic metadata.
## Invariant
`ExprNode::ty` participates in structural equality and hashing by
default. `GlobalVar` and Relax variables retain their symbol identity
rules. `tirx.PrimFunc` compares and hashes authoritative source fields
while excluding its derived type cache until all transformation paths
maintain that cache eagerly. Nested symbolic-shape rendering is isolated
from outer diagnostic configuration so diagnostic context cannot become
script-token content.
## Changes
- include expression types in generic structural equality and hashing
- preserve GlobalVar and Relax variable identity plus definition-safe
SeqExpr traversal
- compare and hash PrimFunc from authoritative fields while excluding
its stale derived type cache
- normalize narrow Relax construction and expected-fixture types exposed
by stricter identity
- isolate nested symbolic-shape token rendering from outer printer
configuration
---
include/tvm/ir/base_expr.h | 5 ++---
include/tvm/ir/expr.h | 5 +++++
include/tvm/relax/expr.h | 19 ++++++++++++++++
include/tvm/tirx/function.h | 25 ++++++++++++++++++++++
python/tvm/relax/op/memory/memory.py | 2 ++
python/tvm/relax/op/vm/vm.py | 2 ++
.../transform/lower_global_view_to_local_view.cc | 7 +++---
src/relax/script/printer/dependent_type.cc | 5 ++++-
.../relax/test_frontend_from_exported_program.py | 6 +++---
tests/python/relax/test_frontend_onnx.py | 2 ++
tests/python/relax/test_frontend_tflite.py | 8 ++++++-
tests/python/relax/test_utils.py | 2 +-
12 files changed, 75 insertions(+), 13 deletions(-)
diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h
index 6d566bd5c9..a836d0410a 100644
--- a/include/tvm/ir/base_expr.h
+++ b/include/tvm/ir/base_expr.h
@@ -299,12 +299,11 @@ class ExprNode : public ffi::Object {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
- // span and ty do not participate in structural equal and hash.
+ // span does not participate in structural equal and hash.
refl::ObjectDef<ExprNode>()
.def_ro("span", &ExprNode::span, refl::DefaultValue(Span()),
refl::AttachFieldFlag::SEqHashIgnore())
- .def_ro("ty", &ExprNode::ty, refl::DefaultValue(Type::Missing()),
- refl::AttachFieldFlag::SEqHashIgnore());
+ .def_ro("ty", &ExprNode::ty, refl::DefaultValue(Type::Missing()));
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index 7f1c84f50f..dd168a2462 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -276,6 +276,11 @@ class GlobalVarNode : public ExprNode {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<GlobalVarNode>().def_ro("name_hint",
&GlobalVarNode::name_hint);
+ // A GlobalVar identifies a module-level symbol. Its type is derived from
the
+ // corresponding function definition and is not part of the symbol
identity.
+ refl::TypeAttrDef<GlobalVarNode>()
+ .def("__s_equal__", &GlobalVarNode::SEqual)
+ .def("__s_hash__", &GlobalVarNode::SHash);
}
bool SEqual(const GlobalVarNode* other,
diff --git a/include/tvm/relax/expr.h b/include/tvm/relax/expr.h
index 83e03c13e7..51b4ed43bd 100644
--- a/include/tvm/relax/expr.h
+++ b/include/tvm/relax/expr.h
@@ -501,6 +501,25 @@ class SeqExprNode : public ExprNode {
refl::ObjectDef<SeqExprNode>()
.def_ro("blocks", &SeqExprNode::blocks)
.def_ro("body", &SeqExprNode::body);
+ refl::TypeAttrDef<SeqExprNode>()
+ .def("__s_equal__", &SeqExprNode::SEqual)
+ .def("__s_hash__", &SeqExprNode::SHash);
+ }
+
+ bool SEqual(const SeqExprNode* other,
+ ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal)
const {
+ // Establish mappings for symbolic variables defined by bindings before
+ // comparing their uses in the SeqExpr result type and body.
+ return equal(blocks, other->blocks, false, "blocks") && equal(ty,
other->ty, false, "ty") &&
+ equal(body, other->body, false, "body");
+ }
+
+ int64_t SHash(int64_t init_hash, ffi::TypedFunction<int64_t(AnyView,
int64_t, bool)> hash) const {
+ int64_t hash_value = init_hash;
+ hash_value = hash(blocks, hash_value, false);
+ hash_value = hash(ty, hash_value, false);
+ hash_value = hash(body, hash_value, false);
+ return hash_value;
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.SeqExpr", SeqExprNode,
ExprNode);
};
diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h
index e4e33f3576..a69b32a433 100644
--- a/include/tvm/tirx/function.h
+++ b/include/tvm/tirx/function.h
@@ -108,6 +108,31 @@ class PrimFuncNode : public BaseFuncNode {
.def_ro("ret_type", &PrimFuncNode::ret_type)
.def_ro("buffer_map", &PrimFuncNode::buffer_map)
.def_ro("body", &PrimFuncNode::body);
+ refl::TypeAttrDef<PrimFuncNode>()
+ .def("__s_equal__", &PrimFuncNode::SEqual)
+ .def("__s_hash__", &PrimFuncNode::SHash);
+ }
+
+ bool SEqual(const PrimFuncNode* other,
+ ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal)
const {
+ // `ty` is derived from the fields below. PrimFunc transformations update
+ // those source fields without maintaining this redundant cache eagerly.
+ // Remove this exception once all PrimFunc mutation paths recompute `ty`.
+ return equal(attrs, other->attrs, false, "attrs") &&
+ equal(params, other->params, true, "params") &&
+ equal(ret_type, other->ret_type, false, "ret_type") &&
+ equal(buffer_map, other->buffer_map, false, "buffer_map") &&
+ equal(body, other->body, false, "body");
+ }
+
+ int64_t SHash(int64_t init_hash, ffi::TypedFunction<int64_t(AnyView,
int64_t, bool)> hash) const {
+ int64_t hash_value = init_hash;
+ hash_value = hash(attrs, hash_value, false);
+ hash_value = hash(params, hash_value, true);
+ hash_value = hash(ret_type, hash_value, false);
+ hash_value = hash(buffer_map, hash_value, false);
+ hash_value = hash(body, hash_value, false);
+ return hash_value;
}
/*!
diff --git a/python/tvm/relax/op/memory/memory.py
b/python/tvm/relax/op/memory/memory.py
index 624d39a133..a1dec11069 100644
--- a/python/tvm/relax/op/memory/memory.py
+++ b/python/tvm/relax/op/memory/memory.py
@@ -98,6 +98,8 @@ def alloc_tensor(
shape = convert_to_expr(shape)
if isinstance(dtype, str):
dtype = DataTypeImm(dtype)
+ if isinstance(runtime_device_ind, int):
+ runtime_device_ind = prim_value(runtime_device_ind)
return _ffi_api.alloc_tensor(storage, offset, shape, dtype,
runtime_device_ind) # type: ignore
diff --git a/python/tvm/relax/op/vm/vm.py b/python/tvm/relax/op/vm/vm.py
index 0fc236d59b..afa9782e3d 100644
--- a/python/tvm/relax/op/vm/vm.py
+++ b/python/tvm/relax/op/vm/vm.py
@@ -98,6 +98,8 @@ def alloc_tensor(
shape = convert_to_expr(shape)
if isinstance(dtype, str):
dtype = DataTypeImm(dtype)
+ if isinstance(runtime_device_ind, int):
+ runtime_device_ind = prim_value(runtime_device_ind)
return _ffi_api.alloc_tensor(storage, offset, shape, dtype,
runtime_device_ind) # type: ignore
diff --git a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
index f37c0d3863..c04970f631 100644
--- a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
+++ b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
@@ -166,10 +166,9 @@ class DistributedBufferCompactor : StmtExprMutator {
}
Stmt new_body = compactor(prim_func->body);
new_body = DistBufferReplacer::BufferReplace(new_body, replace_buffer_map);
- ffi::ObjectPtr<PrimFuncNode> new_func =
ffi::make_object<PrimFuncNode>(*prim_func.get());
- new_func->buffer_map = new_func_buffer_map;
- new_func->body = new_body;
- return std::make_tuple(PrimFunc(new_func), compactor.add_allreduce_kind_);
+ PrimFunc new_func(prim_func->params, new_body, prim_func->ret_type,
new_func_buffer_map,
+ prim_func->attrs, prim_func->span);
+ return std::make_tuple(new_func, compactor.add_allreduce_kind_);
}
private:
diff --git a/src/relax/script/printer/dependent_type.cc
b/src/relax/script/printer/dependent_type.cc
index d5c09cbcc7..c3b77223b5 100644
--- a/src/relax/script/printer/dependent_type.cc
+++ b/src/relax/script/printer/dependent_type.cc
@@ -54,7 +54,10 @@ ExprDoc PrintShapeVar(const PrimExpr& e, const AccessPath&
e_p, const IRDocsifie
}
// Step 3. Stringify the PrimExpr if func var exists
if (func_var_mode) {
- return LiteralDoc::Str(DocToPythonScript(expr_doc, d->cfg), e_p);
+ // This nested render is only converting a shape expression into one token.
+ // Give it an independent configuration so it cannot consume or emit the
+ // enclosing invocation's access-path diagnostics.
+ return LiteralDoc::Str(DocToPythonScript(expr_doc, PrinterConfig()), e_p);
}
return expr_doc;
}
diff --git a/tests/python/relax/test_frontend_from_exported_program.py
b/tests/python/relax/test_frontend_from_exported_program.py
index 1ee88ea846..afd53f1b74 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -6227,9 +6227,9 @@ def test_hamming_window():
with R.dataflow():
lv: R.Tensor((20,), dtype="float32") = R.hamming_window(
R.prim_value(20),
- R.prim_value(1),
- R.prim_value(T.float32(0.54000000000000004)),
- R.prim_value(T.float32(0.46000000000000002)),
+ R.prim_value(True),
+ R.prim_value(T.float64(0.54000000000000004)),
+ R.prim_value(T.float64(0.46000000000000002)),
dtype="float32",
)
gv: R.Tuple(R.Tensor((20,), dtype="float32")) = (lv,)
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index f4dbac634a..9058cdf70a 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -3945,6 +3945,8 @@ def test_dynamic_shape_squeeze(axis):
axes = relax.Var("axes", relax.TensorType([1], "int64"))
gv = relax.Var("gv", tvm.ir.PrimType("int64"))
body = relax.SeqExpr([relax.DataflowBlock([relax.VarBinding(gv, a)])], gv)
+ # Match the importer boundary, where BlockBuilder populates the SeqExpr
result type.
+ body = relax.BlockBuilder().normalize(body)
expected_func = relax.Function([x, axes], body,
tvm.ir.PrimType("int64")).with_attrs(
{"num_input": 1, "global_symbol": "main"}
)
diff --git a/tests/python/relax/test_frontend_tflite.py
b/tests/python/relax/test_frontend_tflite.py
index 3397d08abb..c756227fe0 100644
--- a/tests/python/relax/test_frontend_tflite.py
+++ b/tests/python/relax/test_frontend_tflite.py
@@ -643,7 +643,13 @@ def test_unique():
R.func_attr({"num_input": 1})
with R.dataflow():
lv: R.Tuple(R.Tensor(dtype="int32", ndim=1),
R.Tensor(dtype="int64", ndim=1)) = (
- R.unique(x, R.prim_value(0), R.prim_value(0),
R.prim_value(1), R.prim_value(0))
+ R.unique(
+ x,
+ R.prim_value(False),
+ R.prim_value(False),
+ R.prim_value(True),
+ R.prim_value(False),
+ )
)
lv1: R.Tensor(dtype="int32", ndim=1) = lv[0]
lv2: R.Tensor(dtype="int64", ndim=1) = lv[1]
diff --git a/tests/python/relax/test_utils.py b/tests/python/relax/test_utils.py
index 9ee05b7fdd..b654edfc31 100644
--- a/tests/python/relax/test_utils.py
+++ b/tests/python/relax/test_utils.py
@@ -143,7 +143,7 @@ def test_assert_structural_equal_in_seqexpr():
with pytest.raises(
ValueError,
- match=re.escape("<root>.body.blocks[0].bindings[0].value.op"),
+ match=re.escape("<root>.ty.ret.shape.ty.values[0].value"),
):
assert_structural_equal(func_1, func_2)
