[MLIR][Python] fix operation hashing

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);
 

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) {

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());
 }

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