6 Commits

Author SHA1 Message Date
Boian Petkantchin
abfac563f5
[mlir][mesh] Make sharding propagation and spmdization work on FuncOpInterface (#84415)
Make them more general instead of only supporting `func::FuncOp`.
2024-03-08 08:14:36 -08:00
Boian Petkantchin
adbf21f12b
[mlir][mesh] Add spmdization pass (#80518)
Add a pass that converts a function that has sharding annotations into
SPMD form.
2024-02-06 20:55:14 -08:00
Boian Petkantchin
31fc0a12e1
[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
2024-01-31 07:20:14 -08:00
Chengji Yao
b0d5b4d252
[MLIR][Mesh] Add sharding propagation pass (#71261)
Add a pass that propagates sharding information throughout the graph.
After this pass, each of the operations' operands and results is
annotated with a mesh.shard operation.

The pass is driven by a newly added ShardingInterface, and an implementation
for element-wise and matmul ops in the TOSA dialect is provided.
2023-11-03 21:07:31 -07:00
Mehdi Amini
466abaf152 Revert "[MLIR][Mesh] Add sharding propagation pass (#69665)"
This reverts commit 9d9400d7de9b928e3018af97e8b381a4a6ba5162.
This reverts commit bda763aea0b854178c01eac9f309042d9aaa823b.

The buildbot is broken and tests are failing.
2023-11-03 17:52:41 -07:00
Chengji Yao
9d9400d7de
[MLIR][Mesh] Add sharding propagation pass (#69665)
Add a pass that propagates sharding information throughout the graph.
After this pass, each of the operations' operands and results is
annotated with a `mesh.shard` operation, and the operations themselves
are added with sharding option attributes.

The pass is driven by  a newly added `ShardingInterface`, and an implementation
for element-wise and matmul ops in the TOSA dialect is provided.
2023-11-03 17:12:42 -07:00