[Mlir-commits] [mlir] Add Python bindings for operation structural equivalence and hashing (PR #216357)
Trevor McKay
llvmlistbot at llvm.org
Mon Aug 17 18:55:31 PDT 2026
https://github.com/trmckay updated https://github.com/llvm/llvm-project/pull/216357
>From 8523665086ad76a72b95b557164d269adf4d7b1c Mon Sep 17 00:00:00 2001
From: Trevor McKay <mail at tmckay.com>
Date: Thu, 13 Aug 2026 14:48:15 -0700
Subject: [PATCH 1/4] add Python bindings for structural equivalence and hash
---
mlir/include/mlir/Bindings/Python/IRCore.h | 10 +++
mlir/lib/Bindings/Python/IRCore.cpp | 39 ++++++++++++
mlir/test/python/ir/operation.py | 71 ++++++++++++++++++++++
3 files changed, 120 insertions(+)
diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index 3314e0b2a8fcf..05c7d63ccf766 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -365,6 +365,16 @@ enum class PyWalkOrder : std::underlying_type_t<MlirWalkOrder> {
PostOrder = MlirWalkPostOrder
};
+/// Flags controlling structural operation equivalence and hashing.
+enum class PyOperationEquivalenceFlags : std::underlying_type_t<
+ MlirOperationEquivalenceFlags> {
+ None = MLIR_OPERATION_EQUIVALENCE_NONE,
+ IgnoreLocations = MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS,
+ IgnoreDiscardableAttrs = MLIR_OPERATION_EQUIVALENCE_IGNORE_DISCARDABLE_ATTRS,
+ IgnoreProperties = MLIR_OPERATION_EQUIVALENCE_IGNORE_PROPERTIES,
+ IgnoreCommutativity = MLIR_OPERATION_EQUIVALENCE_IGNORE_COMMUTATIVITY
+};
+
/// Python class mirroring the C MlirDiagnostic struct. Note that these structs
/// are only valid for the duration of a diagnostic callback and attempting
/// to access them outside of that will raise an exception. This applies to
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 75cfd2a0a1c0b..95eb3d857db41 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3291,6 +3291,17 @@ void populateIRCore(nb::module_ &m) {
nb::enum_<PyWalkOrder>(m, "WalkOrder")
.value("PRE_ORDER", PyWalkOrder::PreOrder)
.value("POST_ORDER", PyWalkOrder::PostOrder);
+
+ nb::enum_<PyOperationEquivalenceFlags>(m, "OperationEquivalenceFlags",
+ nb::is_arithmetic(), nb::is_flag())
+ .value("NONE", PyOperationEquivalenceFlags::None)
+ .value("IGNORE_LOCATIONS", PyOperationEquivalenceFlags::IgnoreLocations)
+ .value("IGNORE_DISCARDABLE_ATTRS",
+ PyOperationEquivalenceFlags::IgnoreDiscardableAttrs)
+ .value("IGNORE_PROPERTIES", PyOperationEquivalenceFlags::IgnoreProperties)
+ .value("IGNORE_COMMUTATIVITY",
+ PyOperationEquivalenceFlags::IgnoreCommutativity);
+
nb::enum_<PyWalkResult>(m, "WalkResult")
.value("ADVANCE", PyWalkResult::Advance)
.value("INTERRUPT", PyWalkResult::Interrupt)
@@ -3946,6 +3957,34 @@ void populateIRCore(nb::module_ &m) {
return mlirOperationHashValue(self.getOperation().get());
},
"Returns the hash value of the operation.")
+ .def(
+ "is_structurally_equivalent",
+ [](PyOperationBase &self, PyOperationBase &other,
+ PyOperationEquivalenceFlags flags) {
+ self.getOperation().checkValid();
+ other.getOperation().checkValid();
+ return mlirOperationIsStructurallyEquivalent(
+ self.getOperation().get(), other.getOperation().get(),
+ static_cast<uint32_t>(flags));
+ },
+ "other"_a, "flags"_a = PyOperationEquivalenceFlags::None,
+ R"(Checks whether two operations are structurally equivalent.
+
+The IGNORE_LOCATIONS, IGNORE_DISCARDABLE_ATTRS, IGNORE_PROPERTIES, and
+IGNORE_COMMUTATIVITY flags ignore locations, discardable attributes, properties,
+and commutativity, respectively. The predicate recursively compares regions.)")
+ .def(
+ "structural_hash",
+ [](PyOperationBase &self, PyOperationEquivalenceFlags flags) {
+ self.getOperation().checkValid();
+ return mlirOperationStructuralHashValue(
+ self.getOperation().get(), static_cast<uint32_t>(flags));
+ },
+ "flags"_a = PyOperationEquivalenceFlags::None,
+ R"(Computes a structural hash for the operation.
+
+Use the same flags as is_structurally_equivalent. The hash does not recurse into
+regions, unlike the predicate.)")
.def_prop_ro(
"attributes",
[](PyOperationBase &self) {
diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index 9f1ad01cbcc88..e061d4711c21c 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -1231,6 +1231,77 @@ def testOperationHash():
assert hash(op2) == hash(custom_op2)
+# CHECK-LABEL: TEST: testOperationEquivalence
+ at run
+def testOperationEquivalence():
+ ctx = Context()
+ with ctx:
+ module = Module.parse(
+ r"""
+ func.func @f() -> i32 {
+ %0 = arith.constant 42 : i32
+ %1 = arith.constant 42 : i32
+ %2 = arith.constant 7 : i32
+ %3 = arith.subi %0, %2 : i32
+ %4 = arith.subi %0, %2 : i32
+ %5 = arith.subi %2, %0 : i32
+ %6 = arith.subi %0, %2 {dialect.discardable} : i32
+ %7 = arith.addi %0, %2 : i32
+ %8 = arith.addi %2, %0 : i32
+ return %0 : i32
+ }
+ """,
+ ctx,
+ )
+ operations = list(module.body.operations[0].regions[0].blocks[0].operations)
+ c42a, c42b, c7, sub3, sub4, sub5, sub6, add7, add8 = operations[:-1]
+ flags = OperationEquivalenceFlags
+ ignore_locations = flags.IGNORE_LOCATIONS
+
+ assert c42a.is_structurally_equivalent(c42b, ignore_locations)
+ assert not c42a.is_structurally_equivalent(c42b)
+ assert c42a.is_structurally_equivalent(c42b) == c42a.is_structurally_equivalent(
+ c42b, flags.NONE
+ )
+ assert not c42a.is_structurally_equivalent(c7, ignore_locations)
+ assert c42a.is_structurally_equivalent(c42a, ignore_locations)
+ assert sub3.is_structurally_equivalent(sub4, ignore_locations)
+ assert not sub3.is_structurally_equivalent(sub5, ignore_locations)
+ assert not c42a.is_structurally_equivalent(sub3, ignore_locations)
+
+ discardable_flags = (
+ ignore_locations | flags.IGNORE_DISCARDABLE_ATTRS
+ )
+ assert not sub3.is_structurally_equivalent(sub6, ignore_locations)
+ assert sub3.is_structurally_equivalent(sub6, discardable_flags)
+
+ assert add7.is_structurally_equivalent(add8, ignore_locations)
+ assert not add7.is_structurally_equivalent(
+ add8, ignore_locations | flags.IGNORE_COMMUTATIVITY
+ )
+
+ properties_flags = ignore_locations | flags.IGNORE_PROPERTIES
+ assert c42a.is_structurally_equivalent(c7, properties_flags)
+
+ assert c42a.structural_hash(ignore_locations) == c42b.structural_hash(
+ ignore_locations
+ )
+ assert sub3.structural_hash(ignore_locations) == sub4.structural_hash(
+ ignore_locations
+ )
+ assert sub3.structural_hash(discardable_flags) == sub6.structural_hash(
+ discardable_flags
+ )
+ assert add7.structural_hash(ignore_locations) == add8.structural_hash(
+ ignore_locations
+ )
+ assert c42a.structural_hash(properties_flags) == c7.structural_hash(
+ properties_flags
+ )
+ # CHECK: operation equivalence passed
+ print("operation equivalence passed")
+
+
# CHECK-LABEL: TEST: testOperationParse
@run
def testOperationParse():
>From a67d182ce1e66fe19e5f9531fc0cde18c65cfa8f Mon Sep 17 00:00:00 2001
From: Trevor McKay <mail at tmckay.com>
Date: Fri, 14 Aug 2026 10:06:48 -0700
Subject: [PATCH 2/4] format python
---
mlir/test/python/ir/operation.py | 4 +---
1 file changed, 1 insertion(+), 3 deletions(-)
diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index e061d4711c21c..ffb3694f0821f 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -1269,9 +1269,7 @@ def testOperationEquivalence():
assert not sub3.is_structurally_equivalent(sub5, ignore_locations)
assert not c42a.is_structurally_equivalent(sub3, ignore_locations)
- discardable_flags = (
- ignore_locations | flags.IGNORE_DISCARDABLE_ATTRS
- )
+ discardable_flags = ignore_locations | flags.IGNORE_DISCARDABLE_ATTRS
assert not sub3.is_structurally_equivalent(sub6, ignore_locations)
assert sub3.is_structurally_equivalent(sub6, discardable_flags)
>From f90d0cf8ffdb02bc9cf3396cb5482ecf25f7a516 Mon Sep 17 00:00:00 2001
From: Trevor McKay <mail at tmckay.com>
Date: Mon, 17 Aug 2026 17:19:35 -0700
Subject: [PATCH 3/4] feedback: rewrite docstrings
---
mlir/lib/Bindings/Python/IRCore.cpp | 11 ++---------
1 file changed, 2 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 95eb3d857db41..40633505bd096 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3968,11 +3968,7 @@ void populateIRCore(nb::module_ &m) {
static_cast<uint32_t>(flags));
},
"other"_a, "flags"_a = PyOperationEquivalenceFlags::None,
- R"(Checks whether two operations are structurally equivalent.
-
-The IGNORE_LOCATIONS, IGNORE_DISCARDABLE_ATTRS, IGNORE_PROPERTIES, and
-IGNORE_COMMUTATIVITY flags ignore locations, discardable attributes, properties,
-and commutativity, respectively. The predicate recursively compares regions.)")
+ R"("Checks whether two operations are structurally equivalent. The predicate recursively compares regions.")")
.def(
"structural_hash",
[](PyOperationBase &self, PyOperationEquivalenceFlags flags) {
@@ -3981,10 +3977,7 @@ and commutativity, respectively. The predicate recursively compares regions.)")
self.getOperation().get(), static_cast<uint32_t>(flags));
},
"flags"_a = PyOperationEquivalenceFlags::None,
- R"(Computes a structural hash for the operation.
-
-Use the same flags as is_structurally_equivalent. The hash does not recurse into
-regions, unlike the predicate.)")
+ R"(Computes a structural hash for the operation. The hash does not recurse into regions, unlike the predicate.")")
.def_prop_ro(
"attributes",
[](PyOperationBase &self) {
>From 4110636de498da6e6742f1776ab8e1d7964eca6f Mon Sep 17 00:00:00 2001
From: Trevor McKay <mail at tmckay.com>
Date: Mon, 17 Aug 2026 18:05:46 -0700
Subject: [PATCH 4/4] feedback: refactor tests
---
mlir/test/python/ir/operation.py | 163 +++++++++++++++++++++++--------
1 file changed, 124 insertions(+), 39 deletions(-)
diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index ffb3694f0821f..c6c2fb8219dd2 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -1231,9 +1231,9 @@ def testOperationHash():
assert hash(op2) == hash(custom_op2)
-# CHECK-LABEL: TEST: testOperationEquivalence
+# CHECK-LABEL: TEST: test_structural_equivalence
@run
-def testOperationEquivalence():
+def test_structural_equivalence():
ctx = Context()
with ctx:
module = Module.parse(
@@ -1244,60 +1244,145 @@ def testOperationEquivalence():
%2 = arith.constant 7 : i32
%3 = arith.subi %0, %2 : i32
%4 = arith.subi %0, %2 : i32
- %5 = arith.subi %2, %0 : i32
- %6 = arith.subi %0, %2 {dialect.discardable} : i32
- %7 = arith.addi %0, %2 : i32
- %8 = arith.addi %2, %0 : i32
+ %5 = arith.subi %0, %2 {dialect.discardable} : i32
+ %6 = arith.addi %0, %2 : i32
+ %7 = arith.addi %2, %0 : i32
return %0 : i32
}
""",
ctx,
)
operations = list(module.body.operations[0].regions[0].blocks[0].operations)
- c42a, c42b, c7, sub3, sub4, sub5, sub6, add7, add8 = operations[:-1]
- flags = OperationEquivalenceFlags
- ignore_locations = flags.IGNORE_LOCATIONS
+ c42a, c42b, c7, sub3, sub4, sub5, add6, add7, *_ = operations
- assert c42a.is_structurally_equivalent(c42b, ignore_locations)
assert not c42a.is_structurally_equivalent(c42b)
- assert c42a.is_structurally_equivalent(c42b) == c42a.is_structurally_equivalent(
- c42b, flags.NONE
+ assert c42a.is_structurally_equivalent(
+ c42b, OperationEquivalenceFlags.IGNORE_LOCATIONS
)
- assert not c42a.is_structurally_equivalent(c7, ignore_locations)
- assert c42a.is_structurally_equivalent(c42a, ignore_locations)
- assert sub3.is_structurally_equivalent(sub4, ignore_locations)
- assert not sub3.is_structurally_equivalent(sub5, ignore_locations)
- assert not c42a.is_structurally_equivalent(sub3, ignore_locations)
-
- discardable_flags = ignore_locations | flags.IGNORE_DISCARDABLE_ATTRS
- assert not sub3.is_structurally_equivalent(sub6, ignore_locations)
- assert sub3.is_structurally_equivalent(sub6, discardable_flags)
-
- assert add7.is_structurally_equivalent(add8, ignore_locations)
- assert not add7.is_structurally_equivalent(
- add8, ignore_locations | flags.IGNORE_COMMUTATIVITY
+
+ assert not c42a.is_structurally_equivalent(
+ c7, OperationEquivalenceFlags.IGNORE_LOCATIONS
)
- properties_flags = ignore_locations | flags.IGNORE_PROPERTIES
- assert c42a.is_structurally_equivalent(c7, properties_flags)
+ assert c42a.is_structurally_equivalent(
+ c42a, OperationEquivalenceFlags.IGNORE_LOCATIONS
+ )
+
+ assert sub3.is_structurally_equivalent(
+ sub4, OperationEquivalenceFlags.IGNORE_LOCATIONS
+ )
+
+ assert not c42a.is_structurally_equivalent(
+ sub3, OperationEquivalenceFlags.IGNORE_LOCATIONS
+ )
+
+ discardable_flags = (
+ OperationEquivalenceFlags.IGNORE_LOCATIONS
+ | OperationEquivalenceFlags.IGNORE_DISCARDABLE_ATTRS
+ )
+
+ assert not sub3.is_structurally_equivalent(
+ sub5, OperationEquivalenceFlags.IGNORE_LOCATIONS
+ )
+ assert sub3.is_structurally_equivalent(sub5, discardable_flags)
+
+ assert add6.is_structurally_equivalent(
+ add7, OperationEquivalenceFlags.IGNORE_LOCATIONS
+ )
+ assert not add6.is_structurally_equivalent(
+ add7,
+ OperationEquivalenceFlags.IGNORE_LOCATIONS
+ | OperationEquivalenceFlags.IGNORE_COMMUTATIVITY,
+ )
- assert c42a.structural_hash(ignore_locations) == c42b.structural_hash(
- ignore_locations
+ properties_flags = (
+ OperationEquivalenceFlags.IGNORE_LOCATIONS
+ | OperationEquivalenceFlags.IGNORE_PROPERTIES
)
- assert sub3.structural_hash(ignore_locations) == sub4.structural_hash(
- ignore_locations
+
+ assert c42a.is_structurally_equivalent(c7, properties_flags)
+
+ # CHECK: structural equivalence passed
+ print("structural equivalence passed")
+
+
+# CHECK-LABEL: TEST: test_structural_hash
+ at run
+def test_structural_hash():
+ ctx = Context()
+ with ctx:
+ module = Module.parse(
+ r"""
+ func.func @f() -> i32 {
+ %0 = arith.constant 42 : i32
+ %1 = arith.constant 42 : i32
+ %2 = arith.constant 7 : i32
+ %3 = arith.subi %0, %2 : i32
+ %4 = arith.subi %0, %2 : i32
+ %5 = arith.subi %0, %2 {dialect.discardable} : i32
+ %6 = arith.addi %0, %2 : i32
+ %7 = arith.addi %2, %0 : i32
+ return %0 : i32
+ }
+ """,
+ ctx,
)
- assert sub3.structural_hash(discardable_flags) == sub6.structural_hash(
- discardable_flags
+ operations = list(module.body.operations[0].regions[0].blocks[0].operations)
+ c42a, c42b, c7, sub3, sub4, sub5, add6, add7, *_ = operations
+
+ discardable_flags = (
+ OperationEquivalenceFlags.IGNORE_LOCATIONS
+ | OperationEquivalenceFlags.IGNORE_DISCARDABLE_ATTRS
)
- assert add7.structural_hash(ignore_locations) == add8.structural_hash(
- ignore_locations
+
+ properties_flags = (
+ OperationEquivalenceFlags.IGNORE_LOCATIONS
+ | OperationEquivalenceFlags.IGNORE_PROPERTIES
)
- assert c42a.structural_hash(properties_flags) == c7.structural_hash(
- properties_flags
+
+ c42a_hash = c42a.structural_hash()
+ c42b_hash = c42b.structural_hash()
+ assert c42a_hash != c42b_hash
+
+ c42a_hash = c42a.structural_hash(OperationEquivalenceFlags.IGNORE_LOCATIONS)
+ c42b_hash = c42b.structural_hash(OperationEquivalenceFlags.IGNORE_LOCATIONS)
+ assert c42a_hash == c42b_hash
+
+ sub3_hash = sub3.structural_hash(OperationEquivalenceFlags.IGNORE_LOCATIONS)
+ sub5_hash = sub5.structural_hash(OperationEquivalenceFlags.IGNORE_LOCATIONS)
+ assert sub3_hash != sub5_hash
+
+ sub3_hash = sub3.structural_hash(OperationEquivalenceFlags.IGNORE_LOCATIONS)
+ sub4_hash = sub4.structural_hash(OperationEquivalenceFlags.IGNORE_LOCATIONS)
+ assert sub3_hash == sub4_hash
+
+ sub3_hash = sub3.structural_hash(discardable_flags)
+ sub5_hash = sub5.structural_hash(discardable_flags)
+ assert sub3_hash == sub5_hash
+
+ add6_hash = add6.structural_hash(OperationEquivalenceFlags.IGNORE_LOCATIONS)
+ add7_hash = add7.structural_hash(OperationEquivalenceFlags.IGNORE_LOCATIONS)
+ assert add6_hash == add7_hash
+
+ commutative_flags = (
+ OperationEquivalenceFlags.IGNORE_LOCATIONS
+ | OperationEquivalenceFlags.IGNORE_COMMUTATIVITY
)
- # CHECK: operation equivalence passed
- print("operation equivalence passed")
+
+ add6_hash = add6.structural_hash(commutative_flags)
+ add7_hash = add7.structural_hash(commutative_flags)
+ assert add6_hash != add7_hash
+
+ c42a_hash = c42a.structural_hash(OperationEquivalenceFlags.IGNORE_LOCATIONS)
+ c7_hash = c7.structural_hash(OperationEquivalenceFlags.IGNORE_LOCATIONS)
+ assert c42a_hash != c7_hash
+
+ c42a_hash = c42a.structural_hash(properties_flags)
+ c7_hash = c7.structural_hash(properties_flags)
+ assert c42a_hash == c7_hash
+
+ # CHECK: structural hash passed
+ print("structural hash passed")
# CHECK-LABEL: TEST: testOperationParse
More information about the Mlir-commits
mailing list