[Mlir-commits] [mlir] a6e7d02 - [mlir][python] Add pyi stub files to enable auto completion.

Stella Laurenzo llvmlistbot at llvm.org
Mon Nov 29 20:00:25 PST 2021


Author: Stella Laurenzo
Date: 2021-11-29T19:58:58-08:00
New Revision: a6e7d024a9ebda1564fd78b829c45169add80864

URL: https://github.com/llvm/llvm-project/commit/a6e7d024a9ebda1564fd78b829c45169add80864
DIFF: https://github.com/llvm/llvm-project/commit/a6e7d024a9ebda1564fd78b829c45169add80864.diff

LOG: [mlir][python] Add pyi stub files to enable auto completion.

There is no completely automated facility for generating stubs that are both accurate and comprehensive for native modules. After some experimentation, I found that MyPy's stubgen does the best at generating correct stubs with a few caveats that are relatively easy to fix:
  * Some types resolve to cross module symbols incorrectly.
  * staticmethod and classmethod signatures seem to always be completely generic and need to be manually provided.
  * It does not generate an __all__ which, from testing, causes namespace pollution to be visible to IDE code completion.

As a first step, I did the following:
  * Ran `stubgen` for `_mlir.ir`, `_mlir.passmanager`, and `_mlirExecutionEngine`.
  * Manually looked for all instances where unnamed arguments were being emitted (i.e. as 'arg0', etc) and updated the C++ side to include names (and re-ran stubgen to get a good initial state).
  * Made/noted a few structural changes to each `pyi` file to make it minimally functional.
  * Added the `pyi` files to the CMake rules so they are installed and visible.

To test, I added a `.env` file to the root of the project with `PYTHONPATH=...` set as per instructions. Then reload the developer window (in VsCode) and verify that completion works for various changes to test cases.

There are still a number of overly generic signatures, but I want to check in this low-touch baseline before iterating on more ambiguous changes. This is already a big improvement.

Differential Revision: https://reviews.llvm.org/D114679

Added: 
    mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi
    mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
    mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
    mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi

Modified: 
    mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
    mlir/lib/Bindings/Python/IRAffine.cpp
    mlir/lib/Bindings/Python/IRAttributes.cpp
    mlir/lib/Bindings/Python/IRCore.cpp
    mlir/lib/Bindings/Python/IRModule.h
    mlir/lib/Bindings/Python/IRTypes.cpp
    mlir/lib/Bindings/Python/MainModule.cpp
    mlir/lib/Bindings/Python/Pass.cpp
    mlir/python/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
index c49d9903f8503..b5a0f84d4aba4 100644
--- a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
+++ b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
@@ -105,6 +105,7 @@ PYBIND11_MODULE(_mlirExecutionEngine, m) {
                 mlirStringRefCreate(func.c_str(), func.size()));
             return reinterpret_cast<uintptr_t>(res);
           },
+          py::arg("func_name"),
           "Lookup function `func` in the ExecutionEngine.")
       .def(
           "raw_register_runtime",
@@ -127,5 +128,5 @@ PYBIND11_MODULE(_mlirExecutionEngine, m) {
                 executionEngine.get(),
                 mlirStringRefCreate(fileName.c_str(), fileName.size()));
           },
-          "Dump ExecutionEngine to an object file.");
+          py::arg("file_name"), "Dump ExecutionEngine to an object file.");
 }

diff  --git a/mlir/lib/Bindings/Python/IRAffine.cpp b/mlir/lib/Bindings/Python/IRAffine.cpp
index da80cda9c5823..272de0d7aaaf2 100644
--- a/mlir/lib/Bindings/Python/IRAffine.cpp
+++ b/mlir/lib/Bindings/Python/IRAffine.cpp
@@ -98,10 +98,13 @@ class PyConcreteAffineExpr : public BaseTy {
 
   static void bind(py::module &m) {
     auto cls = ClassTy(m, DerivedTy::pyClassName, py::module_local());
-    cls.def(py::init<PyAffineExpr &>());
-    cls.def_static("isinstance", [](PyAffineExpr &otherAffineExpr) -> bool {
-      return DerivedTy::isaFunction(otherAffineExpr);
-    });
+    cls.def(py::init<PyAffineExpr &>(), py::arg("expr"));
+    cls.def_static(
+        "isinstance",
+        [](PyAffineExpr &otherAffineExpr) -> bool {
+          return DerivedTy::isaFunction(otherAffineExpr);
+        },
+        py::arg("other"));
     DerivedTy::bindDerived(cls);
   }
 
@@ -748,41 +751,50 @@ void mlir::python::populateIRAffine(py::module &m) {
           },
           py::arg("permutation"), py::arg("context") = py::none(),
           "Gets an affine map that permutes its inputs.")
-      .def("get_submap",
-           [](PyAffineMap &self, std::vector<intptr_t> &resultPos) {
-             intptr_t numResults = mlirAffineMapGetNumResults(self);
-             for (intptr_t pos : resultPos) {
-               if (pos < 0 || pos >= numResults)
-                 throw py::value_error("result position out of bounds");
-             }
-             MlirAffineMap affineMap = mlirAffineMapGetSubMap(
-                 self, resultPos.size(), resultPos.data());
-             return PyAffineMap(self.getContext(), affineMap);
-           })
-      .def("get_major_submap",
-           [](PyAffineMap &self, intptr_t nResults) {
-             if (nResults >= mlirAffineMapGetNumResults(self))
-               throw py::value_error("number of results out of bounds");
-             MlirAffineMap affineMap =
-                 mlirAffineMapGetMajorSubMap(self, nResults);
-             return PyAffineMap(self.getContext(), affineMap);
-           })
-      .def("get_minor_submap",
-           [](PyAffineMap &self, intptr_t nResults) {
-             if (nResults >= mlirAffineMapGetNumResults(self))
-               throw py::value_error("number of results out of bounds");
-             MlirAffineMap affineMap =
-                 mlirAffineMapGetMinorSubMap(self, nResults);
-             return PyAffineMap(self.getContext(), affineMap);
-           })
-      .def("replace",
-           [](PyAffineMap &self, PyAffineExpr &expression,
-              PyAffineExpr &replacement, intptr_t numResultDims,
-              intptr_t numResultSyms) {
-             MlirAffineMap affineMap = mlirAffineMapReplace(
-                 self, expression, replacement, numResultDims, numResultSyms);
-             return PyAffineMap(self.getContext(), affineMap);
-           })
+      .def(
+          "get_submap",
+          [](PyAffineMap &self, std::vector<intptr_t> &resultPos) {
+            intptr_t numResults = mlirAffineMapGetNumResults(self);
+            for (intptr_t pos : resultPos) {
+              if (pos < 0 || pos >= numResults)
+                throw py::value_error("result position out of bounds");
+            }
+            MlirAffineMap affineMap = mlirAffineMapGetSubMap(
+                self, resultPos.size(), resultPos.data());
+            return PyAffineMap(self.getContext(), affineMap);
+          },
+          py::arg("result_positions"))
+      .def(
+          "get_major_submap",
+          [](PyAffineMap &self, intptr_t nResults) {
+            if (nResults >= mlirAffineMapGetNumResults(self))
+              throw py::value_error("number of results out of bounds");
+            MlirAffineMap affineMap =
+                mlirAffineMapGetMajorSubMap(self, nResults);
+            return PyAffineMap(self.getContext(), affineMap);
+          },
+          py::arg("n_results"))
+      .def(
+          "get_minor_submap",
+          [](PyAffineMap &self, intptr_t nResults) {
+            if (nResults >= mlirAffineMapGetNumResults(self))
+              throw py::value_error("number of results out of bounds");
+            MlirAffineMap affineMap =
+                mlirAffineMapGetMinorSubMap(self, nResults);
+            return PyAffineMap(self.getContext(), affineMap);
+          },
+          py::arg("n_results"))
+      .def(
+          "replace",
+          [](PyAffineMap &self, PyAffineExpr &expression,
+             PyAffineExpr &replacement, intptr_t numResultDims,
+             intptr_t numResultSyms) {
+            MlirAffineMap affineMap = mlirAffineMapReplace(
+                self, expression, replacement, numResultDims, numResultSyms);
+            return PyAffineMap(self.getContext(), affineMap);
+          },
+          py::arg("expr"), py::arg("replacement"), py::arg("n_result_dims"),
+          py::arg("n_result_syms"))
       .def_property_readonly(
           "is_permutation",
           [](PyAffineMap &self) { return mlirAffineMapIsPermutation(self); })
