5 Commits

Author SHA1 Message Date
Boian Petkantchin
5df2c00af3
[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()`.
2024-01-15 07:39:09 -08:00
Boian Petkantchin
79aa776267
[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.
2024-01-10 07:01:16 -08:00
Boian Petkantchin
7a4c49756d
[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.
2024-01-03 15:47:11 -08:00
Jie Fu
ab43cf26ca [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.
2024-01-03 09:29:14 +08:00
Boian Petkantchin
1a8fb88719
[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.
2024-01-02 15:50:07 -08:00