31 Commits

Author SHA1 Message Date
Maksim Levental
0e4be262f4
[mlir][Python] fix dialect extensions which bind C types (#175405)
Fix some dialect bindings I missed in https://github.com/llvm/llvm-project/pull/174156 so they don't bind C structs (because that leads to multiple registration in the case when multiple packages are used simultaneously).
2026-01-10 21:24:55 -08:00
Maksim Levental
18fc908566
[mlir][Python] move IRTypes and IRAttributes to MLIRPythonSupport (#174118)
This PR continues the work of
https://github.com/llvm/llvm-project/pull/171775 by moving more useful
types/attributes into MLIRPythonSupport.

You can now do 

```c++
struct PyTestIntegerRankedTensorType
    : mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteType<
          PyTestIntegerRankedTensorType,
          mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyRankedTensorType>
struct PyTestTensorValue
    : mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteValue<
          PyTestTensorValue>
```
instead of `mlir_type_subclass` and `mlir_value_subclass`;
**specifically manual registration of the "value caster" via indirection
through the Python interpreter is no longer necessary** . You can also
now freely use all such types at the nanobind API level (e.g., overload
based on `FP*`):

```c++
using mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN;
standaloneM.def("print_fp_type", [](PyF16Type &) { nb::print("this is a fp16 type"); });
standaloneM.def("print_fp_type", [](PyF32Type &) { nb::print("this is a fp32 type"); });
standaloneM.def("print_fp_type", [](PyF64Type &) { nb::print("this is a fp64 type"); });
```

Note, here we only port `PythonTestModuleNanobind` but there is a
follow-up PR that ports **all** in-tree dialect extensions
https://github.com/llvm/llvm-project/pull/174156 to use these. After
that one we can soft deprecate `mlir_pure_subclass`.

Note, depends on https://github.com/llvm/llvm-project/pull/171775
2026-01-05 09:34:58 -08:00
Maksim Levental
f0ef5dba6d
[mlir][Python] create MLIRPythonSupport (#171775)
# What

This PR adds a shared library `MLIRPythonSupport` which contains all of
the CRTP classes ike `PyConcreteValue`, `PyConcreteType`,
`PyConcreteAttribute`, as well as other useful code like `Defaulting*`
and etc enabling their reuse in downstream projects. Downstream projects
can now do

```c++
struct PyTestType : mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteType<PyTestType> {
  ...
};

class PyTestAttr : public mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteAttribute<PyTestAttr> {
  ...
}

NB_MODULE(_mlirPythonTestNanobind, m) {
  PyTestType::bind(m);
  PyTestAttr::bind(m);
}
```

instead of using the discordant alternative
`mlir_type_subclass`/`mlir_attr_subclass` (same goes for
`PyConcreteValue`/`mlir_value_subclass`).

# Why

This PR is mostly code motion (along with CMake) but before I describe
the changes I want to state the goals/benefits:

1. Currently upstream "core" extensions and "dialect" extensions ([all
of the `Dialect*` extensions
here](d7c734b5a1/mlir/lib/Bindings/Python))
are a two-tier system;
**a**. [core
extensions](https://github.com/llvm/llvm-project/blob/main/mlir/lib/Bindings/Python/IRTypes.cpp#L361)
enjoy first class support as far as type inference[^3], type stub
generation, and ease of implementation, while dialect extensions [have
poorer support](https://reviews.llvm.org/D150927), incorrect type stub
generation much more tedious (boilerplate) implementation;
**b**. Crucially, this two-tiered system is reflected in the fact that
**the two sets of types/attributes are not in the same Python object
hierarchy**. To wit: `isinstance(..., Type)` and `isinstance(...,
Attribute)` are not supported for the dialect extensions[^2];
**c**. Since these types are not exposed in public headers, downstream
users (dialect extensions or not) cannot write functions that overload
on e.g. `PyFloat8*Type` - that's quite a [useful
feature](fdbee98df8/cpp_ext/TorchOps.cpp (L29-L69))!
2. The dialect extensions incur a sizeable performance penalty relative
to the core extensions in that every single trip across the wire (either
`python->cpp` or `cpp->python`) requires work in addition to nanobind's
own casting/construction pipeline;
**a**. When going from `python->cpp`, [we extract the capsule object
from the Python
object](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h#L219C24-L219C46)
and then extract from the capsule the `Mlir*` opaque struct/ptr. This
side isn't so onerous;
**b**. When going from `cpp->python` we call long-hand call Python
`import` APIs and construct the Python object using `_CAPICreate`. Note,
there at least 2 `attr` calls incurred in addition to `_CAPICreate`;
this is already much more [efficiently handled by nanobind
itself](4ba51fcf79/src/nb_internals.h (L381-L382))!
3. This division blocks various features: in some configurations[^1] we
trigger a circular import bug because "dialect" types and attributes
perform an [import of the root `_mlir`
module](bd9651bf78/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h (L585))
when they are created (the types themselves, not even instances of those
types). This blocks type stub generation for dialect extensions (i.e.,
the reason we currently only generate type stubs for `_mlir`).

# How

Prior this was not done/possible because of "ODR" issues but I have
resolved those issues; the basic idea for how we solve this is "move
things we want to share into shared libraries":

1. Move IRCore (stuff like `PyConcreteValue`, `PyConcreteType`,
`PyConcreteAttribute`) into `MLIRPythonSupport`;
- Note, we move the rest of the things in `IRModule.h` (renamed to
`IRCore.h`) because `PyConcreteValue`, `PyConcreteType`,
`PyConcreteAttribute` depend on them. This makes for a bigger PR than
one would hope for but ultimately I think we should give people access
to these classes to use as they see fit (specifically inherit from, but
also liberally use in bindings signatures instead of the opaque `Mlir*`
struct wrappers).
2. Put all of this code into a nested namespace
`MLIR_BINDINGS_PYTHON_DOMAIN` which is determined by a compile time
define (and tied to `MLIR_BINDINGS_PYTHON_NB_DOMAIN`). This is necessary
in order to prevent conflicts on both symbol name **and** typeid
(necessary for nanobind to not double register binded types) between
multiple bindings libraries (e.g., `torch-mlir`, and `jax`). Note
[nanobind doesn't support `module_local` like
pybind11](https://nanobind.readthedocs.io/en/latest/porting.html#removed-features).
It does support `NB_DOMAIN` but that is not sufficient for
disambiguating typeids across projects (to wit: we currently define
`NB_DOMAIN` and it was still necessary to move everything to a nested
namespace);
3. Build the [nanobind library itself as a shared
object](https://github.com/wjakob/nanobind/blob/master/cmake/nanobind-config.cmake#L127)
(and link it to both the extensions and `MLIRPythonSupport`).
4. CMake to make this work, in-tree, out-of-tree, downstream, upstream,
etc.

# Testing

Three tests are added here 

1. `PythonTestModuleNanobind` is ported to use
`PyConcreteType<PyTestType>` instead of `mlir_type_subclass` and
`PyConcreteAttribute<PyTestAttr>` instead of `mlir_atrr_subclass`,
verifying this works for non-core extensions in-tree;
2. `StandaloneExtensionNanobind` is ported to use `struct PyCustomType :
mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteType<PyCustomType>`
instead of `mlir_type_subclass` verifying this works for non-core
extensions out-of-tree;
3. `StandaloneExtensionNanobind`'s `smoketest` is extended to also load
another bindings package (namely `mlir`) verifying
`MLIR_BINDINGS_PYTHON_DOMAIN` successfully disambiguates symbols and
typeids.

I have also tested this downstream:
https://github.com/llvm/eudsl/pull/287 as well run the following builder
bots:

mlir-nvidia-gcc7:
https://lab.llvm.org/buildbot/#/buildrequests/6654424?redirect_to_build=true

I have also tested against IREE:
https://github.com/iree-org/iree/pull/21916

# Integration

It is highly recommended to set the CMake var
`MLIR_BINDINGS_PYTHON_NB_DOMAIN` (which will also determine
`MLIR_BINDINGS_PYTHON_DOMAIN`) to something unique for each downstream.
This can also be passed explicitly to `add_mlir_python_modules` if your
project builds multiple bindings packages. I added a `WARNING` to this
effect in `AddMLIRPython.cmake`.

[^3]: Python values being typed correctly when exiting from cpp;
[^1]: Specifically when the modules are imported using `importlib`,
which occurs with nanobind's
[stubgen](https://github.com/wjakob/nanobind/blob/master/src/stubgen.py#L965);
[^2]: The workaround we implemented was a class method for the dialect
bindings called `Class.isinstance(...)`;
2026-01-05 09:08:13 -08:00
Maksim Levental
3d7018c70b
[MLIR][Python] remove pybind11 support (#172581)
This PR removes pybind which has been deprecated for over a year
(https://github.com/llvm/llvm-project/pull/117922).
2025-12-19 09:51:22 -08:00
Maksim Levental
93097b2d47
Revert "[MLIR][Python] use FetchContent_Declare for nanobind and remove pybind (#161230)" (#162309)
This reverts commit 84a214856ad989f37af19f5e8aaa9ec2346dde6f.

This gives us more time to work out the alternative and also people to
migrate
2025-10-07 16:30:10 +00:00
Maksim Levental
84a214856a
[MLIR][Python] use FetchContent_Declare for nanobind and remove pybind (#161230)
Inspired by this comment
https://github.com/llvm/llvm-project/pull/157930#issuecomment-3346634290
(and long-standing issues related to finding nanobind/pybind in the
right place), this PR moves to using `FetchContent_Declare` to get the
nanobind dependency. This is pretty standard (see e.g.,
[IREE](cf60359b74/CMakeLists.txt (L842-L848))).
This PR also removes pybind which has been deprecated for almost a year
(https://github.com/llvm/llvm-project/pull/117922) and which isn't
compatible (for whatever reason) with `FetchContent_Declare`.

---------

Co-authored-by: Jacques Pienaar <jpienaar@google.com>
2025-10-06 17:17:04 +00:00
Maksim Levental
efd96afedf
[MLIR][Python] reland (narrower) type stub generation (#157930)
This a reland of https://github.com/llvm/llvm-project/pull/155741 which
was reverted at https://github.com/llvm/llvm-project/pull/157831. This
version is narrower in scope - it only turns on automatic stub
generation for `MLIRPythonExtension.Core._mlir` and **does not do
anything automatically**. Specifically, the only CMake code added to
`AddMLIRPython.cmake` is the `mlir_generate_type_stubs` function which
is then used only in a manual way. The API for
`mlir_generate_type_stubs` is:

```
Arguments:
  MODULE_NAME: The fully-qualified name of the extension module (used for importing in python).
  DEPENDS_TARGETS: List of targets these type stubs depend on being built; usually corresponding to the
    specific extension module (e.g., something like StandalonePythonModules.extension._standaloneDialectsNanobind.dso)
    and the core bindings extension module (e.g., something like StandalonePythonModules.extension._mlir.dso).
  OUTPUT_DIR: The root output directory to emit the type stubs into.
  OUTPUTS: List of expected outputs.
  DEPENDS_TARGET_SRC_DEPS: List of cpp sources for extension library (for generating a DEPFILE).
  IMPORT_PATHS: List of paths to add to PYTHONPATH for stubgen.
  PATTERN_FILE: (Optional) Pattern file (see https://nanobind.readthedocs.io/en/latest/typing.html#pattern-files).
Outputs:
  NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on.
```

Downstream users should use `mlir_generate_type_stubs` in coordination
with `declare_mlir_python_sources` to turn on stub generation for their
own downstream dialect extensions and upstream dialect extensions if
they so choose. Standalone example shows an example.

Note, downstream will also need to set
`-DMLIR_PYTHON_PACKAGE_PREFIX=...` correctly for their bindings.
2025-09-20 18:47:32 +00:00
Nicholas Junge
6350bb3ed3
[mlir][py] Mark all type caster from_{cpp,python} methods as noexcept (#143866)
This is mentioned as a "must" in
https://nanobind.readthedocs.io/en/latest/porting.html#type-casters when
implementing type casters.

While most of the existing `from_cpp` methods were already marked
noexcept, many of the `from_python` methods were not. This commit adds
the missing noexcept declarations to all type casters found in
`NanobindAdaptors.h`.

---------

Co-authored-by: Maksim Levental <maksim.levental@gmail.com>
2025-07-15 10:58:10 -04:00
Nikhil Kalra
b15ccd436a
[mlir] Better Python diagnostics (#128581)
Updated the Python diagnostics handler to emit notes (in addition to
errors) into the output stream so that users have more context as to
where in the IR the error is occurring.
2025-03-10 15:59:47 -07:00
Michał Górny
047e8e47c1
Reapply "[mlir] Link libraries that aren't included in libMLIR to libMLIR" (#123910)
Use `mlir_target_link_libraries()` to link dependencies of libraries
that are not included in libMLIR, to ensure that they link to the dylib
when they are used in Flang. Otherwise, they implicitly pull in all
their static dependencies, effectively causing Flang binaries to
simultaneously link to the dylib and to static libraries, which is never
a good idea.

I have only covered the libraries that are used by Flang. If you wish, I
can extend this approach to all non-libMLIR libraries in MLIR, making
MLIR itself also link to the dylib consistently.

[v3 with more `-DBUILD_SHARED_LIBS=ON` fixes]
2025-01-22 09:01:50 +00:00
Michał Górny
9decc24c6b Revert "[mlir] Link libraries that aren't included in libMLIR to libMLIR (#123781)"
This reverts commit 4c6242ebf50dde0597df2bace49d534b61122496.  More
BUILD_SHARED_LIBS=ON regressions, sigh.
2025-01-22 09:09:52 +01:00
Michał Górny
4c6242ebf5
[mlir] Link libraries that aren't included in libMLIR to libMLIR (#123781)
Use `mlir_target_link_libraries()` to link dependencies of libraries
that are not included in libMLIR, to ensure that they link to the dylib
when they are used in Flang. Otherwise, they implicitly pull in all
their static dependencies, effectively causing Flang binaries to
simultaneously link to the dylib and to static libraries, which is never
a good idea.

I have only covered the libraries that are used by Flang. If you wish, I
can extend this approach to all non-libMLIR libraries in MLIR, making
MLIR itself also link to the dylib consistently.

[v2 with fixed `-DBUILD_SHARED_LIBS=ON` build]
2025-01-22 07:54:54 +00:00
Michał Górny
8b879d106b Revert "[mlir] Link libraries that aren't included in libMLIR to libMLIR (#123477)"
This reverts commit af6616676fb7f9dd4898290ea684ee0c90f1701d.  It broke
builds with `-DBUILD_SHARED_LIBS=ON`.
2025-01-20 19:33:51 +01:00
Michał Górny
af6616676f
[mlir] Link libraries that aren't included in libMLIR to libMLIR (#123477)
Use `mlir_target_link_libraries()` to link dependencies of libraries
that are not included in libMLIR, to ensure that they link to the dylib
when they are used in Flang. Otherwise, they implicitly pull in all
their static dependencies, effectively causing Flang binaries to
simultaneously link to the dylib and to static libraries, which is never
a good idea.

I have only covered the libraries that are used by Flang. If you wish, I
can extend this approach to all non-libMLIR libraries in MLIR, making
MLIR itself also link to the dylib consistently.
2025-01-20 17:25:20 +00:00
Peter Hawkins
5cd4274772
[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>
2024-12-20 20:32:32 -08:00
Maksim Levental
392622d084
Revert "Revert "[mlir python] Add nanobind support (#119232)
Reverts revert #118517 after (hopefully) fixing builders
(https://github.com/llvm/llvm-zorg/pull/328,
https://github.com/llvm/llvm-zorg/pull/327)

This reverts commit 61bf308cf2fc32452f14861c102ace89f5f36fec.
2024-12-09 16:37:43 -05:00
Maksim Levental
61bf308cf2
Revert "[mlir python] Add nanobind support for standalone dialects." (#118517)
Reverts llvm/llvm-project#117922 because deps aren't met on some of the
post-commit build bots.
2024-12-03 09:26:33 -08:00
Peter Hawkins
afe75b4d5f
[mlir python] Add nanobind support for standalone dialects. (#117922)
This PR allows out-of-tree dialects to write Python dialect modules
using nanobind instead of pybind11.

It may make sense to migrate in-tree dialects and some of the ODS Python
infrastructure to nanobind, but that is a topic for a future change.

This PR makes the following changes:
* adds nanobind to the CMake and Bazel build systems. We also add
robin_map to the Bazel build, which is a dependency of nanobind.
* adds a PYTHON_BINDING_LIBRARY option to various CMake functions, such
as declare_mlir_python_extension, allowing users to select a Python
binding library.
* creates a fork of mlir/include/mlir/Bindings/Python/PybindAdaptors.h
named NanobindAdaptors.h. This plays the same role, using nanobind
instead of pybind11.
* splits CollectDiagnosticsToStringScope out of PybindAdaptors.h and
into a new header mlir/include/mlir/Bindings/Python/Diagnostics.h, since
it is code that is no way related to pybind11 or for that matter,
Python.
* changed the standalone Python extension example to have both pybind11
and nanobind variants.
* changed mlir/python/mlir/dialects/python_test.py to have both pybind11
and nanobind variants.

Notes:
* A slightly unfortunate thing that I needed to do in the CMake
integration was to use FindPython in addition to FindPython3, since
nanobind's CMake integration expects the Python_ names for variables.
Perhaps there's a better way to do this.
2024-12-03 09:13:34 -08:00
Maksim Levental
9315645834
[mlir][python] auto attribute casting (#97786) 2024-07-05 10:43:51 -05:00
Maksim Levental
17ec364b1b
[mlir][python] enable registering dialects with the default Context (#72488) 2023-11-27 19:26:05 -06:00
Maksim Levental
7c850867b9
[mlir][python] value casting (#69644)
This PR adds "value casting", i.e., a mechanism to wrap `ir.Value` in a
proxy class that overloads dunders such as `__add__`, `__sub__`, and
`__mul__` for fun and great profit.

This is thematically similar to
bfb1ba7526
and
9566ee2806.
The example in the test demonstrates the value of the feature (no pun
intended):

```python
    @register_value_caster(F16Type.static_typeid)
    @register_value_caster(F32Type.static_typeid)
    @register_value_caster(F64Type.static_typeid)
    @register_value_caster(IntegerType.static_typeid)
    class ArithValue(Value):
        __add__ = partialmethod(_binary_op, op="add")
        __sub__ = partialmethod(_binary_op, op="sub")
        __mul__ = partialmethod(_binary_op, op="mul")

    a = arith.constant(value=FloatAttr.get(f16_t, 42.42))
    b = a + a
    # CHECK: ArithValue(%0 = arith.addf %cst, %cst : f16)
    print(b)

    a = arith.constant(value=FloatAttr.get(f32_t, 42.42))
    b = a - a
    # CHECK: ArithValue(%1 = arith.subf %cst_0, %cst_0 : f32)
    print(b)

    a = arith.constant(value=FloatAttr.get(f64_t, 42.42))
    b = a * a
    # CHECK: ArithValue(%2 = arith.mulf %cst_1, %cst_1 : f64)
    print(b)
```

**EDIT**: this now goes through the bindings and thus supports automatic
casting of `OpResult` (including as an element of `OpResultList`),
`BlockArgument` (including as an element of `BlockArgumentList`), as
well as `Value`.
2023-11-07 10:49:41 -06:00
Mehdi Amini
863c346209 [MLIR] Switch the default for usePropertiesForAttributes (NFC)
This is adopting properties as storage for attribute by default.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D158581
2023-08-28 20:25:03 -07:00
Mehdi Amini
57f03baf90 Revert "[MLIR] Switch the default for usePropertiesForAttributes (NFC)"
This reverts commit ef3ab3de9787f55b05620e56853909d758b9eb9d.

The revision hasn't actually been approved yet!
2023-08-27 16:30:44 -07:00
Mehdi Amini
ef3ab3de97 [MLIR] Switch the default for usePropertiesForAttributes (NFC)
This is adopting properties as storage for attribute by default.
2023-08-27 16:14:31 -07:00
max
bfb1ba7526 [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
2023-05-26 11:02:05 -05:00
max
d39a784402 [MLIR][python bindings] Expose TypeIDs in python
This diff adds python bindings for `MlirTypeID`. It paves the way for returning accurately typed `Type`s from python APIs (see D150927) and then further along building type "conscious" `Value` APIs (see D150413).

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D150839
2023-05-22 13:19:54 -05:00
Tres Popp
c1fa60b4cd [mlir] Update method cast calls to function calls
The MLIR classes Type/Attribute/Operation/Op/Value support
cast/dyn_cast/isa/dyn_cast_or_null functionality through llvm's doCast
functionality in addition to defining methods with the same name.
This change begins the migration of uses of the method to the
corresponding function call as has been decided as more consistent.

Note that there still exist classes that only define methods directly,
such as AffineExpr, and this does not include work currently to support
a functional cast/isa call.

Context:

* https://mlir.llvm.org/deprecation/ at "Use the free function variants for dyn_cast/cast/isa/…"
* Original discussion at https://discourse.llvm.org/t/preferred-casting-style-going-forward/68443

Implementation:
This follows a previous patch that updated calls
`op.cast<T>()-> cast<T>(op)`. However some cases could not handle an
unprefixed `cast` call due to occurrences of variables named cast, or
occurring inside of class definitions which would resolve to the method.
All C++ files that did not work automatically with `cast<T>()` are
updated here to `llvm::cast` and similar with the intention that they
can be easily updated after the methods are removed through a
find-replace.

See https://github.com/llvm/llvm-project/compare/main...tpopp:llvm-project:tidy-cast-check
for the clang-tidy check that is used and then update printed
occurrences of the function to include `llvm::` before.

One can then run the following:
```
ninja -C $BUILD_DIR clang-tidy

run-clang-tidy -clang-tidy-binary=$BUILD_DIR/bin/clang-tidy -checks='-*,misc-cast-functions'\
                 -export-fixes /tmp/cast/casts.yaml mlir/*\
                 -header-filter=mlir/ -fix

rm -rf $BUILD_DIR/tools/mlir/**/*.inc
```

Differential Revision: https://reviews.llvm.org/D150348
2023-05-12 11:21:30 +02:00
max
69cc3cfb21 [MLIR][python bindings] implement PyValue subclassing to enable operator overloading
Differential Revision: https://reviews.llvm.org/D147758
2023-04-14 14:25:06 -05:00
Stella Laurenzo
5e83a5b475 [mlir] Overhaul C/Python registration APIs to properly scope registration/loading activities.
Since the very first commits, the Python and C MLIR APIs have had mis-placed registration/load functionality for dialects, extensions, etc. This was done pragmatically in order to get bootstrapped and then just grew in. Downstreams largely bypass and do their own thing by providing various APIs to register things they need. Meanwhile, the C++ APIs have stabilized around this and it would make sense to follow suit.

The thing we have observed in canonical usage by downstreams is that each downstream tends to have native entry points that configure its installation to its preferences with one-stop APIs. This patch leans in to this approach with `RegisterEverything.h` and `mlir._mlir_libs._mlirRegisterEverything` being the one-stop entry points for the "upstream packages". The `_mlir_libs.__init__.py` now allows customization of the environment and Context by adding "initialization modules" to the `_mlir_libs` package. If present, `_mlirRegisterEverything` is treated as such a module. Others can be added by downstreams by adding a `_site_initialize_{i}.py` module, where '{i}' is a number starting with zero. The number will be incremented and corresponding module loaded until one is not found. Initialization modules can:

* Perform load time customization to the global environment (i.e. registering passes, hooks, etc).
* Define a `register_dialects(registry: DialectRegistry)` function that can extend the `DialectRegistry` that will be used to bootstrap the `Context`.
* Define a `context_init_hook(context: Context)` function that will be added to a list of callbacks which will be invoked after dialect registration during `Context` initialization.

Note that the `MLIRPythonExtension.RegisterEverything` is not included by default when building a downstream (its corresponding behavior was prior). For downstreams which need the default MLIR initialization to take place, they must add this back in to their Python CMake build just like they add their own components (i.e. to `add_mlir_python_common_capi_library` and `add_mlir_python_modules`). It is perfectly valid to not do this, in which case, only the things explicitly depended on and initialized by downstreams will be built/packaged. If the downstream has not been set up for this, it is recommended to simply add this back for the time being and pay the build time/package size cost.

CMake changes:
* `MLIRCAPIRegistration` -> `MLIRCAPIRegisterEverything` (renamed to signify what it does and force an evaluation: a number of places were incidentally linking this very expensive target)
* `MLIRPythonSoure.Passes` removed (without replacement: just drop)
* `MLIRPythonExtension.AllPassesRegistration` removed (without replacement: just drop)
* `MLIRPythonExtension.Conversions` removed (without replacement: just drop)
* `MLIRPythonExtension.Transforms` removed (without replacement: just drop)

Header changes:
* `mlir-c/Registration.h` is deleted. Dialect registration functionality is now in `IR.h`. Registration of upstream features are in `mlir-c/RegisterEverything.h`. When updating MLIR and a couple of downstreams, I found that proper usage was commingled so required making a choice vs just blind S&R.

Python APIs removed:
  * mlir.transforms and mlir.conversions (previously only had an __init__.py which indirectly triggered `mlirRegisterTransformsPasses()` and `mlirRegisterConversionPasses()` respectively). Downstream impact: Remove these imports if present (they now happen as part of default initialization).
  * mlir._mlir_libs._all_passes_registration, mlir._mlir_libs._mlirTransforms, mlir._mlir_libs._mlirConversions. Downstream impact: None expected (these were internally used).

C-APIs changed:
  * mlirRegisterAllDialects(MlirContext) now takes an MlirDialectRegistry instead. It also used to trigger loading of all dialects, which was already marked with a TODO to remove -- it no longer does, and for direct use, dialects must be explicitly loaded. Downstream impact: Direct C-API users must ensure that needed dialects are loaded or call `mlirContextLoadAllAvailableDialects(MlirContext)` to emulate the prior behavior. Also see the `ir.c` test case (e.g. `  mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("func"));`).
  * mlirDialectHandle* APIs were moved from Registration.h (which now is restricted to just global/upstream registration) to IR.h, arguably where it should have been. Downstream impact: include correct header (likely already doing so).

C-APIs added:
  * mlirContextLoadAllAvailableDialects(MlirContext): Corresponds to C++ API with the same purpose.

Python APIs added:
  * mlir.ir.DialectRegistry: Mapping for an MlirDialectRegistry.
  * mlir.ir.Context.append_dialect_registry(MlirDialectRegistry)
  * mlir.ir.Context.load_all_available_dialects()
  * mlir._mlir_libs._mlirAllRegistration: New native extension that exposes a `register_dialects(MlirDialectRegistry)` entry point and performs all upstream pass/conversion/transforms registration on init. In this first step, we eagerly load this as part of the __init__.py and use it to monkey patch the Context to emulate prior behavior.
  * Type caster and capsule support for MlirDialectRegistry

This should make it possible to build downstream Python dialects that only depend on a subset of MLIR. See: https://github.com/llvm/llvm-project/issues/56037

Here is an example PR, minimally adapting IREE to these changes: https://github.com/iree-org/iree/pull/9638/files In this situation, IREE is opting to not link everything, since it is already configuring the Context to its liking. For projects that would just like to not think about it and pull in everything, add `MLIRPythonExtension.RegisterEverything` to the list of Python sources getting built, and the old behavior will continue.

Reviewed By: mehdi_amini, ftynse

Differential Revision: https://reviews.llvm.org/D128593
2022-07-16 17:27:50 -07:00
Alex Zinenko
89a92fb3ba [mlir] Rework subclass construction in PybindAdaptors.h
The constructor function was being defined without indicating its "__init__"
name, which made it interpret it as a regular fuction rather than a
constructor. When overload resolution failed, Pybind would attempt to print the
arguments actually passed to the function, including "self", which is not
initialized since the constructor couldn't be called. This would result in
"__repr__" being called with "self" referencing an uninitialized MLIR C API
object, which in turn would cause undefined behavior when attempting to print
in C++. Even if the correct name is provided, the mechanism used by
PybindAdaptors.h to bind constructors directly as "__init__" functions taking
"self" is deprecated by Pybind. The new mechanism does not seem to have access
to a fully-constructed "self" object (i.e., the constructor in C++ takes a
`pybind11::detail::value_and_holder` that cannot be forwarded back to Python).

Instead, redefine "__new__" to perform the required checks (there are no
additional initialization needed for attributes and types as they are all
wrappers around a C++ pointer). "__new__" can call its equivalent on a
superclass without needing "self".

Bump pybind11 dependency to 3.8.0, which is the first version that allows one
to redefine "__new__".

Reviewed By: stellaraccident

Differential Revision: https://reviews.llvm.org/D117646
2022-01-19 18:09:05 +01:00
Alex Zinenko
14c9207063 [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
2021-10-25 12:50:42 +02:00