Push shardy unflatenner up past Canonicalizer on reshards.

PiperOrigin-RevId: 908449560

Author: ekayaaslan

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

@@ -42,24 +42,19 @@ void runShardyPartitioner(OpPassManager& pm, int& dumpIndex,
   passOptions.enableFullVersion = options.enableInsertExplicitCollectives;
   passOptions.avoidReshardsOnCalls = options.avoidReshardsOnCalls;
   pm.addNestedPass<func::FuncOp>(createInsertExplicitReshardsPass(passOptions));
+  pm.addPass(createUnflattenCallGraphPass(
+      UnflattenCallGraphPassOptions{options.dedupFunctionsFully}));
+  // Keep a SymbolDCE after UnflattenCallGraph.
+  pm.addPass(createSymbolDCEPass());
   if (options.enableInsertExplicitCollectives) {
     pm.addPass(mlir::sdy::createSaveModuleOpPass(
         options.dumpDirectory, "after_explicit_reshards", dumpIndex++));
     addCanonicalizerPass(pm, kReshardLabel);
-    pm.addPass(createUnflattenCallGraphPass(
-        UnflattenCallGraphPassOptions{options.dedupFunctionsFully}));
-    // Keep a SymbolDCE after UnflattenCallGraph.
-    pm.addPass(createSymbolDCEPass());
     pm.addNestedPass<func::FuncOp>(createReshardToCollectivesPass());
     // NOTE: ReshardToCollectives pass above generates all-slice collectives,
     // which during the canonicalizer below may be converted to reduce scatters
     // by potentially fusing with preceeding all-reduces, which are inserted
     // during InsertExplicitReshards pass.
-  } else {
-    pm.addPass(createUnflattenCallGraphPass(
-        UnflattenCallGraphPassOptions{options.dedupFunctionsFully}));
-    // Keep a SymbolDCE after UnflattenCallGraph.
-    pm.addPass(createSymbolDCEPass());
   }
 
   addCanonicalizerPass(pm, kCollectiveLabel);