[MLIR] [Python] Make Attribute and Type hashable

[jdd]

Summary

Enables putting types and attributes in dicts and sets.

[jdd]

Thanks Stella!

mlir/lib/Bindings/Python/IRCore.cpp

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(
           "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(
           "dump", [](PyType &self) { mlirTypeDump(self); }, kDumpDocstring)
       .def(
           "__str__",
           [](PyType &self) {
             PyPrintAccumulator printAccum;
             mlirTypePrint(self, printAccum.getCallback(),
                           printAccum.getUserData());

mlir/test/python/ir/attributes.py

with Context():
     # CHECK: a1 == a2: False
     print("a1 == a2:", a1 == a2)
     # CHECK: a1 == a3: True
     print("a1 == a3:", a1 == a3)
     # CHECK: a1 == None: False
     print("a1 == None:", a1 == None)
 
 
+# CHECK-LABEL: TEST: testAttrHash
+@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))
+
+
 # CHECK-LABEL: TEST: testAttrCast
 @run
 def testAttrCast():
   with Context():
     a1 = Attribute.parse('"attr1"')
     a2 = Attribute(a1)
     # CHECK: a1 == a2: True
     print("a1 == a2:", a1 == a2)
with Context():
     except RuntimeError as e:
       # CHECK: Error: Invalid attribute (None?) when attempting to create an ArrayAttribute
       print("Error: ", e)
     try:
       ArrayAttr.get([42])
     except RuntimeError as e:
       # CHECK: Error: Invalid attribute when attempting to create an ArrayAttribute
       print("Error: ", e)
-

mlir/test/python/ir/builtin_types.py

def testTypeEq():
   # CHECK: t1 == t2: False
   print("t1 == t2:", t1 == t2)
   # CHECK: t1 == t3: True
   print("t1 == t3:", t1 == t3)
   # CHECK: t1 == None: False
   print("t1 == None:", t1 == None)
 
 
+# CHECK-LABEL: TEST: testTypeHash
+@run
+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))
+
 # CHECK-LABEL: TEST: testTypeCast
 @run
 def testTypeCast():
   ctx = Context()
   t1 = Type.parse("i32", ctx)
   t2 = Type(t1)
   # CHECK: t1 == t2: True
   print("t1 == t2:", t1 == t2)