[Mlir-commits] [mlir] 70a291f - [MLIR][Python] fix operation hashing (#156514)

[llvmlistbot at llvm.org]

Author: Maksim Levental
Date: 2025-09-02T15:12:25-05:00
New Revision: 70a291f3225628d3479452829dc8b72d5d04e034

URL: https://github.com/llvm/llvm-project/commit/70a291f3225628d3479452829dc8b72d5d04e034
DIFF: https://github.com/llvm/llvm-project/commit/70a291f3225628d3479452829dc8b72d5d04e034.diff

LOG: [MLIR][Python] fix operation hashing (#156514)

https://github.com/llvm/llvm-project/pull/155114 broke op hashing
(because the python objects ceased to be reference equivalent). This PR
fixes by binding `OperationEquivalence::computeHash`.

Added: 
    

Modified: 
    mlir/include/mlir-c/IR.h
    mlir/lib/Bindings/Python/IRCore.cpp
    mlir/lib/CAPI/IR/IR.cpp
    mlir/test/python/ir/operation.py

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index e97369778b377..061d7620ba077 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -418,6 +418,9 @@ MLIR_CAPI_EXPORTED MlirModule mlirModuleFromOperation(MlirOperation op);
 /// Checks if two modules are equal.
 MLIR_CAPI_EXPORTED bool mlirModuleEqual(MlirModule lhs, MlirModule rhs);
 
+/// Compute a hash for the given module.
+MLIR_CAPI_EXPORTED size_t mlirModuleHashValue(MlirModule mod);
+
 //===----------------------------------------------------------------------===//
 // Operation state.
 //===----------------------------------------------------------------------===//
@@ -622,6 +625,9 @@ static inline bool mlirOperationIsNull(MlirOperation op) { return !op.ptr; }
 MLIR_CAPI_EXPORTED bool mlirOperationEqual(MlirOperation op,
                                            MlirOperation other);
 
+/// Compute a hash for the given operation.
+MLIR_CAPI_EXPORTED size_t mlirOperationHashValue(MlirOperation op);
+
 /// Gets the context this operation is associated with
 MLIR_CAPI_EXPORTED MlirContext mlirOperationGetContext(MlirOperation op);
 

diff  --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 2df2a73fd88ff..bf4950fc1a070 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3317,7 +3317,9 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           [](PyModule &self, PyModule &other) {
             return mlirModuleEqual(self.get(), other.get());
           },
-          "other"_a);
+          "other"_a)
+      .def("__hash__",
+           [](PyModule &self) { return mlirModuleHashValue(self.get()); });
 
   //----------------------------------------------------------------------------
   // Mapping of Operation.
@@ -3336,7 +3338,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
            [](PyOperationBase &self, nb::object other) { return false; })
       .def("__hash__",
            [](PyOperationBase &self) {
-             return static_cast<size_t>(llvm::hash_value(&self.getOperation()));
+             return mlirOperationHashValue(self.getOperation().get());
            })
       .def_prop_ro("attributes",
                    [](PyOperationBase &self) {

diff  --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index c7069f0017b5d..e9844a7cc1909 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -469,6 +469,10 @@ bool mlirModuleEqual(MlirModule lhs, MlirModule rhs) {
   return unwrap(lhs) == unwrap(rhs);
 }
 
+size_t mlirModuleHashValue(MlirModule mod) {
+  return OperationEquivalence::computeHash(unwrap(mod).getOperation());
+}
+
 //===----------------------------------------------------------------------===//
 // Operation state API.
 //===----------------------------------------------------------------------===//
@@ -640,6 +644,10 @@ bool mlirOperationEqual(MlirOperation op, MlirOperation other) {
   return unwrap(op) == unwrap(other);
 }
 
+size_t mlirOperationHashValue(MlirOperation op) {
+  return OperationEquivalence::computeHash(unwrap(op));
+}
+
 MlirContext mlirOperationGetContext(MlirOperation op) {
   return wrap(unwrap(op)->getContext());
 }

diff  --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index 7759b1797e3c3..4a3625c953d52 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -1051,6 +1051,12 @@ def testOperationHash():
         op = Operation.create("custom.op1")
         assert hash(op) == hash(op.operation)
 
+        module = Module.create()
+        with InsertionPoint(module.body):
+            op2 = Operation.create("custom.op2")
+            custom_op2 = module.body.operations[0]
+            assert hash(op2) == hash(custom_op2)
+
 
 # CHECK-LABEL: TEST: testOperationParse
 @run