Copy UniquifyFunctionInputsOutputsPass to UniquifyAndMergeReturnsPass

This is a preparatory step to modify the new pass to also merge returns.
The original pass and tests are copied and renamed to allow coexistence.

PiperOrigin-RevId: 908184894

Author: petebu

shardy/dialect/mpmd/ir/ops.td

@@ -257,7 +257,8 @@ def FragmentCallOp : Mpmd_Op<"fragment_call",
     Variadic<Mpmd_MeshTensorType>:$tensors,
     TypedArrayAttrBase<Mpmd_UserOrigin, "array of origin infos">:$origin,
     StrAttr:$mesh_name,
-    FlatSymbolRefAttr:$callee);
+    FlatSymbolRefAttr:$callee,
+    OptionalAttr<ArrayAttr>:$inferred_by);
   let results = (outs Variadic<Mpmd_MeshTensorType>);
 
   let extraClassDeclaration = [{

shardy/dialect/mpmd/ir/utils.cc

@@ -379,7 +379,7 @@ SmallVector<MpmdDataflowEdge> GetMpmdDataflowEdges(FuncOp func_op) {
 
 FragmentOp WrapOpWithFragment(
     Operation* op, StringRef mesh_name, RewriterBase& rewriter,
-    std::function<bool(OpOperand&)> should_replace_use) {
+    StringRef inferred_by, std::function<bool(OpOperand&)> should_replace_use) {
   // We set the insertion point right before `op` so assigns of operands will be
   // in the right place regardless of previous insertion point.
   rewriter.setInsertionPoint(op);
@@ -437,6 +437,10 @@ FragmentOp WrapOpWithFragment(
         return block_builder.clone(*op, mapping)->getResults();
       });
 
+  fragment_op->setAttr(
+      kInferredByAttr,
+      rewriter.getArrayAttr({rewriter.getStringAttr(inferred_by)}));
+
   // Unassign all fragment results and replace all uses of `op` with the
   // corresponding unassign op for which `should_replace_use` returns true.
   for (auto [original_result, fragment_result] :

shardy/dialect/mpmd/ir/utils.h

@@ -40,18 +40,16 @@ limitations under the License.
 
 namespace mlir::mpmd {
 
-// Globsl sdy mesh name.
-constexpr StringRef kGlobalMeshName = "mesh";
-
+// Global sdy mesh name.
+inline constexpr StringRef kGlobalMeshName = "mesh";
 // The function attribute that holds the SPMD mesh.
-constexpr StringRef kMeshShapeAttr = "mesh_shape";
+inline constexpr StringRef kMeshShapeAttr = "mesh_shape";
 // The function attribute that holds the MPMD topology.
-constexpr StringRef kTopologyAttr = "topology";
-
+inline constexpr StringRef kTopologyAttr = "topology";
 
 // The suffix of the mesh name for a CPU mesh.
 // LINT.IfChange
-constexpr StringRef kCpuMeshSuffix = "/cpu";
+inline constexpr StringRef kCpuMeshSuffix = "/cpu";
 // LINT.ThenChange(
 //   https://github.com/openxla/shardy/blob/main/shardy/integrations/python/jax/mpmd/types.py
 // )
@@ -76,6 +74,9 @@ inline constexpr StringRef kRematAttributeName = "remat";
 
 inline constexpr StringRef kJaxResultInfoAttr = "jax.result_info";
 
+// The attribute that holds the list of pass names that inferred a fragment.
+inline constexpr StringRef kInferredByAttr = "mpmd.inferred_by";
+
 template <typename... Args>
 std::string StrCat(Args&&... args) {
   std::string result;
@@ -261,6 +262,7 @@ SmallVector<MpmdDataflowEdge> GetMpmdDataflowEdges(func::FuncOp func_op);
 // `should_replace_use` returns true.
 FragmentOp WrapOpWithFragment(
     Operation* op, StringRef mesh_name, RewriterBase& rewriter,
+    StringRef inferred_by,
     std::function<bool(OpOperand&)> should_replace_use = [](OpOperand&) {
       return true;
     });

shardy/dialect/mpmd/transforms/common/BUILD

@@ -42,6 +42,7 @@ cc_library(
         "rule_based_merge.cc",
         "scheduler_preprocess.cc",
         "split_bwd_fragments.cc",
+        "uniquify_and_merge_returns.cc",
         "uniquify_function_inputs_outputs.cc",
         "unroll_for_loops.cc",
     ],

shardy/dialect/mpmd/transforms/common/merge_fragments.cc

@@ -24,6 +24,7 @@ limitations under the License.
 #include <vector>
 
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FormatVariadic.h"
@@ -182,6 +183,31 @@ std::optional<int> MergeCallCounters(FragmentOp producer_op,
   return std::nullopt;
 }
 
+std::optional<ArrayAttr> MergeInferredByAttributes(FragmentOp producer_op,
+                                                   FragmentOp consumer_op) {
+  ArrayAttr producer_inferred_by =
+      producer_op->getAttrOfType<ArrayAttr>(kInferredByAttr);
+  ArrayAttr consumer_inferred_by =
+      consumer_op->getAttrOfType<ArrayAttr>(kInferredByAttr);
+
+  if (!producer_inferred_by && !consumer_inferred_by) {
+    return std::nullopt;
+  }
+
+  llvm::SetVector<Attribute> combined_inferred_by;
+  if (producer_inferred_by) {
+    combined_inferred_by.insert(producer_inferred_by.begin(),
+                                producer_inferred_by.end());
+  }
+  if (consumer_inferred_by) {
+    combined_inferred_by.insert(consumer_inferred_by.begin(),
+                                consumer_inferred_by.end());
+  }
+
+  IRRewriter rewriter(producer_op.getContext());
+  return rewriter.getArrayAttr(combined_inferred_by.takeVector());
+}
+
 // Returns a list of attributes that must be preserved in the merged fragment.
 // Note: origins are preserved by default and require no extra work.
 SmallVector<std::pair<StringRef, Attribute>> MergedAttributes(
@@ -194,6 +220,12 @@ SmallVector<std::pair<StringRef, Attribute>> MergedAttributes(
     attributes.emplace_back(kCallCounterAttrName,
                             rewriter.getUI32IntegerAttr(*merged_call_count));
   }
+
+  if (std::optional<ArrayAttr> merged_inferred_by =
+          MergeInferredByAttributes(producer_op, consumer_op)) {
+    attributes.emplace_back(kInferredByAttr, *merged_inferred_by);
+  }
+
   return attributes;
 }
 

shardy/dialect/mpmd/transforms/common/passes.td

@@ -461,6 +461,36 @@ def UniquifyFunctionInputsOutputsPass :
   let dependentDialects = ["mlir::mpmd::MpmdDialect"];
 }
 
+def UniquifyAndMergeReturnsPass :
+    PassBase<"mpmd-uniquify-and-merge-returns", "DistributedFunctionPass"> {
+  let summary = "Uniquifies any value returned multiple times or any block "
+                "argument directly returned by the function.";
+  let description = [{
+    If a function returns the same value multiple times, creates multiple
+    versions for that value, by creating a fragment assigned to that value's
+    mesh which returns the value multiple times. After this pass, each return
+    operand is unique. This is important to ensure that the respective results
+    are allocated in different buffers, as in the following `jax.jit` example:
+
+    ```python
+    def f(x):
+      y = x + x
+      return y, y
+
+    z1, z2 = f(5)
+    z1 += 1
+    print(z1) ~~> 6
+    print(z2) ~~> 5
+    ```
+
+    Similarly, if a function returns a block argument, this pass creates an
+    identity fragment for that block argument, guaranteeing that values are
+    passed by value to the function, not by reference.
+  }];
+
+  let dependentDialects = ["mlir::mpmd::MpmdDialect"];
+}
+
 def SchedulingUnitVerifierPass :
     PassBase<"mpmd-scheduling-units-verifier", "DistributedFunctionPass"> {
   let summary = "Verifies if the program contains the required scheduling units.";

shardy/dialect/mpmd/transforms/common/test/uniquify_and_merge_returns.mlir

@@ -0,0 +1,142 @@
+// RUN: mpmd_opt %s -mpmd-uniquify-and-merge-returns -split-input-file 2>&1 | FileCheck %s
+
+!mesh_1_tensor = !mpmd.mesh_tensor<"m1", tensor<4xf32>>
+!mesh_2_tensor = !mpmd.mesh_tensor<"m2", tensor<4xf32>>
+
+// CHECK-LABEL: func @no_work_needed
+func.func @no_work_needed(%arg0: !mesh_1_tensor, %arg1: !mesh_2_tensor) -> (!mesh_1_tensor, !mesh_2_tensor) attributes {
+  "topology"=#mpmd.topology<<"m1": <["x"=2]>>, <"m2": <["x"=2]>>>
+} {
+  // CHECK-NEXT: %[[F1:.*]] = mpmd.fragment<mesh="m1", origin=["f1"]>
+  // CHECK:      %[[F2:.*]] = mpmd.fragment<mesh="m2", origin=["f2"]>
+  // CHECK:      return %[[F1]], %[[F2]]
+  %0 = mpmd.fragment<mesh="m1", origin=["f1"]> (%arg0) (%arg2: tensor<4xf32>) {
+    %1 = stablehlo.add %arg2, %arg2 : tensor<4xf32>
+    mpmd.return %1 : tensor<4xf32>
+  } : (!mesh_1_tensor) -> !mesh_1_tensor
+  %1 = mpmd.fragment<mesh="m2", origin=["f2"]> (%arg1) (%arg2: tensor<4xf32>) {
+    %1 = stablehlo.add %arg2, %arg2 : tensor<4xf32>
+    mpmd.return %1 : tensor<4xf32>
+  } : (!mesh_2_tensor) -> !mesh_2_tensor
+  return %0, %1 : !mesh_1_tensor, !mesh_2_tensor
+}
+
+
+// CHECK-LABEL: func @single_mesh_one_return_operand
+func.func @single_mesh_one_return_operand(%arg0: !mesh_1_tensor) -> (!mesh_1_tensor, !mesh_1_tensor, !mesh_1_tensor) attributes {
+  "topology"=#mpmd.topology<<"m1": <["x"=2]>>>
+} {
+  // CHECK-NEXT: %[[F1:.*]]:3 = mpmd.fragment<mesh="m1", origin=["f1"]>
+  // CHECK:      %[[F2:.*]] = mpmd.fragment<mesh="m1", origin=["f2"]> (%[[F1]]#0)
+  // CHECK:      return %[[F2]], %[[F1]]#1, %[[F1]]#2
+  %0 = mpmd.fragment<mesh="m1", origin=["f1"]> (%arg0) (%arg1: tensor<4xf32>) {
+    %1 = stablehlo.add %arg1, %arg1 : tensor<4xf32>
+    mpmd.return %1 : tensor<4xf32>
+  } : (!mesh_1_tensor) -> !mesh_1_tensor
+  %1 = mpmd.fragment<mesh="m1", origin=["f2"]> (%0) (%arg1: tensor<4xf32>) {
+    mpmd.return %arg1 : tensor<4xf32>
+  } : (!mesh_1_tensor) -> !mesh_1_tensor
+  return %1, %0, %0 : !mesh_1_tensor, !mesh_1_tensor, !mesh_1_tensor
+}
+
+// CHECK-LABEL: func @needs_fragment_for_m1_with_many_values
+func.func @needs_fragment_for_m1_with_many_values(%arg0: !mesh_1_tensor, %arg1: !mesh_2_tensor
+) -> (!mesh_2_tensor, !mesh_1_tensor, !mesh_1_tensor, !mesh_1_tensor, !mesh_1_tensor, !mesh_1_tensor) attributes {
+  "topology"=#mpmd.topology<<"m1": <["x"=2]>>, <"m2": <["x"=2]>>>
+} {
+  // CHECK-NEXT: %[[F1:.*]] = mpmd.fragment<mesh="m1", origin=["f1"]>
+  // CHECK:      %[[F2:.*]] = mpmd.fragment<mesh="m2", origin=["f2"]>
+  // CHECK:      %[[F3:.*]]:5 = mpmd.fragment<mesh="m1", origin=["f3"]> (%[[F1]], %arg0)
+  // CHECK:      return %[[F2]], %[[F3]]#0, %[[F3]]#2, %[[F3]]#1, %[[F3]]#3, %[[F3]]#4
+  %0 = mpmd.fragment<mesh="m1", origin=["f1"]> (%arg0) (%arg2: tensor<4xf32>) {
+    mpmd.return %arg2 : tensor<4xf32>
+  } : (!mesh_1_tensor) -> !mesh_1_tensor
+  %1 = mpmd.fragment<mesh="m2", origin=["f2"]> (%arg1) (%arg2: tensor<4xf32>) {
+    mpmd.return %arg2 : tensor<4xf32>
+  } : (!mesh_2_tensor) -> !mesh_2_tensor
+  %2 = mpmd.fragment<mesh="m1", origin=["f3"]> (%arg0) (%arg2: tensor<4xf32>) {
+    mpmd.return %arg2 : tensor<4xf32>
+  } : (!mesh_1_tensor) -> !mesh_1_tensor
+  return %1, %0, %2, %0, %2, %2 : !mesh_2_tensor, !mesh_1_tensor, !mesh_1_tensor, !mesh_1_tensor, !mesh_1_tensor, !mesh_1_tensor
+}
+
+// CHECK-LABEL: func @needs_fragment_for_m1_and_m2
+func.func @needs_fragment_for_m1_and_m2(%arg0: !mesh_1_tensor, %arg1: !mesh_2_tensor
+) -> (!mesh_1_tensor, !mesh_2_tensor, !mesh_2_tensor, !mesh_1_tensor, !mesh_1_tensor, !mesh_1_tensor) attributes {
+  "topology"=#mpmd.topology<<"m1": <["x"=2]>>, <"m2": <["x"=2]>>>
+} {
+  // CHECK: %[[F1:.*]] = mpmd.fragment<mesh="m1", origin=["f1"]>
+  // CHECK: %[[F2:.*]]:2 = mpmd.fragment<mesh="m2", origin=["f2"]>
+  // CHECK: %[[F3:.*]]:4 = mpmd.fragment<mesh="m1", origin=["f3"]> (%[[F1]], %arg0)
+  // CHECK: return %[[F3]]#0, %[[F2]]#0, %[[F2]]#1, %[[F3]]#2, %[[F3]]#1, %[[F3]]#3
+  %0 = mpmd.fragment<mesh="m1", origin=["f1"]> (%arg0) (%arg2: tensor<4xf32>) {
+    mpmd.return %arg2 : tensor<4xf32>
+  } : (!mesh_1_tensor) -> !mesh_1_tensor
+  %1 = mpmd.fragment<mesh="m2", origin=["f2"]> (%arg1) (%arg2: tensor<4xf32>) {
+    mpmd.return %arg2 : tensor<4xf32>
+  } : (!mesh_2_tensor) -> !mesh_2_tensor
+  %2 = mpmd.fragment<mesh="m1", origin=["f3"]> (%arg0) (%arg2: tensor<4xf32>) {
+    mpmd.return %arg2 : tensor<4xf32>
+  } : (!mesh_1_tensor) -> !mesh_1_tensor
+  return %0, %1, %1, %2, %0, %2 : !mesh_1_tensor, !mesh_2_tensor, !mesh_2_tensor, !mesh_1_tensor, !mesh_1_tensor, !mesh_1_tensor
+}
+
+// -----
+
+!dist_mesh_tensor = !mpmd.mesh_tensor<"m1", tensor<4xf32>, sharding=<@mesh, [{"x"}]>>
+
+module {
+
+// CHECK-LABEL: func @single_mesh_one_return_operand
+func.func @single_mesh_one_return_operand_with_global_view(%arg0: !dist_mesh_tensor) -> (!dist_mesh_tensor, !dist_mesh_tensor, !dist_mesh_tensor) attributes {
+  "topology"=#mpmd.topology<<"m1": <["x"=2]>>>
+} {
+  // CHECK-NEXT: %[[F1:.*]]:3 = mpmd.fragment<mesh="m1", origin=["f1"]>
+  // CHECK:      %[[F2:.*]] = mpmd.fragment<mesh="m1", origin=["f2"]> (%[[F1]]#0)
+  // CHECK:      return %[[F2]], %[[F1]]#1, %[[F1]]#2
+  %0 = mpmd.fragment<mesh="m1", origin=["f1"]> (%arg0) (%arg1: tensor<4xf32>) {
+    %1 = stablehlo.add %arg1, %arg1 : tensor<4xf32>
+    mpmd.return %1 : tensor<4xf32>
+  } : (!dist_mesh_tensor) -> !dist_mesh_tensor
+  %1 = mpmd.fragment<mesh="m1", origin=["f2"]> (%0) (%arg1: tensor<4xf32>) {
+    mpmd.return %arg1 : tensor<4xf32>
+  } : (!dist_mesh_tensor) -> !dist_mesh_tensor
+  return %1, %0, %0 : !dist_mesh_tensor, !dist_mesh_tensor, !dist_mesh_tensor
+}
+}
+
+// -----
+
+!mesh_tensor = !mpmd.mesh_tensor<"m", tensor<4xui32>, sharding=<@mesh, [{"x"}]>>
+
+// CHECK-LABEL: func @f
+func.func @f(%arg0: !mesh_tensor) -> (!mesh_tensor, !mesh_tensor, !mesh_tensor)
+  attributes {"topology"=#mpmd.topology<<"m": <["x"=2]>>>}
+{
+  // CHECK-NEXT: %[[F1:.*]] = mpmd.fragment<mesh="m", origin=["f"]> (%arg0) (%arg1: tensor<4xui32>) {
+  // CHECK-NEXT:   return %arg1
+  // CHECK-NEXT: }
+  // CHECK-NEXT: %[[F2:.*]]:2 = mpmd.fragment<mesh="m", origin=[]> (%arg0) {mpmd.inferred_by = ["uniquify"]} (%arg1: tensor<4xui32>) {
+  // CHECK-NEXT:   return %arg1, %arg1
+  // CHECK-NEXT: }
+  // CHECK-NEXT: return %[[F2]]#0, %[[F1]], %[[F2]]#1
+  %0 = mpmd.fragment<mesh="m", origin=["f"]>(%arg0) (%arg1: tensor<4xui32>) {
+    mpmd.return %arg1 : tensor<4xui32>
+  } : (!mesh_tensor) -> !mesh_tensor
+  func.return %arg0, %0, %arg0 : !mesh_tensor, !mesh_tensor, !mesh_tensor
+}
+
+// -----
+
+!mesh_tensor = !mpmd.mesh_tensor<"m", tensor<4xui32>, sharding=<@mesh, [{"x"}]>>
+
+// CHECK-LABEL: func @identity_function
+func.func @identity_function(%arg0: !mesh_tensor) -> !mesh_tensor
+  attributes {"topology"=#mpmd.topology<<"m": <["x"=2]>>>}
+{
+  // CHECK-NEXT: %[[F:.*]] = mpmd.fragment<mesh="m", origin=[]> (%arg0) {mpmd.inferred_by = ["uniquify"]} (%arg1: tensor<4xui32>) {
+  // CHECK-NEXT:   return %arg1
+  // CHECK-NEXT: }
+  // CHECK-NEXT: return %[[F]]
+  func.return %arg0 : !mesh_tensor
+}

shardy/dialect/mpmd/transforms/common/test/uniquify_function_inputs_outputs_with_reshard.mlir

@@ -28,7 +28,7 @@ func.func @single_mesh_one_return_operand(%arg0: !mesh_1_tensor) -> (!mesh_1_ten
 } {
   // CHECK-NEXT: %[[F1:.*]] = mpmd.fragment<mesh="m1", origin=["f1"]>
   // CHECK:      %[[F2:.*]] = mpmd.fragment<mesh="m1", origin=["f2"]>
-  // CHECK:      %[[UF:.*]]:2 = mpmd.fragment<mesh="m1", origin=[]> (%[[F1]]) (%arg1: tensor<4xf32>) {
+  // CHECK:      %[[UF:.*]]:2 = mpmd.fragment<mesh="m1", origin=[]> (%[[F1]]) {mpmd.inferred_by = ["uniquify"]} (%arg1: tensor<4xf32>) {
   // CHECK:         mpmd.return %arg1, %arg1 : tensor<4xf32>, tensor<4xf32>
   // CHECK:      %[[F2]], %[[UF]]#0, %[[UF]]#1
   %0 = mpmd.fragment<mesh="m1", origin=["f1"]> (%arg0) (%arg1: tensor<4xf32>) {
@@ -49,7 +49,7 @@ func.func @needs_fragment_for_m1_with_many_values(%arg0: !mesh_1_tensor, %arg1:
   // CHECK-NEXT: %[[F1:.*]] = mpmd.fragment<mesh="m1", origin=["f1"]>
   // CHECK:      %[[F2:.*]] = mpmd.fragment<mesh="m2", origin=["f2"]>
   // CHECK:      %[[F3:.*]] = mpmd.fragment<mesh="m1", origin=["f3"]>
-  // CHECK:      %[[UF:.*]]:5 = mpmd.fragment<mesh="m1", origin=[]> (%[[F1]], %[[F3]]) (%[[A1:.*]]: tensor<4xf32>, %[[A2:.*]]: tensor<4xf32>)
+  // CHECK:      %[[UF:.*]]:5 = mpmd.fragment<mesh="m1", origin=[]> (%[[F1]], %[[F3]]) {mpmd.inferred_by = ["uniquify"]} (%[[A1:.*]]: tensor<4xf32>, %[[A2:.*]]: tensor<4xf32>)
   // CHECK-NEXT:   mpmd.return %[[A1]], %[[A1]], %[[A2]], %[[A2]], %[[A2]]
   // CHECK-NEXT: }
   // CHECK-NEXT: return %[[F2]], %[[UF]]#0, %[[UF]]#2, %[[UF]]#1, %[[UF]]#3, %[[UF]]#4
@@ -70,8 +70,8 @@ func.func @needs_fragment_for_m1_and_m2(%arg0: !mesh_1_tensor, %arg1: !mesh_2_te
 ) -> (!mesh_1_tensor, !mesh_2_tensor, !mesh_2_tensor, !mesh_1_tensor, !mesh_1_tensor, !mesh_1_tensor) attributes {
   "topology"=#mpmd.topology<<"m1": <["x"=2]>>, <"m2": <["x"=2]>>>
 } {
-  // CHECK: %[[UF1:.*]]:4 = mpmd.fragment<mesh="m1", origin=[]>
-  // CHECK: %[[UF2:.*]]:2 = mpmd.fragment<mesh="m2", origin=[]>
+  // CHECK: %[[UF1:.*]]:4 = mpmd.fragment<mesh="m1", origin=[]> ({{.*}}) {mpmd.inferred_by = ["uniquify"]}
+  // CHECK: %[[UF2:.*]]:2 = mpmd.fragment<mesh="m2", origin=[]> ({{.*}}) {mpmd.inferred_by = ["uniquify"]}
   // CHECK: return %[[UF1]]#0, %[[UF2]]#0, %[[UF2]]#1, %[[UF1]]#2, %[[UF1]]#1, %[[UF1]]#3
   %0 = mpmd.fragment<mesh="m1", origin=["f1"]> (%arg0) (%arg2: tensor<4xf32>) {
     mpmd.return %arg2 : tensor<4xf32>
@@ -97,7 +97,7 @@ func.func @single_mesh_one_return_operand_with_global_view(%arg0: !dist_mesh_ten
 } {
   // CHECK-NEXT: %[[F1:.*]] = mpmd.fragment<mesh="m1", origin=["f1"]>
   // CHECK:      %[[F2:.*]] = mpmd.fragment<mesh="m1", origin=["f2"]>
-  // CHECK:      %[[UF:.*]]:2 = mpmd.fragment<mesh="m1", origin=[]> (%[[F1]]) (%arg1: tensor<4xf32>) {
+  // CHECK:      %[[UF:.*]]:2 = mpmd.fragment<mesh="m1", origin=[]> (%[[F1]]) {mpmd.inferred_by = ["uniquify"]} (%arg1: tensor<4xf32>) {
   // CHECK:         mpmd.return %arg1, %arg1 : tensor<4xf32>, tensor<4xf32>
   // CHECK:      %[[F2]], %[[UF]]#0, %[[UF]]#1
   %0 = mpmd.fragment<mesh="m1", origin=["f1"]> (%arg0) (%arg1: tensor<4xf32>) {
@@ -122,7 +122,7 @@ func.func @f(%arg0: !mesh_tensor) -> (!mesh_tensor, !mesh_tensor, !mesh_tensor)
   // CHECK-NEXT: %[[F1:.*]] = mpmd.fragment<mesh="m", origin=["f"]> (%arg0) (%arg1: tensor<4xui32>) {
   // CHECK-NEXT:   return %arg1
   // CHECK-NEXT: }
-  // CHECK-NEXT: %[[F2:.*]]:2 = mpmd.fragment<mesh="m", origin=[]> (%arg0) (%arg1: tensor<4xui32>) {
+  // CHECK-NEXT: %[[F2:.*]]:2 = mpmd.fragment<mesh="m", origin=[]> (%arg0) {mpmd.inferred_by = ["uniquify"]} (%arg1: tensor<4xui32>) {
   // CHECK-NEXT:   return %arg1, %arg1
   // CHECK-NEXT: }
   // CHECK-NEXT: return %[[F2]]#0, %[[F1]], %[[F2]]#1
@@ -140,7 +140,7 @@ func.func @f(%arg0: !mesh_tensor) -> (!mesh_tensor, !mesh_tensor, !mesh_tensor)
 func.func @identity_function(%arg0: !mesh_tensor) -> !mesh_tensor
   attributes {"topology"=#mpmd.topology<<"m": <["x"=2]>>>}
 {
-  // CHECK-NEXT: %[[F:.*]] = mpmd.fragment<mesh="m", origin=[]> (%arg0) (%arg1: tensor<4xui32>) {
+  // CHECK-NEXT: %[[F:.*]] = mpmd.fragment<mesh="m", origin=[]> (%arg0) {mpmd.inferred_by = ["uniquify"]} (%arg1: tensor<4xui32>) {
   // CHECK-NEXT:   return %arg1
   // CHECK-NEXT: }
   // CHECK-NEXT: return %[[F]]