Drop inlining on shardy.

It adds func data flow edges unconditionally. It propagates on FuncDataFlowEdges between its sources to users. If the operand of FuncDataFlowEdges is:
- BlockArgument: The operand is a func argument. The source is the operand of the caller of the func.
- OpResult: The opearnd is a call result. The source is the terminator of the called func.

PiperOrigin-RevId: 900761579

Author: ekayaaslan

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

@@ -38,9 +38,10 @@ using func::FuncOp;
 
 // Adds func input and output data flow edges. Adds func input data flow edges
 // only for non-main funcs.
-void addFuncDataFlowEdges(ModuleOp moduleOp, const SymbolTable& symbolTable,
-                          IRRewriter& rewriter) {
-  FuncOp mainFuncOp = getMainFuncOrDie(moduleOp, symbolTable);
+void addFuncDataFlowEdgeOps(ModuleOp moduleOp, const SymbolTable& symbolTable,
+                            IRRewriter& rewriter) {
+  FuncOp mainFuncOp =
+      getMainFuncOrDie(moduleOp, symbolTable, /*useSingleFunc=*/true);
   moduleOp.walk([&](FuncOp funcOp) {
     if (funcOp == mainFuncOp) {
       return;
@@ -84,8 +85,8 @@ struct AddDataFlowEdgesPass
       addDataFlowEdges(op.getBlockArgumentEdgeOwners(), rewriter);
       addDataFlowEdges(op.getOpResultEdgeOwners(), rewriter);
     });
-    if (enableNativeNonFlatSupport) {
-      addFuncDataFlowEdges(moduleOp, symbolTable, rewriter);
+    if (addFuncDataFlowEdges) {
+      addFuncDataFlowEdgeOps(moduleOp, symbolTable, rewriter);
     }
   }
 };

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

@@ -43,7 +43,7 @@ void addImportPipeline(OpPassManager& pm, int& dumpIndex,
       options.dumpDirectory, "before_propagation", dumpIndex++));
 
   pm.addPass(createAddDataFlowEdgesPass(
-      AddDataFlowEdgesPassOptions{options.enableNativeNonFlatSupport}));
+      AddDataFlowEdgesPassOptions{/*addFuncDataFlowEdges=*/true}));
   pm.addPass(
       createApplyShardingConstraintsPass(ApplyShardingConstraintsPassOptions{
           options.debugShardingOrigins, options.debugPropagationEdgeSharding}));

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

@@ -92,11 +92,9 @@ def AddDataFlowEdgesPass : Pass<"sdy-add-data-flow-edges", "ModuleOp"> {
   let dependentDialects = ["mlir::sdy::SdyDialect"];
 
   let options = [
-    Option<"enableNativeNonFlatSupport", "enable-native-non-flat-support", "bool",
+    Option<"addFuncDataFlowEdges", "add-func-data-flow-edges", "bool",
            /*default=*/"false",
-           "Whether to propagate shardings directly on a non-flat graph without "
-           "flattening it. The default is false, meaning it will flatten the "
-           "graph and then propagate.">
+           "Whether to add func data flow edges.">
   ];
 }
 

…dges_enable_native_non_flat_support.mlir …edges_add_func_data_flow_edges_true.mlirshardy/dialect/sdy/transforms/import/test/add_data_flow_edges_enable_native_non_flat_support.mlir renamed to shardy/dialect/sdy/transforms/import/test/add_data_flow_edges_add_func_data_flow_edges_true.mlir shardy/dialect/sdy/transforms/import/test/add_data_flow_edges_enable_native_non_flat_support.mlir renamed to shardy/dialect/sdy/transforms/import/test/add_data_flow_edges_add_func_data_flow_edges_true.mlir

