[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] Apply ClangTidy findings.

move constructors should be marked noexcept

Apply clang-tidy fixes for performance-unnecessary-value-param in IRInterfaces.cpp (NFC)

Apply clang-tidy fixes for misc-include-cleaner in IRInterfaces.cpp (NFC)

[MLIR][python bindings] Fix inferReturnTypes + AttrSizedOperandSegments for optional operands

Right now `inferTypeOpInterface.inferReturnTypes` fails because there's a cast in there to `py::sequence

[MLIR][python bindings] Fix inferReturnTypes + AttrSizedOperandSegments for optional operands

Right now `inferTypeOpInterface.inferReturnTypes` fails because there's a cast in there to `py::sequence` which throws a `TypeError` when it tries to cast the `None`s. Note `None`s are inserted into `operands` for omitted operands passed to the generated builder:

```
operands.append(_get_op_result_or_value(start) if start is not None else None)
operands.append(_get_op_result_or_value(stop) if stop is not None else None)
operands.append(_get_op_result_or_value(step) if step is not None else None)
```

Note also that skipping appending to the list operands doesn't work either because [[ https://github.com/llvm/llvm-project/blob/27c37327da67020f938aabf0f6405f57d688441e/mlir/lib/Bindings/Python/IRCore.cpp#L1585 | build generic ]] checks against the number of operand segments expected.

Currently the only way around is to handroll through `ir.Operation.create`.

Reviewed By: rkayaith

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

show more ...

[MLIR][python bindings] Add TypeCaster for returning refined types from python APIs

depends on D150839

