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

Frank Schlimbach llvmlistbot at llvm.org
Mon Nov 4 02:20:57 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:
----------------
fschlimb wrote:

Yes, ideally this would be a 2d array. As we allow dynamic and static values, I didn't intuitively know how to implement the parsing for it. Mechanics exist for 1d arrays, though. We can add this in an extra PR, but at this point I didn't want to spend time on it. Of course it'd be great if someone had a solution.

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


More information about the Mlir-commits mailing list