[mlir][python] Make Operation and Value hashable

[rkayaith]

Summary

This allows operations and values to be used as dict keys

[ftynse]

mlir/lib/Bindings/Python/IRCore.cpp

2175

Let's use llvm::hash_value that has the desirable properties of hash values instead of blindly reusing the pointer. Also, please prefer C++ static_cast to C-style casts in C++ code.

[ftynse]

mlir/test/python/ir/operation.py

754

I don't think hashes are guaranteed to be different. They happen to be because they are currently pointer addresses.

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

2175

Should I update the hash functions for PyType and PyAttribute this way as well?

Also I noticed that there's llvm::hash_value overloads for Value, Type, and Attribute (but not Operation) defined already, should I thread those through the C bindings and use them here? Or is it sufficient to just hash the pointers that are currently being returned?

[rkayaith]

hash pointers instead of returning them directly

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

2175

I went with hashing the pointers in here for now, and also updated PyType and PyAttribute to do this as well.

mlir/test/python/ir/operation.py

754

I removed all the tests that checked for hashes being different, since pointer addresses are no longer being used directly.

[ftynse]

Let's hash pointers for now, my primary concern was the loss of hash function properties that defies the purpose of hash tables. FYI, there's no llvm::hash_value overload for Operation because we never use operations by value, only by pointer, and LLVM already knows how to hash pointers.

[rkayaith]

Thanks, makes sense. Could you land the changes for me?

mlir/lib/Bindings/Python/IRCore.cpp

py::class_<PyOperationBase>(m, "_OperationBase", py::module_local())
                                return self.getOperation().getCapsule();
                              })
       .def("__eq__",
            [](PyOperationBase &self, PyOperationBase &other) {
              return &self.getOperation() == &other.getOperation();
            })
       .def("__eq__",
            [](PyOperationBase &self, py::object other) { return false; })
