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

[Matteo Franciolini]

================
@@ -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.
----------------
mfrancio wrote:

Apologies if i did not provide more context in my previous message, but is there an instance where the input tensor shape is not directly available that warrant the N+1 representation? 

If that's the case, then I would agree with you in that the original implementation (size N, inclusive prefix sum representing the end offset) was probably better, especially because now you have a flattened 2d array of n+1 values which adds additional complexity.

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