This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 6ae00166a [IR] Fix structural equality and hashing of function
attributes (#163)
6ae00166a is described below
commit 6ae00166ab8881f581278736afb398043ca721eb
Author: Lei Wang <[email protected]>
AuthorDate: Tue Nov 3 00:04:08 2026 +0800
[IR] Fix structural equality and hashing of function attributes (#163)
Fix structural equality and hashing for `Function` attributes that
reference the
function's parameters, including parameters used as map keys.
A function with a parameter reference in its attributes currently
compares unequal
to its own JSON round-trip:
```python
import tvm_ffi
from tvm_tirx import ir, tirx
x = ir.Var("x", "int32")
f = tirx.Function([x], [ir.Return(x + 1)], ir.PrimType("int32"))
f = f.with_attr("parameter", x)
g = ir.load_json(ir.save_json(f))
print(g.attrs["parameter"].same_as(g.params[0]))
print(tvm_ffi.structural_equal(f, g))
print(tvm_ffi.structural_hash(f) == tvm_ffi.structural_hash(g))
```
Before this change, the output is `True, False, False`. After the
change, all three
results are `True`. The first result shows that JSON deserialization
preserves the
reference from the attribute to the restored parameter.
The order matters because the `equal` callback carries variable-binding
state
through the recursive comparison. In `FunctionObj::SEqual`, the calls
have two
different roles:
- `equal(params, other->params, true, "params")` compares parameter
definitions.
For matching parameters, it establishes the correspondence between the
variables
on the two sides. This allows equivalent functions to use different
variable
objects and names.
- `equal(attrs, other->attrs, false, "attrs")` compares attributes using
the current
variable correspondence. Under the default comparison policy, distinct
variable
objects without an established correspondence are treated as distinct
free
variables.
Previously, attributes were visited first. When an attribute referred to
a
parameter, its counterpart in the restored function was a different
object, and
the parameter correspondence had not been established yet. The attribute
comparison returned `false`, and the `&&` expression stopped before
visiting the
parameter definitions. Visiting parameters first establishes the
correspondence
that the attribute comparison needs.
`FunctionObj::SHash` has the same ordering requirement. A direct
parameter reference
visited through attributes before its definition can be memoized with an
identity-based free-variable hash. Visiting parameters first assigns
their
definition-based hashes before attributes consume them. Object map keys
also
depend on the existing variable correspondence and hash memo.
This change visits `params` before `attrs` in both `SEqual` and `SHash`.
Attributes
still participate in comparison and hashing, and free variables retain
their identity semantics. Regression tests cover JSON round-trips,
equivalent
functions with renamed parameters, parameters used as map keys, and
negative cases
for different parameter positions and distinct same-name free variables.
---
include/tvm/tirx/function.h | 7 ++--
tests/python/tirx/ir/test_function.py | 64 ++++++++++++++++++++++++++++++++++-
2 files changed, 67 insertions(+), 4 deletions(-)
diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h
index 224b1798f..fadbd7dab 100644
--- a/include/tvm/tirx/function.h
+++ b/include/tvm/tirx/function.h
@@ -72,16 +72,17 @@ class FunctionObj : public BaseFuncObj {
bool SEqual(const FunctionObj* other,
ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal)
const {
// `ty` is derived from the signature; compare its source fields directly.
- return equal(attrs, other->attrs, false, "attrs") &&
- equal(params, other->params, true, "params") &&
+ // Establish parameter bindings before comparing attributes that may
reference them.
+ return equal(params, other->params, true, "params") &&
+ equal(attrs, other->attrs, false, "attrs") &&
equal(ret_type, other->ret_type, false, "ret_type") &&
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(attrs, hash_value, false);
hash_value = hash(ret_type, hash_value, false);
hash_value = hash(body, hash_value, false);
return hash_value;
diff --git a/tests/python/tirx/ir/test_function.py
b/tests/python/tirx/ir/test_function.py
index f5bf5a7b2..fbda5b7d8 100644
--- a/tests/python/tirx/ir/test_function.py
+++ b/tests/python/tirx/ir/test_function.py
@@ -14,10 +14,72 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Test TIRx function structural mutation."""
+"""Test TIRx function structural operations."""
+import pytest
import tvm_ffi
+from tvm_tirx import ir, tirx
+
+
[email protected]("as_map_key", [False, True])
+def test_parameter_attr_json_roundtrip(as_map_key):
+ x = ir.Var("x", "int32")
+ func = tirx.Function([x], [ir.Return(x + 1)], ir.PrimType("int32"))
+ func = func.with_attr("parameter", {x: 7} if as_map_key else x)
+
+ restored = ir.load_json(ir.save_json(func))
+ attr = restored.attrs["parameter"]
+ restored_ref = next(iter(attr)) if as_map_key else attr
+ assert restored_ref.same_as(restored.params[0])
+ assert not restored.params[0].same_as(x)
+ ir.assert_structural_equal(func, restored)
+ ir.assert_structural_equal(restored, func)
+ assert tvm_ffi.structural_hash(func) == tvm_ffi.structural_hash(restored)
+
+
[email protected]("as_map_key", [False, True])
+def test_parameter_attr_alpha_equivalence(as_map_key):
+ def make_function(name):
+ x = ir.Var(name, "int32")
+ func = tirx.Function([x], [ir.Return(x + 1)], ir.PrimType("int32"))
+ return func.with_attr("parameter", {x: 7} if as_map_key else x)
+
+ lhs = make_function("x")
+ rhs = make_function("renamed")
+ ir.assert_structural_equal(lhs, rhs)
+ ir.assert_structural_equal(rhs, lhs)
+ assert tvm_ffi.structural_hash(lhs) == tvm_ffi.structural_hash(rhs)
+
+
[email protected]("as_map_key", [False, True])
+def test_parameter_attr_distinguishes_parameters(as_map_key):
+ x = ir.Var("x", "int32")
+ y = ir.Var("x", "int32")
+ func = tirx.Function([x, y], [ir.Return(x + y)], ir.PrimType("int32"))
+ lhs = func.with_attr("parameter", {x: 7} if as_map_key else x)
+ rhs = func.with_attr("parameter", {y: 7} if as_map_key else y)
+
+ assert not tvm_ffi.structural_equal(lhs, rhs)
+ assert not tvm_ffi.structural_equal(rhs, lhs)
+
+
[email protected]("as_map_key", [False, True])
+def test_function_attr_preserves_free_variable_identity(as_map_key):
+ def make_function(free_var):
+ x = ir.Var("x", "int32")
+ func = tirx.Function([x], [ir.Return(x + 1)], ir.PrimType("int32"))
+ return func.with_attr("free_var", {free_var: 7} if as_map_key else
free_var)
+
+ free_var = ir.Var("x", "int32")
+ lhs = make_function(free_var)
+ rhs = make_function(free_var)
+ other = make_function(ir.Var("x", "int32"))
+ ir.assert_structural_equal(lhs, rhs)
+ assert tvm_ffi.structural_hash(lhs) == tvm_ffi.structural_hash(rhs)
+ assert not tvm_ffi.structural_equal(lhs, other)
+ assert not tvm_ffi.structural_equal(other, lhs)
+
def test_structural_map_specializes_function_signature():
from tvm_tirx.ir import Evaluate, FuncType, IntImm, Return, Var,
assert_structural_equal
