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

Matteo Franciolini llvmlistbot at llvm.org
Thu Oct 31 14:29:08 PDT 2024


================
@@ -429,6 +429,86 @@ tryMoveLastSplitAxisInResharding(ImplicitLocOpBuilder &builder, MeshOp mesh,
   return std::nullopt;
 }
 
+// Detect a change in the halo size (only) and create necessary operations if
+// needed. A changed halo sizes requires copying the "core" of the source tensor
+// into the "core" of the destination tensor followed by an update halo
+// operation.
+static std::optional<std::tuple<TypedValue<ShapedType>, MeshSharding>>
+tryUpdateHaloInResharding(ImplicitLocOpBuilder &builder, MeshOp mesh,
+                          MeshSharding sourceSharding,
+                          MeshSharding targetSharding,
+                          ShapedType sourceUnshardedShape,
+                          TypedValue<ShapedType> sourceShard) {
+  // currently handles only cases where halo sizes differ but everything else
----------------
mfrancio wrote:

Some comments in this file miss proper capitalization and punctuation.

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


More information about the Mlir-commits mailing list