Add a pass to reorder functions in pre-order of the call graph.

shardy/dialect/sdy/ir/utils.cc

@@ -1012,16 +1012,29 @@ bool walkCalls(ModuleOp moduleOp, ProcessCallOpFn processCallOp,
   return true;
 }
 
-void iterateFuncs(ModuleOp moduleOp, ProcessFuncOpFn processFuncOp) {
+void iterateFuncs(ModuleOp moduleOp, ProcessFuncOpFn processFuncOp,
+                  bool preOrder) {
   CallGraph callGraph(moduleOp);
   llvm::ReversePostOrderTraversal<const CallGraph*> rpo(&callGraph);
-  for (CallGraphNode* node : llvm::reverse(rpo)) {
-    if (node->isExternal()) {
-      continue;
+  if (preOrder) {
+    for (CallGraphNode* node : rpo) {
+      if (node->isExternal()) {
+        continue;
+      }
+      mlir::Region* region = node->getCallableRegion();
+      if (FuncOp funcOp = dyn_cast_or_null<FuncOp>(region->getParentOp())) {
+        processFuncOp(funcOp);
+      }
     }
-    mlir::Region* region = node->getCallableRegion();
-    if (FuncOp funcOp = dyn_cast_or_null<FuncOp>(region->getParentOp())) {
-      processFuncOp(funcOp);
+  } else {
+    for (CallGraphNode* node : llvm::reverse(rpo)) {
+      if (node->isExternal()) {
+        continue;
+      }
+      mlir::Region* region = node->getCallableRegion();
+      if (FuncOp funcOp = dyn_cast_or_null<FuncOp>(region->getParentOp())) {
+        processFuncOp(funcOp);
+      }
     }
   }
 }

shardy/dialect/sdy/ir/utils.h

@@ -671,9 +671,11 @@ bool walkCalls(ModuleOp moduleOp, ProcessCallOpFn processCallOp,
 // Iterates on the funcs and performs `processFuncOp` on funcs. Iterates on the
 // funcs and blocks in post order of the call graph by default, that is, the
 // functions are processed before their callers, and child blocks are processed
-// before their parents.
+// before their parents. Iterates funcs and blocks in pre order if `preOrder` is
+// true.
 using ProcessFuncOpFn = std::function<void(func::FuncOp)>;
-void iterateFuncs(ModuleOp moduleOp, ProcessFuncOpFn processFuncOp);
+void iterateFuncs(ModuleOp moduleOp, ProcessFuncOpFn processFuncOp,
+                  bool preOrder = false);
 
 // Returns the reduction operation used in the scatter's update computation if
 // it is a recognized associative and commutative binary op applied to all

shardy/dialect/sdy/transforms/import/BUILD

@@ -42,6 +42,7 @@ cc_library(
         "inline_meshes.cc",
         "lift_inlined_meshes.cc",
         "manual_axes_cleanup.cc",
+        "pre_order_funcs.cc",
         "propagate_sharding_from_func_to_call.cc",
         "remove_size_one_axes.cc",
         "sharding_group_import.cc",

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

@@ -197,6 +197,16 @@ def ConstantOrScalarSplitterPass : Pass<"sdy-constant-or-scalar-splitter", "Modu
   let dependentDialects = ["mlir::sdy::SdyDialect"];
 }
 
+def PreOrderFuncsPass : Pass<"sdy-pre-order-funcs", "ModuleOp"> {
+  let summary = "Reorders functions in the module in pre-order of the call graph.";
+  let description = [{
+    Reorders functions in the module in pre-order of the call graph.
+    This is useful when we inline func/calls, as propagation is top-down on blocks,
+    and a pre-order iteration on funcs will emulate it.
+  }];
+  let dependentDialects = ["mlir::sdy::SdyDialect"];
+}
+
 def ShardingGroupImportPass : Pass<"sdy-sharding-group-import", "ModuleOp"> {
   let summary = "Canonicalization and validation pass for sharding groups.";
   let description = [{

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

@@ -0,0 +1,53 @@
+/* Copyright 2026 The Shardy Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <vector>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Block.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "shardy/dialect/sdy/ir/utils.h"
+#include "shardy/dialect/sdy/transforms/import/passes.h"  // IWYU pragma: keep
+
+namespace mlir {
+namespace sdy {
+
+#define GEN_PASS_DEF_PREORDERFUNCSPASS
+#include "shardy/dialect/sdy/transforms/import/passes.h.inc"
+
+namespace {
+
+struct PreOrderFuncsPass
+    : public impl::PreOrderFuncsPassBase<PreOrderFuncsPass> {
+  using PreOrderFuncsPassBase::PreOrderFuncsPassBase;
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    std::vector<func::FuncOp> funcsInPreOrder;
+    iterateFuncs(
+        moduleOp,
+        [&](func::FuncOp funcOp) { funcsInPreOrder.push_back(funcOp); },
+        /*preOrder=*/true);
+    mlir::Block& body = moduleOp.getBodyRegion().front();
+    for (func::FuncOp funcOp : funcsInPreOrder) {
+      funcOp->moveBefore(&body, body.end());
+    }
+  }
+};
+
+}  // namespace
+
+}  // namespace sdy
+}  // namespace mlir

shardy/dialect/sdy/transforms/import/test/pre_order_funcs.mlir

@@ -0,0 +1,19 @@
+// RUN: sdy_opt %s -sdy-pre-order-funcs | FileCheck %s
+
+// CHECK: func.func @main
+// CHECK: func.func @func1
+// CHECK: func.func @func2
+
+func.func @func2() {
+  return
+}
+
+func.func @func1() {
+  func.call @func2() : () -> ()
+  return
+}
+
+func.func @main() {
+  func.call @func1() : () -> ()
+  return
+}

shardy/dialect/sdy/transforms/propagation/propagation_pipeline.cc

@@ -64,6 +64,7 @@ void addPropagationPipeline(OpPassManager& pm, int& dumpIndex,
     optionsWithKeepShardingRules.keepShardingRules = true;
     // We intentionally don't increment the dump index here, since this pass
     // might dump 0 to multiple files, and will use a nested dump index.
+    pm.addPass(createPreOrderFuncsPass());
     pm.addPass(createUserPriorityPropagationPass(optionsWithKeepShardingRules,
                                                  dumpIndex));
   }