10 Commits

Author SHA1 Message Date
Boian Petkantchin
5e29112719
[mlir][mesh] Add verification and canonicalization for some collectives (#74905)
Add verification and canonicalization for
broadcast, gather, recv, reduce, scatter, send and shift.

The canonicalizations only remove trivial collectives with empty
mesh_axes attrubutes.
2023-12-15 06:41:10 -08:00
Boian Petkantchin
944e031e36
[mlir][mesh] Use tensor shape notation for the shape of a cluster (#73826)
Examle:

substitute
mesh.cluster @mesh0(rank = 2, dim_sizes = [0, 4])

with
mesh.cluster @mesh0(rank = 2, dim_sizes = ?x4)

Same as tensor/memref shapes. The only difference is for 0-rank shapes.
With tensors you would have something like `tensor<f32>`. Here to avoid
matching an empty string a 0-rank shape is denoted by `[]`.
2023-12-08 11:34:44 -08:00
Boian Petkantchin
dff2f59be3
[mlir][mesh] Add TableGen deffinitions of more collective ops (#73842)
Add definitions for
broadcast, gather, receive, reduce, scatter, send and shift.
2023-12-04 09:11:47 -08:00
Boian Petkantchin
5f7c8c1068
[mlir][mesh] Add collective communication operations (#71960)
Add all-gather, all-reduce, all-to-all and reduce-scatter. These
operations have device mesh semantics.
2023-11-21 06:50:24 -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
Jie Fu
bda763aea0 [mlir] Fix -Wreturn-type in MeshOps.cpp (NFC)
/llvm-project/mlir/lib/Dialect/Mesh/IR/MeshOps.cpp:73:1: error: non-void function does not return a value in all control paths [-Werror,-Wreturn-type]
}
^
1 error generated.
2023-11-04 08:47:30 +08: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
Chengji Yao
1a21196b9a
[MLIR] reverse int8 type's printing logic (#69361)
Specializing for 8-bit integers to ensure values are printed as integers

Fixes #69310
2023-10-18 10:30:13 -07:00
Chengji Yao
08545e8516
[MLIR] Add a new Mesh dialect (#68007)
This is the 1st PR of [Mesh sharding
RFC](https://discourse.llvm.org/t/open-mlir-meeting-9-28-2023-rfc-sharding-framework-design-for-device-mesh/73695),
includes

Includes:

- mesh.cluster op
- mesh.shard op (the mesh.annotate op in the RFC slides, the name is
modified a bit from @stellaraccident 's advice, which I think might be a
bit more concise)
- MeshSharding attribute
2023-10-10 11:35:40 -07:00