[Mlir-commits] [mlir] mlir::mesh::shardingOp adding shard-size control (PR #98145)

[Frank Schlimbach]

================
@@ -129,15 +328,28 @@ def Mesh_ShardOp : Mesh_Op<"shard", [
     as an operand in subsequent operations. If not, the sharding applies to the
     operation that defines the tensor value.
 
+    4. `force`: A unit attribute requesting an explicit sharding of the data,
+    therefore not allowing to be optimizied away. This is useful in the presence
+    of halos and inplace semantics.
----------------
fschlimb wrote:

This might well be a lack of understanding on my side.

An `insert_slice` op does not have a result. So I am not sure I fully follow your line of thought.

This was not meant for the sharding propagation, I agree that this does not need this attribute. 

It was meant to serve the spmdization: when spmdization converts `shard` ops into communication ops, e.g. when the input/operand of a `shard` op is another `shard`. If the shardings of the `shard` ops are identical, the algorithm will simply do nothing, e.g. no sharding/communication will happen. However, an in-place operation will require a halo-update even if the sharding stays identical all the time. The elimination of the `shard` op would be incorrect.

```
...
%x = ...
%t = ...
%sharding = mesh.sharding halo_sizes = [2,2] ...
%t1 = mesh.shard %t %sharding
%x1 = mesh.shard %x %sharding
insert_slice %t1 %x1
%t2 = mesh.shard %t1  // <- this must remain and convert to a halo update
%t3 = mesh.shard %t2
<use %t3>
...
```

One solution could be to let the propagation pass insert an explicit `update_halo` op followed by a `shard` op. This would be pretty specific to halos (and not a s generic as a forced resharding) but maybe that is good enough for now. It might also make passes more complicated because now sharding propagation will create IR with sharding ops as well as explicit communication.

Not sure which one is better. Other suggestions?

https://github.com/llvm/llvm-project/pull/98145