@@ -1,4 +1,4 @@
-// RUN: sdy_opt %s -split-input-file -sdy-add-data-flow-edges='enable-native-non-flat-support=true'  | FileCheck %s
+// RUN: sdy_opt %s -split-input-file -sdy-add-data-flow-edges='add-func-data-flow-edges=true'  | FileCheck %s
 
 // CHECK-LABEL: @bar(%arg0: tensor<8xf32>)
 func.func @bar(%arg0: tensor<8xf32>) -> tensor<8xf32> {

shardy/dialect/sdy/transforms/import/test/import_pipeline.mlir

@@ -17,6 +17,7 @@ func.func @main(%arg0: tensor<32x96xf32>) -> tensor<32x96xf32> {
 
 // -----
 
+
 // CHECK-LABEL: func @main
 func.func @main(%arg0: tensor<8x8xf32>, %arg1: tensor<8x8xf32>) {
   // CHECK-DAG: sdy.sharding_group %arg0 group_id=0 : tensor<8x8xf32>
@@ -30,10 +31,11 @@ func.func @main(%arg0: tensor<8x8xf32>, %arg1: tensor<8x8xf32>) {
 
 // -----
 
+
 sdy.mesh @mesh = <["c"=2, "a"=2, "b"=2]>
 
-// CHECK-LABEL: @add_manual_axes_to_replicated
-func.func @add_manual_axes_to_replicated(%arg0: tensor<8xf32>) -> tensor<8xf32> {
+// CHECK-LABEL: func @main
+func.func @main(%arg0: tensor<8xf32>) -> tensor<8xf32> {
   // CHECK-NEXT: sdy.manual_computation(%arg0)
   // CHECK-SAME{LITERAL}: in_shardings=[<@mesh, [{"c", ?}], replicated={"a"}>]
   // CHECK-SAME{LITERAL}:  out_shardings=[<@mesh, [{"c", ?}], replicated={"a"}>]
@@ -46,14 +48,15 @@ func.func @add_manual_axes_to_replicated(%arg0: tensor<8xf32>) -> tensor<8xf32>
 
 // -----
 
+
 sdy.mesh @mesh = <["c"=2, "a"=2, "b"=2]>
 
 // Due to the in_sharding being fully closed, the in_sharding is added to the
 // func arg but with the manual axis added as replicated.
-// CHECK-LABEL: @add_manual_axes_to_replicated_applied_constraint
+// CHECK-LABEL: func @main
 // CHECK-SAME     %arg0: tensor<16x16xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"c"}, {}], replicated={"a"}>}
 // CHECK-SAME     -> tensor<16x16xf32> {
-func.func @add_manual_axes_to_replicated_applied_constraint(%arg0: tensor<8xf32>) -> tensor<8xf32> {
+func.func @main(%arg0: tensor<8xf32>) -> tensor<8xf32> {
   %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"c"}]>] out_shardings=[<@mesh, [{"c"}]>] manual_axes={"c", "a"} (%arg1: tensor<4xf32>) {
     %1 = stablehlo.add %arg1, %arg1 : tensor<4xf32>
     sdy.return %1 : tensor<4xf32>
@@ -63,27 +66,31 @@ func.func @add_manual_axes_to_replicated_applied_constraint(%arg0: tensor<8xf32>
 
 // -----
 
+
 sdy.mesh @mesh = <["a"=2]>
 
 // This test verifies that the manual axes are cleaned up before adding data
 // flow edges.
-func.func @manual_axes_cleanup_before_adding_data_flow_edges(%arg0: tensor<8xf32>) -> tensor<8xf32> {
+func.func @main(%arg0: tensor<8xf32>) -> tensor<8xf32> {
   %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}]>] out_shardings=[<@mesh, [{?}]>] manual_axes={"a"} (%arg1: tensor<8xf32>) {
     %1 = stablehlo.add %arg1, %arg1 : tensor<8xf32>
     sdy.return %1 : tensor<8xf32>
   } : (tensor<8xf32>) -> tensor<8xf32>
-  // CHECK: sdy.data_flow_edge %0 sharding=<@mesh, [{?}], replicated={"a"}> : tensor<8xf32>
+  // CHECK: sdy.data_flow_edge %0 sharding=<@mesh, [{"c"}], replicated={"a"}> : tensor<8xf32>
   return %0 : tensor<8xf32>
 }
 
 // -----
 
+
 sdy.mesh @mesh = <["a"=2]>
 
-// CHECK-LABEL: func @single_call
-func.func @single_call(%arg0: tensor<8xf32>) -> tensor<8xf32> {
-  // CHECK-NEXT: %0 = call @foo(%arg0) : (tensor<8xf32>) -> tensor<8xf32>
-  // CHECK-NEXT: return %0 : tensor<8xf32>
+// test: single_call
+// CHECK-LABEL: func @main
+func.func @main(%arg0: tensor<8xf32>) -> tensor<8xf32> {
+  // CHECK: %[[CALL:.*]] = call @foo(%arg0) : (tensor<8xf32>) -> tensor<8xf32>
+  // CHECK: %[[EDGE:.*]] = sdy.func_data_flow_edge %[[CALL]] : tensor<8xf32>
+  // CHECK: return %[[EDGE]] : tensor<8xf32>
   %0 = call @foo(%arg0) : (tensor<8xf32>) -> tensor<8xf32>
   return %0 : tensor<8xf32>
 }

shardy/dialect/sdy/transforms/propagation/BUILD

@@ -57,6 +57,7 @@ cc_library(
         ":sharding_projection",
         ":utils",
         "//shardy/common:file_utils",
+        "//shardy/common:logging",
         "//shardy/dialect/sdy/ir:dialect",
         "//shardy/dialect/sdy/transforms/common:op_properties",
         "//shardy/dialect/sdy/transforms/common:propagation_options",