[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
diff --git a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
index c49d990..b5a0f84 100644
--- a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
+++ b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
@@ -105,6 +105,7 @@
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 @@
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 da80cda..272de0d 100644
--- a/mlir/lib/Bindings/Python/IRAffine.cpp
+++ b/mlir/lib/Bindings/Python/IRAffine.cpp
@@ -98,10 +98,13 @@
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 @@
},
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 @@
},
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");
+ .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);
- })
+ 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 7db90ec..17b3b34 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -337,7 +337,7 @@
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 c70cfc5..8a110fc 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -1631,10 +1631,13 @@
/// 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 @@
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 @@
}
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 @@
[](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 @@
// 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 @@
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 @@
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 @@
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 @@
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 dc024a2..f0d0cc6 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/mlir/lib/Bindings/Python/IRModule.h
@@ -244,8 +244,7 @@
: 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 @@
: 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 @@
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 @@
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 89fdb1f..380aa36 100644
--- a/mlir/lib/Bindings/Python/IRTypes.cpp
+++ b/mlir/lib/Bindings/Python/IRTypes.cpp
@@ -262,6 +262,7 @@
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 @@
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 @@
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 @@
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 5489a4d..896ee43 100644
--- a/mlir/lib/Bindings/Python/MainModule.cpp
+++ b/mlir/lib/Bindings/Python/MainModule.cpp
@@ -30,14 +30,19 @@
.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 @@
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 @@
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 6aa1c65..2c38a3a 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -79,7 +79,7 @@
[](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 @@
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 cdb2e5a..f7b84b0 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -20,6 +20,11 @@
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 @@
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 0000000..d4aab68
--- /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 0000000..47ebeb2
--- /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 0000000..7003bf0
--- /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 0000000..50ff6c5
--- /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: ...
