[Mlir-commits] [mlir] [mlir][mesh] Handling changed halo region sizes during spmdization (PR #114238)
Frank Schlimbach
llvmlistbot at llvm.org
Mon Nov 4 02:14:29 PST 2024
================
@@ -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:
Yes get your point. We need the last value to able to compute the shard size of the last shard. As said elsewhere, I had to either do it as you suggest and have an array of size N+1 or have the first value given implicitly. I often prefer avoiding redundancy so I did it this way. I have no strong opinion here, if you have a clear preference I can switch to the longer array solution.
https://github.com/llvm/llvm-project/pull/114238
More information about the Mlir-commits
mailing list