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

This change introduces a new pass to reorder the functions within an MLIR module. The functions are arranged such that callers appear before their callees, following a pre-order traversal of the call graph.

This ordering keeps the initial population of the worklist the same on the greedy rewriter driver during the propagation between inlined (named computations) and non-inlined (func/calls).

PiperOrigin-RevId: 901743552

Author: ekayaaslan

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));
   }