Commit Graph

346 Commits

Author SHA1 Message Date
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 4f51b2bfe3.
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
Amy Wang
de7857ab23 [mlir][python] python binding for the affine.store op (#68816)
This PR creates the necessary files to support bindings for operations
in the affine dialect.

This is the first of many PRs which will progressively introduce
affine.load, affine.for, etc operations. I would like to
acknowledge the work by Nelli's author @makslevental :
https://github.com/makslevental/nelli/blob/main/nelli/mlir/affine/affine.py
which jump-starts the work.
2023-10-11 16:37:11 -04:00
Maksim Levental
27c6d55cae [mlir][python] generate value builders (#68308)
This PR adds the additional generation of what I'm calling "value
builders" (a term I'm not married to) that look like this:

```python
def empty(sizes, element_type, *, loc=None, ip=None):
    return get_result_or_results(tensor.EmptyOp(sizes=sizes, element_type=element_type, loc=loc, ip=ip))
```

which instantiates a `tensor.EmptyOp` and then immediately grabs the
result (`OpResult`) and then returns that *instead of a handle to the
op*.

What's the point of adding these when `EmptyOp.result` already exists?
My claim/feeling/intuition is that eDSL users are more comfortable with
a value centric programming model (i.e., passing values as operands) as
opposed to an operator instantiation programming model. Thus this change
enables (or at least goes towards) the bindings supporting such a user
and use case. For example,

```python
i32 = IntegerType.get_signless(32)
...
ten1 = tensor.empty((10, 10), i32)
ten2 = tensor.empty((10, 10), i32)
ten3 = arith.addi(ten1, ten2)
```

Note, in order to present a "pythonic" API and enable "pythonic" eDSLs,
the generated identifiers (op names and operand names) are snake case
instead of camel case and thus `llvm::convertToSnakeFromCamelCase`
needed a small fix. Thus this PR is stacked on top of
https://github.com/llvm/llvm-project/pull/68375.

In addition, as a kind of victory lap, this PR adds a "rangefor" that
looks and acts exactly like python's `range` but emits `scf.for`.
2023-10-09 14:16:28 -07:00
Jack Frankland
e29a253c9e [mlir][tosa][linalg] Apply direct tosa -> linalg Conv2D lowering (#68304)
TOSA defines the filter channel ordering for 2D convolution operation
`tosa.conv2d` as `[OC, KH, KW, IC]`. The LinAlg dialect supports `[F, H,
W, C]` and `[H, W, C, F]` orderings via the `linalg.conv_2d_nhwc_fhwc`
and `linalg.conv_2d_nhwc_hwcf` operations respectively. Where `F == OC`,
`KH == H`, `KW == W` and `C == IC`.

Currently `tosa.conv2d` is lowered to `linalg.conv2d_nhwc_hwcf` meaning
we need to insert a transposition operation to permute the filter
channels before they can be passed as weights to the linalg op, that is
`[F, H, W, C]` -> `[H, W, C, F]`. An analogous transformation needs to
be applied to the quantized operation that lowers to
`linalg.conv_2d_nhwc_hwcf_q`.

This commit updates the TOSA->LinAlg lowering so that `tosa.conv2d` is
lowered to `linalg.conv2d_nhwc_fhwc` removing the need for the
introduction of a transposition operation and making the mapping 1-1. It
also adds a `linalg.conv_2d_nhwc_fhwc_q` quantized operation to the
LinAlg dialect so the same direct 1-1 mapping can be applied to the
quantized variant.

This commit does not add any new lit tests but repurposes the current
TosaToLinalgNamed tests by removing the checks for transpositions and
updating the targeted LinAlg operations from `linalg.conv2d_nhwc_hwcf`
to linalg.conv2d_nhwc_fhwc`.

Signed-off-by: Jack Frankland <jack.frankland@arm.com>
2023-10-06 17:10:39 -07:00