This diff uses `MlirTypeID` to register `TypeCaster`s (i.e., `[](PyType pyType) -> DerivedTy {

[MLIR][python bindings] Add TypeCaster for returning refined types from python APIs

depends on D150839

This diff uses `MlirTypeID` to register `TypeCaster`s (i.e., `[](PyType pyType) -> DerivedTy { return pyType; }`) for all concrete types (i.e., `PyConcrete<...>`) that are then queried for (by `MlirTypeID`) and called in `struct type_caster<MlirType>::cast`. The result is that anywhere an `MlirType mlirType` is returned from a python binding, that `mlirType` is automatically cast to the correct concrete type. For example:

```
c0 = arith.ConstantOp(f32, 0.0)
# CHECK: F32Type(f32)
print(repr(c0.result.type))

unranked_tensor_type = UnrankedTensorType.get(f32)
unranked_tensor = tensor.FromElementsOp(unranked_tensor_type, [c0]).result

# CHECK: UnrankedTensorType
print(type(unranked_tensor.type).__name__)
# CHECK: UnrankedTensorType(tensor<*xf32>)
print(repr(unranked_tensor.type))
```

This functionality immediately extends to typed attributes (i.e., `attr.type`).

The diff also implements similar functionality for `mlir_type_subclass`es but in a slightly different way - for such types (which have no cpp corresponding `class` or `struct`) the user must provide a type caster in python (similar to how `AttrBuilder` works) or in cpp as a `py::cpp_function`.

Reviewed By: ftynse

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

show more ...

[MLIR] Add InferShapedTypeOpInterface bindings

Add C and python bindings for InferShapedTypeOpInterface
and ShapedTypeComponents. This allows users to invoke
InferShapedTypeOpInterface for ops that

[MLIR] Add InferShapedTypeOpInterface bindings

Add C and python bindings for InferShapedTypeOpInterface
and ShapedTypeComponents. This allows users to invoke
InferShapedTypeOpInterface for ops that implement it.

Reviewed By: ftynse

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

show more ...

Introduce MLIR Op Properties

This new features enabled to dedicate custom storage inline within operations.
This storage can be used as an alternative to attributes to store data that is
specific to

Introduce MLIR Op Properties

This new features enabled to dedicate custom storage inline within operations.
This storage can be used as an alternative to attributes to store data that is
specific to an operation. Attribute can also be stored inside the properties
storage if desired, but any kind of data can be present as well. This offers
a way to store and mutate data without uniquing in the Context like Attribute.
See the OpPropertiesTest.cpp for an example where a struct with a
std::vector<> is attached to an operation and mutated in-place:

struct TestProperties {
int a = -1;
float b = -1.;
std::vector<int64_t> array = {-33};
};

More complex scheme (including reference-counting) are also possible.

The only constraint to enable storing a C++ object as "properties" on an
operation is to implement three functions:

- convert from the candidate object to an Attribute
- convert from the Attribute to the candidate object
- hash the object

Optional the parsing and printing can also be customized with 2 extra
functions.

A new options is introduced to ODS to allow dialects to specify:

let usePropertiesForAttributes = 1;

When set to true, the inherent attributes for all the ops in this dialect
will be using properties instead of being stored alongside discardable
attributes.
The TestDialect showcases this feature.

Another change is that we introduce new APIs on the Operation class
to access separately the inherent attributes from the discardable ones.
We envision deprecating and removing the `getAttr()`, `getAttrsDictionary()`,
and other similar method which don't make the distinction explicit, leading
to an entirely separate namespace for discardable attributes.

Recommit d572cd1b067f after fixing python bindings build.

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

show more ...

[mlir][py] Fix unused var

[mlir][py] Fix infer return type invocation for variadics

Previously we only allowed the flattened list passed in, but the same
input provided here as to buildGeneric so flatten accordingly. We have

[mlir][py] Fix infer return type invocation for variadics

Previously we only allowed the flattened list passed in, but the same
input provided here as to buildGeneric so flatten accordingly. We have
less info here than in buildGeneric so the error is more generic if
unpacking fails.

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

show more ...

[mlir] Use std::optional instead of llvm::Optional (NFC)

This patch replaces (llvm::|)Optional< with std::optional<. I'll post
a separate patch to remove #include "llvm/ADT/Optional.h".

This is pa

[mlir] Use std::optional instead of llvm::Optional (NFC)

This patch replaces (llvm::|)Optional< with std::optional<. I'll post
a separate patch to remove #include "llvm/ADT/Optional.h".

This is part of an effort to migrate from llvm::Optional to
std::optional:

https://discourse.llvm.org/t/deprecating-llvm-optional-x-hasvalue-getvalue-getvalueor/63716

show more ...

[mlir] Add #include <optional> (NFC)

This patch adds #include <optional> to those files containing
llvm::Optional<...> or Optional<...>.

I'll post a separate patch to actually replace llvm::Optiona

[mlir] Add #include <optional> (NFC)

This patch adds #include <optional> to those files containing
llvm::Optional<...> or Optional<...>.

I'll post a separate patch to actually replace llvm::Optional with
std::optional.

This is part of an effort to migrate from llvm::Optional to
std::optional:

https://discourse.llvm.org/t/deprecating-llvm-optional-x-hasvalue-getvalue-getvalueor/63716

show more ...

[windows] Remove unused pybind exception params

Resolve MSVC warning C4104 for unreferenced variable

Reviewed By: mehdi_amini

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

Apply clang-tidy fixes for performance-unnecessary-value-param to MLIR (NFC)

Reviewed By: Mogball

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

Fix more clang-tidy cleanups in mlir/ (NFC)

[mlir] support interfaces in Python bindings

Introduce the initial support for operation interfaces in C API and Python
bindings. Interfaces are a key component of MLIR's extensibility and should be

[mlir] support interfaces in Python bindings

Introduce the initial support for operation interfaces in C API and Python
bindings. Interfaces are a key component of MLIR's extensibility and should be
available in bindings to make use of full potential of MLIR.

This initial implementation exposes InferTypeOpInterface all the way to the
Python bindings since it can be later used to simplify the operation
construction methods by inferring their return types instead of requiring the
user to do so. The general infrastructure for binding interfaces is defined and
InferTypeOpInterface can be used as an example for binding other interfaces.

Reviewed By: gysit

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

show more ...