@@ -876,32 +888,35 @@ void mlir::python::populateIRAffine(py::module &m) {
           },
           py::arg("num_dims"), py::arg("num_symbols"),
           py::arg("context") = py::none())
-      .def("get_replaced",
-           [](PyIntegerSet &self, py::list dimExprs, py::list symbolExprs,
-              intptr_t numResultDims, intptr_t numResultSymbols) {
-             if (static_cast<intptr_t>(dimExprs.size()) !=
-                 mlirIntegerSetGetNumDims(self))
-               throw py::value_error(
-                   "Expected the number of dimension replacement expressions "
-                   "to match that of dimensions");
-             if (static_cast<intptr_t>(symbolExprs.size()) !=
-                 mlirIntegerSetGetNumSymbols(self))
-               throw py::value_error(
-                   "Expected the number of symbol replacement expressions "
-                   "to match that of symbols");
-
-             SmallVector<MlirAffineExpr> dimAffineExprs, symbolAffineExprs;
-             pyListToVector<PyAffineExpr>(
-                 dimExprs, dimAffineExprs,
-                 "attempting to create an IntegerSet by replacing dimensions");
-             pyListToVector<PyAffineExpr>(
-                 symbolExprs, symbolAffineExprs,
-                 "attempting to create an IntegerSet by replacing symbols");
-             MlirIntegerSet set = mlirIntegerSetReplaceGet(
-                 self, dimAffineExprs.data(), symbolAffineExprs.data(),
-                 numResultDims, numResultSymbols);
-             return PyIntegerSet(self.getContext(), set);
-           })
+      .def(
+          "get_replaced",
+          [](PyIntegerSet &self, py::list dimExprs, py::list symbolExprs,
+             intptr_t numResultDims, intptr_t numResultSymbols) {
+            if (static_cast<intptr_t>(dimExprs.size()) !=
+                mlirIntegerSetGetNumDims(self))
+              throw py::value_error(
+                  "Expected the number of dimension replacement expressions "
+                  "to match that of dimensions");
+            if (static_cast<intptr_t>(symbolExprs.size()) !=
+                mlirIntegerSetGetNumSymbols(self))
+              throw py::value_error(
+                  "Expected the number of symbol replacement expressions "
+                  "to match that of symbols");
+
+            SmallVector<MlirAffineExpr> dimAffineExprs, symbolAffineExprs;
+            pyListToVector<PyAffineExpr>(
+                dimExprs, dimAffineExprs,
+                "attempting to create an IntegerSet by replacing dimensions");
+            pyListToVector<PyAffineExpr>(
+                symbolExprs, symbolAffineExprs,
+                "attempting to create an IntegerSet by replacing symbols");
+            MlirIntegerSet set = mlirIntegerSetReplaceGet(
+                self, dimAffineExprs.data(), symbolAffineExprs.data(),
+                numResultDims, numResultSymbols);
+            return PyIntegerSet(self.getContext(), set);
+          },
+          py::arg("dim_exprs"), py::arg("symbol_exprs"),
+          py::arg("num_result_dims"), py::arg("num_result_symbols"))
       .def_property_readonly("is_canonical_empty",
                              [](PyIntegerSet &self) {
                                return mlirIntegerSetIsCanonicalEmpty(self);

diff  --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index 7db90ec8dff5d..17b3b34a2ea30 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -337,7 +337,7 @@ class PyStringAttribute : public PyConcreteAttribute<PyStringAttribute> {
               mlirStringAttrTypedGet(type, toMlirStringRef(value));
           return PyStringAttribute(type.getContext(), attr);
         },
-
+        py::arg("type"), py::arg("value"),
         "Gets a uniqued string attribute associated to a type");
     c.def_property_readonly(
         "value",

diff  --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index c70cfc5654f3b..8a110fcc4218b 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -1631,10 +1631,13 @@ class PyConcreteValue : public PyValue {
   /// Binds the Python module objects to functions of this class.
   static void bind(py::module &m) {
     auto cls = ClassTy(m, DerivedTy::pyClassName, py::module_local());
-    cls.def(py::init<PyValue &>(), py::keep_alive<0, 1>());
-    cls.def_static("isinstance", [](PyValue &otherValue) -> bool {
-      return DerivedTy::isaFunction(otherValue);
-    });
+    cls.def(py::init<PyValue &>(), py::keep_alive<0, 1>(), py::arg("value"));
+    cls.def_static(
+        "isinstance",
+        [](PyValue &otherValue) -> bool {
+          return DerivedTy::isaFunction(otherValue);
+        },
+        py::arg("other_value"));
     DerivedTy::bindDerived(cls);
   }
 
@@ -1657,9 +1660,12 @@ class PyBlockArgument : public PyConcreteValue<PyBlockArgument> {
     c.def_property_readonly("arg_number", [](PyBlockArgument &self) {
       return mlirBlockArgumentGetArgNumber(self.get());
     });
-    c.def("set_type", [](PyBlockArgument &self, PyType type) {
-      return mlirBlockArgumentSetType(self.get(), type);
-    });
+    c.def(
+        "set_type",
+        [](PyBlockArgument &self, PyType type) {
+          return mlirBlockArgumentSetType(self.get(), type);
+        },
+        py::arg("type"));
   }
 };
 
@@ -1952,6 +1958,7 @@ void mlir::python::populateIRCore(py::module &m) {
             }
             return PyDialectDescriptor(self.getRef(), dialect);
           },
+          py::arg("dialect_name"),
           "Gets or loads a dialect by name, returning its descriptor object")
       .def_property(
           "allow_unregistered_dialects",
@@ -1961,15 +1968,19 @@ void mlir::python::populateIRCore(py::module &m) {
           [](PyMlirContext &self, bool value) {
             mlirContextSetAllowUnregisteredDialects(self.get(), value);
           })
-      .def("enable_multithreading",
-           [](PyMlirContext &self, bool enable) {
-             mlirContextEnableMultithreading(self.get(), enable);
-           })
-      .def("is_registered_operation",
-           [](PyMlirContext &self, std::string &name) {
-             return mlirContextIsRegisteredOperation(
-                 self.get(), MlirStringRef{name.data(), name.size()});
-           });
+      .def(
+          "enable_multithreading",
+          [](PyMlirContext &self, bool enable) {
+            mlirContextEnableMultithreading(self.get(), enable);
+          },
+          py::arg("enable"))
+      .def(
+          "is_registered_operation",
+          [](PyMlirContext &self, std::string &name) {
+            return mlirContextIsRegisteredOperation(
+                self.get(), MlirStringRef{name.data(), name.size()});
+          },
+          py::arg("operation_name"));
 
   //----------------------------------------------------------------------------
   // Mapping of PyDialectDescriptor
@@ -2013,7 +2024,7 @@ void mlir::python::populateIRCore(py::module &m) {
   // Mapping of PyDialect
   //----------------------------------------------------------------------------
   py::class_<PyDialect>(m, "Dialect", py::module_local())
-      .def(py::init<py::object>(), "descriptor")
+      .def(py::init<py::object>(), py::arg("descriptor"))
       .def_property_readonly(
           "descriptor", [](PyDialect &self) { return self.getDescriptor(); })
       .def("__repr__", [](py::object self) {
@@ -2332,7 +2343,7 @@ void mlir::python::populateIRCore(py::module &m) {
 
   auto opViewClass =
       py::class_<PyOpView, PyOperationBase>(m, "OpView", py::module_local())
-          .def(py::init<py::object>())
+          .def(py::init<py::object>(), py::arg("operation"))
           .def_property_readonly("operation", &PyOpView::getOperationObject)
           .def_property_readonly(
               "context",
@@ -2426,7 +2437,7 @@ void mlir::python::populateIRCore(py::module &m) {
             mlirRegionInsertOwnedBlock(parent, 0, block);
             return PyBlock(parent.getParentOperation(), block);
           },
-          py::arg("parent"), py::arg("pyArgTypes") = py::list(),
+          py::arg("parent"), py::arg("arg_types") = py::list(),
           "Creates and returns a new Block at the beginning of the given "
           "region (with given argument types).")
       .def(
@@ -2499,6 +2510,7 @@ void mlir::python::populateIRCore(py::module &m) {
             operation.getOperation().setAttached(
                 self.getParentOperation().getObject());
           },
+          py::arg("operation"),
           "Appends an operation to this block. If the operation is currently "
           "in another block, it will be moved.");
 
@@ -2758,8 +2770,8 @@ void mlir::python::populateIRCore(py::module &m) {
   py::class_<PySymbolTable>(m, "SymbolTable", py::module_local())
       .def(py::init<PyOperationBase &>())
       .def("__getitem__", &PySymbolTable::dunderGetItem)
-      .def("insert", &PySymbolTable::insert)
-      .def("erase", &PySymbolTable::erase)
+      .def("insert", &PySymbolTable::insert, py::arg("operation"))
+      .def("erase", &PySymbolTable::erase, py::arg("operation"))
       .def("__delitem__", &PySymbolTable::dunderDel)
       .def("__contains__", [](PySymbolTable &table, const std::string &name) {
         return !mlirOperationIsNull(mlirSymbolTableLookup(

diff  --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h
index dc024a24793e0..f0d0cc654eabb 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/mlir/lib/Bindings/Python/IRModule.h
@@ -244,8 +244,7 @@ class DefaultingPyMlirContext
     : public Defaulting<DefaultingPyMlirContext, PyMlirContext> {
 public:
   using Defaulting::Defaulting;
-  static constexpr const char kTypeDescription[] =
-      "[ThreadContextAware] mlir.ir.Context";
+  static constexpr const char kTypeDescription[] = "mlir.ir.Context";
   static PyMlirContext &resolve();
 };
 
@@ -339,8 +338,7 @@ class DefaultingPyLocation
     : public Defaulting<DefaultingPyLocation, PyLocation> {
 public:
   using Defaulting::Defaulting;
-  static constexpr const char kTypeDescription[] =
-      "[ThreadContextAware] mlir.ir.Location";
+  static constexpr const char kTypeDescription[] = "mlir.ir.Location";
   static PyLocation &resolve();
 
   operator MlirLocation() const { return *get(); }
@@ -678,10 +676,14 @@ class PyConcreteType : public BaseTy {
 
   static void bind(pybind11::module &m) {
     auto cls = ClassTy(m, DerivedTy::pyClassName, pybind11::module_local());
-    cls.def(pybind11::init<PyType &>(), pybind11::keep_alive<0, 1>());
-    cls.def_static("isinstance", [](PyType &otherType) -> bool {
-      return DerivedTy::isaFunction(otherType);
-    });
+    cls.def(pybind11::init<PyType &>(), pybind11::keep_alive<0, 1>(),
+            pybind11::arg("cast_from_type"));
+    cls.def_static(
+        "isinstance",
+        [](PyType &otherType) -> bool {
+          return DerivedTy::isaFunction(otherType);
+        },
+        pybind11::arg("other"));
     DerivedTy::bindDerived(cls);
   }
 
@@ -768,10 +770,14 @@ class PyConcreteAttribute : public BaseTy {
   static void bind(pybind11::module &m) {
     auto cls = ClassTy(m, DerivedTy::pyClassName, pybind11::buffer_protocol(),
                        pybind11::module_local());
-    cls.def(pybind11::init<PyAttribute &>(), pybind11::keep_alive<0, 1>());
-    cls.def_static("isinstance", [](PyAttribute &otherAttr) -> bool {
-      return DerivedTy::isaFunction(otherAttr);
-    });
+    cls.def(pybind11::init<PyAttribute &>(), pybind11::keep_alive<0, 1>(),
+            pybind11::arg("cast_from_attr"));
+    cls.def_static(
+        "isinstance",
+        [](PyAttribute &otherAttr) -> bool {
+          return DerivedTy::isaFunction(otherAttr);
+        },
+        pybind11::arg("other"));
     cls.def_property_readonly("type", [](PyAttribute &attr) {
       return PyType(attr.getContext(), mlirAttributeGetType(attr));
     });

diff  --git a/mlir/lib/Bindings/Python/IRTypes.cpp b/mlir/lib/Bindings/Python/IRTypes.cpp
index 89fdb1f06a91a..380aa36d77824 100644
--- a/mlir/lib/Bindings/Python/IRTypes.cpp
+++ b/mlir/lib/Bindings/Python/IRTypes.cpp
@@ -262,6 +262,7 @@ class PyShapedType : public PyConcreteType<PyShapedType> {
           self.requireHasRank();
           return mlirShapedTypeIsDynamicDim(self, dim);
         },
+        py::arg("dim"),
         "Returns whether the dim-th dimension of the given shaped type is "
         "dynamic.");
     c.def(
@@ -270,10 +271,12 @@ class PyShapedType : public PyConcreteType<PyShapedType> {
           self.requireHasRank();
           return mlirShapedTypeGetDimSize(self, dim);
         },
+        py::arg("dim"),
         "Returns the dim-th dimension of the given ranked shaped type.");
     c.def_static(
         "is_dynamic_size",
         [](int64_t size) -> bool { return mlirShapedTypeIsDynamicSize(size); },
+        py::arg("dim_size"),
         "Returns whether the given dimension size indicates a dynamic "
         "dimension.");
     c.def(
@@ -282,6 +285,7 @@ class PyShapedType : public PyConcreteType<PyShapedType> {
           self.requireHasRank();
           return mlirShapedTypeIsDynamicStrideOrOffset(val);
         },
+        py::arg("dim_size"),
         "Returns whether the given value is used as a placeholder for dynamic "
         "strides and offsets in shaped types.");
     c.def_property_readonly(
@@ -544,7 +548,7 @@ class PyTupleType : public PyConcreteType<PyTupleType> {
           MlirType t = mlirTupleTypeGetType(self, pos);
           return PyType(self.getContext(), t);
         },
-        "Returns the pos-th type in the tuple type.");
+        py::arg("pos"), "Returns the pos-th type in the tuple type.");
     c.def_property_readonly(
         "num_types",
         [](PyTupleType &self) -> intptr_t {

diff  --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp
index 5489a4d3e3810..896ee432ecdad 100644
--- a/mlir/lib/Bindings/Python/MainModule.cpp
+++ b/mlir/lib/Bindings/Python/MainModule.cpp
@@ -30,14 +30,19 @@ PYBIND11_MODULE(_mlir, m) {
       .def_property("dialect_search_modules",
                     &PyGlobals::getDialectSearchPrefixes,
                     &PyGlobals::setDialectSearchPrefixes)
-      .def("append_dialect_search_prefix",
-           [](PyGlobals &self, std::string moduleName) {
-             self.getDialectSearchPrefixes().push_back(std::move(moduleName));
-             self.clearImportCache();
-           })
+      .def(
+          "append_dialect_search_prefix",
+          [](PyGlobals &self, std::string moduleName) {
+            self.getDialectSearchPrefixes().push_back(std::move(moduleName));
+            self.clearImportCache();
+          },
+          py::arg("module_name"))
       .def("_register_dialect_impl", &PyGlobals::registerDialectImpl,
+           py::arg("dialect_namespace"), py::arg("dialect_class"),
            "Testing hook for directly registering a dialect")
       .def("_register_operation_impl", &PyGlobals::registerOperationImpl,
+           py::arg("operation_name"), py::arg("operation_class"),
+           py::arg("raw_opview_class"),
            "Testing hook for directly registering an operation");
 
   // Aside from making the globals accessible to python, having python manage
@@ -55,6 +60,7 @@ PYBIND11_MODULE(_mlir, m) {
         PyGlobals::get().registerDialectImpl(dialectNamespace, pyClass);
         return pyClass;
       },
+      py::arg("dialect_class"),
       "Class decorator for registering a custom Dialect wrapper");
   m.def(
       "register_operation",
@@ -78,7 +84,9 @@ PYBIND11_MODULE(_mlir, m) {
               return opClass;
             });
       },
-      "Class decorator for registering a custom Operation wrapper");
+      py::arg("dialect_class"),
+      "Produce a class decorator for registering an Operation class as part of "
+      "a dialect");
 
   // Define and populate IR submodule.
   auto irModule = m.def_submodule("ir", "MLIR IR Bindings");

diff  --git a/mlir/lib/Bindings/Python/Pass.cpp b/mlir/lib/Bindings/Python/Pass.cpp
index 6aa1c651c9a06..2c38a3a25c6c8 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -79,7 +79,7 @@ void mlir::python::populatePassManagerSubmodule(py::module &m) {
           [](PyPassManager &passManager, bool enable) {
             mlirPassManagerEnableVerifier(passManager.get(), enable);
           },
-          "Enable / disable verify-each.")
+          py::arg("enable"), "Enable / disable verify-each.")
       .def_static(
           "parse",
           [](const std::string pipeline, DefaultingPyMlirContext context) {
@@ -106,6 +106,7 @@ void mlir::python::populatePassManagerSubmodule(py::module &m) {
               throw SetPyError(PyExc_RuntimeError,
                                "Failure while executing pass pipeline.");
           },
+          py::arg("module"),
           "Run the pass manager on the provided module, throw a RuntimeError "
           "on failure.")
       .def(

diff  --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index cdb2e5a3c0c49..f7b84b0337111 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -20,6 +20,11 @@ declare_mlir_python_sources(MLIRPythonSources.Core
     ir.py
     passmanager.py
     dialects/_ods_common.py
+
+    # The main _mlir module has submodules: include stubs from each.
+    _mlir_libs/_mlir/__init__.pyi
+    _mlir_libs/_mlir/ir.pyi
+    _mlir_libs/_mlir/passmanager.pyi
 )
 
 declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
@@ -27,6 +32,7 @@ declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
   ADD_TO_PARENT MLIRPythonSources
   SOURCES
     execution_engine.py
+    _mlir_libs/_mlirExecutionEngine.pyi
   SOURCES_GLOB
     runtime/*.py
 )

diff  --git a/mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi b/mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi
new file mode 100644
index 0000000000000..d4aab68060625
--- /dev/null
+++ b/mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi
@@ -0,0 +1,13 @@
+from typing import List
+
+globals: _Globals
+
+class _Globals:
+    dialect_search_modules: List[str]
+    def __init__(self, *args, **kwargs) -> None: ...
+    def _register_dialect_impl(self, dialect_namespace: str, dialect_class: object) -> None: ...
+    def _register_operation_impl(self, operation_name: str, operation_class: object, raw_opview_class: object) -> None: ...
+    def append_dialect_search_prefix(self, module_name: str) -> None: ...
+
+def register_dialect(dialect_class: object) -> object: ...
+def register_operation(dialect_class: object) -> object: ...

diff  --git a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
new file mode 100644
index 0000000000000..47ebeb291c35f
--- /dev/null
+++ b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
@@ -0,0 +1,863 @@
+# Originally imported via:
+#   stubgen {...} -m mlir._mlir_libs._mlir.ir
+# Local modifications:
+#   * Rewrite references to 'mlir.ir.' to local types
+#   * Add __all__ with the following incantation:
+#       egrep '^class ' ir.pyi | awk -F ' |:|\\(' '{print "    \"" $2 "\","}'
+#   * Local edits to signatures and types that MyPy did not auto detect (or
+#     detected incorrectly).
+
+from typing import Any, ClassVar, List, Optional
+
+from typing import overload
+
+__all__ = [
+    "AffineAddExpr",
+    "AffineBinaryExpr",
+    "AffineCeilDivExpr",
+    "AffineConstantExpr",
+    "AffineDimExpr",
+    "AffineExpr",
+    "AffineExprList",
+    "AffineFloorDivExpr",
+    "AffineMap",
+    "AffineMapAttr",
+    "AffineModExpr",
+    "AffineMulExpr",
+    "AffineSymbolExpr",
+    "ArrayAttr",
+    "ArrayAttributeIterator",
+    "Attribute",
+    "BF16Type",
+    "Block",
+    "BlockArgument",
+    "BlockArgumentList",
+    "BlockIterator",
+    "BlockList",
+    "BoolAttr",
+    "ComplexType",
+    "Context",
+    "DenseElementsAttr",
+    "DenseFPElementsAttr",
+    "DenseIntElementsAttr",
+    "Dialect",
+    "DialectDescriptor",
+    "Dialects",
+    "DictAttr",
+    "F16Type",
+    "F32Type",
+    "F64Type",
+    "FlatSymbolRefAttr",
+    "FloatAttr",
+    "FunctionType",
+    "IndexType",
+    "InferTypeOpInterface",
+    "InsertionPoint",
+    "IntegerAttr",
+    "IntegerSet",
+    "IntegerSetConstraint",
+    "IntegerSetConstraintList",
+    "IntegerType",
+    "Location",
+    "MemRefType",
+    "Module",
+    "NamedAttribute",
+    "NoneType",
+    "OpAttributeMap",
+    "OpOperandList",
+    "OpResult",
+    "OpResultList",
+    "OpView",
+    "Operation",
+    "OperationIterator",
+    "OperationList",
+    "RankedTensorType",
+    "Region",
+    "RegionIterator",
+    "RegionSequence",
+    "ShapedType",
+    "StringAttr",
+    "SymbolTable",
+    "TupleType",
+    "Type",
+    "TypeAttr",
+    "UnitAttr",
+    "UnrankedMemRefType",
+    "UnrankedTensorType",
+    "Value",
+    "VectorType",
+    "_GlobalDebug",
+    "_OperationBase",
+]
+
+
+class AffineAddExpr(AffineBinaryExpr):
+    def __init__(self, expr: AffineExpr) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+
+class AffineBinaryExpr(AffineExpr):
+    def __init__(self, expr: AffineExpr) -> None: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def lhs(self) -> AffineExpr: ...
+    @property
+    def rhs(self) -> AffineExpr: ...
+
+class AffineCeilDivExpr(AffineBinaryExpr):
+    def __init__(self, expr: AffineExpr) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+
+class AffineConstantExpr(AffineExpr):
+    def __init__(self, expr: AffineExpr) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def value(self) -> int: ...
+
+class AffineDimExpr(AffineExpr):
+    def __init__(self, expr: AffineExpr) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def position(self) -> int: ...
+
+class AffineExpr:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def _CAPICreate(self) -> AffineExpr: ...
+    def compose(self, arg0) -> AffineExpr: ...
+    def dump(self) -> None: ...
+    def get_add(self, *args, **kwargs) -> Any: ...
+    def get_ceil_div(self, *args, **kwargs) -> Any: ...
+    def get_constant(self, *args, **kwargs) -> Any: ...
+    def get_dim(self, *args, **kwargs) -> Any: ...
+    def get_floor_div(self, *args, **kwargs) -> Any: ...
+    def get_mod(self, *args, **kwargs) -> Any: ...
+    def get_mul(self, *args, **kwargs) -> Any: ...
+    def get_symbol(self, *args, **kwargs) -> Any: ...
+    def __add__(self, other) -> Any: ...
+    @overload
+    def __eq__(self, arg0: AffineExpr) -> bool: ...
+    @overload
+    def __eq__(self, arg0: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    def __mod__(self, other) -> Any: ...
+    def __mul__(self, other) -> Any: ...
+    def __radd__(self, other) -> Any: ...
+    def __rmod__(self, other) -> Any: ...
+    def __rmul__(self, other) -> Any: ...
+    def __rsub__(self, other) -> Any: ...
+    def __sub__(self, other) -> Any: ...
+    @property
+    def _CAPIPtr(self) -> object: ...
+    @property
+    def context(self) -> object: ...
+
+class AffineExprList:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __add__(self, arg0: AffineExprList) -> List[AffineExpr]: ...
+    @overload
+    def __getitem__(self, arg0: int) -> AffineExpr: ...
+    @overload
+    def __getitem__(self, arg0: slice) -> AffineExprList: ...
+    def __len__(self) -> int: ...
+
+class AffineFloorDivExpr(AffineBinaryExpr):
+    def __init__(self, expr: AffineExpr) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+
+class AffineMap:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def _CAPICreate(self) -> AffineMap: ...
+    def compress_unused_symbols(self, *args, **kwargs) -> Any: ...
+    def dump(self) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def get_constant(self, *args, **kwargs) -> Any: ...
+    def get_empty(self, *args, **kwargs) -> Any: ...
+    def get_identity(self, *args, **kwargs) -> Any: ...
+    def get_major_submap(self, n_results: int) -> AffineMap: ...
+    def get_minor_identity(self, *args, **kwargs) -> Any: ...
+    def get_minor_submap(self, n_results: int) -> AffineMap: ...
+    def get_permutation(self, *args, **kwargs) -> Any: ...
+    def get_submap(self, result_positions: List[int]) -> AffineMap: ...
+    def replace(self, expr: AffineExpr, replacement: AffineExpr, n_result_dims: int, n_result_syms: int) -> AffineMap: ...
+    @overload
+    def __eq__(self, arg0: AffineMap) -> bool: ...
+    @overload
+    def __eq__(self, arg0: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    @property
+    def _CAPIPtr(self) -> object: ...
+    @property
+    def context(self) -> object: ...
+    @property
+    def is_permutation(self) -> bool: ...
+    @property
+    def is_projected_permutation(self) -> bool: ...
+    @property
+    def n_dims(self) -> int: ...
+    @property
+    def n_inputs(self) -> int: ...
+    @property
+    def n_symbols(self) -> int: ...
+    @property
+    def results(self) -> Any: ...
+
+class AffineMapAttr(Attribute):
+    def __init__(self, cast_from_attr: Attribute) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def type(self) -> Type: ...
+
+class AffineModExpr(AffineBinaryExpr):
+    def __init__(self, expr: AffineExpr) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+
+class AffineMulExpr(AffineBinaryExpr):
+    def __init__(self, expr: AffineExpr) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+
+class AffineSymbolExpr(AffineExpr):
+    def __init__(self, expr: AffineExpr) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def position(self) -> int: ...
+
+class ArrayAttr(Attribute):
+    def __init__(self, cast_from_attr: Attribute) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    def __add__(self, arg0: list) -> ArrayAttr: ...
+    def __getitem__(self, arg0: int) -> Attribute: ...
+    def __iter__(self) -> Any: ...
+    def __len__(self) -> int: ...
+    @property
+    def type(self) -> Type: ...
+
+class ArrayAttributeIterator:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __iter__(self) -> ArrayAttributeIterator: ...
+    def __next__(self) -> Attribute: ...
+
+class Attribute:
+    def __init__(self, cast_from_type: Attribute) -> None: ...
+    def _CAPICreate(self) -> Attribute: ...
+    def dump(self) -> None: ...
+    def get_named(self, *args, **kwargs) -> Any: ...
+    def parse(self, *args, **kwargs) -> Any: ...
+    @overload
+    def __eq__(self, arg0: Attribute) -> bool: ...
+    @overload
+    def __eq__(self, arg0: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    @property
+    def _CAPIPtr(self) -> object: ...
+    @property
+    def context(self) -> object: ...
+    @property
+    def type(self) -> Any: ...
+
+class BF16Type(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+
+class Block:
+    __hash__: ClassVar[None] = ...
+    def __init__(self, *args, **kwargs) -> None: ...
+    def append(self, operation: _OperationBase) -> None: ...
+    def create_after(self, *args) -> Block: ...
+    def create_at_start(self, *args, **kwargs) -> Any: ...
+    def create_before(self, *args) -> Block: ...
+    @overload
+    def __eq__(self, arg0: Block) -> bool: ...
+    @overload
+    def __eq__(self, arg0: object) -> bool: ...
+    def __iter__(self) -> Any: ...
+    @property
+    def arguments(self) -> Any: ...
+    @property
+    def operations(self) -> Any: ...
+    @property
+    def owner(self) -> object: ...
+    @property
+    def region(self) -> Region: ...
+
+class BlockArgument(Value):
+    def __init__(self, value: Value) -> None: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    def set_type(self, type: Type) -> None: ...
+    @property
+    def arg_number(self) -> int: ...
+    @property
+    def owner(self) -> Block: ...
+
+class BlockArgumentList:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __add__(self, arg0: BlockArgumentList) -> List[BlockArgument]: ...
+    @overload
+    def __getitem__(self, arg0: int) -> BlockArgument: ...
+    @overload
+    def __getitem__(self, arg0: slice) -> BlockArgumentList: ...
+    def __len__(self) -> int: ...
+    @property
+    def types(self) -> List[Type]: ...
+
+class BlockIterator:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __iter__(self) -> BlockIterator: ...
+    def __next__(self) -> Block: ...
+
+class BlockList:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def append(self, *args) -> Block: ...
+    def __getitem__(self, arg0: int) -> Block: ...
+    def __iter__(self) -> BlockIterator: ...
+    def __len__(self) -> int: ...
+
+class BoolAttr(Attribute):
+    def __init__(self, cast_from_attr: Attribute) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def type(self) -> Type: ...
+    @property
+    def value(self) -> bool: ...
+
+class ComplexType(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def element_type(self) -> Type: ...
+
+class Context:
+    current: ClassVar[Context] = ...  # read-only
+    allow_unregistered_dialects: bool
+    def __init__(self) -> None: ...
+    def _CAPICreate(self) -> object: ...
+    def _get_context_again(self) -> object: ...
+    def _get_live_count(self, *args, **kwargs) -> Any: ...
+    def _get_live_module_count(self) -> int: ...
+    def _get_live_operation_count(self) -> int: ...
+    def enable_multithreading(self, enable: bool) -> None: ...
+    def get_dialect_descriptor(self, *args, **kwargs) -> Any: ...
+    def is_registered_operation(self, operation_name: str) -> bool: ...
+    def __enter__(self) -> object: ...
+    def __exit__(self, arg0: object, arg1: object, arg2: object) -> None: ...
+    @property
+    def _CAPIPtr(self) -> object: ...
+    @property
+    def d(self) -> Any: ...
+    @property
+    def dialects(self) -> Any: ...
+
+class DenseElementsAttr(Attribute):
+    def __init__(self, cast_from_attr: Attribute) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def get_splat(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    def __len__(self) -> int: ...
+    @property
+    def is_splat(self) -> bool: ...
+    @property
+    def type(self) -> Type: ...
+
+class DenseFPElementsAttr(DenseElementsAttr):
+    def __init__(self, cast_from_attr: Attribute) -> None: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    def __getitem__(self, arg0: int) -> float: ...
+    @property
+    def type(self) -> Type: ...
+
+class DenseIntElementsAttr(DenseElementsAttr):
+    def __init__(self, cast_from_attr: Attribute) -> None: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    def __getitem__(self, arg0: int) -> int: ...
+    @property
+    def type(self) -> Type: ...
+
+class Dialect:
+    def __init__(self, descriptor: object) -> None: ...
+    @property
+    def descriptor(self) -> object: ...
+
+class DialectDescriptor:
+    def __init__(self, *args, **kwargs) -> None: ...
+    @property
+    def namespace(self) -> str: ...
+
+class Dialects:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __getattr__(self, arg0: str) -> object: ...
+    def __getitem__(self, arg0: str) -> object: ...
+
+class DictAttr(Attribute):
+    def __init__(self, cast_from_attr: Attribute) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    def __contains__(self, arg0: str) -> bool: ...
+    @overload
+    def __getitem__(self, arg0: str) -> Attribute: ...
+    @overload
+    def __getitem__(self, arg0: int) -> NamedAttribute: ...
+    def __len__(self) -> int: ...
+    @property
+    def type(self) -> Type: ...
+
+class F16Type(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+
+class F32Type(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+
+class F64Type(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+
+class FlatSymbolRefAttr(Attribute):
+    def __init__(self, cast_from_attr: Attribute) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def type(self) -> Type: ...
+    @property
+    def value(self) -> str: ...
+
+class FloatAttr(Attribute):
+    def __init__(self, cast_from_attr: Attribute) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def get_f32(self, *args, **kwargs) -> Any: ...
+    def get_f64(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def type(self) -> Type: ...
+    @property
+    def value(self) -> float: ...
+
+class FunctionType(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def inputs(self) -> list: ...
+    @property
+    def results(self) -> list: ...
+
+class IndexType(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+
+class InferTypeOpInterface:
+    def __init__(self, object: object, context: Context = ...) -> None: ...
+    def inferReturnTypes(self, operands: Optional[List[Value]] = ..., attributes: Optional[Attribute] = ..., regions: Optional[List[Region]] = ..., context: Context = ..., loc: Location = ...) -> List[Type]: ...
+    @property
+    def operation(self) -> object: ...
+    @property
+    def opview(self) -> object: ...
+
+class InsertionPoint:
+    current: ClassVar[InsertionPoint] = ...  # read-only
+    @overload
+    def __init__(self, block: Block) -> None: ...
+    @overload
+    def __init__(self, beforeOperation: _OperationBase) -> None: ...
+    def at_block_begin(self, *args, **kwargs) -> Any: ...
+    def at_block_terminator(self, *args, **kwargs) -> Any: ...
+    def insert(self, operation: _OperationBase) -> None: ...
+    def __enter__(self) -> object: ...
+    def __exit__(self, arg0: object, arg1: object, arg2: object) -> None: ...
+    @property
+    def block(self) -> Block: ...
+
+class IntegerAttr(Attribute):
+    def __init__(self, cast_from_attr: Attribute) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def type(self) -> Type: ...
+    @property
+    def value(self) -> int: ...
+
+class IntegerSet:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def _CAPICreate(self) -> IntegerSet: ...
+    def dump(self) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def get_empty(self, *args, **kwargs) -> Any: ...
+    def get_replaced(self, dim_exprs: list, symbol_exprs: list, num_result_dims: int, num_result_symbols: int) -> IntegerSet: ...
+    @overload
+    def __eq__(self, arg0: IntegerSet) -> bool: ...
+    @overload
+    def __eq__(self, arg0: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    @property
+    def _CAPIPtr(self) -> object: ...
+    @property
+    def constraints(self) -> Any: ...
+    @property
+    def context(self) -> object: ...
+    @property
+    def is_canonical_empty(self) -> bool: ...
+    @property
+    def n_dims(self) -> int: ...
+    @property
+    def n_equalities(self) -> int: ...
+    @property
+    def n_inequalities(self) -> int: ...
+    @property
+    def n_inputs(self) -> int: ...
+    @property
+    def n_symbols(self) -> int: ...
+
+class IntegerSetConstraint:
+    def __init__(self, *args, **kwargs) -> None: ...
+    @property
+    def expr(self) -> AffineExpr: ...
+    @property
+    def is_eq(self) -> bool: ...
+
+class IntegerSetConstraintList:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __add__(self, arg0: IntegerSetConstraintList) -> List[IntegerSetConstraint]: ...
+    @overload
+    def __getitem__(self, arg0: int) -> IntegerSetConstraint: ...
+    @overload
+    def __getitem__(self, arg0: slice) -> IntegerSetConstraintList: ...
+    def __len__(self) -> int: ...
+
+class IntegerType(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get_signed(self, *args, **kwargs) -> Any: ...
+    def get_signless(self, *args, **kwargs) -> Any: ...
+    def get_unsigned(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def is_signed(self) -> bool: ...
+    @property
+    def is_signless(self) -> bool: ...
+    @property
+    def is_unsigned(self) -> bool: ...
+    @property
+    def width(self) -> int: ...
+
+class Location:
+    current: ClassVar[Location] = ...  # read-only
+    __hash__: ClassVar[None] = ...
+    def __init__(self, *args, **kwargs) -> None: ...
+    def _CAPICreate(self) -> Location: ...
+    def callsite(self, *args, **kwargs) -> Any: ...
+    def file(self, *args, **kwargs) -> Any: ...
+    def name(self, *args, **kwargs) -> Any: ...
+    def unknown(self, *args, **kwargs) -> Any: ...
+    def __enter__(self) -> object: ...
+    @overload
+    def __eq__(self, arg0: Location) -> bool: ...
+    @overload
+    def __eq__(self, arg0: object) -> bool: ...
+    def __exit__(self, arg0: object, arg1: object, arg2: object) -> None: ...
+    @property
+    def _CAPIPtr(self) -> object: ...
+    @property
+    def context(self) -> object: ...
+
+class MemRefType(ShapedType):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def affine_map(self) -> AffineMap: ...
+    @property
+    def layout(self) -> Attribute: ...
+    @property
+    def memory_space(self) -> Attribute: ...
+
+class Module:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def _CAPICreate(self) -> object: ...
+    def create(self, *args, **kwargs) -> Any: ...
+    def dump(self) -> None: ...
+    def parse(self, *args, **kwargs) -> Any: ...
+    @property
+    def _CAPIPtr(self) -> object: ...
+    @property
+    def body(self) -> Any: ...
+    @property
+    def context(self) -> object: ...
+    @property
+    def operation(self) -> object: ...
+
+class NamedAttribute:
+    def __init__(self, *args, **kwargs) -> None: ...
+    @property
+    def attr(self) -> Attribute: ...
+    @property
+    def name(self) -> str: ...
+
+class NoneType(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+
+class OpAttributeMap:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __contains__(self, arg0: str) -> bool: ...
+    def __delitem__(self, arg0: str) -> None: ...
+    @overload
+    def __getitem__(self, arg0: str) -> Attribute: ...
+    @overload
+    def __getitem__(self, arg0: int) -> NamedAttribute: ...
+    def __len__(self) -> int: ...
+    def __setitem__(self, arg0: str, arg1: Attribute) -> None: ...
+
+class OpOperandList:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __add__(self, arg0: OpOperandList) -> List[Value]: ...
+    @overload
+    def __getitem__(self, arg0: int) -> Value: ...
+    @overload
+    def __getitem__(self, arg0: slice) -> OpOperandList: ...
+    def __len__(self) -> int: ...
+    def __setitem__(self, arg0: int, arg1: Value) -> None: ...
+
+class OpResult(Value):
+    def __init__(self, value: Value) -> None: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def owner(self) -> object: ...
+    @property
+    def result_number(self) -> int: ...
+
+class OpResultList:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __add__(self, arg0: OpResultList) -> List[OpResult]: ...
+    @overload
+    def __getitem__(self, arg0: int) -> OpResult: ...
+    @overload
+    def __getitem__(self, arg0: slice) -> OpResultList: ...
+    def __len__(self) -> int: ...
+    @property
+    def types(self) -> List[Type]: ...
+
+class OpView(_OperationBase):
+    _ODS_OPERAND_SEGMENTS: ClassVar[None] = ...
+    _ODS_REGIONS: ClassVar[tuple] = ...
+    _ODS_RESULT_SEGMENTS: ClassVar[None] = ...
+    def __init__(self, operation: object) -> None: ...
+    @classmethod
+    def build_generic(self, *args, **kwargs) -> Any: ...
+    @property
+    def context(self) -> object: ...
+    @property
+    def operation(self) -> object: ...
+
+class Operation(_OperationBase):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def _CAPICreate(self) -> object: ...
+    def create(self, *args, **kwargs) -> Any: ...
+    def erase(self) -> None: ...
+    @property
+    def _CAPIPtr(self) -> object: ...
+    @property
+    def context(self) -> object: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def opview(self) -> object: ...
+    @property
+    def parent(self) -> object: ...
+
+class OperationIterator:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __iter__(self) -> OperationIterator: ...
+    def __next__(self) -> object: ...
+
+class OperationList:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __getitem__(self, arg0: int) -> object: ...
+    def __iter__(self) -> OperationIterator: ...
+    def __len__(self) -> int: ...
+
+class RankedTensorType(ShapedType):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def encoding(self) -> Optional[Attribute]: ...
+
+class Region:
+    __hash__: ClassVar[None] = ...
+    def __init__(self, *args, **kwargs) -> None: ...
+    @overload
+    def __eq__(self, arg0: Region) -> bool: ...
+    @overload
+    def __eq__(self, arg0: object) -> bool: ...
+    def __iter__(self) -> Any: ...
+    @property
+    def blocks(self) -> Any: ...
+    @property
+    def owner(self) -> object: ...
+
+class RegionIterator:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __iter__(self) -> RegionIterator: ...
+    def __next__(self) -> Region: ...
+
+class RegionSequence:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __getitem__(self, arg0: int) -> Region: ...
+    def __len__(self) -> int: ...
+
+class ShapedType(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get_dim_size(self, dim: int) -> int: ...
+    def is_dynamic_dim(self, dim: int) -> bool: ...
+    def is_dynamic_size(self, *args, **kwargs) -> Any: ...
+    def is_dynamic_stride_or_offset(self, dim_size: int) -> bool: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def element_type(self) -> Type: ...
+    @property
+    def has_rank(self) -> bool: ...
+    @property
+    def has_static_shape(self) -> bool: ...
+    @property
+    def rank(self) -> int: ...
+    @property
+    def shape(self) -> List[int]: ...
+
+class StringAttr(Attribute):
+    def __init__(self, cast_from_attr: Attribute) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def get_typed(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def type(self) -> Type: ...
+    @property
+    def value(self) -> str: ...
+
+class SymbolTable:
+    def __init__(self, arg0: _OperationBase) -> None: ...
+    def erase(self, operation: _OperationBase) -> None: ...
+    def insert(self, operation: _OperationBase) -> Attribute: ...
+    def __contains__(self, arg0: str) -> bool: ...
+    def __delitem__(self, arg0: str) -> None: ...
+    def __getitem__(self, arg0: str) -> object: ...
+
+class TupleType(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get_tuple(self, *args, **kwargs) -> Any: ...
+    def get_type(self, pos: int) -> Type: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def num_types(self) -> int: ...
+
+class Type:
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def _CAPICreate(self) -> Type: ...
+    def dump(self) -> None: ...
+    def parse(self, *args, **kwargs) -> Any: ...
+    @overload
+    def __eq__(self, arg0: Type) -> bool: ...
+    @overload
+    def __eq__(self, arg0: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    @property
+    def _CAPIPtr(self) -> object: ...
+    @property
+    def context(self) -> object: ...
+
+class TypeAttr(Attribute):
+    def __init__(self, cast_from_attr: Attribute) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def type(self) -> Type: ...
+    @property
+    def value(self) -> Type: ...
+
+class UnitAttr(Attribute):
+    def __init__(self, cast_from_attr: Attribute) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def type(self) -> Type: ...
+
+class UnrankedMemRefType(ShapedType):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+    @property
+    def memory_space(self) -> Attribute: ...
+
+class UnrankedTensorType(ShapedType):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+
+class Value:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def _CAPICreate(self) -> Value: ...
+    def dump(self) -> None: ...
+    @overload
+    def __eq__(self, arg0: Value) -> bool: ...
+    @overload
+    def __eq__(self, arg0: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    @property
+    def _CAPIPtr(self) -> object: ...
+    @property
+    def context(self) -> Any: ...
+    @property
+    def owner(self) -> object: ...
+    @property
+    def type(self) -> Type: ...
+
+class VectorType(ShapedType):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    def get(self, *args, **kwargs) -> Any: ...
+    def isinstance(self, *args, **kwargs) -> Any: ...
+
+class _GlobalDebug:
+    flag: ClassVar[bool] = ...
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class _OperationBase:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def detach_from_parent(self) -> object: ...
+    def get_asm(self, binary: bool = ..., large_elements_limit: Optional[int] = ..., enable_debug_info: bool = ..., pretty_debug_info: bool = ..., print_generic_op_form: bool = ..., use_local_scope: bool = ...) -> object: ...
+    def move_after(self, other: _OperationBase) -> None: ...
+    def move_before(self, other: _OperationBase) -> None: ...
+    def print(self, file: object = ..., binary: bool = ..., large_elements_limit: Optional[int] = ..., enable_debug_info: bool = ..., pretty_debug_info: bool = ..., print_generic_op_form: bool = ..., use_local_scope: bool = ...) -> None: ...
+    def verify(self) -> bool: ...
+    @overload
+    def __eq__(self, arg0: _OperationBase) -> bool: ...
+    @overload
+    def __eq__(self, arg0: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    @property
+    def _CAPIPtr(self) -> object: ...
+    @property
+    def attributes(self) -> Any: ...
+    @property
+    def location(self) -> Location: ...
+    @property
+    def operands(self) -> Any: ...
+    @property
+    def regions(self) -> Any: ...
+    @property
+    def result(self) -> Any: ...
+    @property
+    def results(self) -> Any: ...

diff  --git a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
new file mode 100644
index 0000000000000..7003bf0f04d1e
--- /dev/null
+++ b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
@@ -0,0 +1,24 @@
+# Originally imported via:
+#   stubgen {...} -m mlir._mlir_libs._mlir.passmanager
+# Local modifications:
+#   * Relative imports for cross-module references.
+#   * Add __all__
+
+from typing import Any
+
+from . import ir as _ir
+
+__all__ = [
+    "PassManager",
+]
+
+class PassManager:
+    def __init__(self, context: _ir.Context = ...) -> None: ...
+    def _CAPICreate(self) -> object: ...
+    def _testing_release(self) -> None: ...
+    def enable_ir_printing(self) -> None: ...
+    def enable_verifier(self, enable: bool) -> None: ...
+    def parse(self, *args, **kwargs) -> Any: ...
+    def run(self, module: _ir.Module) -> None: ...
+    @property
+    def _CAPIPtr(self) -> object: ...

diff  --git a/mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi b/mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi
new file mode 100644
index 0000000000000..50ff6c5b1f139
--- /dev/null
+++ b/mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi
@@ -0,0 +1,23 @@
+# Originally imported via:
+#   stubgen {...} -m mlir._mlir_libs._mlirExecutionEngine
+# Local modifications:
+#   * Relative imports for cross-module references.
+#   * Add __all__
+
+from typing import List
+
+from ._mlir import ir as _ir
+
+__all__ = [
+    "ExecutionEngine",
+]
+
+class ExecutionEngine:
+    def __init__(self, module: _ir.Module, opt_level: int = ..., shared_libs: List[str] = ...) -> None: ...
+    def _CAPICreate(self) -> object: ...
+    def _testing_release(self) -> None: ...
+    def dump_to_object_file(self, file_name: str) -> None: ...
+    def raw_lookup(self, func_name: str) -> int: ...
+    def raw_register_runtime(self, name: str, callback: object) -> None: ...
+    @property
+    def _CAPIPtr(self) -> object: ...


        


More information about the Mlir-commits mailing list