Check for overflow axes regardless of being default/minimal or full version.

Also refactor to run compatibility check also on the default/minimal version. It prepares to unify/simplify the logic to find reduction axes.

PiperOrigin-RevId: 812733006

Author: Google-ML-Automation

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

@@ -71,15 +71,8 @@ bool hasShardedPermutationFactors(
                                !factorSharding.axisRefs.empty();
                       });
 }
+}  // namespace
 
-// Checks if factor sharding is compatible, that is, it satisfies:
-// 1. Factors are sharded the same way across operands and results.
-// 2. Factors that need replication are unsharded.
-//
-// Returns the common axes per factor if the factor sharding is compatible.
-// Otherwise, returns empty AxesPerFactor.
-//
-// Assumes factor shardings do not have overflow axes.
 // TODO(enver): Handle the case when some factor shardings have overflow axes.
 AxesPerFactor getCompatibleFactorShardings(
     const ShardingProjection& shardingProjection,
@@ -112,6 +105,8 @@ AxesPerFactor getCompatibleFactorShardings(
   return commonAxesPerFactor;
 }
 
+namespace {
+
 void insertExplicitReshardsOnOperand(
     Operation* op, const int64_t operandIndex,
     const ShardingProjection& shardingProjection,
@@ -758,22 +753,12 @@ void distributeAxisRefsToBatchingFactors(
 }
 }  // namespace
 
-// Assumes there are no overflow axes.
-//
-// Guarantees to return a non-empty AxesPerFactor.
 AxesPerFactor findCommonAxes(ArrayRef<TensorShardingAttr> inShardings,
                              ArrayRef<TensorShardingAttr> outShardings,
                              const ShardingProjection& shardingProjection,
                              OpShardingRuleAttr shardingRule,
                              ArrayRef<int64_t> tensorSizes,
                              const SymbolTable& symbolTable, const Mesh& mesh) {
-  // Checks if factors are sharded the same way across operands and results.
-  if (AxesPerFactor commonAxesPerFactor =
-          getCompatibleFactorShardings(shardingProjection, shardingRule);
-      !commonAxesPerFactor.empty()) {
-    return commonAxesPerFactor;
-  }
-
   // Handle the special case of unary operations without factors that need
   // replication. Reshard only one of the tensors.
   if (shardingRule.getNonScalarTensorIndices().size() == 2 &&

shardy/dialect/sdy/transforms/export/explicit_reshards_util.h

@@ -89,6 +89,18 @@ SmallVector<AxisRefAttr> getReductionAxes(const AxesPerFactor& axesPerFactor,
 // Returns true iff any tensor factor sharding has non-empty overflow axes.
 bool hasOverflowAxes(const ShardingProjection& shardingProjection);
 
+// Checks if factor sharding is compatible, that is, it satisfies:
+// 1. Factors are sharded the same way across operands and results.
+// 2. Factors that need replication are unsharded.
+//
+// Returns the common axes per factor if the factor sharding is compatible.
+// Otherwise, returns empty AxesPerFactor.
+//
+// Assumes factor shardings do not have overflow axes.
+AxesPerFactor getCompatibleFactorShardings(
+    const ShardingProjection& shardingProjection,
+    OpShardingRuleAttr shardingRule);
+
 // Insert explicit reshards for operands and results that change by
 // the given `shardingProjection` for a given `op`. The reshards are inserted
 // only to make the given operation compatible.
@@ -154,6 +166,7 @@ void insertAllReducesForReductionFactors(Operation* op,
 // - If the op has no results, none of the operands has unreduced axes.
 // - Operand and result meshes are the same ignoring device id order.
 // - There are no overflow axes.
+// - Some shardings are not compatible.
 //
 // Guarantees to return a non-empty AxesPerFactor.
 AxesPerFactor findCommonAxes(ArrayRef<TensorShardingAttr> inShardings,

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

@@ -401,7 +401,8 @@ bool isOnFullVersion(Operation* op, const bool enableFullVersion) {
 // - If the op has no results, none of the operands has unreduced axes.
 // - Operand and result meshes are the same ignoring device id order.
 //
-// Returns the union of common reduction axes which may not be canonicalized.
+// Returns the union of axes along all the reduction factors which may not be
+// canonicalized.
 SmallVector<AxisRefAttr> processOp(Operation* op,
                                    ArrayRef<TensorShardingAttr> inShardings,
                                    ArrayRef<TensorShardingAttr> outShardings,
@@ -413,17 +414,24 @@ SmallVector<AxisRefAttr> processOp(Operation* op,
       inShardings, outShardings, shardingRule, mesh.attr(),
       /*closedIfMissing=*/true);
 
+  // Return without inserting reshards if any factor sharding has overflow
+  // axes. This case is not handled yet.
+  // TODO(enver): Handle the case when factor shardings have overflow axes.
+  if (hasOverflowAxes(shardingProjection)) {
+    return {};
+  }
+
   if (onFullVersion) {
-    // Return without inserting reshards if any factor sharding has overflow
-    // axes. This case is not handled yet.
-    // TODO(b/446833985): Handle the case when factor shardings have overflow
-    // axes.
-    if (hasOverflowAxes(shardingProjection)) {
-      return {};
-    }
+    // Checks if factors are sharded the same way across operands and results.
     AxesPerFactor commonAxesPerFactor =
-        findCommonAxes(inShardings, outShardings, shardingProjection,
-                       shardingRule, getTensorSizes(op), symbolTable, mesh);
+        getCompatibleFactorShardings(shardingProjection, shardingRule);
+    // Find compatible shardings if it is not already compatible.
+    if (commonAxesPerFactor.empty()) {
+      commonAxesPerFactor =
+          findCommonAxes(inShardings, outShardings, shardingProjection,
+                         shardingRule, getTensorSizes(op), symbolTable, mesh);
+    }
+
     UpdateTensorShardings updateTensorShardings(shardingRule.getNumOperands(),
                                                 shardingRule.getNumResults());
     for (const auto& [index, axes] : llvm::enumerate(commonAxesPerFactor)) {
@@ -453,8 +461,8 @@ SmallVector<AxisRefAttr> processOp(Operation* op,
     return {};
   }
 
-  // TODO(enver): Factor out finding common axes per factor. Share logic with
-  // getCompatibleFactorShardings.
+  // TODO(enver): Repurpose getCompatibleFactorShardings to return compatible
+  // factors, and simplify the following logic.
   SmallVector<AxisRefAttr> axesAlongAllReductionFactors;
   for (int64_t reductionFactor : shardingRule.getReductionFactors()) {
     // We only iterate operands since reduction factors are not in results.