comments/docs

Author: fschlimb

mlir/include/mlir/Dialect/Mesh/IR/MeshOps.td

@@ -1066,18 +1066,20 @@ def Mesh_UpdateHaloOp : Mesh_Op<"update_halo", [
   let summary = "Update halo data.";
   let description = [{
     This operation updates halo regions of shards, e.g. if their sharding
-    specified halos and the actual tensor data might have changed
+    specified halos and the actual tensor/memref 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.
 
+    Source and destination might have different halo sizes.
+
     Assumes all devices hold tensors with same-sized halo data as specified
     by `source_halo_sizes/static_source_halo_sizes` and
-    `destination_halo_sizes/static_destination_halo_sizes`
+    `destination_halo_sizes/static_destination_halo_sizes` in source shard
+    and destination/result shard.
 
     `split_axes` specifies for each tensor axis along which mesh axes its halo
     data is updated.
 
-    Source and destination might have different halo sizes.
   }];
   let arguments = (ins
     AnyTypeOf<[AnyNon0RankedMemRef, AnyNon0RankedTensor]>:$source,

mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp

@@ -429,12 +429,18 @@ 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
+  // stays the same (from source to destination sharding)
   if (sourceSharding.equalSplitAndPartialAxes(targetSharding) &&
       sourceSharding.getPartialAxes().empty() &&
       targetSharding.getPartialAxes().empty() &&
@@ -454,6 +460,9 @@ tryUpdateHaloInResharding(ImplicitLocOpBuilder &builder, MeshOp mesh,
     SmallVector<int64_t> srcCoreOffs(rank, 0), tgtCoreOffs(rank, 0),
         strides(rank, 1), outShape(sourceShard.getType().getShape()),
         coreShape(sourceShard.getType().getShape());
+
+    // determine "core" of source and destination
+    // the core is the local part of the shard excluding halo regions
     for (auto i = 0u; i < rank; ++i) {
       if (i < splitAxes.size() && !splitAxes[i].empty()) {
         if (!srcHaloSizes.empty()) {
@@ -465,6 +474,8 @@ tryUpdateHaloInResharding(ImplicitLocOpBuilder &builder, MeshOp mesh,
             coreShape[i] + tgtHaloSizes[i * 2] + tgtHaloSizes[i * 2 + 1];
       }
     }
+
+    // extract core from source and copy into destination core
     auto noVals = ValueRange{};
     auto initVal = builder.create<tensor::EmptyOp>(
         sourceShard.getLoc(), outShape, sourceShard.getType().getElementType());
@@ -476,6 +487,8 @@ tryUpdateHaloInResharding(ImplicitLocOpBuilder &builder, MeshOp mesh,
     auto initOprnd = builder.create<tensor::InsertSliceOp>(
         sourceShard.getLoc(), core, initVal, noVals, noVals, noVals,
         tgtCoreOffs, coreShape, strides);
+
+    // finally update the halo
     auto updateHaloResult = builder.create<UpdateHaloOp>(
         sourceShard.getLoc(),
         RankedTensorType::get(outShape, sourceShard.getType().getElementType()),
@@ -546,10 +559,13 @@ TypedValue<ShapedType> reshard(ImplicitLocOpBuilder &builder, MeshOp mesh,
                                MeshSharding targetSharding,
                                TypedValue<ShapedType> sourceUnshardedValue,
                                TypedValue<ShapedType> sourceShard) {
+  // If source and destination sharding are the same, no need to do anything.
   if (sourceSharding == targetSharding) {
     return sourceShard;
   }
 
+  // tries to handle the case where the resharding is needed because the halo
+  // sizes are different. Supports arbitrary mesh dimensionality.
   if (auto tryRes = tryUpdateHaloInResharding(
           builder, mesh, sourceSharding, targetSharding,
           sourceUnshardedValue.getType(), sourceShard)) {