[mlir python] Port in-tree dialects to nanobind. (#119924)

This is a companion to #118583, although it can be landed independently
because since #117922 dialects do not have to use the same Python

[mlir python] Port in-tree dialects to nanobind. (#119924)

This is a companion to #118583, although it can be landed independently
because since #117922 dialects do not have to use the same Python
binding framework as the Python core code.

This PR ports all of the in-tree dialect and pass extensions to
nanobind, with the exception of those that remain for testing pybind11
support.

This PR also:
* removes CollectDiagnosticsToStringScope from NanobindAdaptors.h. This
was overlooked in a previous PR and it is duplicated in Diagnostics.h.

---------

Co-authored-by: Jacques Pienaar <jpienaar@google.com>

show more ...

[mlir python] Port Python core code to nanobind. (#120473)

Relands #118583, with a fix for Python 3.8 compatibility. It was not
possible to set the buffer protocol accessers via slots in Python 3.8

[mlir python] Port Python core code to nanobind. (#120473)

Relands #118583, with a fix for Python 3.8 compatibility. It was not
possible to set the buffer protocol accessers via slots in Python 3.8.

Why? https://nanobind.readthedocs.io/en/latest/why.html says it better
than I can, but my primary motivation for this change is to improve MLIR
IR construction time from JAX.

For a complicated Google-internal LLM model in JAX, this change improves
the MLIR
lowering time by around 5s (out of around 30s), which is a significant
speedup for simply switching binding frameworks.

To a large extent, this is a mechanical change, for instance changing
`pybind11::` to `nanobind::`.

Notes:
* this PR needs Nanobind 2.4.0, because it needs a bug fix
(https://github.com/wjakob/nanobind/pull/806) that landed in that
release.
* this PR does not port the in-tree dialect extension modules. They can
be ported in a future PR.
* I removed the py::sibling() annotations from def_static and def_class
in `PybindAdapters.h`. These ask pybind11 to try to form an overload
with an existing method, but it's not possible to form mixed
pybind11/nanobind overloads this ways and the parent class is now
defined in nanobind. Better solutions may be possible here.
* nanobind does not contain an exact equivalent of pybind11's buffer
protocol support. It was not hard to add a nanobind implementation of a
similar API.
* nanobind is pickier about casting to std::vector<bool>, expecting that
the input is a sequence of bool types, not truthy values. In a couple of
places I added code to support truthy values during casting.
* nanobind distinguishes bytes (`nb::bytes`) from strings (e.g.,
`std::string`). This required nb::bytes overloads in a few places.

show more ...

Revert "[mlir python] Port Python core code to nanobind. (#118583)"

This reverts commit 41bd35b58bb482fd466aa4b13aa44a810ad6470f.

Breakage detected, rolling back.

[mlir python] Port Python core code to nanobind. (#118583)

Why? https://nanobind.readthedocs.io/en/latest/why.html says it better
than I can, but my primary motivation for this change is to improve

[mlir python] Port Python core code to nanobind. (#118583)

Why? https://nanobind.readthedocs.io/en/latest/why.html says it better
than I can, but my primary motivation for this change is to improve MLIR
IR construction time from JAX.

For a complicated Google-internal LLM model in JAX, this change improves
the MLIR
lowering time by around 5s (out of around 30s), which is a significant
speedup for simply switching binding frameworks.

To a large extent, this is a mechanical change, for instance changing
`pybind11::`
to `nanobind::`.

Notes:
* this PR needs Nanobind 2.4.0, because it needs a bug fix
(https://github.com/wjakob/nanobind/pull/806) that landed in that
release.
* this PR does not port the in-tree dialect extension modules. They can
be ported in a future PR.
* I removed the py::sibling() annotations from def_static and def_class
in `PybindAdapters.h`. These ask pybind11 to try to form an overload
with an existing method, but it's not possible to form mixed
pybind11/nanobind overloads this ways and the parent class is now
defined in nanobind. Better solutions may be possible here.
* nanobind does not contain an exact equivalent of pybind11's buffer
protocol support. It was not hard to add a nanobind implementation of a
similar API.
* nanobind is pickier about casting to std::vector<bool>, expecting that
the input is a sequence of bool types, not truthy values. In a couple of
places I added code to support truthy values during casting.
* nanobind distinguishes bytes (`nb::bytes`) from strings (e.g.,
`std::string`). This required nb::bytes overloads in a few places.

show more ...

[MLIR][Python] enhance python ir printing with pringing flags (#117836)

Close https://github.com/llvm/llvm-project/pull/65854

[MLIR][Python] Add the `--mlir-print-ir-tree-dir` to the C and Python API (#117339)

[MLIR] [Python] align python ir printing with mlir-print-ir-after-all (#107522)

When using the `enable_ir_printing` API from Python, it invokes IR
printing with default args, printing the IR before

[MLIR] [Python] align python ir printing with mlir-print-ir-after-all (#107522)

When using the `enable_ir_printing` API from Python, it invokes IR
printing with default args, printing the IR before each pass and
printing IR after pass only if there have been changes. This PR attempts
to align the `enable_ir_printing` API with the documentation

show more ...

[mlir][Python] Apply ClangTidy findings.

move constructors should be marked noexcept

[mlir][python] Clear PyOperations instead of invalidating them. (#70044)

`PyOperations` are Python-level handles to `Operation *` instances. When
the latter are modified by C++, the former need to

[mlir][python] Clear PyOperations instead of invalidating them. (#70044)

`PyOperations` are Python-level handles to `Operation *` instances. When
the latter are modified by C++, the former need to be invalidated.
#69746 implements such invalidation mechanism by setting all
`PyReferences` to `invalid`. However, that is not enough: they also need
to be removed from the `liveOperations` map since other parts of the
code (such as `PyOperation::createDetached`) assume that that map only
contains valid refs.

This is required to actually solve the issue in #69730.

show more ...

[MLIR][python bindings] invalidate ops after PassManager run (#69746)

Fixes https://github.com/llvm/llvm-project/issues/69730 (also see
https://reviews.llvm.org/D155543).

There are two things o

[MLIR][python bindings] invalidate ops after PassManager run (#69746)

Fixes https://github.com/llvm/llvm-project/issues/69730 (also see
https://reviews.llvm.org/D155543).

There are two things outstanding (why I didn't land before):

1. add some C API tests for `mlirOperationWalk`;
2. potentially refactor how the invalidation in `run` works; the first
version of the code looked like this:
```cpp
if (invalidateOps) {
auto *context = op.getOperation().getContext().get();
MlirOperationWalkCallback invalidatingCallback =
[](MlirOperation op, void *userData) {
PyMlirContext *context =
static_cast<PyMlirContext *>(userData);
context->setOperationInvalid(op);
};
auto numRegions =
mlirOperationGetNumRegions(op.getOperation().get());
for (int i = 0; i < numRegions; ++i) {
MlirRegion region =
mlirOperationGetRegion(op.getOperation().get(), i);
for (MlirBlock block = mlirRegionGetFirstBlock(region);
!mlirBlockIsNull(block);
block = mlirBlockGetNextInRegion(block))
for (MlirOperation childOp =
mlirBlockGetFirstOperation(block);
!mlirOperationIsNull(childOp);
childOp = mlirOperationGetNextInBlock(childOp))
mlirOperationWalk(childOp, invalidatingCallback, context,
MlirWalkPostOrder);
}
}
```
This is verbose and ugly but it has the important benefit of not
executing `mlirOperationEqual(rootOp->get(), op)` for every op
underneath the root op.

Supposing there's no desire for the slightly more efficient but highly
convoluted approach, I can land this "posthaste".
But, since we have eyes on this now, any suggestions or approaches (or
needs/concerns) are welcome.

show more ...

[MLIR][python bindings] use pybind C++ APIs for throwing python errors.

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

[mlir][python] Capture error diagnostics in exceptions

This updates most (all?) error-diagnostic-emitting python APIs to
capture error diagnostics and include them in the raised exception's
message:

[mlir][python] Capture error diagnostics in exceptions

This updates most (all?) error-diagnostic-emitting python APIs to
capture error diagnostics and include them in the raised exception's
message:
```
>>> Operation.parse('"arith.addi"() : () -> ()'))
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
mlir._mlir_libs.MLIRError: Unable to parse operation assembly:
error: "-":1:1: 'arith.addi' op requires one result
note: "-":1:1: see current operation: "arith.addi"() : () -> ()
```

The diagnostic information is available on the exception for users who
may want to customize the error message:
```
>>> try:
... Operation.parse('"arith.addi"() : () -> ()')
... except MLIRError as e:
... print(e.message)
... print(e.error_diagnostics)
... print(e.error_diagnostics[0].message)
...
Unable to parse operation assembly
[<mlir._mlir_libs._mlir.ir.DiagnosticInfo object at 0x7fed32bd6b70>]
'arith.addi' op requires one result
```

Error diagnostics captured in exceptions aren't propagated to diagnostic
handlers, to avoid double-reporting of errors. The context-level
`emit_error_diagnostics` option can be used to revert to the old
behaviour, causing error diagnostics to be reported to handlers instead
of as part of exceptions.

API changes:
- `Operation.verify` now raises an exception on verification failure,
instead of returning `false`
- The exception raised by the following methods has been changed to
`MLIRError`:
- `PassManager.run`
- `{Module,Operation,Type,Attribute}.parse`
- `{RankedTensorType,UnrankedTensorType}.get`
- `{MemRefType,UnrankedMemRefType}.get`
- `VectorType.get`
- `FloatAttr.get`

closes #60595

depends on D144804, D143830

Reviewed By: stellaraccident

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

show more ...

[mlir][python] Allow running pass manager on any operation

`PassManager.run` is currently restricted to running on `builtin.module`
ops, but this restriction doesn't exist on the C++ side. This upda

[mlir][python] Allow running pass manager on any operation

`PassManager.run` is currently restricted to running on `builtin.module`
ops, but this restriction doesn't exist on the C++ side. This updates it
to take `ir.Operation/OpView` instead of `ir.Module`.

Depends on D143354

Reviewed By: mehdi_amini

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

show more ...

[mlir][CAPI] Allow running pass manager on any operation

`mlirPassManagerRun` is currently restricted to running on
`builtin.module` ops, but this restriction doesn't exist on the C++
side. This ren

[mlir][CAPI] Allow running pass manager on any operation

`mlirPassManagerRun` is currently restricted to running on
`builtin.module` ops, but this restriction doesn't exist on the C++
side. This renames it to `mlirPassManagerRunOnOp` and updates it to take
`MlirOperation` instead of `MlirModule`.

Depends on D143352

Reviewed By: mehdi_amini

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

show more ...

[mlir][python] Allow adding to existing pass manager

This adds a `PassManager.add` method which adds pipeline elements to the
pass manager. This allows for progressively building up a pipeline from

[mlir][python] Allow adding to existing pass manager

This adds a `PassManager.add` method which adds pipeline elements to the
pass manager. This allows for progressively building up a pipeline from
python without string manipulation.

Reviewed By: mehdi_amini

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

show more ...

[mlir][python] Include anchor op in PassManager constructor

This adds an extra argument for specifying the pass manager's anchor op,
with a default of `any`. Previously the anchor was always default

[mlir][python] Include anchor op in PassManager constructor

This adds an extra argument for specifying the pass manager's anchor op,
with a default of `any`. Previously the anchor was always defaulted to
`builtin.module`.

Reviewed By: mehdi_amini

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

show more ...

[mlir][python] Include anchor op in PassManager.parse

The pipeline string must now include the pass manager's anchor op. This
makes the parse API properly roundtrip the printed form of a pass
manage

[mlir][python] Include anchor op in PassManager.parse

The pipeline string must now include the pass manager's anchor op. This
makes the parse API properly roundtrip the printed form of a pass
manager.

Reviewed By: mehdi_amini

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

show more ...

[mlir][python] Include pipeline parse errors in exception message

Currently any errors during pipeline parsing are reported to stderr.
This adds a new pipeline parsing function to the C api that rep

[mlir][python] Include pipeline parse errors in exception message

Currently any errors during pipeline parsing are reported to stderr.
This adds a new pipeline parsing function to the C api that reports
errors through a callback, and updates the python bindings to use it.

Reviewed By: mehdi_amini

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

show more ...

[mlir] Prefix pass manager options with `mlir-`

With this change, there's going to be a clear distinction between LLVM
and MLIR pass maanger options (e.g. `-mlir-print-after-all` vs
`-print-after-al

[mlir] Prefix pass manager options with `mlir-`

With this change, there's going to be a clear distinction between LLVM
and MLIR pass maanger options (e.g. `-mlir-print-after-all` vs
`-print-after-all`). This change is desirable from the point of view of
projects that depend on both LLVM and MLIR, e.g. Flang.

For consistency, all pass manager options in MLIR are prefixed with
`mlir-`, even options that don't have equivalents in LLVM .

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

show more ...

Adjust "end namespace" comment in MLIR to match new agree'd coding style

See D115115 and this mailing list discussion:
https://lists.llvm.org/pipermail/llvm-dev/2021-December/154199.html

Differenti

Adjust "end namespace" comment in MLIR to match new agree'd coding style

See D115115 and this mailing list discussion:
https://lists.llvm.org/pipermail/llvm-dev/2021-December/154199.html

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

show more ...

[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 ex

[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

show more ...

[mlir][python] Apply py::module_local() to all classes.

* This allows multiple MLIR-API embedding downstreams to co-exist in the same process.
* I believe this is the last thing needed to enable iso

[mlir][python] Apply py::module_local() to all classes.

* This allows multiple MLIR-API embedding downstreams to co-exist in the same process.
* I believe this is the last thing needed to enable isolated embedding.

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

show more ...

[mlir][python] Add simple debugging and printing helpers

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

NFC: Break up the mlir python bindings into individual sources.

* IRModules.cpp -> (IRCore.cpp, IRAffine.cpp, IRAttributes.cpp, IRTypes.cpp).
* The individual pieces now compile in the 5-15s range w

NFC: Break up the mlir python bindings into individual sources.

* IRModules.cpp -> (IRCore.cpp, IRAffine.cpp, IRAttributes.cpp, IRTypes.cpp).
* The individual pieces now compile in the 5-15s range whereas IRModules.cpp was starting to approach a minute (didn't capture a before time).
* More fine grained splitting is possible, but this represents the most obvious.

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

show more ...

[mlir][Python] Allow PassManager to interop with the capsule APIs.

* Used in npcomp to cast Python objects via the C-API.

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