349 Commits

Author SHA1 Message Date
Oleksandr "Alex" Zinenko
73140daebb
[mlir] expose transform dialect symbol merge to python (#87690)
This functionality is available in C++, make it available in Python
directly to operate on transform modules.
2024-04-17 15:01:59 +02:00
Kirill Podoprigora
4cd7bb07c7
[mlir] Remove `dataclasses package from mlir requirements.txt` (#87223)
The ``dataclasses`` package makes sense for Python 3.6, becauses
``dataclasses`` is only included in the standard library with 3.7
version. Now, 3.6 has reached EOL, so all current supported versions of
Python (3.8, 3.9, 3.10, 3.11, 3.12) have this feature in their standard
libraries.

Therefore there's no need to install the ``dataclasses`` package now.
2024-04-01 12:07:10 -07:00
Steven Varoumas
eb861acd49
[mlir][python] Enable python bindings for Index dialect (#85827)
This small patch enables python bindings for the index dialect.

---------

Co-authored-by: Steven Varoumas <steven.varoumas1@huawei.com>
2024-03-20 16:56:22 +01:00
Oleksandr "Alex" Zinenko
5d59fa90ce
Reapply "[mlir][py] better support for arith.constant construction" (#84142)
Arithmetic constants for vector types can be constructed from objects
implementing Python buffer protocol such as `array.array`. Note that
until Python 3.12, there is no typing support for buffer protocol
implementers, so the annotations use array explicitly.

Reverts llvm/llvm-project#84103
2024-03-07 17:14:08 +01:00
Mehdi Amini
96fc54828a
Revert "[mlir][py] better support for arith.constant construction" (#84103)
Reverts llvm/llvm-project#83259

This broke an integration test on Windows
2024-03-05 18:57:45 -08:00
Oleksandr "Alex" Zinenko
a691f65a84
[mlir][py] better support for arith.constant construction (#83259)
Arithmetic constants for vector types can be constructed from objects
implementing Python buffer protocol such as `array.array`. Note that
until Python 3.12, there is no typing support for buffer protocol
implementers, so the annotations use array explicitly.
2024-03-05 16:09:59 +01:00
Sergei Lebedev
6ce5159945
[MLIR][Python] Use ir.Value directly instead of _SubClassValueT (#82341)
_SubClassValueT is only useful when it is has >1 usage in a signature.
This was not true for the signatures produced by tblgen.

For example

def call(result, callee, operands_, *, loc=None, ip=None) ->
_SubClassValueT:
        ...

here a type checker does not have enough information to infer a type
argument for _SubClassValueT, and thus effectively treats it as Any.
2024-02-21 12:59:23 +01: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
Oleksandr "Alex" Zinenko
bd8fcf75df
[mlir][python] expose LLVMStructType API (#81672)
Expose the API for constructing and inspecting StructTypes from the LLVM
dialect. Separate constructor methods are used instead of overloads for
better readability, similarly to IntegerType.
2024-02-14 15:03:04 +01:00
Sergei Lebedev
82f3cbc860
[MLIR][Python] Added a base class to all builtin floating point types (#81720)
This allows to

* check if a given ir.Type is a floating point type via isinstance() or
issubclass()
* get the bitwidth of a floating point type

See motivation and discussion in
https://discourse.llvm.org/t/add-floattype-to-mlir-python-bindings/76959.
2024-02-14 13:02:49 +01:00
Rolf Morel
4c654b7b91
[MLIR][Python] Add missing peel_front argument to LoopPeelOp's extension class (#81424) 2024-02-12 11:35:43 -06: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
Maksim Levental
404af14f92
[mlir][python] enable memref.subview (#79393) 2024-01-30 16:21:56 -06:00
Rageking8
5f4c89edd0
Fix unsigned typos (#76670) 2024-01-27 22:20:08 -08:00
Quinn Dawkins
5caab8bbc0
[mlir][transform] Add transform.get_operand op (#78397)
Similar to `transform.get_result`, except it returns a handle to the
operand indicated by a positional specification, same as is defined for
the linalg match ops.

Additionally updates `get_result` to take the same positional specification.
This makes the use case of wanting to get all of the results of an
operation easier by no longer requiring the user to reconstruct the list
of results one-by-one.
2024-01-18 09:33:14 -05:00
martin-luecke
06e3abcb54
[MLIR][transform][python] Introduce abstractions for handles to values and parameters (#77305)
In addition to the existing `OpHandle` which provides an abstraction to
emit transform ops targeting operations this introduces a similar
concept for _values_ and _parameters_ in form of `ValueHandle` and
`ParamHandle`.

New core transform abstractions:
- `constant_param`
- `OpHandle.get_result`
- `OpHandle.print`
- `ValueHandle.get_defining_op`
2024-01-15 10:31:22 +01:00
Maksim Levental
83be8a7400
[mlir][python] add MemRefTypeAttr attr builder (#76371) 2024-01-06 16:42:14 -06:00
Sergei Lebedev
3737712dae
Slightly improved ir.pyi type annotations (#76728)
* Replaced `Any` with static types where appropriate
* Removed undocumented `__str__` and `__repr__` -- these are always
defined via `object`
2024-01-04 09:49:57 +01:00
Jungwook Park
2292fd0129
[mlir][spirv] Add support for C-API/python binding to SPIR-V dialect (#76055)
Enable bindings.

---------

Co-authored-by: jungpark-mlir <jungwook@jungwook-22.04>
2024-01-02 08:11:44 -08:00
Maksim Levental
537b2aa264
[mlir][python] meta region_op (#75673) 2023-12-21 11:20:29 -06:00
Maksim Levental
acaff70841
[mlir][python] move transform extras (#76102) 2023-12-20 17:29:11 -06:00
martin-luecke
681eacc1b6
[MLIR][transform][python] add sugared python abstractions for transform dialect (#75073)
This adds Python abstractions for the different handle types of the
transform dialect

The abstractions allow for straightforward chaining of transforms by
calling their member functions.
As an initial PR for this infrastructure, only a single transform is
included: `transform.structured.match`.
With a future `tile` transform abstraction an example of the usage is: 
```Python
def script(module: OpHandle):
    module.match_ops(MatchInterfaceEnum.TilingInterface).tile(tile_sizes=[32,32])
```
to generate the following IR:
```mlir
%0 = transform.structured.match interface{TilingInterface} in %arg0
%tiled_op, %loops = transform.structured.tile_using_for %0 [32, 32]
```

These abstractions are intended to enhance the usability and flexibility
of the transform dialect by providing an accessible interface that
allows for easy assembly of complex transformation chains.
2023-12-15 13:04:43 +01:00
Maksim Levental
d36b483f4f
[mlir][python] update type stubs (#75099) 2023-12-11 18:35:02 -06:00
Maksim Levental
db3bc49487
[mlir][python] fix up affine for (#74495) 2023-12-07 10:55:55 -06:00
Mehdi Amini
6b0ed49c8e [mlir] Fix missing cmake dependency causing non-deterministic build failure (NFC)
Fixes #74611
2023-12-07 03:22:45 -08:00
Peter Hawkins
45e7b410c0
[mlir:python] Fail immediately if importing an initializer module raises ImportError (#74595) 2023-12-06 12:42:11 -06:00
Amy Wang
543589af49
[mlir][python] python binding wrapper for the affine.AffineForOp (#74408)
This PR creates the wrapper class AffineForOp and adds a testcase for
it. A testcase for the AffineLoadOp is also added.
2023-12-05 02:33:56 -05:00
Felix Schneider
ed22bf6991
[mlir][linalg] Fix weight dimension ordering in 2D grouped conv (#73855)
The `conv_2d_ngchw_fgchw` Op implements 2d grouped convolution with
dimensions ordered as given in the name. However, the current
implementation orders weights as `gfchw` instead of `fgchw`. This was
already pointed out in an old phabricator revision which never landed:
https://reviews.llvm.org/D150064

This patch
1) Adds a new op `conv_2d_ngchw_gfchw`
2) Fixes the dimension ordering of the old op `conv_2d_ngchw_fgchw`
3) Adds tests with non-dynamic dimensions so that it's easier to
  understand.
2023-12-01 22:16:00 +01:00
Maksim Levental
17ec364b1b
[mlir][python] enable registering dialects with the default Context (#72488) 2023-11-27 19:26:05 -06:00
Maksim Levental
225648e91c
[mlir][python] add type wrappers (#71218) 2023-11-27 15:58:00 -06:00
Maksim Levental
4eaf3a9c12
[mlir][python] reformat transform ext (#71136) 2023-11-16 13:37:52 -06:00
Nicolas Vasilache
1451411e64 [mlir][python] NFC - Expose LoopExtensionOps to SCFLoopTransformOps.td 2023-11-15 16:25:22 +00:00
Maksim Levental
e9453f3c3c
[mlir][python] fix scf.for_ convenience builder (#72170) 2023-11-13 20:25:41 -06:00
Stella Laurenzo
0677e54653
[mlir][python] Allow contexts to be created with a custom thread pool. (#72042)
The existing initialization sequence always enables multi-threading at
MLIRContext construction time, making it impractical to provide a
customized thread pool.

Here, this is changed to always create the context with threading
disabled, process all site-specific init hooks (which can set thread
pools) and ultimately enable multi-threading unless if site-configured
to not do so.

This should preserve the existing user-visible initialization behavior
while also letting downstreams ensure that contexts are always created
with a shared thread pool. This was tested with IREE, which has such a
concept. Using site-specific thread tuning produced up to 2x single
compilation job behavior and customization of batch compilation (i.e. as
part of a build system) to utilize half the memory and run the entire
test suite ~2x faster. Given this, I believe that the additional
configurability can well pay for itself for implementations that use it.
We may also want to present user-level Python APIs for controlling
threading configuration in the future.
2023-11-11 21:41:56 -08:00
Nicolas Vasilache
5967375fcf [mlir][python] Add support for arg_attrs and other attrs to NamedSequenceOp 2023-11-08 13:42:16 +00:00
Nicolas Vasilache
af3d856944 [mlir][python] Reland - Add sugared builder for transform.named_sequence
Address issues with #71597 post-revert and and reland
2023-11-08 09:34:29 +00:00
Nicolas Vasilache
cb3880515f Revert "[mlir][python]Add sugared buider for transform.named_sequence (#71597)"
This reverts commit 4f51b2bfe3ec11b597272be6caa00efb575bc59f.
2023-11-08 09:34:29 +00:00
Nicolas Vasilache
4f51b2bfe3
[mlir][python]Add sugared buider for transform.named_sequence (#71597) 2023-11-08 09:49:57 +01:00
Maksim Levental
c86d35a5f4
[mlir][python] factor out pure python core sources (#71592)
I'd like to be able to install just the Python core sources (without
building/including the pybind sources).
2023-11-07 19:52:43 -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
Maksim Levental
5192e299cf
[mlir][python] remove various caching mechanisms (#70831)
This PR removes the various caching mechanisms currently in the python
bindings - both positive caching and negative caching.
2023-11-03 13:28:20 -05:00
Maksim Levental
2ab14dff43
[mlir][python] fix python_test dialect and I32/I64ElementsBuilder (#70871)
This PR fixes the `I32ElementsAttr` and `I64ElementsAttr` builders and
tests them through the `python_test` dialect.
2023-10-31 19:55:42 -05:00
Mehdi Amini
b2bdc45580 [mlir][python] Fix possible use of variable use before set
The _mlirRegisterEverything symbol may not be built by some customers. The
code here was intended to support this, but didn't properly initialize the
init_module variable.
This would break JAX with:

NameError: free variable 'init_module' referenced before assignment in enclosing scope
2023-10-31 10:31:20 -07:00
Matthias Springer
04736c7f7a
[mlir][SCF] Use transform.get_parent_op instead of transform.loop.get_parent_for (#70757)
Add a new attribute to `get_parent_op` to get the n-th parent. Remove
`transform.loop.get_parent_for`, which is no longer needed.
2023-10-31 18:36:40 +09:00
Jungwook Park
6995183e17
[mlir][python] Register LLVM translations in the RegisterEverything for python (#70428)
Added missing register_translations in python to replicate the same in
the C-API
Cleaned up the current calls to register passes where the other calls
are already embedded in the mlirRegisterAllPasses.
found here,
https://discourse.llvm.org/t/opencl-example/74187
2023-10-30 14:46:21 -07:00
bjacob
8c8336fcad
Add missing linalg.batch_vecmat named op (#70218)
Linalg currently has these named ops:
* `matmul`
* `matvec`
* `vecmat`
* `batch_matmul`
* `batch_matvec`

But it does not have:
* `batch_vecmat`

This PRs adds that for consistency, and I have a short-term need for it
( https://github.com/openxla/iree/issues/15158 ), so not having this
would cause some contortion on my end.
2023-10-25 11:41:24 -04:00
Maksim Levental
a9694043c9
[mlir][linalg] regionBuilder for transpose, broadcast (#69742)
Currently, `linalg.transpose` and `linalg.broadcast` can't be emitted
through either the C API or the python bindings (which of course go
through the C API). See
https://discourse.llvm.org/t/how-to-build-linalg-transposeop-in-mlir-pybind/73989/10.

The reason is even though they're named ops, there is no opdsl
`@linalg_structured_op` for them and thus while they can be instantiated
they cannot be passed to
[`mlirLinalgFillBuiltinNamedOpRegion`](a7cccb9cbb/mlir/lib/CAPI/Dialect/Linalg.cpp (L18)).
I believe the issue is they both take a `IndexAttrDef` but
`IndexAttrDef` cannot represent dynamic rank. Note, if I'm mistaken and
there is a way to write the `@linalg_structured_op` let me know.

The solution here simply implements the `regionBuilder` interface which
is then picked up by
[`LinalgDialect::addNamedOpBuilders`](7557530f42/mlir/lib/Dialect/Linalg/IR/LinalgDialect.cpp (L116)).

Extension classes are added "by hand" that mirror the API of the
`@linalg_structured_op`s. Note, the extension classes are added to to
`dialects/linalg/__init__.py` instead of
`dialects/linalg/opdsl/ops/core_named_ops.py` in order that they're not
confused for opdsl generators/emitters.
2023-10-20 16:14:46 -05:00
Maksim Levental
dd473f1dd1
[mlir][python] simplify extensions (#69642)
https://github.com/llvm/llvm-project/pull/68853 enabled a lot of nice
cleanup. Note, I made sure each of the touched extensions had tests.
2023-10-19 18:07:06 -05:00
Maksim Levental
a2288a8944
[mlir][python] remove mixins (#68853)
This PR replaces the mixin `OpView` extension mechanism with the
standard inheritance mechanism.

Why? Firstly, mixins are not very pythonic (inheritance is usually used
for this), a little convoluted, and too "tight" (can only be used in the
immediately adjacent `_ext.py`). Secondly, it (mixins) are now blocking
are correct implementation of "value builders" (see
[here](https://github.com/llvm/llvm-project/pull/68764)) where the
problem becomes how to choose the correct base class that the value
builder should call.

This PR looks big/complicated but appearances are deceiving; 4 things
were needed to make this work:

1. Drop `skipDefaultBuilders` in
`OpPythonBindingGen::emitDefaultOpBuilders`
2. Former mixin extension classes are converted to inherit from the
generated `OpView` instead of being "mixins"
a. extension classes that simply were calling into an already generated
`super().__init__` continue to do so
b. (almost all) extension classes that were calling `self.build_generic`
because of a lack of default builder being generated can now also just
call `super().__init__`
3. To handle the [lone single
use-case](https://sourcegraph.com/search?q=context%3Aglobal+select_opview_mixin&patternType=standard&sm=1&groupBy=repo)
of `select_opview_mixin`, namely
[linalg](https://github.com/llvm/llvm-project/blob/main/mlir/python/mlir/dialects/_linalg_ops_ext.py#L38),
only a small change was necessary in `opdsl/lang/emitter.py` (thanks to
the emission/generation of default builders/`__init__`s)
4. since the `extend_opview_class` decorator is removed, we need a way
to register extension classes as the desired `OpView` that `op.opview`
conjures into existence; so we do the standard thing and just enable
replacing the existing registered `OpView` i.e.,
`register_operation(_Dialect, replace=True)`.

Note, the upgrade path for the common case is to change an extension to
inherit from the generated builder and decorate it with
`register_operation(_Dialect, replace=True)`. In the slightly more
complicated case where `super().__init(self.build_generic(...))` is
called in the extension's `__init__`, this needs to be updated to call
`__init__` in `OpView`, i.e., the grandparent (see updated docs). 
Note, also `<DIALECT>_ext.py` files/modules will no longer be automatically loaded.

Note, the PR has 3 base commits that look funny but this was done for
the purpose of tracking the line history of moving the
`<DIALECT>_ops_ext.py` class into `<DIALECT>.py` and updating (commit
labeled "fix").
2023-10-19 16:20:14 -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