[mlir][python] Downcast attributes in more places

[rkayaith]

Summary

Update remaining PyAttribute-returning APIs to return MlirAttribute instead,
so that they go through the downcasting mechanism.

[makslevental]

LGTM

mlir/lib/Bindings/Python/IRAttributes.cpp

static constexpr GetTypeIDFunctionTy getTypeIdFunction =
       mlirArrayAttrGetTypeID;
 
   class PyArrayAttributeIterator {
   public:
     PyArrayAttributeIterator(PyAttribute attr) : attr(std::move(attr)) {}
 
     PyArrayAttributeIterator &dunderIter() { return *this; }
 
-    PyAttribute dunderNext() {
+    MlirAttribute dunderNext() {
       // TODO: Throw is an inefficient way to stop iteration.
       if (nextIndex >= mlirArrayAttrGetNumElements(attr.get()))
         throw py::stop_iteration();
-      return PyAttribute(attr.getContext(),
-                         mlirArrayAttrGetElement(attr.get(), nextIndex++));
+      return mlirArrayAttrGetElement(attr.get(), nextIndex++);
     }
 
     static void bind(py::module &m) {
       py::class_<PyArrayAttributeIterator>(m, "ArrayAttributeIterator",
                                            py::module_local())
           .def("__iter__", &PyArrayAttributeIterator::dunderIter)
           .def("__next__", &PyArrayAttributeIterator::dunderNext);
     }
 
   private:
     PyAttribute attr;
     int nextIndex = 0;
   };
 
-  PyAttribute getItem(intptr_t i) {
-    return PyAttribute(getContext(), mlirArrayAttrGetElement(*this, i));
+  MlirAttribute getItem(intptr_t i) {
+    return mlirArrayAttrGetElement(*this, i);
   }
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](py::list attributes, DefaultingPyMlirContext context) {
           SmallVector<MlirAttribute> mlirAttributes;
           mlirAttributes.reserve(py::len(attributes));
c.def("__len__", &PyDenseElementsAttribute::dunderLen)
         .def_static("get_splat", PyDenseElementsAttribute::getSplat,
                     py::arg("shaped_type"), py::arg("element_attr"),
                     "Gets a DenseElementsAttr where all values are the same")
         .def_property_readonly("is_splat",
                                [](PyDenseElementsAttribute &self) -> bool {
                                  return mlirDenseElementsAttrIsSplat(self);
                                })
         .def("get_splat_value",
-             [](PyDenseElementsAttribute &self) -> PyAttribute {
-               if (!mlirDenseElementsAttrIsSplat(self)) {
+             [](PyDenseElementsAttribute &self) {
+               if (!mlirDenseElementsAttrIsSplat(self))
                  throw py::value_error(
                      "get_splat_value called on a non-splat attribute");
-               }
-               return PyAttribute(self.getContext(),
-                                  mlirDenseElementsAttrGetSplatValue(self));
+               return mlirDenseElementsAttrGetSplatValue(self);
              })
         .def_buffer(&PyDenseElementsAttribute::accessBuffer);
   }
 
 private:
   static bool isUnsignedIntegerFormat(const std::string &format) {
     if (format.empty())
       return false;
c.def_static(
                                     mlirNamedAttributes.data());
           return PyDictAttribute(context->getRef(), attr);
         },
         py::arg("value") = py::dict(), py::arg("context") = py::none(),
         "Gets an uniqued dict attribute");
     c.def("__getitem__", [](PyDictAttribute &self, const std::string &name) {
       MlirAttribute attr =
           mlirDictionaryAttrGetElementByName(self, toMlirStringRef(name));
-      if (mlirAttributeIsNull(attr)) {
+      if (mlirAttributeIsNull(attr))
         throw py::key_error("attempt to access a non-existent attribute");
-      }
-      return PyAttribute(self.getContext(), attr);
+      return attr;
     });
     c.def("__getitem__", [](PyDictAttribute &self, intptr_t index) {
       if (index < 0 || index >= self.dunderLen()) {
         throw py::index_error("attempt to access out of bounds attribute");
       }
       MlirNamedAttribute namedAttr = mlirDictionaryAttrGetElement(self, index);
       return PyNamedAttribute(
           namedAttr.attribute,

mlir/lib/Bindings/Python/IRCore.cpp

void PySymbolTable::erase(PyOperationBase &symbol) {
   symbol.getOperation().valid = false;
 }
 
 void PySymbolTable::dunderDel(const std::string &name) {
   py::object operation = dunderGetItem(name);
   erase(py::cast<PyOperationBase &>(operation));
 }
 
-PyAttribute PySymbolTable::insert(PyOperationBase &symbol) {
+MlirAttribute PySymbolTable::insert(PyOperationBase &symbol) {
   operation->checkValid();
   symbol.getOperation().checkValid();
   MlirAttribute symbolAttr = mlirOperationGetAttributeByName(
       symbol.getOperation().get(), mlirSymbolTableGetSymbolAttributeName());
   if (mlirAttributeIsNull(symbolAttr))
     throw py::value_error("Expected operation to have a symbol name.");
-  return PyAttribute(
-      symbol.getOperation().getContext(),
-      mlirSymbolTableInsert(symbolTable, symbol.getOperation().get()));
+  return mlirSymbolTableInsert(symbolTable, symbol.getOperation().get());
 }
 
-PyAttribute PySymbolTable::getSymbolName(PyOperationBase &symbol) {
+MlirAttribute PySymbolTable::getSymbolName(PyOperationBase &symbol) {
   // Op must already be a symbol.
   PyOperation &operation = symbol.getOperation();
   operation.checkValid();
   MlirStringRef attrName = mlirSymbolTableGetSymbolAttributeName();
   MlirAttribute existingNameAttr =
       mlirOperationGetAttributeByName(operation.get(), attrName);
   if (mlirAttributeIsNull(existingNameAttr))
     throw py::value_error("Expected operation to have a symbol name.");
-  return PyAttribute(symbol.getOperation().getContext(), existingNameAttr);
+  return existingNameAttr;
 }
 
 void PySymbolTable::setSymbolName(PyOperationBase &symbol,
                                   const std::string &name) {
   // Op must already be a symbol.
   PyOperation &operation = symbol.getOperation();
   operation.checkValid();
   MlirStringRef attrName = mlirSymbolTableGetSymbolAttributeName();
   MlirAttribute existingNameAttr =
       mlirOperationGetAttributeByName(operation.get(), attrName);
   if (mlirAttributeIsNull(existingNameAttr))
     throw py::value_error("Expected operation to have a symbol name.");
   MlirAttribute newNameAttr =
       mlirStringAttrGet(operation.getContext()->get(), toMlirStringRef(name));
   mlirOperationSetAttributeByName(operation.get(), attrName, newNameAttr);
 }
 
-PyAttribute PySymbolTable::getVisibility(PyOperationBase &symbol) {
+MlirAttribute PySymbolTable::getVisibility(PyOperationBase &symbol) {
   PyOperation &operation = symbol.getOperation();
   operation.checkValid();
   MlirStringRef attrName = mlirSymbolTableGetVisibilityAttributeName();
   MlirAttribute existingVisAttr =
       mlirOperationGetAttributeByName(operation.get(), attrName);
   if (mlirAttributeIsNull(existingVisAttr))
     throw py::value_error("Expected operation to have a symbol visibility.");
-  return PyAttribute(symbol.getOperation().getContext(), existingVisAttr);
+  return existingVisAttr;
 }
 
 void PySymbolTable::setVisibility(PyOperationBase &symbol,
                                   const std::string &visibility) {
   if (visibility != "public" && visibility != "private" &&
       visibility != "nested")
     throw py::value_error(
         "Expected visibility to be 'public', 'private' or 'nested'");

 
 /// A list of operation attributes. Can be indexed by name, producing
 /// attributes, or by index, producing named attributes.
 class PyOpAttributeMap {
 public:
   PyOpAttributeMap(PyOperationRef operation)
       : operation(std::move(operation)) {}
 
-  PyAttribute dunderGetItemNamed(const std::string &name) {
+  MlirAttribute dunderGetItemNamed(const std::string &name) {
     MlirAttribute attr = mlirOperationGetAttributeByName(operation->get(),
                                                          toMlirStringRef(name));
     if (mlirAttributeIsNull(attr)) {
       throw py::key_error("attempt to access a non-existent attribute");
     }
-    return PyAttribute(operation->getContext(), attr);
+    return attr;
   }
 
   PyNamedAttribute dunderGetItemIndexed(intptr_t index) {
     if (index < 0 || index >= dunderLen()) {
       throw py::index_error("attempt to access out of bounds attribute");
     }
     MlirNamedAttribute namedAttr =
         mlirOperationGetAttribute(operation->get(), index);

mlir/lib/Bindings/Python/IRModule.h

public:
   void erase(PyOperationBase &symbol);
 
   /// Removes the operation with the given name from the symbol table and erases
   /// it, throws if there is no such symbol in the table.
   void dunderDel(const std::string &name);
 
   /// Inserts the given operation into the symbol table. The operation must have
   /// the symbol trait.
-  PyAttribute insert(PyOperationBase &symbol);
+  MlirAttribute insert(PyOperationBase &symbol);
 
   /// Gets and sets the name of a symbol op.
-  static PyAttribute getSymbolName(PyOperationBase &symbol);
+  static MlirAttribute getSymbolName(PyOperationBase &symbol);
   static void setSymbolName(PyOperationBase &symbol, const std::string &name);
 
   /// Gets and sets the visibility of a symbol op.
-  static PyAttribute getVisibility(PyOperationBase &symbol);
+  static MlirAttribute getVisibility(PyOperationBase &symbol);
   static void setVisibility(PyOperationBase &symbol,
                             const std::string &visibility);
 
   /// Replaces all symbol uses within an operation. See the API
   /// mlirSymbolTableReplaceAllSymbolUses for all caveats.
   static void replaceAllSymbolUses(const std::string &oldSymbol,
                                    const std::string &newSymbol,
                                    PyOperationBase &from);

mlir/lib/Bindings/Python/IRTypes.cpp

c.def_static(
           if (mlirTypeIsNull(t))
             throw MLIRError("Invalid type", errors.take());
           return PyRankedTensorType(elementType.getContext(), t);
         },
         py::arg("shape"), py::arg("element_type"),
         py::arg("encoding") = py::none(), py::arg("loc") = py::none(),
         "Create a ranked tensor type");
     c.def_property_readonly(
-        "encoding", [](PyRankedTensorType &self) -> std::optional<PyAttribute> {
+        "encoding",
+        [](PyRankedTensorType &self) -> std::optional<MlirAttribute> {
           MlirAttribute encoding = mlirRankedTensorTypeGetEncoding(self.get());
           if (mlirAttributeIsNull(encoding))
             return std::nullopt;
-          return PyAttribute(self.getContext(), encoding);
+          return encoding;
         });
   }
 };
 
 /// Unranked Tensor Type subclass - UnrankedTensorType.
 class PyUnrankedTensorType
     : public PyConcreteType<PyUnrankedTensorType, PyShapedType> {
 public:
c.def_static(
              throw MLIRError("Invalid type", errors.take());
            return PyMemRefType(elementType.getContext(), t);
          },
          py::arg("shape"), py::arg("element_type"),
          py::arg("layout") = py::none(), py::arg("memory_space") = py::none(),
          py::arg("loc") = py::none(), "Create a memref type")
         .def_property_readonly(
             "layout",
-            [](PyMemRefType &self) -> PyAttribute {
-              MlirAttribute layout = mlirMemRefTypeGetLayout(self);
-              return PyAttribute(self.getContext(), layout);
+            [](PyMemRefType &self) -> MlirAttribute {
+              return mlirMemRefTypeGetLayout(self);
             },
             "The layout of the MemRef type.")
         .def_property_readonly(
             "affine_map",
             [](PyMemRefType &self) -> PyAffineMap {
               MlirAffineMap map = mlirMemRefTypeGetAffineMap(self);
               return PyAffineMap(self.getContext(), map);
             },
             "The layout of the MemRef type as an affine map.")
         .def_property_readonly(
             "memory_space",
-            [](PyMemRefType &self) -> PyAttribute {
+            [](PyMemRefType &self) -> std::optional<MlirAttribute> {
               MlirAttribute a = mlirMemRefTypeGetMemorySpace(self);
-              return PyAttribute(self.getContext(), a);
+              if (mlirAttributeIsNull(a))
+                return std::nullopt;
+              return a;
             },
             "Returns the memory space of the given MemRef type.");
   }
 };
 
 /// Unranked MemRef Type subclass - UnrankedMemRefType.
 class PyUnrankedMemRefType
     : public PyConcreteType<PyUnrankedMemRefType, PyShapedType> {
c.def_static(
            if (mlirTypeIsNull(t))
              throw MLIRError("Invalid type", errors.take());
            return PyUnrankedMemRefType(elementType.getContext(), t);
          },
          py::arg("element_type"), py::arg("memory_space"),
          py::arg("loc") = py::none(), "Create a unranked memref type")
         .def_property_readonly(
             "memory_space",
-            [](PyUnrankedMemRefType &self) -> PyAttribute {
-              MlirAttribute a = mlirMemRefTypeGetMemorySpace(self);
-              return PyAttribute(self.getContext(), a);
+            [](PyUnrankedMemRefType &self) -> std::optional<MlirAttribute> {
+              MlirAttribute a = mlirUnrankedMemrefGetMemorySpace(self);
+              if (mlirAttributeIsNull(a))
+                return std::nullopt;
+              return a;
             },
             "Returns the memory space of the given Unranked MemRef type.");
   }
 };
 
 /// Tuple Type subclass - TupleType.
 class PyTupleType : public PyConcreteType<PyTupleType> {
 public:

mlir/test/python/ir/array_attributes.py

with Context(), Location.unknown():
         t = IntegerType.get_signless(32)
         element = IntegerAttr.get(t, 555)
         shaped_type = RankedTensorType.get((2, 3, 4), t)
         attr = DenseElementsAttr.get_splat(shaped_type, element)
         # CHECK: dense<555> : tensor<2x3x4xi32>
         print(attr)
         # CHECK: is_splat: True
         print("is_splat:", attr.is_splat)
-        assert attr.get_splat_value() == element
+
+        # CHECK: splat_value: IntegerAttr(555 : i32)
+        splat_value = attr.get_splat_value()
+        print("splat_value:", repr(splat_value))
+        assert splat_value == element
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsSplatFloat
 @run
 def testGetDenseElementsSplatFloat():
     with Context(), Location.unknown():
         t = F32Type.get()
         element = FloatAttr.get(t, 1.2)

mlir/test/python/ir/attributes.py

with Context():
 
         a = DictAttr.get(dict_attr)
 
         # CHECK attr: {integerattr = 42 : i32, stringattr = "string"}
         print("attr:", a)
 
         assert len(a) == 2
 
-        # CHECK: 42 : i32
-        print(a["integerattr"])
+        # CHECK: integerattr: IntegerAttr(42 : i32)
+        print("integerattr:", repr(a["integerattr"]))
 
-        # CHECK: "string"
-        print(a["stringattr"])
+        # CHECK: stringattr: StringAttr("string")
+        print("stringattr:", repr(a["stringattr"]))
 
         # CHECK: True
         print("stringattr" in a)
 
         # CHECK: False
         print("not_in_dict" in a)
 
         # Check that exceptions are raised as expected.
with Context():
         # CHECK: f32
         print(ShapedType(type_attr.value).element_type)
 
 
 # CHECK-LABEL: TEST: testArrayAttr
 @run
 def testArrayAttr():
     with Context():
-        raw = Attribute.parse("[42, true, vector<4xf32>]")
-    # CHECK: attr: [42, true, vector<4xf32>]
-    print("raw attr:", raw)
-    # CHECK: - 42
-    # CHECK: - true
-    # CHECK: - vector<4xf32>
-    for attr in ArrayAttr(raw):
-        print("- ", attr)
+        arr = Attribute.parse("[42, true, vector<4xf32>]")
+    # CHECK: arr: [42, true, vector<4xf32>]
+    print("arr:", arr)
+    # CHECK: - IntegerAttr(42 : i64)
+    # CHECK: - BoolAttr(true)
+    # CHECK: - TypeAttr(vector<4xf32>)
+    for attr in arr:
+        print("- ", repr(attr))
 
     with Context():
         intAttr = Attribute.parse("42")
         vecAttr = Attribute.parse("vector<4xf32>")
         boolAttr = BoolAttr.get(True)
         raw = ArrayAttr.get([vecAttr, boolAttr, intAttr])
     # CHECK: attr: [vector<4xf32>, true, 42]
     print("raw attr:", raw)
-    # CHECK: - vector<4xf32>
-    # CHECK: - true
-    # CHECK: - 42
-    arr = ArrayAttr(raw)
+    # CHECK: - TypeAttr(vector<4xf32>)
+    # CHECK: - BoolAttr(true
+    # CHECK: - IntegerAttr(42 : i64)
+    arr = raw
     for attr in arr:
-        print("- ", attr)
-    # CHECK: attr[0]: vector<4xf32>
-    print("attr[0]:", arr[0])
-    # CHECK: attr[1]: true
-    print("attr[1]:", arr[1])
-    # CHECK: attr[2]: 42
-    print("attr[2]:", arr[2])
+        print("- ", repr(attr))
+    # CHECK: attr[0]: TypeAttr(vector<4xf32>)
+    print("attr[0]:", repr(arr[0]))
+    # CHECK: attr[1]: BoolAttr(true)
+    print("attr[1]:", repr(arr[1]))
+    # CHECK: attr[2]: IntegerAttr(42 : i64)
+    print("attr[2]:", repr(arr[2]))
     try:
         print("attr[3]:", arr[3])
     except IndexError as e:
         # CHECK: Error: ArrayAttribute index out of range
         print("Error: ", e)
     with Context():
         try:
             ArrayAttr.get([None])

mlir/test/python/ir/builtin_types.py

with Context(), Location.unknown():
             tensor_invalid = RankedTensorType.get(shape, none)
         except MLIRError as e:
             # CHECK: Invalid type:
             # CHECK: error: unknown: invalid tensor element type: 'none'
             print(e)
         else:
             print("Exception not produced")
 
+        tensor = RankedTensorType.get(shape, f32, StringAttr.get("encoding"))
+        assert tensor.shape == shape
+        assert tensor.encoding.value == "encoding"
+
         # Encoding should be None.
         assert RankedTensorType.get(shape, f32).encoding is None
 
-        tensor = RankedTensorType.get(shape, f32)
-        assert tensor.shape == shape
 
 
 # CHECK-LABEL: TEST: testUnrankedTensorType
 @run
 def testUnrankedTensorType():
     with Context(), Location.unknown():
         f32 = F32Type.get()
         loc = Location.unknown()

 def testMemRefType():
     with Context(), Location.unknown():
         f32 = F32Type.get()
         shape = [2, 3]
         loc = Location.unknown()
         memref_f32 = MemRefType.get(shape, f32, memory_space=Attribute.parse("2"))
         # CHECK: memref type: memref<2x3xf32, 2>
         print("memref type:", memref_f32)
-        # CHECK: memref layout: affine_map<(d0, d1) -> (d0, d1)>
-        print("memref layout:", memref_f32.layout)
+        # CHECK: memref layout: AffineMapAttr(affine_map<(d0, d1) -> (d0, d1)>)
+        print("memref layout:", repr(memref_f32.layout))
         # CHECK: memref affine map: (d0, d1) -> (d0, d1)
         print("memref affine map:", memref_f32.affine_map)
-        # CHECK: memory space: 2
-        print("memory space:", memref_f32.memory_space)
+        # CHECK: memory space: IntegerAttr(2 : i64)
+        print("memory space:", repr(memref_f32.memory_space))
 
         layout = AffineMapAttr.get(AffineMap.get_permutation([1, 0]))
         memref_layout = MemRefType.get(shape, f32, layout=layout)
         # CHECK: memref type: memref<2x3xf32, affine_map<(d0, d1) -> (d1, d0)>>
         print("memref type:", memref_layout)
         # CHECK: memref layout: affine_map<(d0, d1) -> (d1, d0)>
         print("memref layout:", memref_layout.layout)
         # CHECK: memref affine map: (d0, d1) -> (d1, d0)
         print("memref affine map:", memref_layout.affine_map)
-        # CHECK: memory space: <<NULL ATTRIBUTE>>
+        # CHECK: memory space: None
         print("memory space:", memref_layout.memory_space)
 
         none = NoneType.get()
         try:
             memref_invalid = MemRefType.get(shape, none)
         except MLIRError as e:
             # CHECK: Invalid type:
             # CHECK: error: unknown: invalid memref element type
             print(e)
         else:
             print("Exception not produced")
 
     assert memref_f32.shape == shape
 
 
 # CHECK-LABEL: TEST: testUnrankedMemRefType
 @run
 def testUnrankedMemRefType():
     with Context(), Location.unknown():
         f32 = F32Type.get()
         loc = Location.unknown()
         unranked_memref = UnrankedMemRefType.get(f32, Attribute.parse("2"))
         # CHECK: unranked memref type: memref<*xf32, 2>
         print("unranked memref type:", unranked_memref)
+        # CHECK: memory space: IntegerAttr(2 : i64)
+        print("memory space:", repr(unranked_memref.memory_space))
         try:
             invalid_rank = unranked_memref.rank
         except ValueError as e:
             # CHECK: calling this method requires that the type has a rank.
             print(e)
         else:
             print("Exception not produced")
         try:

mlir/test/python/ir/operation.py

def testOperationAttributes():
     "some.op"() { some.attribute = 1 : i8,
                   other.attribute = 3.0,
                   dependent = "text" } : () -> ()
   """,
         ctx,
     )
     op = module.body.operations[0]
     assert len(op.attributes) == 3
-    iattr = IntegerAttr(op.attributes["some.attribute"])
-    fattr = FloatAttr(op.attributes["other.attribute"])
-    sattr = StringAttr(op.attributes["dependent"])
+    iattr = op.attributes["some.attribute"]
+    fattr = op.attributes["other.attribute"]
+    sattr = op.attributes["dependent"]
     # CHECK: Attribute type i8, value 1
     print(f"Attribute type {iattr.type}, value {iattr.value}")
     # CHECK: Attribute type f64, value 3.0
     print(f"Attribute type {fattr.type}, value {fattr.value}")
     # CHECK: Attribute value text
     print(f"Attribute value {sattr.value}")
     # CHECK: Attribute value b'text'
     print(f"Attribute value {sattr.value_bytes}")

mlir/test/python/ir/symbol_table.py

with Context() as ctx:
         assert "qux" in symbol_table
 
         # Check that insertion actually renames this symbol in the symbol table.
         foo2 = m2.body.operations[0]
         m1.body.append(foo2)
         updated_name = symbol_table.insert(foo2)
         assert foo2.name.value != "foo"
         assert foo2.name == updated_name
+        assert isinstance(updated_name, StringAttr)
 
         # CHECK: module
         # CHECK:   func private @foo()
         # CHECK:   func private @qux()
         # CHECK:   func private @foo{{.*}}
         print(m1)
 
         try:
with Context() as ctx:
         foo, bar = list(m.operation.regions[0].blocks[0].operations)[0:2]
         SymbolTable.set_symbol_name(bar, "bam")
         # Note that module.operation counts as a "nested symbol table" which won't
         # be traversed into, so it is necessary to traverse its children.
         SymbolTable.replace_all_symbol_uses("bar", "bam", foo)
         # CHECK: call @bam()
         # CHECK: func private @bam
         print(m)
-        # CHECK: Foo symbol: "foo"
-        # CHECK: Bar symbol: "bam"
-        print(f"Foo symbol: {SymbolTable.get_symbol_name(foo)}")
-        print(f"Bar symbol: {SymbolTable.get_symbol_name(bar)}")
+        # CHECK: Foo symbol: StringAttr("foo")
+        # CHECK: Bar symbol: StringAttr("bam")
+        print(f"Foo symbol: {repr(SymbolTable.get_symbol_name(foo))}")
+        print(f"Bar symbol: {repr(SymbolTable.get_symbol_name(bar))}")
 
 
 # CHECK-LABEL: testSymbolTableVisibility
 @run
 def testSymbolTableVisibility():
     with Context() as ctx:
         m = Module.parse(
             """
       func.func private @foo() {
         return
       }
       """
         )
         foo = m.operation.regions[0].blocks[0].operations[0]
-        # CHECK: Existing visibility: "private"
-        print(f"Existing visibility: {SymbolTable.get_visibility(foo)}")
+        # CHECK: Existing visibility: StringAttr("private")
+        print(f"Existing visibility: {repr(SymbolTable.get_visibility(foo))}")
         SymbolTable.set_visibility(foo, "public")
         # CHECK: func public @foo
         print(m)
 
 
 # CHECK: testWalkSymbolTables
 @run
 def testWalkSymbolTables():