Confine logics to find reduction axes and to check for unreduced axes inside the method that inserts all reduces.

It is a refactoring.

PiperOrigin-RevId: 812766518

Author: Google-ML-Automation

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

@@ -804,14 +804,60 @@ SmallVector<int64_t> getTensorSizes(Operation* op) {
   return tensorSizes;
 }
 
-SmallVector<AxisRefAttr> getReductionAxes(const AxesPerFactor& axesPerFactor,
-                                          OpShardingRuleAttr shardingRule) {
-  SmallVector<AxisRefAttr> reductionAxes;
+namespace {
+
+// Returns reduction axes that are the union of all axes on reduction factors.
+// The result axes are not necessarilly canonicalized.
+//
+// Returns empty axes if not `onFullVersion` and op results do not have
+// unreduced axes.
+//
+// Assumes `commonAxesPerFactor` is non-empty if `onFullVersion` is true.
+//
+// Hard fails if some reduction factors do not have compatible shardings.
+SmallVector<AxisRefAttr> getReductionAxes(
+    Operation* op, const ShardingProjection& shardingProjection,
+    const AxesPerFactor& commonAxesPerFactor, OpShardingRuleAttr shardingRule,
+    const bool onFullVersion) {
+  if (onFullVersion) {
+    SmallVector<AxisRefAttr> reductionAxes;
+    for (int64_t reductionFactor : shardingRule.getReductionFactors()) {
+      reductionAxes.append(commonAxesPerFactor[reductionFactor]);
+    }
+    return reductionAxes;
+  }
+
+  if (getUnreducedAxes(op->getResult(0)).empty()) {
+    return {};
+  }
+
+  // TODO(enver): Repurpose getCompatibleFactorShardings to return compatible
+  // factors, and simplify the following logic.
+  SmallVector<AxisRefAttr> axesAlongAllReductionFactors;
   for (int64_t reductionFactor : shardingRule.getReductionFactors()) {
-    reductionAxes.append(axesPerFactor[reductionFactor]);
+    // We only iterate operands since reduction factors are not in results.
+    bool seen = false;
+    SmallVector<AxisRefAttr> axesAlongCurrentReductionFactor;
+    for (const TensorFactorShardings& tensorFactorSharding :
+         shardingProjection.getOperands()) {
+      if (std::optional<ArrayRef<AxisRefAttr>> factorSharding =
+              getFactorSharding(tensorFactorSharding, reductionFactor)) {
+        if (seen) {
+          SDY_CHECK(axesAlongCurrentReductionFactor == *factorSharding)
+              << "For the operation " << op
+              << ", the result has unreduced axes while the operand has "
+                 "incompatible sharding along reduction factors.";
+        } else {
+          axesAlongCurrentReductionFactor = llvm::to_vector(*factorSharding);
+          seen = true;
+        }
+      }
+    }
+    axesAlongAllReductionFactors.append(axesAlongCurrentReductionFactor);
   }
-  return reductionAxes;
+  return axesAlongAllReductionFactors;
 }
+}  // namespace
 
 TensorShardingAttr insertAllReduceIfUnreducedToReplicated(
     OpOperand& use, TensorShardingAttr sourceSharding,
@@ -869,11 +915,16 @@ ArrayRef<AxisRefAttr> getUnreducedAxes(Value value) {
   return getUnreducedAxes(getSharding(value));
 }
 
-void insertAllReducesForReductionFactors(Operation* op,
-                                         ArrayRef<AxisRefAttr> reductionAxes,
-                                         const Mesh& mesh,
-                                         IRRewriter& rewriter) {
-  if (reductionAxes.empty() || op->getResults().empty()) {
+void insertAllReducesForReductionFactors(
+    Operation* op, const ShardingProjection& shardingProjection,
+    AxesPerFactor& commonAxesPerFactor, OpShardingRuleAttr shardingRule,
+    const Mesh& mesh, IRRewriter& rewriter, const bool onFullVersion) {
+  if (op->getResults().empty()) {
+    return;
+  }
+  SmallVector<AxisRefAttr> reductionAxes = getReductionAxes(
+      op, shardingProjection, commonAxesPerFactor, shardingRule, onFullVersion);
+  if (reductionAxes.empty()) {
     return;
   }
 

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

@@ -81,11 +81,6 @@ ArrayRef<AxisRefAttr> getUnreducedAxes(Value value);
 // Returns a concatenated array of operand and result tensor sizes.
 SmallVector<int64_t> getTensorSizes(Operation* op);
 
-// Returns reduction axes that are the union of all axes on reduction factors.
-// The result axes are not necessarilly canonicalized.
-SmallVector<AxisRefAttr> getReductionAxes(const AxesPerFactor& axesPerFactor,
-                                          OpShardingRuleAttr shardingRule);
-
 // Returns true iff any tensor factor sharding has non-empty overflow axes.
 bool hasOverflowAxes(const ShardingProjection& shardingProjection);
 
@@ -147,14 +142,22 @@ void insertExplicitReshards(Operation* op,
                             OpShardingRuleAttr shardingRule,
                             const SymbolTable& symbolTable, const Mesh& mesh);
 
-// Inserts an `sdy.all-reduce` for each result of `op` if `reductionAxes`
-// is non-empty. Assume the followings:
+// Inserts an `sdy.all-reduce` for each result of `op`.
+//
+// Assumes the followings:
 // - All op results have the same unreduced axes.
 // - All op results have the same mesh as `mesh` ignoring device id orders.
-void insertAllReducesForReductionFactors(Operation* op,
-                                         ArrayRef<AxisRefAttr> reductionAxes,
-                                         const Mesh& mesh,
-                                         IRRewriter& rewriter);
+// - `commonAxesPerFactor` is non-empty if `onFullVersion` is true.
+//
+// In case `onFullVersion` is false, it inserts all reduces only if op results
+// have some unreduced axes.
+//
+// Hard fails if the reduction factors do not have compatible shardings, and op
+// results have unreduced axes.
+void insertAllReducesForReductionFactors(
+    Operation* op, const ShardingProjection& shardingProjection,
+    AxesPerFactor& commonAxesPerFactor, OpShardingRuleAttr shardingRule,
+    const Mesh& mesh, IRRewriter& rewriter, bool onFullVersion);
 
 // Finds common factor axes on the operands and results of `op` so that the
 // sharding of `op` is compatible with its sharding rule.

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

@@ -158,18 +158,14 @@ void insertExplicitReshardsOnDataFlowOp(ShardableDataFlowOpInterface& op,
 // return %reshard  : tensor<4x8xf32>
 // ```
 template <class OpTy>
-void processDot(OpTy op, ArrayRef<TensorShardingAttr> inShardings,
+void processDot(OpTy op, ShardingProjection& shardingProjection,
                 ArrayRef<TensorShardingAttr> outShardings, IRRewriter& rewriter,
                 const SymbolTable& symbolTable, OpShardingRuleAttr shardingRule,
                 const Mesh& mesh) {
   if (outShardings.empty()) {
     // Result doesn't have a sharding.
     return;
   }
-  ShardingProjection shardingProjection =
-      ShardingProjection::build(inShardings, outShardings, shardingRule,
-                                mesh.attr(), /*closedIfMissing=*/true);
-
   const TensorFactorShardings& lhsSharding = shardingProjection.getOperand(0);
   const TensorFactorShardings& rhsSharding = shardingProjection.getOperand(1);
   TensorFactorShardings& resultSharding =
@@ -400,31 +396,25 @@ bool isOnFullVersion(Operation* op, const bool enableFullVersion) {
 // - All op results have the same unreduced axes.
 // - 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.
 //
 // 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,
-                                   IRRewriter& rewriter,
-                                   const SymbolTable& symbolTable,
-                                   OpShardingRuleAttr shardingRule,
-                                   const Mesh& mesh, const bool onFullVersion) {
-  ShardingProjection shardingProjection = ShardingProjection::build(
-      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 {};
-  }
-
+//
+// Guarantees to return non-empty `AxesPerFactor` if `onFullVersion` is true.
+AxesPerFactor processOp(Operation* op, ShardingProjection& shardingProjection,
+                        ArrayRef<TensorShardingAttr> inShardings,
+                        ArrayRef<TensorShardingAttr> outShardings,
+                        IRRewriter& rewriter, const SymbolTable& symbolTable,
+                        OpShardingRuleAttr shardingRule, const Mesh& mesh,
+                        const bool onFullVersion) {
+  // Checks if factors are sharded the same way across operands and results.
+  AxesPerFactor commonAxesPerFactor =
+      getCompatibleFactorShardings(shardingProjection, shardingRule);
+
+  // TODO(b/446833985): Return common axes per factor also when the sharding
+  // projection have overflow axes.
   if (onFullVersion) {
-    // Checks if factors are sharded the same way across operands and results.
-    AxesPerFactor commonAxesPerFactor =
-        getCompatibleFactorShardings(shardingProjection, shardingRule);
     // Find compatible shardings if it is not already compatible.
     if (commonAxesPerFactor.empty()) {
       commonAxesPerFactor =
@@ -443,49 +433,19 @@ SmallVector<AxisRefAttr> processOp(Operation* op,
     insertExplicitReshards(op, inShardings, outShardings, shardingProjection,
                            updateTensorShardings, rewriter, shardingRule,
                            symbolTable, mesh);
-
-    return getReductionAxes(commonAxesPerFactor, shardingRule);
-  }
-
-  TypeSwitch<Operation*>(op)
-      .Case<stablehlo::DotOp>([&](stablehlo::DotOp dotOp) {
-        processDot(dotOp, inShardings, outShardings, rewriter, symbolTable,
-                   shardingRule, mesh);
-      })
-      .Case<stablehlo::DotGeneralOp>([&](stablehlo::DotGeneralOp dotGeneralOp) {
-        processDot(dotGeneralOp, inShardings, outShardings, rewriter,
-                   symbolTable, shardingRule, mesh);
-      });
-
-  if (outShardings.empty() || getUnreducedAxes(outShardings[0]).empty()) {
-    return {};
+  } else {
+    TypeSwitch<Operation*>(op)
+        .Case<stablehlo::DotOp>([&](stablehlo::DotOp dotOp) {
+          processDot(dotOp, shardingProjection, outShardings, rewriter,
+                     symbolTable, shardingRule, mesh);
+        })
+        .Case<stablehlo::DotGeneralOp>(
+            [&](stablehlo::DotGeneralOp dotGeneralOp) {
+              processDot(dotGeneralOp, shardingProjection, outShardings,
+                         rewriter, symbolTable, shardingRule, mesh);
+            });
   }
-
-  // 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.
-    bool seen = false;
-    SmallVector<AxisRefAttr> axesAlongCurrentReductionFactor;
-    for (const TensorFactorShardings& tensorFactorSharding :
-         shardingProjection.getOperands()) {
-      if (std::optional<ArrayRef<AxisRefAttr>> factorSharding =
-              getFactorSharding(tensorFactorSharding, reductionFactor)) {
-        if (seen) {
-          SDY_CHECK(axesAlongCurrentReductionFactor == *factorSharding)
-              << "For the operation " << op
-              << ", the result has unreduced axes while the operand has "
-                 "incompatible sharding along reduction factors.";
-        } else {
-          axesAlongCurrentReductionFactor = llvm::to_vector(*factorSharding);
-          seen = true;
-        }
-      }
-    }
-    axesAlongAllReductionFactors.append(axesAlongCurrentReductionFactor);
-  }
-  return axesAlongAllReductionFactors;
+  return commonAxesPerFactor;
 }
 
 struct InsertExplicitReshardsPass
@@ -544,11 +504,22 @@ struct InsertExplicitReshardsPass
         return;
       }
 
-      SmallVector<AxisRefAttr> reductionAxes =
-          processOp(op, inShardings, outShardings, rewriter, symbolTable,
-                    shardingRule, *mesh, onFullVersion);
+      ShardingProjection shardingProjection = ShardingProjection::build(
+          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;
+      }
+      AxesPerFactor commonAxesPerFactor =
+          processOp(op, shardingProjection, inShardings, outShardings, rewriter,
+                    symbolTable, shardingRule, *mesh, onFullVersion);
       // TODO(b/440055868): Insert a reshard from unreduced to replicated axes.
-      insertAllReducesForReductionFactors(op, reductionAxes, *mesh, rewriter);
+      insertAllReducesForReductionFactors(op, shardingProjection,
+                                          commonAxesPerFactor, shardingRule,
+                                          *mesh, rewriter, onFullVersion);
 
       // TODO(enver): Remove sharding rules from ops.
     });