removing ProcessMultiIndexOp canonicalizer, fixing MeshSharding.equalSplitAndPartialAxes

Author: fschlimb

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

@@ -132,7 +132,6 @@ def Mesh_ProcessMultiIndexOp : Mesh_Op<"process_multi_index", [
     OpBuilder<(ins "::mlir::mesh::MeshOp":$mesh)>,
     OpBuilder<(ins "StringRef":$mesh, "ArrayRef<MeshAxis>":$axes)>
   ];
-  let hasCanonicalizer = 1;
 }
 
 def Mesh_ProcessLinearIndexOp : Mesh_Op<"process_linear_index", [
@@ -1061,7 +1060,8 @@ def Mesh_UpdateHaloOp : Mesh_Op<"update_halo", [
   TypesMatchWith<
     "result has same type as destination",
     "result", "destination", "$_self">,
-  DeclareOpInterfaceMethods<SymbolUserOpInterface>
+  DeclareOpInterfaceMethods<SymbolUserOpInterface>,
+  AttrSizedOperandSegments
 ]> {
   let summary = "Update halo data.";
   let description = [{
@@ -1071,21 +1071,23 @@ def Mesh_UpdateHaloOp : Mesh_Op<"update_halo", [
     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 `source_halo_sizes/static_source_halo_sizes`.
+    by `source_halo_sizes/static_source_halo_sizes` and
+    `destination_halo_sizes/static_destination_halo_sizes`
 
     `split_axes` specifies for each tensor axis along which mesh axes its halo
     data is updated.
 
-    The destination halo sizes are allowed differ from the source sizes. The sizes
-    of the inner (local) shard is inferred from the destination shape and source sharding.
+    Source and destination might have different halo sizes.
   }];
   let arguments = (ins
     AnyTypeOf<[AnyNon0RankedMemRef, AnyNon0RankedTensor]>:$source,
     AnyTypeOf<[AnyNon0RankedMemRef, AnyNon0RankedTensor]>:$destination,
     FlatSymbolRefAttr:$mesh,
     Mesh_MeshAxesArrayAttr:$split_axes,
     Variadic<I64>:$source_halo_sizes,
-    DefaultValuedAttr<DenseI64ArrayAttr, "{}">:$static_source_halo_sizes
+    DefaultValuedAttr<DenseI64ArrayAttr, "{}">:$static_source_halo_sizes,
+    Variadic<I64>:$destination_halo_sizes,
+    DefaultValuedAttr<DenseI64ArrayAttr, "{}">:$static_destination_halo_sizes
   );
   let results = (outs
     AnyTypeOf<[AnyNon0RankedMemRef, AnyNon0RankedTensor]>:$result
@@ -1095,6 +1097,7 @@ def Mesh_UpdateHaloOp : Mesh_Op<"update_halo", [
     `on` $mesh
     `split_axes` `=` $split_axes
     (`source_halo_sizes` `=` custom<DynamicIndexList>($source_halo_sizes, $static_source_halo_sizes)^)?
+    (`destination_halo_sizes` `=` custom<DynamicIndexList>($destination_halo_sizes, $static_destination_halo_sizes)^)?
     attr-dict `:` type($source) `->` type($result)
   }];
   let extraClassDeclaration = [{

mlir/lib/Dialect/Mesh/IR/MeshOps.cpp

@@ -592,7 +592,12 @@ bool MeshSharding::equalSplitAndPartialAxes(const MeshSharding &rhs) const {
     return false;
   }
 
-  if (!getPartialAxes().empty() && getPartialType() != rhs.getPartialType()) {
+  if (getPartialAxes().size() != rhs.getPartialAxes().size() ||
+      (!getPartialAxes().empty() && getPartialType() != rhs.getPartialType()) ||
+      !llvm::equal(
+          llvm::make_range(getPartialAxes().begin(), getPartialAxes().end()),
+          llvm::make_range(rhs.getPartialAxes().begin(),
+                           rhs.getPartialAxes().end()))) {
     return false;
   }
 
@@ -776,62 +781,6 @@ void ProcessMultiIndexOp::getAsmResultNames(
   setNameFn(getResults()[0], "proc_linear_idx");
 }
 
-namespace {
-#ifndef NDEBUG
-static std::vector<int> convertStringToVector(const std::string &str) {
-  std::vector<int> result;
-  std::stringstream ss(str);
-  std::string item;
-  while (std::getline(ss, item, ',')) {
-    result.push_back(std::stoi(item));
-  }
-  return result;
-}
-#endif // NDEBUG
-
-std::optional<SmallVector<Value>> getMyMultiIndex(OpBuilder &b,
-                                                  ::mlir::mesh::MeshOp mesh) {
-#ifndef NDEBUG
-  if (auto envStr = getenv("DEBUG_MESH_INDEX")) {
-    auto myIdx = convertStringToVector(envStr);
-    if (myIdx.size() == mesh.getShape().size()) {
-      SmallVector<Value> idxs;
-      for (auto i : myIdx) {
-        idxs.push_back(b.create<::mlir::arith::ConstantOp>(mesh->getLoc(),
-                                                           b.getIndexAttr(i)));
-      }
-      return idxs;
-    } else {
-      mesh->emitError() << "DEBUG_MESH_INDEX has wrong size";
-    }
-  }
-#endif // NDEBUG
-  return std::nullopt;
-}
-
-class FoldStaticIndex final : public OpRewritePattern<ProcessMultiIndexOp> {
-public:
-  using OpRewritePattern<ProcessMultiIndexOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(ProcessMultiIndexOp op,
-                                PatternRewriter &b) const override {
-#ifndef NDEBUG
-    SymbolTableCollection tmp;
-    if (auto idxs = getMyMultiIndex(b, getMesh(op, tmp))) {
-      b.replaceOp(op, idxs.value());
-      return success();
-    }
-#endif // NDEBUG
-    return failure();
-  }
-};
-} // namespace
-
-void ProcessMultiIndexOp::getCanonicalizationPatterns(
-    mlir::RewritePatternSet &results, mlir::MLIRContext *context) {
-  results.add<FoldStaticIndex>(context);
-}
-
 //===----------------------------------------------------------------------===//
 // mesh.process_linear_index op
 //===----------------------------------------------------------------------===//

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

@@ -483,7 +483,9 @@ tryUpdateHaloInResharding(ImplicitLocOpBuilder &builder, MeshOp mesh,
         MeshAxesArrayAttr::get(builder.getContext(),
                                sourceSharding.getSplitAxes()),
         sourceSharding.getDynamicHaloSizes(),
-        sourceSharding.getStaticHaloSizes());
+        sourceSharding.getStaticHaloSizes(),
+        targetSharding.getDynamicHaloSizes(),
+        targetSharding.getStaticHaloSizes());
     return std::make_tuple(
         cast<TypedValue<ShapedType>>(targetShard.getResult()), targetSharding);
   }
@@ -568,10 +570,9 @@ TypedValue<ShapedType> reshard(OpBuilder &builder, MeshOp mesh, ShardOp source,
   auto sourceSharding = source.getSharding();
   auto targetSharding = target.getSharding();
   ImplicitLocOpBuilder implicitLocOpBuilder(target->getLoc(), builder);
-  auto shard =
-      reshard(implicitLocOpBuilder, mesh, sourceSharding, targetSharding,
-              cast<TypedValue<ShapedType>>(source.getSrc()), sourceShardValue);
-  return shard;
+  return reshard(implicitLocOpBuilder, mesh, sourceSharding, targetSharding,
+                 cast<TypedValue<ShapedType>>(source.getSrc()),
+                 sourceShardValue);
 }
 
 TypedValue<ShapedType> reshard(OpBuilder &builder, ShardOp source,