[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