+      .def("__hash__",
+           [](PyOperationBase &self) {
+             return static_cast<size_t>(llvm::hash_value(&self.getOperation()));
+           })
       .def_property_readonly("attributes",
                              [](PyOperationBase &self) {
                                return PyOpAttributeMap(
                                    self.getOperation().getRef());
                              })
       .def_property_readonly("operands",
                              [](PyOperationBase &self) {
                                return PyOpOperandList(
py::class_<PyAttribute>(m, "Attribute", py::module_local())
           "get_named",
           [](PyAttribute &self, std::string name) {
             return PyNamedAttribute(self, std::move(name));
           },
           py::keep_alive<0, 1>(), "Binds a name to the attribute")
       .def("__eq__",
            [](PyAttribute &self, PyAttribute &other) { return self == other; })
       .def("__eq__", [](PyAttribute &self, py::object &other) { return false; })
-      .def("__hash__", [](PyAttribute &self) { return (size_t)self.get().ptr; })
+      .def("__hash__",
+           [](PyAttribute &self) {
+             return static_cast<size_t>(llvm::hash_value(self.get().ptr));
+           })
       .def(
           "dump", [](PyAttribute &self) { mlirAttributeDump(self); },
           kDumpDocstring)
       .def(
           "__str__",
           [](PyAttribute &self) {
             PyPrintAccumulator printAccum;
             mlirAttributePrint(self, printAccum.getCallback(),
py::class_<PyType>(m, "Type", py::module_local())
           },
           py::arg("asm"), py::arg("context") = py::none(),
           kContextParseTypeDocstring)
       .def_property_readonly(
           "context", [](PyType &self) { return self.getContext().getObject(); },
           "Context that owns the Type")
       .def("__eq__", [](PyType &self, PyType &other) { return self == other; })
       .def("__eq__", [](PyType &self, py::object &other) { return false; })
-      .def("__hash__", [](PyType &self) { return (size_t)self.get().ptr; })
+      .def("__hash__",
+           [](PyType &self) {
+             return static_cast<size_t>(llvm::hash_value(self.get().ptr));
+           })
       .def(
           "dump", [](PyType &self) { mlirTypeDump(self); }, kDumpDocstring)
       .def(
           "__str__",
           [](PyType &self) {
             PyPrintAccumulator printAccum;
             mlirTypePrint(self, printAccum.getCallback(),
                           printAccum.getUserData());
py::class_<PyValue>(m, "Value", py::module_local())
                    "the IR");
             return self.getParentOperation().getObject();
           })
       .def("__eq__",
            [](PyValue &self, PyValue &other) {
              return self.get().ptr == other.get().ptr;
            })
       .def("__eq__", [](PyValue &self, py::object other) { return false; })
+      .def("__hash__",
+           [](PyValue &self) {
+             return static_cast<size_t>(llvm::hash_value(self.get().ptr));
+           })
       .def(
           "__str__",
           [](PyValue &self) {
             PyPrintAccumulator printAccum;
             printAccum.parts.append("Value(");
             mlirValuePrint(self.get(), printAccum.getCallback(),
                            printAccum.getUserData());
             printAccum.parts.append(")");

mlir/test/python/ir/attributes.py

 @run
 def testAttrHash():
   with Context():
     a1 = Attribute.parse('"attr1"')
     a2 = Attribute.parse('"attr2"')
     a3 = Attribute.parse('"attr1"')
     # CHECK: hash(a1) == hash(a3): True
     print("hash(a1) == hash(a3):", a1.__hash__() == a3.__hash__())
-    # In general, hashes don't have to be unique. In this case, however, the
-    # hash is just the underlying pointer so it will be.
-    # CHECK: hash(a1) == hash(a2): False
-    print("hash(a1) == hash(a2):", a1.__hash__() == a2.__hash__())
 
     s = set()
     s.add(a1)
     s.add(a2)
     s.add(a3)
     # CHECK: len(s): 2
     print("len(s): ", len(s))
 

mlir/test/python/ir/builtin_types.py

 def testTypeHash():
   ctx = Context()
   t1 = Type.parse("i32", ctx)
   t2 = Type.parse("f32", ctx)
   t3 = Type.parse("i32", ctx)
 
   # CHECK: hash(t1) == hash(t3): True
   print("hash(t1) == hash(t3):", t1.__hash__() == t3.__hash__())
-  # In general, hashes don't have to be unique. In this case, however, the
-  # hash is just the underlying pointer so it will be.
-  # CHECK: hash(t1) == hash(t2): False
-  print("hash(t1) == hash(t2):", t1.__hash__() == t2.__hash__())
 
   s = set()
   s.add(t1)
   s.add(t2)
   s.add(t3)
   # CHECK: len(s): 2
   print("len(s): ", len(s))
 

mlir/test/python/ir/operation.py

def testOperationLoc():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with ctx:
     loc = Location.name("loc")
     op = Operation.create("custom.op", loc=loc)
     assert op.location == loc
     assert op.operation.location == loc
 
+
 # CHECK-LABEL: TEST: testModuleMerge
 @run
 def testModuleMerge():
   with Context():
     m1 = Module.parse("func private @foo()")
     m2 = Module.parse("""
       func private @bar()
       func private @qux()
     """)
     foo = m1.body.operations[0]
     bar = m2.body.operations[0]
     qux = m2.body.operations[1]
     bar.move_before(foo)
     qux.move_after(foo)
 
     # CHECK: module
     # CHECK: func private @bar
     # CHECK: func private @foo
with Context() as ctx:
 
     try:
       symbol_table.insert(m2.body.operations[0])
     except ValueError as e:
       if "Expected operation to have a symbol name" not in str(e):
         raise
     else:
       assert False, "exepcted ValueError when adding a non-symbol"
+
+
+# CHECK-LABEL: TEST: testOperationHash
+@run
+def testOperationHash():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  with ctx, Location.unknown():
+    op = Operation.create("custom.op1")
+    assert hash(op) == hash(op.operation)

mlir/test/python/ir/value.py

module = Module.parse(
     }""", ctx)
   func = module.body.operations[0]
   assert BlockArgument.isinstance(func.regions[0].blocks[0].arguments[0])
   assert not OpResult.isinstance(func.regions[0].blocks[0].arguments[0])
 
   op = func.regions[0].blocks[0].operations[0]
   assert not BlockArgument.isinstance(op.results[0])
   assert OpResult.isinstance(op.results[0])
+
+
+# CHECK-LABEL: TEST: testValueHash
+@run
+def testValueHash():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  module = Module.parse(
+      r"""
+    func @foo(%arg0: f32) -> f32 {
+      %0 = "some_dialect.some_op"(%arg0) : (f32) -> f32
+      return %0 : f32
+    }""", ctx)
+
+  [func] = module.body.operations
+  block = func.entry_block
+  op, ret = block.operations
+  assert hash(block.arguments[0]) == hash(op.operands[0])
+  assert hash(op.result) == hash(ret.operands[0])