[mpmd] Merge sideways in export pipeline.

Adds sideways merging of inferred fragments to the export pipeline to allow merging fragments on the same mesh that are separated by operations on other meshes. This ensures that inferred fragments created in the uniquify pass are merged correctly.

PiperOrigin-RevId: 899501008

Author: petebu

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

@@ -90,6 +90,12 @@ void addExportPipeline(OpPassManager& pm, const ExportOptions& options) {
   // fragments to compile) and may cause performance regressions. Thus, we merge
   // them with other fragments.
   pm.addNestedPass<FuncOp>(createMergeInferredFragmentsPass());
+  {
+    MergeInferredFragmentsPassOptions mergeInferredOptions;
+    mergeInferredOptions.mergeSideways = true;
+    pm.addNestedPass<FuncOp>(
+        createMergeInferredFragmentsPass(std::move(mergeInferredOptions)));
+  }
 
   // Mark each fragment with the inputs and outputs which are offloaded to host
   // memory.

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

@@ -1,4 +1,4 @@
-// RUN: mpmd_opt %s -mpmd-export-pipeline 2>&1 | FileCheck %s
+// RUN: mpmd_opt %s -mpmd-export-pipeline -split-input-file 2>&1 | FileCheck %s
 
 !mesh_1_tensor_4_8_f32 = !mpmd.mesh_tensor<"m1", tensor<4x8xf32>>
 
@@ -17,3 +17,36 @@ func.func @main(%arg0: !mesh_1_tensor_4_8_f32 {tf.aliasing_output = 0: i32}, %ar
   } : (!mesh_1_tensor_4_8_f32, !mesh_1_tensor_4_8_f32) -> (!mesh_1_tensor_4_8_f32)
   func.return %0 : !mesh_1_tensor_4_8_f32
 }
+
+// -----
+
+!mesh_1_tensor_4_8_f32 = !mpmd.mesh_tensor<"m1", tensor<4x8xf32>>
+!mesh_2_tensor_4_8_f32 = !mpmd.mesh_tensor<"m2", tensor<4x8xf32>>
+
+// This test verifies that an explicit fragment and an inferred fragment
+// (created by the UniquifyFunctionInputsOutputsPass for the duplicated return
+// of the transfer result) are merged sideways. Without sideways merge, the
+// transfer result would produce a separate inferred fragment call on m1.
+// The function-level returns remain unique SSA values (%[[RES]]#0, #1, #2),
+// preserving the invariant established by the uniquify pass, even though the
+// fragment body internally returns the same value in multiple positions.
+// CHECK-LABEL: func.func @test_sideways_merge
+func.func @test_sideways_merge(%arg0: !mesh_1_tensor_4_8_f32, %arg1: !mesh_2_tensor_4_8_f32)
+  -> (!mesh_1_tensor_4_8_f32, !mesh_1_tensor_4_8_f32, !mesh_1_tensor_4_8_f32) attributes {
+      "topology"=#mpmd.topology<
+      <"m1": <["x"=2]>>,
+      <"m2": <["x"=2]>>
+    >} {
+  // CHECK: %[[RES:.*]]:3 = mpmd.fragment_call<mesh="m1", origin=["f1"]> @[[CALLEE_M1:.*]]
+  // CHECK-NOT: mpmd.fragment_call<mesh="m1"
+  // CHECK: return %[[RES]]#0, %[[RES]]#1, %[[RES]]#2
+
+  %0 = mpmd.fragment<mesh="m1", origin=["f1"]> (%arg0) (%arg2: tensor<4x8xf32>) {
+    %4 = stablehlo.add %arg2, %arg2 : tensor<4x8xf32>
+    mpmd.return %4 : tensor<4x8xf32>
+  } : (!mesh_1_tensor_4_8_f32) -> !mesh_1_tensor_4_8_f32
+
+  %1 = mpmd.transfer %arg1 : (!mesh_2_tensor_4_8_f32) -> !mesh_1_tensor_4_8_f32
+
+  func.return %0, %1, %1 : !mesh_1_tensor_4_8_f32, !mesh_1_tensor_4_8_f32, !mesh_1_tensor_4_8_f32
+}