Push shardy outliner up out of export pipeline.

PiperOrigin-RevId: 900660386

Author: ekayaaslan

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

@@ -70,7 +70,6 @@ void runShardyPartitioner(OpPassManager& pm, int& dumpIndex,
 
 void addExportPipeline(OpPassManager& pm, int& dumpIndex,
                        const ExportOptions& options) {
-  pm.addPass(createExportNamedComputationsPass());
   pm.addNestedPass<func::FuncOp>(createConstantOrScalarMergerPass());
   if (!options.avoidExportForPartitioning) {
     pm.addPass(createRemoveShardingGroupsPass());

shardy/dialect/sdy/transforms/export/test/export_pipeline.mlir

@@ -64,7 +64,7 @@ func.func @main(
 // -----
 sdy.mesh @mesh = <["a"=2, "b"=2, "c"=2]>
 
-// Test: named_computation_with_shardings
+// Test: call_with_shardings
 // CHECK-LABEL: func @main
 func.func @main(%arg0: tensor<8x2xi32>, %arg1: tensor<4x2xi32>) -> tensor<12x2xi32> {
   // CHECK-NEXT: %0 = sdy.reshard %arg0 <@mesh, [{"a"}, {}]> : tensor<8x2xi32>
@@ -73,15 +73,13 @@ func.func @main(%arg0: tensor<8x2xi32>, %arg1: tensor<4x2xi32>) -> tensor<12x2xi
   // CHECK-NEXT: %3 = sdy.reshard %1#1 <@mesh, [{}, {"a"}]> : tensor<4x2xi32>
   // CHECK-NEXT: %4 = stablehlo.concatenate %2, %3, dim = 0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"a"}]>]>}
   // CHECK-NEXT: return %4 : tensor<12x2xi32>
-  %0:2 = sdy.named_computation<"foo">(%arg0, %arg1) in_shardings=[<@mesh, [{"a"}, {}]>, <@mesh, [{?}, {}]>] out_shardings=[<@mesh, [{"a"}, {}]>, <@mesh, [{?}, {}]>] (%arg2: tensor<8x2xi32>, %arg3: tensor<4x2xi32>)  {
-    sdy.return %arg2, %arg3 : tensor<8x2xi32>, tensor<4x2xi32>
-  } : (tensor<8x2xi32>, tensor<4x2xi32>) -> (tensor<8x2xi32>, tensor<4x2xi32>)
+  %0:2 = call @foo(%arg0, %arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>, <@mesh, [{}, {}]>]>} : (tensor<8x2xi32>, tensor<4x2xi32>) -> (tensor<8x2xi32>, tensor<4x2xi32>)
   %1 = stablehlo.concatenate %0#0, %0#1, dim=0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"a"}]>]>} : (tensor<8x2xi32>, tensor<4x2xi32>) -> tensor<12x2xi32>
   return %1 : tensor<12x2xi32>
 }
 
-// CHECK-LABEL: func private @foo(%arg0: tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}, %arg1: tensor<4x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>})
-// CHECK-SAME:       -> (tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}, tensor<4x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>})
-// CHECK-SAME:       attributes {sdy.original_func_name = "foo"} {
-// CHECK-NEXT:    return %arg0, %arg1 : tensor<8x2xi32>, tensor<4x2xi32>
-// CHECK-NEXT:  }
+func.func private @foo(%arg0: tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}, %arg1: tensor<4x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>})
+    -> (tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}, tensor<4x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>})
+    attributes {sdy.original_func_name = "foo"} {
+  return %arg0, %arg1 : tensor<8x2xi32>, tensor<4x2xi32>
+}

shardy/dialect/sdy/transforms/propagation/propagation_pipeline.cc

@@ -74,6 +74,7 @@ void addPropagationPipeline(OpPassManager& pm, int& dumpIndex,
                                       dumpIndex++));
     AutoPartitionerRegistry::addPasses(pm);
   }
+  pm.addPass(createExportNamedComputationsPass());
   ExportOptions exportOptions;
   populateExportOptions(exportOptions, options);
   addExportPipeline(pm, dumpIndex, exportOptions);