62 Commits

Author SHA1 Message Date
Peter Hawkins
41bd35b58b
[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.
2024-12-18 11:16:11 -08:00
Jonas Rickert
abad8455ab
[mlir] Expose skipRegions option for Op printing in the C and Python bindings (#96150)
The MLIR C and Python Bindings expose various methods from
`mlir::OpPrintingFlags` . This PR adds a binding for the `skipRegions`
method, which allows to skip the printing of Regions when printing Ops.
It also exposes this option as parameter in the python `get_asm` and
`print` methods
2024-06-20 10:15:08 -05:00
Jacques Pienaar
18cf1cd92b
[mlir] Add PDL C & Python usage (#94714)
Following a rather direct approach to expose PDL usage from C and then
Python. This doesn't yes plumb through adding support for custom
matchers through this interface, so constrained to basics initially.

This also exposes greedy rewrite driver. Only way currently to define
patterns is via PDL (just to keep small). The creation of the PDL
pattern module could be improved to avoid folks potentially accessing
the module used to construct it post construction. No ergonomic work
done yet.

---------

Signed-off-by: Jacques Pienaar <jpienaar@google.com>
2024-06-11 07:45:12 -07:00
Oleksandr "Alex" Zinenko
67897d77ed
[mlir][py] invalidate nested operations when parent is deleted (#93339)
When an operation is erased in Python, its children may still be in the
"live" list inside Python bindings. After this, if some of the newly
allocated operations happen to reuse the same pointer address, this will
trigger an assertion in the bindings. This assertion would be incorrect
because the operations aren't actually live. Make sure we remove the
children operations from the "live" list when erasing the parent.

This also concentrates responsibility over the removal from the "live"
list and invalidation in a single place.

Note that this requires the IR to be sufficiently structurally valid so
a walk through it can succeed. If this invariant was broken by, e.g, C++
pass called from Python, there isn't much we can do.
2024-05-30 10:06:02 +02:00
Hideto Ueno
47148832d4
[mlir][python] Add walk method to PyOperationBase (#87962)
This commit adds `walk` method to PyOperationBase that uses a python
object as a callback, e.g. `op.walk(callback)`. Currently callback must
return a walk result explicitly.

We(SiFive) have implemented walk method with python in our internal
python tool for a while. However the overhead of python is expensive and
it didn't scale well for large MLIR files. Just replacing walk with this
version reduced the entire execution time of the tool by 30~40% and
there are a few configs that the tool takes several hours to finish so
this commit significantly improves tool performance.
2024-04-17 15:09:47 +09:00
Oleksandr "Alex" Zinenko
91f1161133
[mlir] expose transform interpreter to Python (#82365)
Transform interpreter functionality can be used standalone without going
through the interpreter pass, make it available in Python.
2024-02-21 11:01:00 +01:00
John Demme
d1fdb41629
[MLIR][Python] Add method for getting the live operation objects (#78663)
Currently, a method exists to get the count of the operation objects
which are still alive. This helps for sanity checking, but isn't
terribly useful for debugging. This new method returns the actual
operation objects which are still alive.

This allows Python code like the following:

```
    gc.collect()
    live_ops = ir.Context.current._get_live_operation_objects()
    for op in live_ops:
      print(f"Warning: {op} is still live. Referrers:")
      for referrer in gc.get_referrers(op)[0]:
        print(f"  {referrer}")
```
2024-02-08 11:39:06 -08:00
Alex Zinenko
78bd124649 Revert "[mlir][python] Make the Context/Operation capsule creation methods work as documented. (#76010)"
This reverts commit bbc29768683b394b34600347f46be2b8245ddb30.

This change seems to be at odds with the non-owning part semantics of
MlirOperation in C API. Since downstream clients can only take and
return MlirOperation, it does not sound correct to force all returns of
MlirOperation transfer ownership. Specifically, this makes it impossible
for downstreams to implement IR-traversing functions that, e.g., look at
neighbors of an operation.

The following patch triggers the exception, and there does not seem to
be an alternative way for a downstream binding writer to express this:

```
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 39757dfad5be..2ce640674245 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3071,6 +3071,11 @@ void mlir::python::populateIRCore(py::module &m) {
                   py::arg("successors") = py::none(), py::arg("regions") = 0,
                   py::arg("loc") = py::none(), py::arg("ip") = py::none(),
                   py::arg("infer_type") = false, kOperationCreateDocstring)
+      .def("_get_first_in_block", [](PyOperation &self) -> MlirOperation {
+        MlirBlock block = mlirOperationGetBlock(self.get());
+        MlirOperation first = mlirBlockGetFirstOperation(block);
+        return first;
+      })
       .def_static(
           "parse",
           [](const std::string &sourceStr, const std::string &sourceName,
diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index f59b1a26ba48..6b12b8da5c24 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -24,6 +24,25 @@ def expect_index_error(callback):
     except IndexError:
         pass

+@run
+def testCustomBind():
+    ctx = Context()
+    ctx.allow_unregistered_dialects = True
+    module = Module.parse(
+        r"""
+    func.func @f1(%arg0: i32) -> i32 {
+      %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
+      return %1 : i32
+    }
+  """,
+        ctx,
+    )
+    add = module.body.operations[0].regions[0].blocks[0].operations[0]
+    op = add.operation
+    # This will get a reference to itself.
+    f1 = op._get_first_in_block()
+
+

 # Verify iterator based traversal of the op/region/block hierarchy.
 # CHECK-LABEL: TEST: testTraverseOpRegionBlockIterators
```
2023-12-21 10:06:44 +00:00
Stella Laurenzo
bbc2976868
[mlir][python] Make the Context/Operation capsule creation methods work as documented. (#76010)
This fixes a longstanding bug in the `Context._CAPICreate` method
whereby it was not taking ownership of the PyMlirContext wrapper when
casting to a Python object. The result was minimally that all such
contexts transferred in that way would leak. In addition, counter to the
documentation for the `_CAPICreate` helper (see
`mlir-c/Bindings/Python/Interop.h`) and the `forContext` /
`forOperation` methods, we were silently upgrading any unknown
context/operation pointer to steal-ownership semantics. This is
dangerous and was causing some subtle bugs downstream where this
facility is getting the most use.

This patch corrects the semantics and will only do an ownership transfer
for `_CAPICreate`, and it will further require that it is an ownership
transfer (if already transferred, it was just silently succeeding).
Removing the mis-aligned behavior made it clear where the downstream was
doing the wrong thing.

It also adds some `_testing_` functions to create unowned context and
operation capsules so that this can be fully tested upstream, reworking
the tests to verify the behavior.

In some torture testing downstream, I was not able to trigger any memory
corruption with the newly enforced semantics. When getting it wrong, a
regular exception is raised.
2023-12-20 12:18:58 -08:00
Adrian Kuegel
ea2e83af55 [mlir][Python] Apply ClangTidy findings.
move constructors should be marked noexcept
2023-12-11 09:43:08 +00:00
Jacques Pienaar
204acc5c10
[mlir][py] Overload print with state. (#72064)
Enables reusing the AsmState when printing from Python. Also moves the
fileObject and binary to the end (pybind11::object was resulting in the
overload not working unless `state=` was specified).

---------

Co-authored-by: Maksim Levental <maksim.levental@gmail.com>
2023-11-13 10:21:21 -08: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
Ingo Müller
fa19ef7a10
[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.
2023-10-25 07:17:56 +02:00
Maksim Levental
bdc3e6cb45
[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.
2023-10-20 20:28:32 -05:00
Tomás Longeri
5a600c23f9
[mlir][python] Expose PyInsertionPoint's reference operation (#69082)
The reason I want this is that I am writing my own Python bindings and
would like to use the insertion point from
`PyThreadContextEntry::getDefaultInsertionPoint()` to call C++ functions
that take an `OpBuilder` (I don't need to expose it in Python but it
also seems appropriate). AFAICT, there is currently no way to translate
a `PyInsertionPoint` into an `OpBuilder` because the operation is
inaccessible.
2023-10-18 16:53:18 +02:00
Jacques Pienaar
a677a17327 [mlir][py] Enable AsmState overload for operation. 2023-09-25 12:25:08 -07:00
Jacques Pienaar
75453714f0
[mlir][python] Expose AsmState python side. (#66819)
This does basic plumbing, ideally want a context approach to reduce
needing to thread these manually, but the current is useful even in that
state.

Made Value.get_name change backwards compatible, so one could either set
a field or create a state to pass in.
2023-09-20 15:12:06 -07:00
Jacques Pienaar
f573bc24d4 [mlir][py] Reuse more of CAPI build time inference.
This reduces code generated for type inference and instead reuses
facilities CAPI side that performed same role.

Differential Revision: https://reviews.llvm.org/D156041t
2023-07-23 21:26:52 -07:00
Adam Paszke
c83318e3e0 [MLIR][Python] Implement pybind adapters for MlirBlock
Reviewed By: jpienaar

Differential Revision: https://reviews.llvm.org/D155092
2023-07-12 22:27:01 -07:00
Rahul Kayaith
974c1596ab [mlir][python] Downcast attributes in more places
Update remaining `PyAttribute`-returning APIs to return `MlirAttribute` instead,
so that they go through the downcasting mechanism.

Reviewed By: makslevental

Differential Revision: https://reviews.llvm.org/D154462
2023-07-10 22:01:34 -04:00
max
9566ee2806 [MLIR][python bindings] TypeCasters for Attributes
Differential Revision: https://reviews.llvm.org/D151840
2023-06-07 12:01:00 -05: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
4811270bac [MLIR][python bindings] use pybind C++ APIs for throwing python errors.
Differential Revision: https://reviews.llvm.org/D151167
2023-05-23 11:31:16 -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
Jacques Pienaar
5c90e1ffb0 [mlir][bytecode] Return error instead of min version
Can't return a well-formed IR output while enabling version to be bumped
up during emission. Previously it would return min version but
potentially invalid IR which was confusing, instead make it return
error and abort immediately instead.

Differential Revision: https://reviews.llvm.org/D149569
2023-04-30 22:11:02 -07:00
Jacques Pienaar
0610e2f6a2 [mlir][bytecode] Allow client to specify a desired version.
Add method to set a desired bytecode file format to generate. Change
write method to be able to return status including the minimum bytecode
version needed by reader. This enables generating an older version of
the bytecode (not dialect ops, attributes or types). But this does not
guarantee that an older version can always be generated, e.g., if a
dialect uses a new encoding only available at later bytecode version.
This clamps setting to at most current version.

Differential Revision: https://reviews.llvm.org/D146555
2023-04-29 05:35:53 -07:00
Rahul Kayaith
e6d738e0c7 [mlir][python] Mark operator== overloads as const
This resolves some warnings when building with C++20, e.g.:
```
llvm-project/mlir/lib/Bindings/Python/IRAffine.cpp:545:60: warning: ISO C++20 considers use of overloaded operator '==' (with operand types 'mlir::python::PyAffineExpr' and 'mlir::python::PyAffineExpr') to be ambiguous despite there being a unique best viable function [-Wambiguous-reversed-operator]
                        PyAffineExpr &other) { return self == other; })
                                                      ~~~~ ^  ~~~~~
llvm-project/mlir/lib/Bindings/Python/IRAffine.cpp:350:20: note: ambiguity is between a regular call to this operator and a call with the argument order reversed
bool PyAffineExpr::operator==(const PyAffineExpr &other) {
                   ^
```

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D147018
2023-03-28 11:06:13 -04:00
Rahul Kayaith
3ea4c5014d [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
2023-03-07 14:59:22 -05:00
Rahul Kayaith
a7f8b7cd8e [mlir][python] Remove "Raw" OpView classes
The raw `OpView` classes are used to bypass the constructors of `OpView`
subclasses, but having a separate class can create some confusing
behaviour, e.g.:
```
op = MyOp(...)
# fails, lhs is 'MyOp', rhs is '_MyOp'
assert type(op) == type(op.operation.opview)
```

Instead we can use `__new__` to achieve the same thing without a
separate class:
```
my_op = MyOp.__new__(MyOp)
OpView.__init__(my_op, op)
```

Reviewed By: stellaraccident

Differential Revision: https://reviews.llvm.org/D143830
2023-03-01 18:17:14 -05:00
rkayaith
37107e177e [mlir][python] Add generic operation parse APIs
Currently the bindings only allow for parsing IR with a top-level
`builtin.module` op, since the parse APIs insert an implicit module op.
This change adds `Operation.parse`, which returns whatever top-level op
is actually in the source.

To simplify parsing of specific operations, `OpView.parse` is also
added, which handles the error checking for `OpView` subclasses.

Reviewed By: ftynse, stellaraccident

Differential Revision: https://reviews.llvm.org/D143352
2023-03-01 18:17:12 -05:00
Kazu Hirata
2140260d91 [mlir] Remove remaining uses of llvm::Optional (NFC)
This patch removes one "using" declaration and #include
"llvm/ADT/Optional.h".  It keeps several "using" declarations in
headers for downstream users.

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
2023-01-14 01:34:49 -08:00
Kazu Hirata
0a81ace004 [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
2023-01-14 01:25:58 -08:00
Kazu Hirata
a1fe1f5f77 [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
2023-01-13 21:05:06 -08:00
Fangrui Song
14ce58f3bf [mlir][python] llvm::Optional::value => operator*
And convert it to std::optional while updating.
2022-12-19 04:28:55 +00:00
Mehdi Amini
30c7c42341 Apply clang-tidy fixes for performance-unnecessary-value-param in IRCore.cpp (NFC) 2022-10-08 18:18:13 +00:00
Mehdi Amini
89418ddcb5 Plumb write_bytecode to the Python API
This adds a `write_bytecode` method to the Operation class.
The method takes a file handle and writes the binary blob to it.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D133210
2022-09-05 12:02:06 +00:00
Jacques Pienaar
10de551297 [mlir][python] Address deprecation warning for hasValue 2022-08-07 15:28:18 -07: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
John Demme
6b0bed7ea5 [MLIR] [Python] Add a method to clear live operations map
Introduce a method on PyMlirContext (and plumb it through to Python) to
invalidate all of the operations in the live operations map and clear
it. Since Python has no notion of private data, an end-developer could
reach into some 3rd party API which uses the MLIR Python API (that is
behaving correctly with regard to holding references) and grab a
reference to an MLIR Python Operation, preventing it from being
deconstructed out of the live operations map. This allows the API
developer to clear the map when it calls C++ code which could delete
operations, protecting itself from its users.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D123895
2022-04-19 15:14:09 -07:00
Dominik Grewe
774818c09c Expose MlirOperationClone in Python bindings.
Expose MlirOperationClone in Python bindings.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D122526
2022-03-28 15:58:22 +02:00
Mehdi Amini
e8d073951b Apply clang-tidy fixes for performance-unnecessary-value-param to MLIR (NFC) 2022-01-14 02:26:27 +00:00
Mehdi Amini
bd87241c44 Apply clang-tidy fixes for modernize-use-override to MLIR (NFC) 2022-01-14 02:26:27 +00:00
Mehdi Amini
9940dcfa4a Apply clang-tidy fixes for modernize-use-equals-default to MLIR (NFC) 2022-01-14 02:26:27 +00:00
Stella Laurenzo
7ee25bc56f [mlir][python] Add bindings for diagnostic handler.
I considered multiple approaches for this but settled on this one because I could make the lifetime management work in a reasonably easy way (others had issues with not being able to cast to a Python reference from a C++ constructor). We could stand to have more formatting helpers, but best to get the core mechanism in first.

Differential Revision: https://reviews.llvm.org/D116568
2022-01-04 11:04:37 -08:00
Mehdi Amini
4f415216ca Apply clang-tidy fixes for performance-unnecessary-value-param to MLIR (NFC) 2022-01-02 22:37:13 +00:00
Mehdi Amini
1fc096af1e Apply clang-tidy fixes for performance-unnecessary-value-param to MLIR (NFC)
Reviewed By: Mogball

Differential Revision: https://reviews.llvm.org/D116250
2022-01-02 01:45:18 +00:00
Stella Laurenzo
bdc3183742 [mlir][python] Implement more SymbolTable methods.
* set_symbol_name, get_symbol_name, set_visibility, get_visibility, replace_all_symbol_uses, walk_symbol_tables
* In integrations I've been doing, I've been reaching for all of these to do both general IR manipulation and module merging.
* I don't love the replace_all_symbol_uses underlying APIs since they necessitate SYMBOL_COUNT walks and have various sharp edges. I'm hoping that whatever emerges eventually for this can still retain this simple API as a one-shot.

Differential Revision: https://reviews.llvm.org/D114687
2021-11-29 20:31:13 -08:00
Stella Laurenzo
a6e7d024a9 [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
2021-11-29 19:58:58 -08:00
Stella Laurenzo
ace1d0ad3d [mlir][python] Normalize asm-printing IR behavior.
While working on an integration, I found a lot of inconsistencies on IR printing and verification. It turns out that we were:
  * Only doing "soft fail" verification on IR printing of Operation, not of a Module.
  * Failed verification was interacting badly with binary=True IR printing (causing a TypeError trying to pass an `str` to a `bytes` based handle).
  * For systematic integrations, it is often desirable to control verification yourself so that you can explicitly handle errors.

This patch:
  * Trues up the "soft fail" semantics by having `Module.__str__` delegate to `Operation.__str__` vs having a shortcut implementation.
  * Fixes soft fail in the presence of binary=True (and adds an additional happy path test case to make sure the binary functionality works).
  * Adds an `assume_verified` boolean flag to the `print`/`get_asm` methods which disables internal verification, presupposing that the caller has taken care of it.

It turns out that we had a number of tests which were generating illegal IR but it wasn't being caught because they were doing a print on the `Module` vs operation. All except two were trivially fixed:
  * linalg/ops.py : Had two tests for direct constructing a Matmul incorrectly. Fixing them made them just like the next two tests so just deleted (no need to test the verifier only at this level).
  * linalg/opdsl/emit_structured_generic.py : Hand coded conv and pooling tests appear to be using illegal shaped inputs/outputs, causing a verification failure. I just used the `assume_verified=` flag to restore the original behavior and left a TODO. Will get someone who owns that to fix it properly in a followup (would also be nice to break this file up into multiple test modules as it is hard to tell exactly what is failing).

Notes to downstreams:
  * If, like some of our tests, you get verification failures after this patch, it is likely that your IR was always invalid and you will need to fix the root cause. To temporarily revert to prior (broken) behavior, replace calls like `print(module)` with `print(module.operation.get_asm(assume_verified=True))`.

Differential Revision: https://reviews.llvm.org/D114680
2021-11-28 18:02:01 -08:00
Alex Zinenko
30d61893fb [mlir] provide C API and Python bindings for symbol tables
Symbol tables are a largely useful top-level IR construct, for example, they
make it easy to access functions in a module by name instead of traversing the
list of module's operations to find the corresponding function.

Depends On D112886

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D112821
2021-11-02 14:22:58 +01:00