From 2df299d47d94e94a538bfb43773fd9e285da93db Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental@gmail.com>
Date: Tue, 2 Sep 2025 11:41:09 -0700
Subject: [PATCH] [MLIR][Python] fix operation hashing

---
 mlir/include/mlir-c/IR.h            | 6 ++++++
 mlir/lib/Bindings/Python/IRCore.cpp | 6 ++++--
 mlir/lib/CAPI/IR/IR.cpp             | 8 ++++++++
 mlir/test/python/ir/operation.py    | 6 ++++++
 4 files changed, 24 insertions(+), 2 deletions(-)

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