Push shardy inliner down past ApplyShardingConstraints pass for early inliner.

PiperOrigin-RevId: 904958581

Author: ekayaaslan

shardy/dialect/sdy/transforms/import/import_pipeline.cc

@@ -32,12 +32,6 @@ void addImportPipeline(OpPassManager& pm, int& dumpIndex,
   pm.addPass(createRemoveSizeOneAxesPass());
   pm.addPass(createPropagateShardingFromFuncToCallPass());
   pm.addPass(createConstantOrScalarSplitterPass());
-  if (!options.enableLateInlining) {
-    pm.addPass(createImportFuncCallsPass(ImportFuncCallsPassOptions{
-        /*addDataFlowEdgesOnNamedComputations=*/false}));
-    // Keep SymbolDCEPass after ImportFuncCallsPass.
-    pm.addPass(createSymbolDCEPass());
-  }
   pm.addPass(createSymbolDCEPass());
   pm.addPass(createManualAxesCleanupPass());
 
@@ -53,6 +47,11 @@ void addImportPipeline(OpPassManager& pm, int& dumpIndex,
   pm.addPass(
       createApplyShardingConstraintsPass(ApplyShardingConstraintsPassOptions{
           options.debugShardingOrigins, options.debugPropagationEdgeSharding}));
+  if (!options.enableLateInlining) {
+    pm.addPass(createImportFuncCallsPass());
+    // Keep SymbolDCEPass after ImportFuncCallsPass.
+    pm.addPass(createSymbolDCEPass());
+  }
   // The sharding group import pass must run after applying sharding
   // constraints. This ensures we can detect sharding conflicts between group
   // members which have pre-propagation shardings due to sharding constraints.