[mlir][mesh, MPI] Mesh2mpi (#104566)

Pass for lowering `Mesh` to `MPI`.
Initial commit lowers `UpdateHaloOp` only.

[mlir][mesh] Handling changed halo region sizes during spmdization (#114238)

* Changed `MeshSharding::sharded_dims_sizes` from representing sizes per
shard to offsets to origin per shard.
- Lo

[mlir][mesh] Handling changed halo region sizes during spmdization (#114238)

* Changed `MeshSharding::sharded_dims_sizes` from representing sizes per
shard to offsets to origin per shard.
- Local shard size are now a simple subtraction
- Offsets are now readily available without a reduction operation
- Enables constant value/shape propagation through standard
canonicalization
- Renamed to `sharded_dims_offsets` accordingly.
* First spmdization pattern for halo regions.
- Triggers when source and destination shardings differ only in their
halo sizes
- Copies local data from source into a new tensor and calls update_halo
- Supports arbitrary mesh dimensions (unlike the other patterns which
work on 1d meshes only)
* `UpdateHaloOp` implements `DestinationStyleOpInterface` and accepts
tensors and memrefs
- also accepts target and source halo sizes; both are required for
proper lowering
* minor refactoring for testing partial MeshSharding equality
* Canonicalization for ShardingOp folding constant values into
respective `static_*` attributes

show more ...

[mlir][mesh] Shardingcontrol (#102598)

This is a fixed copy of #98145 (necessary after it got reverted).

@sogartar @yaochengji
This PR adds the following to #98145:
- `UpdateHaloOp` accepts a `

[mlir][mesh] Shardingcontrol (#102598)

This is a fixed copy of #98145 (necessary after it got reverted).

@sogartar @yaochengji
This PR adds the following to #98145:
- `UpdateHaloOp` accepts a `memref` (instead of a tensor) and not
returning a result to clarify its inplace-semantics
- `UpdateHaloOp` accepts `split_axis` to allow multiple mesh-axes per
tensor/memref-axis (similar to `mesh.sharding`)
- The implementation of `Shardinginterface` for tensor operation
(`tensor.empty` for now) moved from the tensor library to the mesh
interface library. `spmdize` uses features from `mesh` dialect.
@rengolin agreed that `tensor` should not depend on `mesh` so this
functionality cannot live in a `tensor`s lib. The unfulfilled dependency
caused the issues leading to reverting #98145. Such cases are generally
possible and might lead to re-considering the current structure (like
for tosa ops).
- rebased onto latest main
--------------------------
Replacing `#mesh.sharding` attribute with operation `mesh.sharding`
- extended semantics now allow providing optional `halo_sizes` and
`sharded_dims_sizes`
- internally a sharding is represented as a non-IR class
`mesh::MeshSharding`

What previously was
```mlir
%sharded0 = mesh.shard %arg0 <@mesh0, [[0]]> : tensor<4x8xf32>
%sharded1 = mesh.shard %arg1 <@mesh0, [[0]]> annotate_for_users : tensor<16x8xf32>
```
is now
```mlir
%sharding = mesh.sharding @mesh0, [[0]] : !mesh.sharding
%0 = mesh.shard %arg0 to %sharding : tensor<4x8xf32>
%1 = mesh.shard %arg1 to %sharding annotate_for_users : tensor<16x8xf32>
```
and allows additional annotations to control the shard sizes:
```mlir
mesh.mesh @mesh0 (shape = 4)
%sharding0 = mesh.sharding @mesh0, [[0]] halo_sizes = [1, 2] : !mesh.sharding
%0 = mesh.shard %arg0 to %sharding0 : tensor<4x8xf32>
%sharding1 = mesh.sharding @mesh0, [[0]] sharded_dims_sizes = [3, 5, 5, 3] : !mesh.sharding
%1 = mesh.shard %arg1 to %sharding1 annotate_for_users : tensor<16x8xf32>
```
- `mesh.shard` op accepts additional optional attribute `force`, useful
for halo updates
- Some initial spmdization support for the new semantics
- Support for `tensor.empty` reacting on `sharded_dims_sizes` and
`halo_sizes` in the sharding
- New collective operation `mesh.update_halo` as a spmdized target for
shardings with `halo_sizes`

---------

Co-authored-by: frank.schlimbach <fschlimb@smtp.igk.intel.com>
Co-authored-by: Jie Fu <jiefu@tencent.com>

show more ...

Revert "[mlir][mesh] adding shard-size control (#98145)"

This reverts commit fca69838caf19854769ada21a71da91fcfcbde73.

Also reverts the fixup: "[mlir] Fix -Wunused-variable in MeshOps.cpp (NFC)"

T

Revert "[mlir][mesh] adding shard-size control (#98145)"

This reverts commit fca69838caf19854769ada21a71da91fcfcbde73.

Also reverts the fixup: "[mlir] Fix -Wunused-variable in MeshOps.cpp (NFC)"

This reverts commit fc737368fe6e27d6ecf76e522cb43a32aaca992a.

show more ...

[mlir][mesh] adding shard-size control (#98145)

- Replacing `#mesh.sharding` attribute with operation `mesh.sharding`
- extended semantics now allow providing optional `halo_sizes` and
`sharded_di

[mlir][mesh] adding shard-size control (#98145)

- Replacing `#mesh.sharding` attribute with operation `mesh.sharding`
- extended semantics now allow providing optional `halo_sizes` and
`sharded_dims_sizes`
- internally a sharding is represented as a non-IR class
`mesh::MeshSharding`

What previously was
```mlir
%sharded0 = mesh.shard %arg0 <@mesh0, [[0]]> : tensor<4x8xf32>
%sharded1 = mesh.shard %arg1 <@mesh0, [[0]]> annotate_for_users : tensor<16x8xf32>
```
is now
```mlir
%sharding = mesh.sharding @mesh0, [[0]] : !mesh.sharding
%0 = mesh.shard %arg0 to %sharding : tensor<4x8xf32>
%1 = mesh.shard %arg1 to %sharding annotate_for_users : tensor<16x8xf32>
```
and allows additional annotations to control the shard sizes:
```mlir
mesh.mesh @mesh0 (shape = 4)
%sharding0 = mesh.sharding @mesh0, [[0]] halo_sizes = [1, 2] : !mesh.sharding
%0 = mesh.shard %arg0 to %sharding0 : tensor<4x8xf32>
%sharding1 = mesh.sharding @mesh0, [[0]] sharded_dims_sizes = [3, 5, 5, 3] : !mesh.sharding
%1 = mesh.shard %arg1 to %sharding1 annotate_for_users : tensor<16x8xf32>
```
- `mesh.shard` op accepts additional optional attribute `force`, useful
for halo updates
- Some initial spmdization support for the new semantics
- Support for `tensor.empty` reacting on `sharded_dims_sizes` and
`halo_sizes` in the sharding
- New collective operation `mesh.update_halo` as a spmdized target for
shardings with `halo_sizes`

@sogartar @yaochengji

show more ...

mlir/LogicalResult: move into llvm (#97309)

This patch is part of a project to move the Presburger library into
LLVM.

[mlir][mesh] Fix wrong argument passed to targetShardingInUnsplitLastAxis (#95059)

In unsplitLastAxisInResharding, wrong argument was passed when calling
targetShardingInUnsplitLastAxis.There weren

[mlir][mesh] Fix wrong argument passed to targetShardingInUnsplitLastAxis (#95059)

In unsplitLastAxisInResharding, wrong argument was passed when calling
targetShardingInUnsplitLastAxis.There weren't any tests to uncover this.
I added one in mesh-spmdization.mlir for Linalg and one in
resharding-spmdization.mlir for Mesh dialects.

show more ...

[mlir][mesh] Insert resharding during sharding propagation (#84514)

If there are conflicts between the sharding annotations of some op,
insert resharding.
Make the Spmdization pass more forgiving

[mlir][mesh] Insert resharding during sharding propagation (#84514)

If there are conflicts between the sharding annotations of some op,
insert resharding.
Make the Spmdization pass more forgiving to allow for more than 2
chained `mesh.shard` ops.

Implement `getReductionLoopIteratorKinds` in ShardingInterface for
linalg ops.

show more ...

Switch member calls to `isa/dyn_cast/cast/...` to free function calls. (#89356)

This change cleans up call sites. Next step is to mark the member
functions deprecated.

See https://mlir.llvm.org/

Switch member calls to `isa/dyn_cast/cast/...` to free function calls. (#89356)

This change cleans up call sites. Next step is to mark the member
functions deprecated.

See https://mlir.llvm.org/deprecation and
https://discourse.llvm.org/t/preferred-casting-style-going-forward.

show more ...

[mlir][mesh] Make sharding propagation and spmdization work on FuncOpInterface (#84415)

Make them more general instead of only supporting `func::FuncOp`.

[mlir][mesh] add support in spmdization for incomplete sharding annotations (#82442)

Don't require that `mesh.shard` operations come in pairs. If there is
only a single `mesh.shard` operation we as

[mlir][mesh] add support in spmdization for incomplete sharding annotations (#82442)

Don't require that `mesh.shard` operations come in pairs. If there is
only a single `mesh.shard` operation we assume that the producer result
and consumer operand have the same sharding.

show more ...

[mlir][mesh] Add all-slice operation (#81218)

This op is the inverse of all-gather. It is useful to have an explicit
concise representation instead of having a blob of slicing logic.

Add lowerin

[mlir][mesh] Add all-slice operation (#81218)

This op is the inverse of all-gather. It is useful to have an explicit
concise representation instead of having a blob of slicing logic.

Add lowering for the op that slices from the tensor based on the
in-group process index.

Make resharding generate an all-slice instead of inserting the slicing
logic directly.

show more ...

[mlir][mesh] Add spmdization pass (#80518)

Add a pass that converts a function that has sharding annotations into
SPMD form.

[mlir][mesh] Refactoring code organization, tests and docs (#79606)

* Split out `MeshDialect.h` form `MeshOps.h` that defines the dialect
class. Reduces include clutter if you care only about the d

[mlir][mesh] Refactoring code organization, tests and docs (#79606)

* Split out `MeshDialect.h` form `MeshOps.h` that defines the dialect
class. Reduces include clutter if you care only about the dialect and
not the ops.

* Expose functions `getMesh` and `collectiveProcessGroupSize`. There
functions are useful for outside users of the dialect.

* Remove unused code.

* Remove examples and tests of mesh.shard attribute in tensor encoding.
Per the decision that Spmdization would be performed on sharding
annotations and there will be no tensors with sharding specified in the
type.
For more info see this RFC comment:
https://discourse.llvm.org/t/rfc-sharding-framework-design-for-device-mesh/73533/81

show more ...

[mlir][mesh] Rename cluster to mesh (#79484)

Rename
* Op mesh.cluster -> mesh.mesh
* Op mesh.cluster_shape -> mesh.mesh_shape
* variables and attributes.

The name `mesh` is more specific to wh

[mlir][mesh] Rename cluster to mesh (#79484)

Rename
* Op mesh.cluster -> mesh.mesh
* Op mesh.cluster_shape -> mesh.mesh_shape
* variables and attributes.

The name `mesh` is more specific to what it really represents. It is a
mesh of devices.
The name `cluster` implies a broader posibility of device
configurations. When just the word `mesh` is used the meaning can often
be inferred from the context whether it refers to the mesh dialect or a
device mesh. The full name can be used when needed.

show more ...

[mlir][mesh] Remove rank attribute and rename dim_sizes to shape in ClusterOp (#77838)

Remove the somewhat redundant rank attribute.
Before this change
```
mesh.cluster @mesh(rank = 3, dim_sizes

[mlir][mesh] Remove rank attribute and rename dim_sizes to shape in ClusterOp (#77838)

Remove the somewhat redundant rank attribute.
Before this change
```
mesh.cluster @mesh(rank = 3, dim_sizes = 2x3)
```
After
```
mesh.cluster @mesh(shape = 2x3x?)
```

The rank is instead determined by the provided shape. With this change
no longer `getDimSizes()` can be wrongly assumed to have size equal to
the cluster rank.
Now `getShape().size()` will always equal `getRank()`.

show more ...

[mlir][mesh] Add lowering of process multi-index op (#77490)

* Rename mesh.process_index -> mesh.process_multi_index.
* Add mesh.process_linear_index op.
* Add lowering of mesh.process_multi_index

[mlir][mesh] Add lowering of process multi-index op (#77490)

* Rename mesh.process_index -> mesh.process_multi_index.
* Add mesh.process_linear_index op.
* Add lowering of mesh.process_multi_index into an expression using
mesh.process_linear_index, mesh.cluster_shape and
affine.delinearize_index.

This is useful to lower mesh ops and prepare them for further lowering
where the runtime may have only the linear index of a device/process.
For example in MPI we have a rank (linear index) in a communicator.

show more ...

[mlir][mesh] Use one type for mesh axis (#76830)

Make all ops and attributes use the types MeshAxis and MeshAxesAttr
instead of int16_t, int32_t, DenseI16ArrayAttr and DenseI32ArrayAttr.

[mlir][mesh] Fix -Wunused-variable in Spmdization.cpp (NFC)

llvm-project/mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp:573:14:
error: unused variable 'targetShardType' [-Werror,-Wunused-variable

[mlir][mesh] Fix -Wunused-variable in Spmdization.cpp (NFC)

llvm-project/mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp:573:14:
error: unused variable 'targetShardType' [-Werror,-Wunused-variable]
ShapedType targetShardType =
^
1 error generated.

show more ...

[mlir][mesh] Add resharding spmdization on a 1D device mesh (#76179)

The current implementation supports only sharding of tensor axes that
have size divisible by the mesh axis size.