Push shardy unflatenner up past SinkDataFlowEdges pass.

In the process, split out SinkFuncDataFlowEdges from SinkDataFlowEdges. It is because func arguments needs to get their shardings from func data flow edges in order for the shardy unflattener to deduplicate on funcs based on their input/output shardings properly.

PiperOrigin-RevId: 908391477

Author: ekayaaslan

shardy/dialect/sdy/transforms/export/BUILD

@@ -48,6 +48,7 @@ cc_library(
         "resolve_permutation_factors.cc",
         "sharding_constraint_to_reshard.cc",
         "sink_data_flow_edges.cc",
+        "sink_func_data_flow_edges.cc",
         "unflatten_call_graph.cc",
         "update_non_divisible_input_output_shardings.cc",
     ],

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

@@ -40,26 +40,16 @@ void runShardyPartitioner(OpPassManager& pm, int& dumpIndex,
                           const ExportOptions& options) {
   InsertExplicitReshardsPassOptions passOptions;
   passOptions.enableFullVersion = options.enableInsertExplicitCollectives;
-  passOptions.avoidReshardsOnCalls = options.avoidReshardsOnCalls;
   pm.addNestedPass<func::FuncOp>(createInsertExplicitReshardsPass(passOptions));
   if (options.enableInsertExplicitCollectives) {
     pm.addPass(mlir::sdy::createSaveModuleOpPass(
         options.dumpDirectory, "after_explicit_reshards", dumpIndex++));
     addCanonicalizerPass(pm, kReshardLabel);
-    pm.addPass(createUnflattenCallGraphPass(
-        UnflattenCallGraphPassOptions{options.dedupFunctionsFully}));
-    // Keep a SymbolDCE after UnflattenCallGraph.
-    pm.addPass(createSymbolDCEPass());
     pm.addNestedPass<func::FuncOp>(createReshardToCollectivesPass());
     // NOTE: ReshardToCollectives pass above generates all-slice collectives,
     // which during the canonicalizer below may be converted to reduce scatters
     // by potentially fusing with preceeding all-reduces, which are inserted
     // during InsertExplicitReshards pass.
-  } else {
-    pm.addPass(createUnflattenCallGraphPass(
-        UnflattenCallGraphPassOptions{options.dedupFunctionsFully}));
-    // Keep a SymbolDCE after UnflattenCallGraph.
-    pm.addPass(createSymbolDCEPass());
   }
 
   addCanonicalizerPass(pm, kCollectiveLabel);
@@ -82,16 +72,22 @@ void runShardyPartitioner(OpPassManager& pm, int& dumpIndex,
 void addExportPipeline(OpPassManager& pm, int& dumpIndex,
                        const ExportOptions& options) {
   pm.addNestedPass<func::FuncOp>(createConstantOrScalarMergerPass());
-  if (!options.avoidExportForPartitioning) {
-    pm.addPass(createRemoveShardingGroupsPass());
-    pm.addNestedPass<func::FuncOp>(createShardingConstraintToReshardPass());
-  }
+  pm.addNestedPass<func::FuncOp>(createSinkFuncDataFlowEdgesPass());
+  pm.addPass(createUnflattenCallGraphPass(
+      UnflattenCallGraphPassOptions{options.dedupFunctionsFully}));
+  // Keep a SymbolDCE after UnflattenCallGraph.
+  pm.addPass(createSymbolDCEPass());
   pm.addNestedPass<func::FuncOp>(
       createSinkDataFlowEdgesPass(SinkDataFlowEdgesPassOptions{
           /*sinkDebugShardingOrigins=*/options.dumpShardingOrigins,
           /*sinkDebugPropagationEdgeSharding=*/options.dumpPropagationEdges}));
+  if (!options.avoidExportForPartitioning) {
+    pm.addPass(createRemoveShardingGroupsPass());
+    pm.addNestedPass<func::FuncOp>(createShardingConstraintToReshardPass());
+  }
   if (options.updateNonDivisibleInputOutputShardings) {
     pm.addPass(createUpdateNonDivisibleInputOutputShardingsPass());
+
     pm.addPass(createRemoveSubAxesInInputOutputShardingsPass());
   }
   pm.addPass(createCloseShardingsPass());
@@ -106,11 +102,6 @@ void addExportPipeline(OpPassManager& pm, int& dumpIndex,
   // reshards/collectives.
   if (!options.avoidExportForPartitioning) {
     runShardyPartitioner(pm, dumpIndex, options);
-  } else {
-    pm.addPass(createUnflattenCallGraphPass(
-        UnflattenCallGraphPassOptions{options.dedupFunctionsFully}));
-    // Keep a SymbolDCE after UnflattenCallGraph.
-    pm.addPass(createSymbolDCEPass());
   }
   if (options.dumpPropagationEdges || options.dumpShardingOrigins) {
     pm.addPass(createRemovePropagationDebugInfoPass());

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

@@ -495,10 +495,8 @@ struct InsertExplicitReshardsPass
       }
 
       if (CallOp callOp = dyn_cast<CallOp>(op)) {
-        if (!avoidReshardsOnCalls) {
           insertExplicitReshardsOnCallOp(callOp, rewriter, symbolTable,
                                          onFullVersion);
-        }
         return;
       }
 

shardy/dialect/sdy/transforms/export/passes.td

@@ -43,6 +43,15 @@ def SinkDataFlowEdgesPass : Pass<"sdy-sink-data-flow-edges", "func::FuncOp"> {
   ];
 }
 
+def SinkFuncDataFlowEdgesPass : Pass<"sdy-sink-func-data-flow-edges", "func::FuncOp"> {
+  let summary = "Sinks all `FuncDataFlowEdgeOp` into their input.";
+  let description = [{
+    Moves the sharding of each `FuncDataFlowEdgeOp` to its input and replaces
+    the op with its input.
+  }];
+  let dependentDialects = ["mlir::sdy::SdyDialect"];
+}
+
 def UpdateNonDivisibleInputOutputShardingsPass : Pass<"sdy-update-non-divisible-input-output-shardings", "ModuleOp"> {
   let summary = "Makes FuncOp inputs/outputs evenly sharded, removing any need for padding due to non-divisible shardings.";
   let description = [{
@@ -123,10 +132,6 @@ def InsertExplicitReshardsPass : Pass<"sdy-insert-explicit-reshards", "func::Fun
       Option<"enableFullVersion", "enable-full-version",
             "bool", /*default=*/"false",
             "Enable full version.">,
-      Option<"avoidReshardsOnCalls",
-            "avoid-reshards-on-calls",
-            "bool", /*default=*/"false",
-            "Avoid explicit reshards/collectives on calls.">
     ];
 }
 

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

@@ -67,20 +67,6 @@ struct SinkDataFlowEdgesPass
         rewriter.replaceOp(dataFlowEdgeOp, input);
         return WalkResult::skip();
       }
-      if (isa<FuncDataFlowEdgeOp>(op)) {
-        FuncDataFlowEdgeOp funcEdgeOp = cast<FuncDataFlowEdgeOp>(op);
-        Value operand = funcEdgeOp.getOperand();
-        Value result = funcEdgeOp.getResult();
-        TensorShardingAttr operandSharding = getSharding(operand);
-        if (TensorShardingAttr sharding = getSharding(result)) {
-          setSharding(operand, sharding);
-        } else if (operandSharding) {
-          setSharding(operand,
-                      TensorShardingAttr::getFullyOpenLike(operandSharding));
-        }
-        rewriter.replaceOp(funcEdgeOp, operand);
-        return WalkResult::skip();
-      }
       auto shardableDataFlowOp = dyn_cast<ShardableDataFlowOpInterface>(op);
       if (!shardableDataFlowOp) {
         return WalkResult::advance();

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

@@ -0,0 +1,86 @@
+/* Copyright 2024 The Shardy Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cassert>
+#include <memory>  // IWYU pragma: keep
+#include <tuple>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+#include "mlir/IR/ValueRange.h"
+#include "mlir/IR/Visitors.h"
+#include "mlir/Pass/Pass.h"  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"
+#include "shardy/dialect/sdy/ir/constants.h"
+#include "shardy/dialect/sdy/ir/dialect.h"
+#include "shardy/dialect/sdy/ir/utils.h"
+#include "shardy/dialect/sdy/transforms/export/passes.h"  // IWYU pragma: keep
+#include "shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.h"
+#include "shardy/dialect/sdy/transforms/propagation/utils.h"
+
+namespace mlir {
+namespace sdy {
+
+#define GEN_PASS_DEF_SINKFUNCDATAFLOWEDGESPASS
+#include "shardy/dialect/sdy/transforms/export/passes.h.inc"
+
+namespace {
+
+struct SinkFuncDataFlowEdgesPass
+    : public impl::SinkFuncDataFlowEdgesPassBase<SinkFuncDataFlowEdgesPass> {
+  using SinkFuncDataFlowEdgesPassBase::SinkFuncDataFlowEdgesPassBase;
+
+  void runOnOperation() final {
+    func::FuncOp funcOp = getOperation();
+    IRRewriter rewriter(funcOp);
+    // Copy the sharding from data flow edges to the data flow ops.
+    funcOp.walk<WalkOrder::PreOrder>([&](Operation* op) {
+      // Since we are doing the walk in preorder with a forward iterator, ops
+      // are walked before their users and regions. Since `DataFlowEdgeOp` can
+      // only appear inside the data flow op's region or as its user, we always
+      // encounter the data flow op before their data flow edges. This means it
+      // is safe to erase the `FuncDataFlowEdgeOp` at this point. We need the
+      // skip at the end because it's a condition to erase the op. See the
+      // documentation for `Operation::walk` for more details.
+      if (isa<FuncDataFlowEdgeOp>(op)) {
+        FuncDataFlowEdgeOp funcEdgeOp = cast<FuncDataFlowEdgeOp>(op);
+        Value operand = funcEdgeOp.getOperand();
+        Value result = funcEdgeOp.getResult();
+        TensorShardingAttr operandSharding = getSharding(operand);
+        if (TensorShardingAttr sharding = getSharding(result)) {
+          setSharding(operand, sharding);
+        } else if (operandSharding) {
+          setSharding(operand,
+                      TensorShardingAttr::getFullyOpenLike(operandSharding));
+        }
+        rewriter.replaceOp(funcEdgeOp, operand);
+        return WalkResult::skip();
+      }
+      return WalkResult::advance();
+    });
+  }
+};
+
+}  // namespace
+
+}  // namespace sdy
+}  // namespace mlir