Push shardy outliner up past UpdateNonDivisibleInputOutputShardings pass.

UpdateNonDivisibleInputOutputShardings updates non-divisible input output shardings for all func input/outputs, all shardable data flow ops including named-computations in/out shardings. Similarly it should also update them for call ops.

PiperOrigin-RevId: 900135043

Author: ekayaaslan

shardy/dialect/sdy/ir/constants.h

@@ -86,6 +86,9 @@ inline const std::string kEmptyMeshSymbol = "empty_mesh";
 // Attribute name for the original name of the func before flattening.
 inline constexpr llvm::StringRef kOriginalFuncName = "sdy.original_func_name";
 
+// Attribute name of the main func.
+inline constexpr llvm::StringRef kMainFuncName = "main";
+
 }  // namespace sdy
 }  // namespace mlir
 

shardy/dialect/sdy/ir/utils.cc

@@ -1108,5 +1108,25 @@ FuncOp cloneFuncRecursively(FuncOp funcOp, SymbolTable& symbolTable) {
   return clonedFuncOp;
 }
 
+TensorShardingPerValueAttr getFullyClosedLike(mlir::ValueRange values,
+                                              Attribute meshOrRef) {
+  SmallVector<TensorShardingAttr> resultShardings;
+  resultShardings.reserve(values.size());
+  for (mlir::Value value : values) {
+    resultShardings.push_back(TensorShardingAttr::getFullyReplicated(
+        meshOrRef.getContext(), mlir::sdy::getTensorRank(value), meshOrRef,
+        /*isClosed=*/true));
+  }
+  return TensorShardingPerValueAttr::get(meshOrRef.getContext(),
+                                         resultShardings);
+}
+
+// Returns the main func. Dies if there is no main func.
+FuncOp getMainFuncOrDie(ModuleOp moduleOp, SymbolTable& symbolTable) {
+  FuncOp funcOp = symbolTable.lookup<FuncOp>(kMainFuncName);
+  SDY_CHECK(funcOp) << "Failed to lookup function: " << kMainFuncName.str();
+  return funcOp;
+}
+
 }  // namespace sdy
 }  // namespace mlir

shardy/dialect/sdy/ir/utils.h

@@ -678,6 +678,16 @@ Operation* getCommonSupportedReductionOp(stablehlo::ScatterOp scatter);
 mlir::func::FuncOp cloneFuncRecursively(func::FuncOp funcOp,
                                         SymbolTable& symbolTable);
 
+// Returns a `TensorShardingPerValueAttr` on the shardings of the `values`. If
+// the sharding of a value is null, it creates a fully closed sharding for it on
+// the given `meshOrRef` and the rank of the tensor corresponding to the value.
+TensorShardingPerValueAttr getFullyClosedLike(mlir::ValueRange values,
+                                              Attribute meshOrRef);
+
+// Returns the main func. Dies if there is no main func.
+mlir::func::FuncOp getMainFuncOrDie(ModuleOp moduleOp,
+                                    SymbolTable& symbolTable);
+
 }  // namespace sdy
 }  // namespace mlir
 

shardy/dialect/sdy/transforms/common/propagation_options.h

