[IFRT] Apply pass that merges multiple reshards into a single one whe…

…n they have the same source and destination. PiperOrigin-RevId: 717613221
Intel-tensorflow · Jan 20, 2025 · 03260eb · 03260eb
1 parent 8784fe8
commit 03260eb
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/xla/python/ifrt/ir/transforms/passes.cc b/xla/python/ifrt/ir/transforms/passes.cc
@@ -47,6 +47,8 @@ void CreateIfrtToOutlinedAtomProgramsPipeline(
 
   if (!options.propagate_shardings) {
     pm.addPass(CreateIfrtVerifyShardingSpecifiedPass());
+    pm.addNestedPass<mlir::func::FuncOp>(
+        xla::ifrt::CreateIfrtMergeReshardsPass());
     // We can split ifrt.Reshard to ifrt.CopyArrays because all the shardings
     // are specified.
     pm.addPass(CreateIfrtReshardToCopyArraysPass());