Keep reshards that has non-equivalent input and output meshes.

InsertExplicitReshards pass do not insert reshards with non-equivalent input and output meshes. Still it is possible for reshard to collectives pass to have them from user sharding constraints.

PiperOrigin-RevId: 813793497

Author: Google-ML-Automation

shardy/dialect/sdy/transforms/export/reshard_to_collectives.cc

@@ -1314,6 +1314,19 @@ class CollectiveInserter {
   AxisToDimAndIndex outAxisToDimAndIndex;
 };
 
+bool isEquivalentOnMesh(TensorShardingAttr inSharding,
+                        TensorShardingAttr outSharding, ReshardOp reshardOp) {
+  if (isFullyReplicated(inSharding) || isFullyReplicated(outSharding)) {
+    return true;
+  }
+  if (inSharding.getMeshName() == outSharding.getMeshName()) {
+    return true;
+  }
+  MeshAttr inMesh = inSharding.getMesh(reshardOp);
+  MeshAttr outMesh = outSharding.getMesh(reshardOp);
+  return inMesh.equals(outMesh, /*ignoreDeviceOrder=*/true);
+}
+
 class ReshardPattern : public OpConversionPattern<ReshardOp> {
  public:
   using OpConversionPattern::OpConversionPattern;
@@ -1334,20 +1347,19 @@ class ReshardPattern : public OpConversionPattern<ReshardOp> {
       rewriter.replaceOp(op, adaptor.getInput());
       return success();
     }
-    MeshAttr inMesh = inSharding.getMesh(op);
     if (inSharding.getMeshName() != outSharding.getMeshName()) {
-       MeshAttr outMesh = outSharding.getMesh(op);
-       // TODO(enver): Use MeshAttr::equals method instead.
-       if (outMesh.getAxes() != inMesh.getAxes() ||
-           inMesh.getDeviceIds() == outMesh.getDeviceIds() ||
-           (inSharding.isFullyReplicated() &&
-            outSharding.isFullyReplicated())) {
-         // We currently only support a reshard between different meshes if
-         // they have the same axes and different device ids, and at least one
-         // of the sharding isn't fully replicated.
-         return rewriter.notifyMatchFailure(
-             op, [](Diagnostic& diag) { diag << "Incompatible meshes"; });
-       }
+      if (outSharding.isFullyReplicated()) {
+        // TODO(enver): Hard fail if out sharding has unreduced axes.
+        outSharding = TensorShardingAttr::getFullyClosedLike(inSharding);
+      }
+    } else {
+      if (!isEquivalentOnMesh(inSharding, outSharding, op)) {
+        // We currently only support a reshard between different meshes if
+        // they have the same axes and different device ids, and at least one
+        // of the sharding isn't fully replicated.
+        return rewriter.notifyMatchFailure(
+            op, [](Diagnostic& diag) { diag << "Incompatible meshes"; });
+      }
     }
 
     // TODO(tomnatan): we should verify that the operand of ReshardOp has a
@@ -1370,13 +1382,12 @@ struct ReshardToCollectivesPass
     target = std::make_shared<ConversionTarget>(*context);
     target->addLegalOp<AllGatherOp, AllSliceOp, AllToAllOp,
                        CollectivePermuteOp>();
-    if (keepRedundantReshards) {
-      target->addDynamicallyLegalOp<ReshardOp>([](ReshardOp op) {
-        return isEquivalent(getSharding(op.getInput()), op.getSharding());
-      });
-    } else {
-      target->addIllegalOp<ReshardOp>();
-    }
+    target->addDynamicallyLegalOp<ReshardOp>([&](ReshardOp op) {
+      TensorShardingAttr inSharding = getSharding(op.getInput());
+      TensorShardingAttr outSharding = op.getSharding();
+      return (keepRedundantReshards && isEquivalent(inSharding, outSharding)) ||
+             !isEquivalentOnMesh(inSharding, outSharding, op);
+    });
 
     RewritePatternSet patternsInternal(context);
     patternsInternal.add<ReshardPattern>(context);

shardy/dialect/sdy/transforms/export/test/reshard_to_collectives.mlir

@@ -21,6 +21,7 @@ sdy.mesh @empty_mesh = <[]>
 sdy.mesh @empty_mesh_another = <[]>
 
 
+
 // CHECK-LABEL: func @redundant_reshard_fully_replicated
 func.func @redundant_reshard_fully_replicated(%arg0 : tensor<16x8xf32> {sdy.sharding=#sdy.sharding<@mesh2d, [{}, {}]>}) -> tensor<16x8xf32> {
   // CHECK-NEXT: return %arg0
@@ -70,6 +71,38 @@ func.func @redundant_reshard(%arg0 : tensor<16x8xf32> {sdy.sharding=#sdy.shardin
   return %0 : tensor<16x8xf32>
 }
 
+// CHECK-LABEL: func @reshard_from_sharded_to_fully_replicated_same_meshes
+func.func @reshard_from_sharded_to_fully_replicated_same_meshes(%arg0 : tensor<24x8xf32> {sdy.sharding=#sdy.sharding<@mesh1d_6, [{"x"}, {}]>}) -> tensor<24x8xf32> {
+  // CHECK-NEXT: %[[ALL_GATHER:.*]] = sdy.all_gather [{"x"}, {}] %arg0 out_sharding=<@mesh1d_6, [{}, {}]>
+  // CHECK-NEXT: return %[[ALL_GATHER]]
+  %0 = sdy.reshard %arg0 <@mesh1d_6, [{}, {}]> : tensor<24x8xf32>
+  return %0 : tensor<24x8xf32>
+}
+
+// CHECK-LABEL: func @reshard_from_sharded_to_fully_replicated_different_meshes
+func.func @reshard_from_sharded_to_fully_replicated_different_meshes(%arg0 : tensor<24x8xf32> {sdy.sharding=#sdy.sharding<@mesh1d_6, [{"x"}, {}]>}) -> tensor<24x8xf32> {
+  // CHECK-NEXT: %[[ALL_GATHER:.*]] = sdy.all_gather [{"x"}, {}] %arg0 out_sharding=<@mesh1d_6, [{}, {}]>
+  // CHECK-NEXT: return %[[ALL_GATHER]]
+  %0 = sdy.reshard %arg0 <@mesh2d_2x3, [{}, {}]> : tensor<24x8xf32>
+  return %0 : tensor<24x8xf32>
+}
+
+// CHECK-LABEL: func @reshard_from_sharded_to_fully_replicated_different_meshes_with_different_device_counts
+func.func @reshard_from_sharded_to_fully_replicated_different_meshes_with_different_device_counts(%arg0 : tensor<24x8xf32> {sdy.sharding=#sdy.sharding<@mesh1d_6, [{"x"}, {}]>}) -> tensor<24x8xf32> {
+  // CHECK-NEXT: %[[ALL_GATHER:.*]] = sdy.all_gather [{"x"}, {}] %arg0 out_sharding=<@mesh1d_6, [{}, {}]>
+  // CHECK-NEXT: return %[[ALL_GATHER]]
+  %0 = sdy.reshard %arg0 <@mesh2d, [{}, {}]> : tensor<24x8xf32>
+  return %0 : tensor<24x8xf32>
+}
+
+// CHECK-LABEL: func @reshard_from_sharded_to_sharded_different_meshes
+func.func @reshard_from_sharded_to_sharded_different_meshes(%arg0 : tensor<24x8xf32> {sdy.sharding=#sdy.sharding<@mesh1d_6, [{"x"}, {}]>}) -> (tensor<24x8xf32> {sdy.sharding=#sdy.sharding<@mesh2d_2x3, [{"x"}, {}]>}) {
+  // CHECK-NEXT: %[[RESHARD:.*]] = sdy.reshard %arg0 <@mesh2d_2x3, [{"x"}, {}]>
+  // CHECK-NEXT: return %[[RESHARD]]
+  %0 = sdy.reshard %arg0 <@mesh2d_2x3, [{"x"}, {}]> : tensor<24x8xf32>
+  return %0 : tensor<24x8xf32>
+}
+
 // CHECK-LABEL: func @all_gather_single_axis
 func.func @all_gather_single_axis(%arg0 : tensor<16x8xf32> {sdy.sharding=#sdy.sharding<@mesh2d, [{"y"}, {"x"}]>}) -> tensor<16x8xf32> {
   // CHECK-NEXT: sdy.all_gather [{}, {"x"}] %arg0 out_sharding=<@mesh2d, [{"y"}, {}]>