@@ -48,7 +48,8 @@ struct PropagationOptions {
   // auto-partitioner will be invoked after propagation of user-specified
   // shardings.
   bool enableAutoPartitioning = false;
-  // Whether to avoid explicit reshards/collectives on named computations.
+  // Whether to avoid explicit reshards/collectives on named computations/calls.
+  // TODO(enver): Rename to avoidReshardsOnCalls.
   bool avoidReshardsOnNamedComputations = false;
   // Whether to update axes with non-divisible input/output shardings.
   bool updateNonDivisibleInputOutputShardings = true;

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

@@ -40,10 +40,8 @@ void runShardyPartitioner(OpPassManager& pm, int& dumpIndex,
                           const ExportOptions& options) {
   InsertExplicitReshardsPassOptions passOptions;
   passOptions.enableFullVersion = options.enableInsertExplicitCollectives;
-  passOptions.avoidReshardsOnNamedComputations =
-      options.avoidReshardsOnNamedComputations;
+  passOptions.avoidReshardsOnCalls = options.avoidReshardsOnCalls;
   pm.addNestedPass<func::FuncOp>(createInsertExplicitReshardsPass(passOptions));
-  pm.addPass(createExportNamedComputationsPass());
   if (options.enableInsertExplicitCollectives) {
     pm.addPass(mlir::sdy::createSaveModuleOpPass(
         options.dumpDirectory, "after_explicit_reshards", dumpIndex++));
@@ -82,6 +80,7 @@ void addExportPipeline(OpPassManager& pm, int& dumpIndex,
       /*sinkDebugShardingOrigins=*/options.dumpShardingOrigins,
       /*sinkDebugPropagationEdgeSharding=*/options.dumpPropagationEdges,
       /*sinkEnableNativeNonFlatSupport=*/options.enableNativeNonFlatSupport}));
+  pm.addPass(createExportNamedComputationsPass());
   if (options.updateNonDivisibleInputOutputShardings) {
     pm.addPass(createUpdateNonDivisibleInputOutputShardingsPass());
     pm.addPass(createRemoveSubAxesInInputOutputShardingsPass());
@@ -98,8 +97,6 @@ void addExportPipeline(OpPassManager& pm, int& dumpIndex,
   // reshards/collectives.
   if (!options.avoidExportForPartitioning) {
     runShardyPartitioner(pm, dumpIndex, options);
-  } else {
-    pm.addPass(createExportNamedComputationsPass());
   }
   if (options.dumpPropagationEdges || options.dumpShardingOrigins) {
     pm.addPass(createRemovePropagationDebugInfoPass());

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

@@ -51,6 +51,9 @@ namespace sdy {
 
 namespace {
 
+using func::CallOp;
+using func::FuncOp;
+
 void insertExplicitReshardsToTargetSharding(OpOperand& opOperand,
                                             TensorShardingAttr targetSharding,
                                             IRRewriter& rewriter,
@@ -102,22 +105,10 @@ void insertExplicitReshardsOnFuncReturn(Operation* op, func::FuncOp& funcOp,
   }
 }
 
-void insertExplicitReshardsOnDataFlowOp(
-    ShardableDataFlowOpInterface& op, IRRewriter& rewriter,
-    const SymbolTable& symbolTable, const bool onFullVersion,
-    const bool avoidReshardsOnNamedComputations) {
-  if (isa<NamedComputationOp>(op) && avoidReshardsOnNamedComputations) {
-    for (Value owner : op.getOpResultEdgeOwners()) {
-      for (OpOperand* sourceOpOperand : op.getEdgeSources(owner)) {
-        insertExplicitReshardsToTargetSharding(
-            *sourceOpOperand,
-            /*targetSharding=*/op.getEdgeOwnerSharding(owner), rewriter,
-            symbolTable,
-            /*insertAfterOperand=*/true, onFullVersion);
-      }
-    }
-    return;
-  }
+void insertExplicitReshardsOnDataFlowOp(ShardableDataFlowOpInterface& op,
+                                        IRRewriter& rewriter,
+                                        const SymbolTable& symbolTable,
+                                        const bool onFullVersion) {
   for (Value owner : llvm::concat<Value>(op.getOpResultEdgeOwners(),
                                          op.getBlockArgumentEdgeOwners())) {
     TensorShardingAttr ownerSharding = op.transformTargetSharding(
@@ -132,6 +123,33 @@ void insertExplicitReshardsOnDataFlowOp(
   }
 }
 
+void insertExplicitReshardsOnCallOp(CallOp callOp, IRRewriter& rewriter,
+                                    const SymbolTable& symbolTable,
+                                    const bool onFullVersion) {
+  FuncOp funcOp = symbolTable.lookup<FuncOp>(callOp.getCallee());
+  TensorShardingPerValueAttr funcArgShardings =
+      mlir::sdy::getFuncArgShardings(funcOp, symbolTable);
+  if (!funcArgShardings) {
+    mlir::Attribute meshOrRef = getMeshOrRef(
+        callOp.getNumOperands(), symbolTable,
+        [&](int64_t i) { return getSharding(callOp.getOperand(i)); });
+    // Return without inserting reshards as neither func arguments nor call
+    // operands have a sharding with non-maximal mesh.
+    if (!meshOrRef) {
+      return;
+    }
+    funcArgShardings = getFullyClosedLike(callOp.getOperands(), meshOrRef);
+  }
+  rewriter.setInsertionPoint(callOp);
+  for (auto [funcArgSharding, sourceOpOperand] : llvm::zip_equal(
+           funcArgShardings.getShardings(), callOp->getOpOperands())) {
+    insertExplicitReshardsToTargetSharding(
+        sourceOpOperand,
+        /*targetSharding=*/funcArgSharding, rewriter, symbolTable,
+        /*insertAfterOperand=*/true, onFullVersion);
+  }
+}
+
 // Reshard the result of a dot operation if all the following hold:
 //
 // 1. LHS and RHS have fully compatible shardings.
@@ -382,7 +400,7 @@ bool isOnFullVersion(Operation* op, const bool enableFullVersion) {
   }
   // To avoid copies of the same functions with mismatching shardings on the
   // arguments onto multiple callsites.
-  if (isa<NamedComputationOp>(op)) {
+  if (isa<func::CallOp>(op)) {
     return true;
   }
 
@@ -472,8 +490,15 @@ struct InsertExplicitReshardsPass
         // TODO(enver): Prefer resharding the owner when multiple sources are
         // sharded in the same way.
         insertExplicitReshardsOnDataFlowOp(shardableDataFlowOp, rewriter,
-                                           symbolTable, onFullVersion,
-                                           avoidReshardsOnNamedComputations);
+                                           symbolTable, onFullVersion);
+        return;
+      }
+
+      if (CallOp callOp = dyn_cast<CallOp>(op)) {
+        if (!avoidReshardsOnCalls) {
+          insertExplicitReshardsOnCallOp(callOp, rewriter, symbolTable,
+                                         onFullVersion);
+        }
         return;
       }
 

shardy/dialect/sdy/transforms/export/passes.h

@@ -76,10 +76,10 @@ struct ExportOptions : public PassPipelineOptions<ExportOptions> {
       llvm::cl::desc("Sink sdy.propagation_edges attr."),
       llvm::cl::init(false)};
 
-  Option<bool> avoidReshardsOnNamedComputations{
-      *this, "avoid-reshards-on-named-computations",
-      llvm::cl::desc("Avoid inserting explicit reshards/collectives for named "
-                     "computations."),
+  Option<bool> avoidReshardsOnCalls{
+      *this, "avoid-reshards-on-calls",
+      llvm::cl::desc(
+          "Avoid inserting explicit reshards/collectives for calls."),
       llvm::cl::init(false)};
 
   Option<bool> updateNonDivisibleInputOutputShardings{

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

@@ -128,10 +128,10 @@ def InsertExplicitReshardsPass : Pass<"sdy-insert-explicit-reshards", "func::Fun
       Option<"enableFullVersion", "enable-full-version",
             "bool", /*default=*/"false",
             "Enable full version.">,
-      Option<"avoidReshardsOnNamedComputations",
-            "avoid-reshards-on-named-computations",
+      Option<"avoidReshardsOnCalls",
+            "avoid-reshards-on-calls",
             "bool", /*default=*/"false",
-            "Avoid explicit reshards/collectives on named computations.">
+            "Avoid explicit reshards/collectives on calls.">
     ];
 }
 

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

@@ -23,6 +23,7 @@ limitations under the License.
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/SymbolTable.h"
 #include "mlir/Pass/Pass.h"  // IWYU pragma: keep
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/DialectConversion.h"
@@ -82,22 +83,23 @@ struct RemoveSubAxesInInputOutputShardingsPass
       RemoveSubAxesInInputOutputShardingsPassBase;
 
   void runOnOperation() final {
-    for (auto funcOp : getOperation().getOps<func::FuncOp>()) {
-      // Update arguments.
-      updateValueShardings(
-          funcOp.getNumArguments(),
-          [&](int64_t index) { return getSharding(funcOp.getArgument(index)); },
-          [&](int64_t index, TensorShardingAttr sharding) {
-            setSharding(funcOp.getArgument(index), sharding);
-          });
-      // Update results.
-      updateValueShardings(
-          funcOp.getNumResults(),
-          [&](int64_t index) { return getFuncResultSharding(funcOp, index); },
-          [&](int64_t index, TensorShardingAttr sharding) {
-            setFuncResultSharding(funcOp, index, sharding);
-          });
-    }
+    ModuleOp moduleOp = getOperation();
+    SymbolTable symbolTable(moduleOp);
+    func::FuncOp funcOp = getMainFuncOrDie(moduleOp, symbolTable);
+    // Update arguments.
+    updateValueShardings(
+        funcOp.getNumArguments(),
+        [&](int64_t index) { return getSharding(funcOp.getArgument(index)); },
+        [&](int64_t index, TensorShardingAttr sharding) {
+          setSharding(funcOp.getArgument(index), sharding);
+        });
+    // Update results.
+    updateValueShardings(
+        funcOp.getNumResults(),
+        [&](int64_t index) { return getFuncResultSharding(funcOp, index); },
+        [&](int64_t index, TensorShardingAttr sharding) {
+          setFuncResultSharding(funcOp, index, sharding);
+        });
   }
 };
 

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

@@ -0,0 +1,126 @@
+// RUN: sdy_opt %s -split-input-file -sdy-insert-explicit-reshards='enable-full-version=true avoid-reshards-on-calls=true' | FileCheck %s
+
+sdy.mesh @mesh = <["x"=2, "y"=2, "z"=4]>
+
+// CHECK-LABEL: func @call
+func.func @call(%arg0: tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}]>}) -> (tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"z"}]>}) {
+  // CHECK-NEXT: %[[CALL:.*]] = call @foo(%arg0)
+  // CHECK-NEXT: %[[NEGATE:.*]] = stablehlo.negate %[[CALL]]
+  // CHECK-NEXT: return %[[NEGATE]]
+  %0 = call @foo(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"z"}]>]>} : (tensor<210xf32>) -> (tensor<210xf32>)
+  %1 = stablehlo.negate %0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"z"}]>]>} : tensor<210xf32>
+  return %1 : tensor<210xf32>
+}
+
+// CHECK-LABEL: func private @foo
+func.func private @foo(%arg0: tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}]>}) -> (tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"z"}]>}) {
+  // CHECK-NEXT: %[[ABS:.*]] = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>}
+  // CHECK-NEXT: %[[RESHARD:.*]] = sdy.reshard %[[ABS]] <@mesh, [{"z"}]>
+  // CHECK-NEXT: return %[[RESHARD]]
+  %0 = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : tensor<210xf32>
+  return %0 : tensor<210xf32>
+}
+
+// -----
+sdy.mesh @mesh = <["x"=4, "y"=2]>
+
+// CHECK-LABEL: func @call_empty_block
+func.func @call_empty_block(%arg0: tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}]>}) -> (tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}]>}) {
+  // CHECK-NEXT: %[[CALL:.*]] = call @foo(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>}
+  // CHECK-NEXT: %[[NEGATE:.*]] = stablehlo.negate %0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>}
+  %0 = call @foo(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : (tensor<210xf32>) -> (tensor<210xf32>)
+  %1 = stablehlo.negate %0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : tensor<210xf32>
+  return %1 : tensor<210xf32>
+}
+
+// CHECK-LABEL: func private @foo
+func.func private @foo(%arg0: tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}]>}) -> (tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}]>}) {
+  // CHECK-NEXT: %[[RESHARD:.*]] = sdy.reshard %arg0 <@mesh, [{"y"}]>
+  // CHECK-NEXT: return %[[RESHARD]]
+  return %arg0 : tensor<210xf32>
+}
+
+// -----
+sdy.mesh @mesh = <["a"=2, "b"=2, "c"=2]>
+
+// CHECK-LABEL: func @call_with_shardings
+func.func @call_with_shardings(%arg0: tensor<8x2xi32>, %arg1: tensor<4x2xi32>) -> tensor<12x2xi32> {
+  // CHECK-NEXT: %[[CALL:.*]]:2 = call @foo(%arg0, %arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>, <@mesh, [{}, {}]>]>}
+  // CHECK-NEXT: %[[RESHARD1:.*]] = sdy.reshard %[[CALL]]#0 <@mesh, [{}, {"a"}]>
+  // CHECK-NEXT: %[[RESHARD2:.*]] = sdy.reshard %[[CALL]]#1 <@mesh, [{}, {"a"}]>
+  // CHECK-NEXT: %[[CONCAT:.*]] = stablehlo.concatenate %[[RESHARD1]], %[[RESHARD2]], dim = 0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"a"}]>]>}
+  // CHECK-NEXT: %[[RESHARD3:.*]] = sdy.reshard %[[CONCAT]] <@mesh, [{}, {}]>
+  %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
+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, [{}, {}]>}) {
+  // CHECK-NEXT: %[[ABS:.*]] = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>}
+  // CHECK-NEXT: %[[RESHARD0:.*]] = sdy.reshard %[[ABS]] <@mesh, [{}, {"a"}]>
+  // CHECK-NEXT: %[[RESHARD1:.*]] = sdy.reshard %[[RESHARD0]] <@mesh, [{"a"}, {}]>
+  // CHECK-NEXT: return %[[RESHARD1]], %arg1
+  %0 = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"a"}]>]>} : tensor<8x2xi32>
+  return %0, %arg1 : tensor<8x2xi32>, tensor<4x2xi32>
+}
+
+// -----
+sdy.mesh @mesh = <["x"=2, "y"=2, "z"=4]>
+
+// CHECK-LABEL: func @one_argument_to_multiple_calls(
+func.func @one_argument_to_multiple_calls(%arg0: tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}]>}) -> (tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"z"}]>}) {
+  // CHECK-NEXT: %[[CALL0:.*]] = call @foo(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>}
+  // CHECK-NEXT: %[[CALL1:.*]] = call @bar(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"z"}]>]>}
+  // CHECK-NEXT: %[[RESHARD:.*]] = sdy.reshard %[[CALL0]] <@mesh, [{"z"}]>
+  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[RESHARD]], %[[CALL1]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"z"}]>]>}
+  // CHECK-NEXT: return %[[ADD]]
+  %0 = call @foo(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : (tensor<210xf32>) -> (tensor<210xf32>)
+  %1 = call @bar(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"z"}]>]>} : (tensor<210xf32>) -> (tensor<210xf32>)
+  %3 = stablehlo.add %0, %1 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"z"}]>]>} : tensor<210xf32>
+  return %3 : tensor<210xf32>
+}
+
+// CHECK-LABEL: func private @foo
+func.func private @foo(%arg0: tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}]>}) -> (tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}]>}) {
+  // CHECK-NEXT: %[[ABS:.*]] = stablehlo.abs %arg0
+  // CHECK-NEXT: return %[[ABS]]
+  %0 = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : tensor<210xf32>
+  return %0 : tensor<210xf32>
+}
+
+// CHECK-LABEL: func private @bar
+func.func private @bar(%arg0: tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"z"}]>}) -> (tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"z"}]>}) {
+  // CHECK-NEXT: %[[ABS:.*]] = stablehlo.abs %arg0
+  // CHECK-NEXT: return %[[ABS]]
+  %0 = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"z"}]>]>} : tensor<210xf32>
+  return %0 : tensor<210xf32>
+}
+
+// -----
+sdy.mesh @mesh = <["x"=2, "y"=2, "z"=4]>
+
+
+// CHECK-LABEL: func @different_arguments_to_multiple_calls_with_same_input_output_shardings
+func.func @different_arguments_to_multiple_calls_with_same_input_output_shardings(%arg0: tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}]>}) -> (tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}]>}) {
+  // CHECK-NEXT: %[[CALL0:.*]] = call @foo(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>}
+  // CHECK-NEXT: %[[NEGATE:.*]] = stablehlo.negate %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}]>]>}
+  // CHECK-NEXT: %[[RESHARD:.*]] = sdy.reshard %[[NEGATE]] <@mesh, [{"y"}]>
+  // CHECK-NEXT: %[[CALL1:.*]] = call @foo(%[[RESHARD]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>}
+  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[CALL0]], %[[CALL1]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>}
+  // CHECK-NEXT: return %[[ADD]]
+  %0 = call @foo(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : (tensor<210xf32>) -> (tensor<210xf32>)
+  %1 = stablehlo.negate %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : tensor<210xf32>
+  %2 = call @foo(%1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : (tensor<210xf32>) -> (tensor<210xf32>)
+  %4 = stablehlo.add %0, %2 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : tensor<210xf32>
+  return %4 : tensor<210xf32>
+}
+
+// CHECK-LABEL: func private @foo
+func.func private @foo(%arg0: tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}]>}) -> (tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}]>}) {
+  // CHECK-NEXT: %[[ABS:.*]] = stablehlo.abs %arg0
+  // CHECK-NEXT: return %[[ABS]]
+  %3 = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : tensor<210xf32>
+  return %3 : tensor<210xf32>
+}