[all-commits] [llvm/llvm-project] fca698: [mlir][mesh] adding shard-size control (#98145)

[Frank Schlimbach via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: fca69838caf19854769ada21a71da91fcfcbde73
      https://github.com/llvm/llvm-project/commit/fca69838caf19854769ada21a71da91fcfcbde73
  Author: Frank Schlimbach <frank.schlimbach at intel.com>
  Date:   2024-08-07 (Wed, 07 Aug 2024)

  Changed paths:
    M mlir/include/mlir/Dialect/Mesh/IR/CMakeLists.txt
    M mlir/include/mlir/Dialect/Mesh/IR/MeshBase.td
    M mlir/include/mlir/Dialect/Mesh/IR/MeshOps.h
    M mlir/include/mlir/Dialect/Mesh/IR/MeshOps.td
    M mlir/include/mlir/Dialect/Mesh/Interfaces/ShardingInterface.h
    M mlir/include/mlir/Dialect/Mesh/Interfaces/ShardingInterface.td
    M mlir/include/mlir/Dialect/Mesh/Interfaces/ShardingInterfaceImpl.h
    A mlir/include/mlir/Dialect/Tensor/IR/ShardingInterfaceImpl.h
    M mlir/include/mlir/InitAllDialects.h
    M mlir/include/mlir/Interfaces/InferTypeOpInterface.h
    M mlir/lib/Dialect/Linalg/Transforms/MeshShardingInterfaceImpl.cpp
    M mlir/lib/Dialect/Mesh/IR/MeshOps.cpp
    M mlir/lib/Dialect/Mesh/Interfaces/ShardingInterface.cpp
    M mlir/lib/Dialect/Mesh/Transforms/ShardingPropagation.cpp
    M mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp
    M mlir/lib/Dialect/Tensor/IR/CMakeLists.txt
    A mlir/lib/Dialect/Tensor/IR/ShardingInterfaceImpl.cpp
    M mlir/test/Dialect/Linalg/mesh-sharding-propagation.mlir
    M mlir/test/Dialect/Linalg/mesh-spmdization.mlir
    M mlir/test/Dialect/Mesh/canonicalization.mlir
    M mlir/test/Dialect/Mesh/invalid.mlir
    M mlir/test/Dialect/Mesh/ops.mlir
    M mlir/test/Dialect/Mesh/resharding-spmdization.mlir
    M mlir/test/Dialect/Mesh/sharding-propagation.mlir
    M mlir/test/Dialect/Mesh/simplifications.mlir
    M mlir/test/Dialect/Mesh/spmdization.mlir
    A mlir/test/Dialect/Tensor/mesh-spmdization.mlir
    M mlir/test/lib/Dialect/Mesh/TestReshardingSpmdization.cpp

  Log Message:
  -----------
  [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



To unsubscribe from these emails, change your notification settings at https://github.com/llvm/llvm-project/settings/notifications