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

[Boian Petkantchin]

================
@@ -879,4 +1056,39 @@ def Mesh_ShiftOp : Mesh_CollectiveCommunicationOpBase<"shift", [
   let hasCanonicalizer = 1;
 }
 
+def Mesh_UpdateHaloOp : Mesh_CollectiveCommunicationOpBase<"update_halo", [
+    AllShapesMatch<["input", "result"]>,
+    AllElementTypesMatch<["input", "result"]>
+  ]> {
+  let summary = "Update halo data.";
+  let description = [{
+    This operation updates halo regions of shards, e.g. if their sharding
+    specified `halo_sizes` and the actual tensor data might have changed
+    on the remote devices. Changes might be caused by mutating operations
+    and/or if the new halo regions are larger than the existing ones.
+
+    Assumes all devices hold tensors with same-sized halo data as specified
+    by `halo_sizes`.
+
+    `mesh_axes` specifies the tensor axes along which the halo data is updated.
+    Currently each tensor dim can be sharded along a single mesh axis only.
+
+    Optionally resizes to new halo sizes `target_halo_sizes`.
+  }];
+  let arguments = !con(commonArgs, (ins
+    AnyNon0RankedTensor:$input,
+    DenseI64ArrayAttr:$dynamic_halo_sizes,
+    OptionalAttr<DenseI64ArrayAttr>:$target_halo_sizes
+  ));
+  let results = (outs
+    AnyRankedTensor:$result
+  );
+  let assemblyFormat = [{
+    $input `on` $mesh
+    (`mesh_axes` `=` $mesh_axes^)?
+    `halo_sizes` `=` $dynamic_halo_sizes
----------------
sogartar wrote:

Why is there a discrepancy in the naming (`dynamic_`)?

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