4 Commits

Author SHA1 Message Date
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
Jie Fu
046dffce23 Fix -Wunused-variable in TestSimplifications.cpp (NFC)
llvm-project/mlir/test/lib/Dialect/Mesh/TestSimplifications.cpp:36:17:
 error: unused variable 'status' [-Werror,-Wunused-variable]
  LogicalResult status =
                ^
1 error generated.
2024-01-10 07:59:19 +08:00
Boian Petkantchin
ab590377a3
[mlir][mesh] Add folding of ClusterShapeOp (#77033)
If the mesh has static size on some of the requested axes, the result is
substituted with a constant.
2024-01-09 13:42:56 -08:00
Boian Petkantchin
4b3446771f
[mlir][mesh] Add endomorphism simplification for all-reduce (#73150)
Does transformations like
all_reduce(x) + all_reduce(y) -> all_reduce(x + y)

max(all_reduce(x), all_reduce(y)) -> all_reduce(max(x, y))
when the all_reduce element-wise op is max.

Added general rewrite pattern HomomorphismSimplification and
EndomorphismSimplification that encapsulate the general algorithm.
Made specialization for all-reduce with respect to
addf, addi, minsi, maxsi, minimumf and maximumf
in the Arithmetic dialect.
2023-12-12 10:21:52 -08:00