[MLIR][Python] reland (narrower) type stub generation (#157930)

This a reland of llvm/llvm-project#155741 which
was reverted at llvm/llvm-project#157831. This
version is narrower in scope - it only turns on automatic stub
generation for `MLIRPythonExtension.Core._mlir` and **does not do
anything automatically**. Specifically, the only CMake code added to
`AddMLIRPython.cmake` is the `mlir_generate_type_stubs` function which
is then used only in a manual way. The API for
`mlir_generate_type_stubs` is:

```
Arguments:
  MODULE_NAME: The fully-qualified name of the extension module (used for importing in python).
  DEPENDS_TARGETS: List of targets these type stubs depend on being built; usually corresponding to the
    specific extension module (e.g., something like StandalonePythonModules.extension._standaloneDialectsNanobind.dso)
    and the core bindings extension module (e.g., something like StandalonePythonModules.extension._mlir.dso).
  OUTPUT_DIR: The root output directory to emit the type stubs into.
  OUTPUTS: List of expected outputs.
  DEPENDS_TARGET_SRC_DEPS: List of cpp sources for extension library (for generating a DEPFILE).
  IMPORT_PATHS: List of paths to add to PYTHONPATH for stubgen.
  PATTERN_FILE: (Optional) Pattern file (see https://nanobind.readthedocs.io/en/latest/typing.html#pattern-files).
Outputs:
  NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on.
```

Downstream users should use `mlir_generate_type_stubs` in coordination
with `declare_mlir_python_sources` to turn on stub generation for their
own downstream dialect extensions and upstream dialect extensions if
they so choose. Standalone example shows an example.

Note, downstream will also need to set
`-DMLIR_PYTHON_PACKAGE_PREFIX=...` correctly for their bindings.

Author: makslevental

mlir/include/mlir/Bindings/Python/NanobindAdaptors.h

@@ -512,8 +512,13 @@ class mlir_attribute_subclass : public pure_subclass {
           .attr("replace")(superCls.attr("__name__"), captureTypeName);
     });
     if (getTypeIDFunction) {
-      def_staticmethod("get_static_typeid",
-                       [getTypeIDFunction]() { return getTypeIDFunction(); });
+      def_staticmethod(
+          "get_static_typeid",
+          [getTypeIDFunction]() { return getTypeIDFunction(); },
+          // clang-format off
+          nanobind::sig("def get_static_typeid() -> " MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID"))
+          // clang-format on
+      );
       nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
           .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
               getTypeIDFunction())(nanobind::cpp_function(
@@ -582,8 +587,9 @@ class mlir_type_subclass : public pure_subclass {
 
     // 'isinstance' method.
     static const char kIsinstanceSig[] =
-        "def isinstance(other_type: " MAKE_MLIR_PYTHON_QUALNAME(
-            "ir") ".Type) -> bool";
+        // clang-format off
+        "def isinstance(other_type: " MAKE_MLIR_PYTHON_QUALNAME("ir.Type") ") -> bool";
+    // clang-format on
     def_staticmethod(
         "isinstance",
         [isaFunction](MlirType other) { return isaFunction(other); },
@@ -599,8 +605,13 @@ class mlir_type_subclass : public pure_subclass {
       // `def_property_readonly_static` is not available in `pure_subclass` and
       // we do not want to introduce the complexity that pybind uses to
       // implement it.
-      def_staticmethod("get_static_typeid",
-                       [getTypeIDFunction]() { return getTypeIDFunction(); });
+      def_staticmethod(
+          "get_static_typeid",
+          [getTypeIDFunction]() { return getTypeIDFunction(); },
+          // clang-format off
+          nanobind::sig("def get_static_typeid() -> " MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID"))
+          // clang-format on
+      );
       nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
           .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
               getTypeIDFunction())(nanobind::cpp_function(
@@ -665,8 +676,9 @@ class mlir_value_subclass : public pure_subclass {
 
     // 'isinstance' method.
     static const char kIsinstanceSig[] =
-        "def isinstance(other_value: " MAKE_MLIR_PYTHON_QUALNAME(
-            "ir") ".Value) -> bool";
+        // clang-format off
+        "def isinstance(other_value: " MAKE_MLIR_PYTHON_QUALNAME("ir.Value") ") -> bool";
+    // clang-format on
     def_staticmethod(
         "isinstance",
         [isaFunction](MlirValue other) { return isaFunction(other); },

mlir/lib/Bindings/Python/DialectPDL.cpp

@@ -68,6 +68,8 @@ static void populateDialectPDLSubmodule(const nanobind::module_ &m) {
   rangeType.def_property_readonly(
       "element_type",
       [](MlirType type) { return mlirPDLRangeTypeGetElementType(type); },
+      nb::sig(
+          "def element_type(self) -> " MAKE_MLIR_PYTHON_QUALNAME("ir.Type")),
       "Get the element type.");
 
   //===-------------------------------------------------------------------===//

mlir/lib/Bindings/Python/IRAttributes.cpp

@@ -167,7 +167,7 @@ struct nb_buffer_info {
 };
 
 class nb_buffer : public nb::object {
-  NB_OBJECT_DEFAULT(nb_buffer, object, "buffer", PyObject_CheckBuffer);
+  NB_OBJECT_DEFAULT(nb_buffer, object, "Buffer", PyObject_CheckBuffer);
 
   nb_buffer_info request() const {
     int flags = PyBUF_STRIDES | PyBUF_FORMAT;
@@ -252,8 +252,13 @@ class PyAffineMapAttribute : public PyConcreteAttribute<PyAffineMapAttribute> {
           return PyAffineMapAttribute(affineMap.getContext(), attr);
         },
         nb::arg("affine_map"), "Gets an attribute wrapping an AffineMap.");
-    c.def_prop_ro("value", mlirAffineMapAttrGetValue,
-                  "Returns the value of the AffineMap attribute");
+    c.def_prop_ro(
+        "value",
+        [](PyAffineMapAttribute &self) {
+          return PyAffineMap(self.getContext(),
+                             mlirAffineMapAttrGetValue(self));
+        },
+        "Returns the value of the AffineMap attribute");
   }
 };
 
@@ -480,11 +485,13 @@ class PyArrayAttribute : public PyConcreteAttribute<PyArrayAttribute> {
 
     PyArrayAttributeIterator &dunderIter() { return *this; }
 
-    MlirAttribute dunderNext() {
+    nb::typed<nb::object, PyAttribute> dunderNext() {
       // TODO: Throw is an inefficient way to stop iteration.
       if (nextIndex >= mlirArrayAttrGetNumElements(attr.get()))
         throw nb::stop_iteration();
-      return mlirArrayAttrGetElement(attr.get(), nextIndex++);
+      return PyAttribute(this->attr.getContext(),
+                         mlirArrayAttrGetElement(attr.get(), nextIndex++))
+          .maybeDownCast();
     }
 
     static void bind(nb::module_ &m) {
@@ -517,12 +524,13 @@ class PyArrayAttribute : public PyConcreteAttribute<PyArrayAttribute> {
         },
         nb::arg("attributes"), nb::arg("context") = nb::none(),
         "Gets a uniqued Array attribute");
-    c.def("__getitem__",
-          [](PyArrayAttribute &arr, intptr_t i) {
-            if (i >= mlirArrayAttrGetNumElements(arr))
-              throw nb::index_error("ArrayAttribute index out of range");
-            return arr.getItem(i);
-          })
+    c.def(
+         "__getitem__",
+         [](PyArrayAttribute &arr, intptr_t i) {
+           if (i >= mlirArrayAttrGetNumElements(arr))
+             throw nb::index_error("ArrayAttribute index out of range");
+           return PyAttribute(arr.getContext(), arr.getItem(i)).maybeDownCast();
+         })
         .def("__len__",
              [](const PyArrayAttribute &arr) {
                return mlirArrayAttrGetNumElements(arr);
@@ -611,10 +619,12 @@ class PyIntegerAttribute : public PyConcreteAttribute<PyIntegerAttribute> {
                   "Returns the value of the integer attribute");
     c.def("__int__", toPyInt,
           "Converts the value of the integer attribute to a Python int");
-    c.def_prop_ro_static("static_typeid",
-                         [](nb::object & /*class*/) -> MlirTypeID {
-                           return mlirIntegerAttrGetTypeID();
-                         });
+    c.def_prop_ro_static(
+        "static_typeid",
+        [](nb::object & /*class*/) {
+          return PyTypeID(mlirIntegerAttrGetTypeID());
+        },
+        nanobind::sig("def static_typeid(/) -> TypeID"));
   }
 
 private:
@@ -657,8 +667,8 @@ class PySymbolRefAttribute : public PyConcreteAttribute<PySymbolRefAttribute> {
   static constexpr const char *pyClassName = "SymbolRefAttr";
   using PyConcreteAttribute::PyConcreteAttribute;
 
-  static MlirAttribute fromList(const std::vector<std::string> &symbols,
-                                PyMlirContext &context) {
+  static PySymbolRefAttribute fromList(const std::vector<std::string> &symbols,
+                                       PyMlirContext &context) {
     if (symbols.empty())
       throw std::runtime_error("SymbolRefAttr must be composed of at least "
                                "one symbol.");
@@ -668,8 +678,10 @@ class PySymbolRefAttribute : public PyConcreteAttribute<PySymbolRefAttribute> {
       referenceAttrs.push_back(
           mlirFlatSymbolRefAttrGet(context.get(), toMlirStringRef(symbols[i])));
     }
-    return mlirSymbolRefAttrGet(context.get(), rootSymbol,
-                                referenceAttrs.size(), referenceAttrs.data());
+    return PySymbolRefAttribute(context.getRef(),
+                                mlirSymbolRefAttrGet(context.get(), rootSymbol,
+                                                     referenceAttrs.size(),
+                                                     referenceAttrs.data()));
   }
 
   static void bindDerived(ClassTy &c) {
@@ -746,7 +758,11 @@ class PyOpaqueAttribute : public PyConcreteAttribute<PyOpaqueAttribute> {
           return PyOpaqueAttribute(context->getRef(), attr);
         },
         nb::arg("dialect_namespace"), nb::arg("buffer"), nb::arg("type"),
-        nb::arg("context") = nb::none(), "Gets an Opaque attribute.");
+        nb::arg("context") = nb::none(),
+        // clang-format off
+        nb::sig("def get(dialect_namespace: str, buffer: typing_extensions.Buffer, type: Type, context: Context | None = None) -> OpaqueAttr"),
+        // clang-format on
+        "Gets an Opaque attribute.");
     c.def_prop_ro(
         "dialect_namespace",
         [](PyOpaqueAttribute &self) {
@@ -764,59 +780,6 @@ class PyOpaqueAttribute : public PyConcreteAttribute<PyOpaqueAttribute> {
   }
 };
 
-class PyStringAttribute : public PyConcreteAttribute<PyStringAttribute> {
-public:
-  static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAString;
-  static constexpr const char *pyClassName = "StringAttr";
-  using PyConcreteAttribute::PyConcreteAttribute;
-  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
-      mlirStringAttrGetTypeID;
-
-  static void bindDerived(ClassTy &c) {
-    c.def_static(
-        "get",
-        [](const std::string &value, DefaultingPyMlirContext context) {
-          MlirAttribute attr =
-              mlirStringAttrGet(context->get(), toMlirStringRef(value));
-          return PyStringAttribute(context->getRef(), attr);
-        },
-        nb::arg("value"), nb::arg("context") = nb::none(),
-        "Gets a uniqued string attribute");
-    c.def_static(
-        "get",
-        [](const nb::bytes &value, DefaultingPyMlirContext context) {
-          MlirAttribute attr =
-              mlirStringAttrGet(context->get(), toMlirStringRef(value));
-          return PyStringAttribute(context->getRef(), attr);
-        },
-        nb::arg("value"), nb::arg("context") = nb::none(),
-        "Gets a uniqued string attribute");
-    c.def_static(
-        "get_typed",
-        [](PyType &type, const std::string &value) {
-          MlirAttribute attr =
-              mlirStringAttrTypedGet(type, toMlirStringRef(value));
-          return PyStringAttribute(type.getContext(), attr);
-        },
-        nb::arg("type"), nb::arg("value"),
-        "Gets a uniqued string attribute associated to a type");
-    c.def_prop_ro(
-        "value",
-        [](PyStringAttribute &self) {
-          MlirStringRef stringRef = mlirStringAttrGetValue(self);
-          return nb::str(stringRef.data, stringRef.length);
-        },
-        "Returns the value of the string attribute");
-    c.def_prop_ro(
-        "value_bytes",
-        [](PyStringAttribute &self) {
-          MlirStringRef stringRef = mlirStringAttrGetValue(self);
-          return nb::bytes(stringRef.data, stringRef.length);
-        },
-        "Returns the value of the string attribute as `bytes`");
-  }
-};
-
 // TODO: Support construction of string elements.
 class PyDenseElementsAttribute
     : public PyConcreteAttribute<PyDenseElementsAttribute> {
@@ -1028,11 +991,14 @@ class PyDenseElementsAttribute
         PyDenseElementsAttribute::bf_releasebuffer;
 #endif
     c.def("__len__", &PyDenseElementsAttribute::dunderLen)
-        .def_static("get", PyDenseElementsAttribute::getFromBuffer,
-                    nb::arg("array"), nb::arg("signless") = true,
-                    nb::arg("type") = nb::none(), nb::arg("shape") = nb::none(),
-                    nb::arg("context") = nb::none(),
-                    kDenseElementsAttrGetDocstring)
+        .def_static(
+            "get", PyDenseElementsAttribute::getFromBuffer, nb::arg("array"),
+            nb::arg("signless") = true, nb::arg("type") = nb::none(),
+            nb::arg("shape") = nb::none(), nb::arg("context") = nb::none(),
+            // clang-format off
+            nb::sig("def get(array: typing_extensions.Buffer, signless: bool = True, type: Type | None = None, shape: Sequence[int] | None = None, context: Context | None = None) -> DenseElementsAttr"),
+            // clang-format on
+            kDenseElementsAttrGetDocstring)
         .def_static("get", PyDenseElementsAttribute::getFromList,
                     nb::arg("attrs"), nb::arg("type") = nb::none(),
                     nb::arg("context") = nb::none(),
@@ -1048,7 +1014,9 @@ class PyDenseElementsAttribute
           if (!mlirDenseElementsAttrIsSplat(self))
             throw nb::value_error(
                 "get_splat_value called on a non-splat attribute");
-          return mlirDenseElementsAttrGetSplatValue(self);
+          return PyAttribute(self.getContext(),
+                             mlirDenseElementsAttrGetSplatValue(self))
+              .maybeDownCast();
         });
   }
 
@@ -1509,6 +1477,9 @@ class PyDenseResourceElementsAttribute
                  nb::arg("array"), nb::arg("name"), nb::arg("type"),
                  nb::arg("alignment") = nb::none(),
                  nb::arg("is_mutable") = false, nb::arg("context") = nb::none(),
+                 // clang-format off
+                 nb::sig("def get_from_buffer(array: typing_extensions.Buffer, name: str, type: Type, alignment: int | None = None, is_mutable: bool = False, context: Context | None = None) -> DenseResourceElementsAttr"),
+                 // clang-format on
                  kDenseResourceElementsAttrGetFromBufferDocstring);
   }
 };
@@ -1556,7 +1527,7 @@ class PyDictAttribute : public PyConcreteAttribute<PyDictAttribute> {
           mlirDictionaryAttrGetElementByName(self, toMlirStringRef(name));
       if (mlirAttributeIsNull(attr))
         throw nb::key_error("attempt to access a non-existent attribute");
-      return attr;
+      return PyAttribute(self.getContext(), attr).maybeDownCast();
     });
     c.def("__getitem__", [](PyDictAttribute &self, intptr_t index) {
       if (index < 0 || index >= self.dunderLen()) {
@@ -1624,7 +1595,8 @@ class PyTypeAttribute : public PyConcreteAttribute<PyTypeAttribute> {
         nb::arg("value"), nb::arg("context") = nb::none(),
         "Gets a uniqued Type attribute");
     c.def_prop_ro("value", [](PyTypeAttribute &self) {
-      return mlirTypeAttrGetValue(self.get());
+      return PyType(self.getContext(), mlirTypeAttrGetValue(self.get()))
+          .maybeDownCast();
     });
   }
 };
@@ -1761,6 +1733,50 @@ nb::object symbolRefOrFlatSymbolRefAttributeCaster(PyAttribute &pyAttribute) {
 
 } // namespace
 
+void PyStringAttribute::bindDerived(ClassTy &c) {
+  c.def_static(
+      "get",
+      [](const std::string &value, DefaultingPyMlirContext context) {
+        MlirAttribute attr =
+            mlirStringAttrGet(context->get(), toMlirStringRef(value));
+        return PyStringAttribute(context->getRef(), attr);
+      },
+      nb::arg("value"), nb::arg("context") = nb::none(),
+      "Gets a uniqued string attribute");
+  c.def_static(
+      "get",
+      [](const nb::bytes &value, DefaultingPyMlirContext context) {
+        MlirAttribute attr =
+            mlirStringAttrGet(context->get(), toMlirStringRef(value));
+        return PyStringAttribute(context->getRef(), attr);
+      },
+      nb::arg("value"), nb::arg("context") = nb::none(),
+      "Gets a uniqued string attribute");
+  c.def_static(
+      "get_typed",
+      [](PyType &type, const std::string &value) {
+        MlirAttribute attr =
+            mlirStringAttrTypedGet(type, toMlirStringRef(value));
+        return PyStringAttribute(type.getContext(), attr);
+      },
+      nb::arg("type"), nb::arg("value"),
+      "Gets a uniqued string attribute associated to a type");
+  c.def_prop_ro(
+      "value",
+      [](PyStringAttribute &self) {
+        MlirStringRef stringRef = mlirStringAttrGetValue(self);
+        return nb::str(stringRef.data, stringRef.length);
+      },
+      "Returns the value of the string attribute");
+  c.def_prop_ro(
+      "value_bytes",
+      [](PyStringAttribute &self) {
+        MlirStringRef stringRef = mlirStringAttrGetValue(self);
+        return nb::bytes(stringRef.data, stringRef.length);
+      },
+      "Returns the value of the string attribute as `bytes`");
+}
+
 void mlir::python::populateIRAttributes(nb::module_ &m) {
   PyAffineMapAttribute::bind(m);
   PyDenseBoolArrayAttribute::bind(m);