Commit Graph

254 Commits

Author SHA1 Message Date
srcarroll
b79db39659 [mlir][linalg] Support ParamType in vector_sizes option of VectorizeOp transform (#87557) 2024-04-09 15:52:40 -05: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
Matthias Gehre
8ec28af8ea Reapply "[mlir][PDL] Add support for native constraints with results (#82760)"
with a small stack-use-after-scope fix in getConstraintPredicates()

This reverts commit c80e6edba4.
2024-03-02 20:57:30 +01:00
Matthias Gehre
c80e6edba4 Revert "[mlir][PDL] Add support for native constraints with results (#82760)"
Due to buildbot failure https://lab.llvm.org/buildbot/#/builders/88/builds/72130

This reverts commit dca32a3b59.
2024-03-01 07:44:30 +01:00
Matthias Gehre
dca32a3b59 [mlir][PDL] Add support for native constraints with results (#82760)
From https://reviews.llvm.org/D153245

This adds support for native PDL (and PDLL) C++ constraints to return
results.

This is useful for situations where a pattern checks for certain
constraints of multiple interdependent attributes and computes a new
attribute value based on them. Currently, for such an example it is
required to escape to C++ during matching to perform the check and after
a successful match again escape to native C++ to perform the computation
during the rewriting part of the pattern. With this work we can do the
computation in C++ during matching and use the result in the rewriting
part of the pattern. Effectively this enables a choice in the trade-off
of memory consumption during matching vs recomputation of values.

This is an example of a situation where this is useful: We have two
operations with certain attributes that have interdependent constraints.
For instance `attr_foo: one_of [0, 2, 4, 8], attr_bar: one_of [0, 2, 4,
8]` and `attr_foo == attr_bar`. The pattern should only match if all
conditions are true. The new operation should be created with a new
attribute which is computed from the two matched attributes e.g.
`attr_baz = attr_foo * attr_bar`. For the check we already escape to
native C++ and have all values at hand so it makes sense to directly
compute the new attribute value as well:

```
Constraint checkAndCompute(attr0: Attr, attr1: Attr) -> Attr;

Pattern example with benefit(1) {
    let foo = op<test.foo>() {attr = attr_foo : Attr};
    let bar = op<test.bar>(foo) {attr = attr_bar : Attr};
    let attr_baz = checkAndCompute(attr_foo, attr_bar);
    rewrite bar with {
        let baz = op<test.baz> {attr=attr_baz};
        replace bar with baz;
    };
}
```
To achieve this the following notable changes were necessary:
PDLL:
- Remove check in PDLL parser that prevented native constraints from
returning results

PDL:
- Change PDL definition of pdl.apply_native_constraint to allow variadic
results

PDL_interp:
- Change PDL_interp definition of pdl_interp.apply_constraint to allow
variadic results

PDLToPDLInterp Pass:
The input to the pass is an arbitrary number of PDL patterns. The pass
collects the predicates that are required to match all of the pdl
patterns and establishes an ordering that allows creation of a single
efficient matcher function to match all of them. Values that are matched
and possibly used in the rewriting part of a pattern are represented as
positions. This allows fusion and thus reusing a single position for
multiple matching patterns. Accordingly, we introduce
ConstraintPosition, which records the type and index of the result of
the constraint. The problem is for the corresponding value to be used in
the rewriting part of a pattern it has to be an input to the
pdl_interp.record_match operation, which is generated early during the
pass such that its surrounding block can be referred to by branching
operations. In consequence the value has to be materialized after the
original pdl.apply_native_constraint has been deleted but before we get
the chance to generate the corresponding pdl_interp.apply_constraint
operation. We solve this by emitting a placeholder value when a
ConstraintPosition is evaluated. These placeholder values (due to fusion
there may be multiple for one constraint result) are replaced later when
the actual pdl_interp.apply_constraint operation is created.

Changes since the phabricator review:
- Addressed all comments
- In particular, removed registerConstraintFunctionWithResults and
instead changed registerConstraintFunction so that contraint functions
always have results (empty by default)
- Thus we don't need to reuse `rewriteFunctions` to store constraint
functions with results anymore, and can instead use
`constraintFunctions`
- Perform a stable sort of ConstraintQuestion, so that
ConstraintQuestion appear before other ConstraintQuestion that use their
results.
- Don't create placeholders for pdl_interp::ApplyConstraintOp. Instead
generate the `pdl_interp::ApplyConstraintOp` before generating the
successor block.
- Fixed a test failure in the pdl python bindings


Original code by @martin-luecke

Co-authored-by: martin-luecke <martinpaul.luecke@amd.com>
2024-03-01 07:29:49 +01:00
Alexander Belyaev
bfcd3fa825 [mlir] Add result name for gpu.block_id and gpu.thread_id ops. (#83393)
expand-arith-ops.mlir fails on windows, but this is unrelated to this PR
2024-02-29 10:57:09 +01:00
Peiming Liu
56d58295dd [mlir][sparse] Introduce batch level format. (#83082) 2024-02-26 16:08:28 -08: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
Peiming Liu
429919e328 [mlir][sparse][pybind][CAPI] remove LevelType enum from CAPI, constru… (#81682)
…ct LevelType from LevelFormat and properties instead.

**Rationale**
We used to explicitly declare every possible combination between
`LevelFormat` and `LevelProperties`, and it now becomes difficult to
scale as more properties/level formats are going to be introduced.
2024-02-13 16:45:22 -08: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
Yinying Li
2a6b521b36 [mlir][sparse] Add more tests and verification for n:m (#81186)
1. Add python test for n out of m
2. Add more methods for python binding
3. Add verification for n:m and invalid encoding tests
4. Add e2e test for n:m

Previous PRs for n:m #80501 #79935
2024-02-09 14:34:36 -05:00
Yinying Li
e5924d6499 [mlir][sparse] Implement parsing n out of m (#79935)
1. Add parsing methods for block[n, m].
2. Encode n and m with the newly extended 64-bit LevelType enum.
3. Update 2:4 methods names/comments to n:m.
2024-02-08 14:38:42 -05:00
Yinying Li
cd481fa827 [mlir][sparse] Change LevelType enum to 64 bit (#80501)
1. C++ enum is set through enum class LevelType : uint_64.
2. C enum is set through typedef uint_64 level_type. It is due to the
limitations in Windows build: setting enum width to ui64 is not
supported in C.
2024-02-05 17:00:52 -05:00
Jacques Pienaar
59eadcd28f [mlir][py] Reduce size of allocation memrefs in test. 2024-02-01 16:27:20 -08:00
Maksim Levental
404af14f92 [mlir][python] enable memref.subview (#79393) 2024-01-30 16:21:56 -06:00
Guray Ozen
12c241b365 [MLIR][NVVM] Explicit Data Type for Output in wgmma.mma_async (#78713)
The current implementation of `nvvm.wgmma.mma_async` Op deduces the data
type of the output matrix from the data type of struct member, which can be
non-intuitive, especially in cases where types like `2xf16` are packed
into `i32`.

This PR addresses this issue by improving the Op to include an explicit
data type for the output matrix.

The modified Op now includes an explicit data type for Matrix-D (<f16>),
and looks as follows:

```
%result = llvm.mlir.undef : !llvm.struct<(struct<(i32, i32, ...
nvvm.wgmma.mma_async
    %descA, %descB, %result,
    #nvvm.shape<m = 64, n = 32, k = 16>,
    D [<f16>, #nvvm.wgmma_scale_out<zero>],
    A [<f16>, #nvvm.wgmma_scale_in<neg>, <col>],
    B [<f16>, #nvvm.wgmma_scale_in<neg>, <col>]
```
2024-01-22 08:37:20 +01:00
Jacques Pienaar
8934b10642 [mlir][arith] Add overflow flags support to arith ops (#78376)
Add overflow flags support to the following ops:
* `arith.addi`
* `arith.subi`
* `arith.muli`

Example of new syntax:
```
%res = arith.addi %arg1, %arg2 overflow<nsw> : i64
```
Similar to existing LLVM dialect syntax
```
%res = llvm.add %arg1, %arg2 overflow<nsw> : i64
```

Tablegen canonicalization patterns updated to always drop flags, proper
support with tests will be added later.

Updated LLVMIR translation as part of this commit as it currenly written
in a way that it will crash when new attributes added to arith ops
otherwise.

Also lower `arith` overflow flags to corresponding SPIR-V op decorations

Discussion

https://discourse.llvm.org/t/rfc-integer-overflow-flags-support-in-arith-dialect/76025

This effectively rolls forward #77211, #77700 and #77714 while adding a
test to ensure the Python usage is not broken. More follow up needed but
unrelated to the core change here. The changes here are minimal and just
correspond to "textual namespacing" ODS side, no C++ or Python changes
were needed.

---------

---------

Co-authored-by: Ivan Butygin <ivan.butygin@gmail.com>, Yi Wu <yi.wu2@arm.com>
2024-01-17 06:12:23 +03: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
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
max
4a6ed4a90d [mlir][python] fix affine test 2023-12-07 16:21:57 -06:00
Maksim Levental
98d8dce6e9 [mlir][affine] implement inferType for delinearize (#74644) 2023-12-07 15:59:52 -06:00
Maksim Levental
db3bc49487 [mlir][python] fix up affine for (#74495) 2023-12-07 10:55:55 -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
Maksim Levental
17ec364b1b [mlir][python] enable registering dialects with the default Context (#72488) 2023-11-27 19:26:05 -06:00
Aart Bik
1944c4f76b [mlir][sparse] rename DimLevelType to LevelType (#73561)
The "Dim" prefix is a legacy left-over that no longer makes sense, since
we have a very strict "Dimension" vs. "Level" definition for sparse
tensor types and their storage.
2023-11-27 14:27:52 -08:00
Maksim Levental
e9453f3c3c [mlir][python] fix scf.for_ convenience builder (#72170) 2023-11-13 20:25:41 -06: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
2859a9cdfa Revert "[mlir][python] NFC - Lint fix"
This reverts commit 8c014e5949.
2023-11-08 09:34:28 +00:00
Nicolas Vasilache
8c014e5949 [mlir][python] NFC - Lint fix
Followup from #71597
2023-11-08 09:02:30 +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
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
Christian Ulmann
7ed96b1c0d [MLIR][LLVM] Remove last typed pointer remnants from tests (#71232)
This commit removes all LLVM dialect typed pointers from the lit tests.
Typed pointers have been deprecated for a while now and it's planned to
soon remove them from the LLVM dialect.

Related PSA:
https://discourse.llvm.org/t/psa-removal-of-typed-pointers-from-the-llvm-dialect/74502
2023-11-04 14:13:31 +01: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
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
Maksim Levental
b0e00ca6a6 [mlir][python] fix replace=True for register_operation and register_type_caster (#70264)
<img
src="https://github.com/llvm/llvm-project/assets/5657668/443852b6-ac25-45bb-a38b-5dfbda09d5a7"
height="400" />
<p></p>


So turns out that none of the `replace=True` things actually work
because of the map caches (except for
`register_attribute_builder(replace=True)`, which doesn't use such a
cache). This was hidden by a series of unfortunate events:

1. `register_type_caster` failure was hidden because it was the same
`TestIntegerRankedTensorType` being replaced with itself (d'oh).
2. `register_operation` failure was hidden behind the "order of events"
in the lifecycle of typical extension import/use. Since extensions are
loaded/registered almost immediately after generated builders are
registered, there is no opportunity for the `operationClassMapCache` to
be populated (through e.g., `module.body.operations[2]` or
`module.body.operations[2].opview` or something). Of course as soon as
you as actually do "late-bind/late-register" the extension, you see it's
not successfully replacing the stale one in `operationClassMapCache`.

I'll take this opportunity to propose we ditch the caches all together.
I've been cargo-culting them but I really don't understand how they
work. There's this comment above `operationClassMapCache`

```cpp
  /// Cache of operation name to external operation class object. This is
  /// maintained on lookup as a shadow of operationClassMap in order for repeat
  /// lookups of the classes to only incur the cost of one hashtable lookup.
  llvm::StringMap<pybind11::object> operationClassMapCache;
```

But I don't understand how that's true given that the canonical thing
`operationClassMap` is already a map:

```cpp
  /// Map of full operation name to external operation class object.
  llvm::StringMap<pybind11::object> operationClassMap;
```

Maybe it wasn't always the case? Anyway things work now but it seems
like an unnecessary layer of complexity for not much gain? But maybe I'm
wrong.
2023-10-30 20:22:27 -05: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
Yinying Li
b165650aee [mlir][sparse] Return actual identity map instead of null map (#70365)
Changes:

1. For both dimToLvl and lvlToDim, always returns the actual map instead
of AffineMap() for identity map.
2. Updated custom builder for encoding to have default values.
3. Non-inferable lvlToDim will still return AffineMap() during
inference, so it will be caught by verifier.
2023-10-26 18:30:34 -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