Push shardy inliner down past ShardingGroupImport pass for late inlining.

shardy/dialect/sdy/transforms/import/BUILD

@@ -58,6 +58,7 @@ cc_library(
         "//shardy/dialect/sdy/transforms/common:sharding_walker",
         "//shardy/dialect/sdy/transforms/propagation:op_sharding_rule_registry",
         "//shardy/dialect/sdy/transforms/propagation:sharding_projection",
+        "//shardy/dialect/sdy/transforms/propagation:utils",
         "//shardy/dialect/sdy/transforms/propagation/debugging:source_sharding",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:Analysis",

shardy/dialect/sdy/transforms/import/add_data_flow_edges.cc

@@ -21,6 +21,7 @@ limitations under the License.
 #include "shardy/dialect/sdy/ir/dialect.h"
 #include "shardy/dialect/sdy/ir/utils.h"
 #include "shardy/dialect/sdy/transforms/import/passes.h"  // IWYU pragma: keep
+#include "shardy/dialect/sdy/transforms/propagation/utils.h"
 
 namespace mlir {
 namespace sdy {
@@ -34,22 +35,6 @@ struct AddDataFlowEdgesPass
     : public impl::AddDataFlowEdgesPassBase<AddDataFlowEdgesPass> {
   using AddDataFlowEdgesPassBase::AddDataFlowEdgesPassBase;
 
-  void addDataFlowEdges(ValueRange edgeOwners, IRRewriter& rewriter) {
-    // We are iterating the owners in a reversed order because we set the
-    // insertion point after each value and we would like to keep the data flow
-    // edges for the arguments/results in the same order as they appear.
-    for (Value edgeOwner : llvm::reverse(edgeOwners)) {
-      rewriter.setInsertionPointAfterValue(edgeOwner);
-      if (!isStaticShapedType(edgeOwner.getType())) {
-        // Skip non-static-shaped tensors, e.g., tokens.
-        continue;
-      }
-      auto dataFlowEdge = DataFlowEdgeOp::create(
-          rewriter, edgeOwner.getLoc(), edgeOwner, getSharding(edgeOwner));
-      rewriter.replaceAllUsesExcept(edgeOwner, dataFlowEdge, dataFlowEdge);
-    }
-  }
-
   void runOnOperation() final {
     func::FuncOp funcOp = getOperation();
     IRRewriter rewriter(funcOp);

shardy/dialect/sdy/transforms/import/import_func_calls.cc

@@ -35,6 +35,7 @@ limitations under the License.
 #include "shardy/dialect/sdy/ir/dialect.h"
 #include "shardy/dialect/sdy/ir/utils.h"
 #include "shardy/dialect/sdy/transforms/import/passes.h"  // IWYU pragma: keep
+#include "shardy/dialect/sdy/transforms/propagation/utils.h"
 
 namespace mlir {
 namespace sdy {
@@ -117,6 +118,18 @@ struct ImportFuncCallsPass
     for (auto [calleeName, _] : calleeNameToMovedRegion) {
       symbolTable.erase(symbolTable.lookup(calleeName));
     }
+
+    // Required for cases that AddDataFlowEdges runs before this pass.
+    // TODO(enver): Drop after the late inlining drops ImportFuncCalls pass
+    // altogether as long as early inlining is before AddDataFlowEdges pass.
+    if (addDataFlowEdgesOnNamedComputations) {
+      moduleOp.walk([&](NamedComputationOp namedComputationOp) {
+        // Add the data flow edges for result owners and block argument owners.
+        addDataFlowEdges(namedComputationOp.getBlockArgumentEdgeOwners(),
+                         rewriter);
+        addDataFlowEdges(namedComputationOp.getOpResultEdgeOwners(), rewriter);
+      });
+    }
   }
 };
 

shardy/dialect/sdy/transforms/import/import_pipeline.cc

@@ -31,9 +31,12 @@ void addImportPipeline(OpPassManager& pm, int& dumpIndex,
   pm.addPass(createLiftInlinedMeshesPass());
   pm.addPass(createRemoveSizeOneAxesPass());
   pm.addPass(createPropagateShardingFromFuncToCallPass());
-  pm.addPass(createImportFuncCallsPass());
-  // Keep SymbolDCEPass after ImportFuncCallsPass.
-  pm.addPass(createSymbolDCEPass());
+  if (!options.enableLateInlining) {
+    pm.addPass(createImportFuncCallsPass(ImportFuncCallsPassOptions{
+        /*addDataFlowEdgesOnNamedComputations=*/false}));
+    // Keep SymbolDCEPass after ImportFuncCallsPass.
+    pm.addPass(createSymbolDCEPass());
+  }
   pm.addPass(createConstantOrScalarSplitterPass());
   pm.addPass(createSymbolDCEPass());
   pm.addPass(createManualAxesCleanupPass());
@@ -54,21 +57,35 @@ void addImportPipeline(OpPassManager& pm, int& dumpIndex,
   // constraints. This ensures we can detect sharding conflicts between group
   // members which have pre-propagation shardings due to sharding constraints.
   pm.addPass(createShardingGroupImportPass());
+  if (options.enableLateInlining) {
+    pm.addPass(createImportFuncCallsPass());
+    // Keep SymbolDCEPass after ImportFuncCallsPass.
+    pm.addPass(createSymbolDCEPass());
+  }
 }
 
 void addImportPipeline(OpPassManager& pm, const PropagationOptions& options) {
   int dumpIndex = 0;
   addImportPipeline(pm, dumpIndex, options);
 }
 
+struct ImportPipelineOptions
+    : public PassPipelineOptions<ImportPipelineOptions> {
+  Option<bool> enableLateInlining{*this, "enable-late-inlining",
+                                  llvm::cl::desc("Whether to late inline."),
+                                  llvm::cl::init(true)};
+};
+
 void registerImportPipeline() {
-  PassPipelineRegistration<>(
+  PassPipelineRegistration<ImportPipelineOptions>(
       "sdy-import-pipeline",
       "Run a sequence of import passes needed as a pre-processing step for "
       "Shardy propagation",
-      [](OpPassManager& pm) {
+      [](OpPassManager& pm, const ImportPipelineOptions& options) {
         int dumpIndex = 0;
-        addImportPipeline(pm, dumpIndex, PropagationOptions());
+        addImportPipeline(pm, dumpIndex,
+                          PropagationOptions{.enableLateInlining =
+                                                 options.enableLateInlining});
       });
 }
 

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

@@ -26,6 +26,11 @@ def ImportFuncCallsPass : Pass<"sdy-import-func-calls", "ModuleOp"> {
      function body for each call op and emit a warning.
   }];
   let dependentDialects = ["mlir::sdy::SdyDialect"];
+  let options = [
+    Option<"addDataFlowEdgesOnNamedComputations", "add-data-flow-edges-on-named-computations", "bool",
+           /*default=*/"true",
+           "Whether to add data flow edges on named computations.">
+  ];
 }
 
 def PropagateShardingFromFuncToCallPass : Pass<"sdy-propagate-sharding-from-func-to-call", "ModuleOp"> {