[Mlir-commits] [mlir] [mlir][mesh] Handling changed halo region sizes during spmdization (PR #114238)

[Frank Schlimbach]

================
@@ -196,16 +197,18 @@ def Mesh_ShardingOp : Mesh_Op<"sharding", [
     e.g. the first sharded dimension gets [1,2] halos and the seconds gets [2,3] halos.
     `?` indicates dynamic halo sizes.
     
-    6. [Optional] Sizes of sharded dimensions of each shard.
-    `sharded_dims_sizes`is provided as a flattened 1d array of i64s: for each device of the
-    device-mesh one value for each sharded tensor dimension.
+    6. [Optional] Offsets for each shard and sharded tensor dimension.
+    `sharded_dims_offsets` is provided as a flattened 1d array of i64s:
+    For each sharded tensor dimension the offsets (starting index) of all shards in that dimension.
+    The offset of each first shard is omitted and is implicitly assumed to be 0.
----------------
fschlimb wrote:

We will need an additional last element, otherwise we cannot deduce the size of of the last shard. 
That is because there is no way to implicitly define the global tensor shape. It should be possible to deduce the sizes without pulling in a dependency to the input tensor.  The start will always be 0, so that is redundant information and can be avoided to keep the array to size N. Not so the end.
I will change to having offset for shard `i` as `sharded_dims_offsets[i]` and append the offset to the end.

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