[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/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);
