Add utility to walk funcs.

ekayaaslan · copybara-github · commit 2dd4be972c1a · 2026-04-15T09:23:08.000-07:00
Also simplify the utility to walk calls.

Add few more utilities:
- getOriginalFuncName
- getFuncOpOrDie

PiperOrigin-RevId: 900202519
diff --git a/shardy/dialect/sdy/ir/utils.cc b/shardy/dialect/sdy/ir/utils.cc
@@ -912,6 +912,11 @@ StringAttr getOriginalFuncName(FuncOp funcOp) {
   return funcOp.getSymNameAttr();
 }
 
+StringAttr getOriginalFuncName(CallOp callOp, const SymbolTable& symbolTable) {
+  FuncOp funcOp = getFuncOpOrDie(callOp.getCallee(), symbolTable);
+  return getOriginalFuncName(funcOp);
+}
+
 mlir::Attribute getMeshOrRef(
     int64_t numElements, const SymbolTable& symbolTable,
     std::function<TensorShardingAttr(int64_t)> getSharding) {
@@ -974,9 +979,8 @@ TensorShardingPerValueAttr getFuncResultShardings(
   return TensorShardingPerValueAttr::get(funcOp.getContext(), resultShardings);
 }
 
-std::optional<FuncOp> walkCalls(ModuleOp moduleOp,
-                                ProcessCallOpFn processCallOp, bool preOrder) {
-  FuncOp mainFuncOp;
+bool walkCalls(ModuleOp moduleOp, ProcessCallOpFn processCallOp,
+               bool preOrder) {
   CallGraph callGraph(moduleOp);
   llvm::ReversePostOrderTraversal<const CallGraph*> rpo(&callGraph);
   if (preOrder) {  // Iterate pre-order.
@@ -985,16 +989,11 @@ std::optional<FuncOp> walkCalls(ModuleOp moduleOp,
         continue;
       }
       mlir::Region* region = node->getCallableRegion();
-      // The first func is the main one as it is in pre-order.
-      if (auto funcOp = dyn_cast_or_null<FuncOp>(region->getParentOp());
-          !mainFuncOp && funcOp) {
-        mainFuncOp = funcOp;
-      }
       if (region
               ->walk<WalkOrder::PreOrder>(
                   [&](CallOp callOp) { return processCallOp(callOp); })
               .wasInterrupted()) {
-        return std::nullopt;
+        return false;
       }
     }
   } else {
@@ -1004,25 +1003,27 @@ std::optional<FuncOp> walkCalls(ModuleOp moduleOp,
         continue;
       }
       mlir::Region* region = node->getCallableRegion();
-      // The last func is the main one as it is in post-order.
-      if (auto funcOp = dyn_cast_or_null<FuncOp>(region->getParentOp());
-          funcOp) {
-        mainFuncOp = funcOp;
-      }
       if (region->walk([&](CallOp callOp) { return processCallOp(callOp); })
               .wasInterrupted()) {
-        return std::nullopt;
+        return false;
       }
     }
   }
-  return mainFuncOp;
+  return true;
 }
 
-FuncOp walkCallsOrDie(ModuleOp moduleOp, ProcessCallOpFn processCallOp,
-                      bool preOrder) {
-  auto mainFuncOp = walkCalls(moduleOp, processCallOp, preOrder);
-  SDY_CHECK(mainFuncOp);
-  return *mainFuncOp;
+void iterateFuncs(ModuleOp moduleOp, ProcessFuncOpFn processFuncOp) {
+  CallGraph callGraph(moduleOp);
+  llvm::ReversePostOrderTraversal<const CallGraph*> rpo(&callGraph);
+  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);
+    }
+  }
 }
 
 Operation* getCommonSupportedReductionOp(stablehlo::ScatterOp scatter) {
@@ -1108,5 +1109,11 @@ FuncOp cloneFuncRecursively(FuncOp funcOp, SymbolTable& symbolTable) {
   return clonedFuncOp;
 }
 
+FuncOp getFuncOpOrDie(StringRef funcSymName, const SymbolTable& symbolTable) {
+  FuncOp funcOp = symbolTable.lookup<FuncOp>(funcSymName);
+  SDY_CHECK(funcOp) << "Failed to lookup function: " << funcSymName.str();
+  return funcOp;
+}
+
 }  // namespace sdy
 }  // namespace mlir
diff --git a/shardy/dialect/sdy/ir/utils.h b/shardy/dialect/sdy/ir/utils.h
@@ -642,6 +642,11 @@ DenseIntElementsAttr getReplicaGroups(AxisRefListAttr reductionAxesAttr,
 // case there is no such attribute attached, create one on the name of `funcOp`.
 StringAttr getOriginalFuncName(func::FuncOp funcOp);
 
+// Gets `kOriginalFuncName` attribute attached to the func of `callOp`. In case
+// there is no such attribute attached, create one on the name of the func.
+StringAttr getOriginalFuncName(func::CallOp callOp,
+                               const SymbolTable& symbolTable);
+
 // Returns the shardings for the arguments of `funcOp`, with fully replicated
 // shardings for empty shardings on `funcOp`.
 TensorShardingPerValueAttr getFuncArgShardings(func::FuncOp funcOp,
@@ -657,16 +662,18 @@ TensorShardingPerValueAttr getFuncResultShardings(
 // functions are processed before their callers, and child blocks are processed
 // before their parents. Iterates calls and blocks in pre order if `preOrder` is
 // true, that is, the functions are processed after their callers, and child
-// blocks are processed after their parents. Returns nullopt if the walk was
-// interrupted, returns the main func otherwise.
+// blocks are processed after their parents. Returns false if the walk was
+// interrupted, returns true otherwise.
 using ProcessCallOpFn = std::function<mlir::WalkResult(func::CallOp)>;
-std::optional<func::FuncOp> walkCalls(ModuleOp moduleOp,
-                                      ProcessCallOpFn processCallOp,
-                                      bool preOrder = false);
-// Walks calls as in `walkCalls` and returns the main func. Dies if the the walk
-// is interrupted, or otherwise it can not identify the main function.
-func::FuncOp walkCallsOrDie(ModuleOp moduleOp, ProcessCallOpFn processCallOp,
-                            bool preOrder = false);
+bool walkCalls(ModuleOp moduleOp, ProcessCallOpFn processCallOp,
+               bool preOrder = false);
+
+// 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.
+using ProcessFuncOpFn = std::function<void(func::FuncOp)>;
+void iterateFuncs(ModuleOp moduleOp, ProcessFuncOpFn processFuncOp);
 
 // Returns the reduction operation used in the scatter's update computation if
 // it is a recognized associative and commutative binary op applied to all
@@ -678,6 +685,11 @@ Operation* getCommonSupportedReductionOp(stablehlo::ScatterOp scatter);
 mlir::func::FuncOp cloneFuncRecursively(func::FuncOp funcOp,
                                         SymbolTable& symbolTable);
 
+// Returns the funcOp on `funcSymName`. Dies if the func does not exist on the
+// `symbolTable`.
+func::FuncOp getFuncOpOrDie(StringRef funcSymName,
+                            const SymbolTable& symbolTable);
+
 }  // namespace sdy
 }  // namespace mlir
 
diff --git a/shardy/dialect/sdy/ir/utils_test.cc b/shardy/dialect/sdy/ir/utils_test.cc
@@ -647,47 +647,59 @@ TEST_F(UtilsTest, WalkCalls_Interrupted) {
   EXPECT_THAT(calledFuncs, ElementsAre("baz"));
 }
 
-TEST_F(UtilsTest, WalkCallsPostOrder_ReturnsMainFunc) {
+TEST_F(UtilsTest, IterateFuncs_ThreeCalls) {
   auto localModule = mlir::parseSourceString<ModuleOp>(
       "module {\n"
-      "  func.func private @bar()\n"
-      "  func.func private @foo() {\n"
-      "    call @bar() : () -> ()\n"
+      "  func.func private @bar() {\n"
+      "    call @baz() : () -> ()\n"
+      "    return\n"
+      "  }\n"
+      "  func.func private @baz() {\n"
       "    return\n"
       "  }\n"
-      "  func.func @mymain() {\n"
+      "  func.func @main() {\n"
       "    call @foo() : () -> ()\n"
       "    return\n"
       "  }\n"
+      "  func.func private @foo() {\n"
+      "    call @bar() : () -> ()\n"
+      "    return\n"
+      "  }\n"
       "}",
       &context);
-  EXPECT_THAT(
-      walkCalls(localModule.get(),
-                [&](func::CallOp callOp) { return WalkResult::advance(); })
-          ->getName(),
-      "mymain");
+  std::vector<std::string> iteratedFuncs;
+  iterateFuncs(localModule.get(), [&](func::FuncOp funcOp) {
+    iteratedFuncs.push_back(funcOp.getName().str());
+  });
+  EXPECT_THAT(iteratedFuncs, ElementsAre("baz", "bar", "foo", "main"));
 }
 
-TEST_F(UtilsTest, WalkCallsPreOrder_ReturnsMainFunc) {
+TEST_F(UtilsTest, IterateFuncs_Triangle) {
   auto localModule = mlir::parseSourceString<ModuleOp>(
       "module {\n"
-      "  func.func private @bar()\n"
-      "  func.func private @foo() {\n"
-      "    call @bar() : () -> ()\n"
+      "  func.func private @bar() {\n"
+      "    call @baz() : () -> ()\n"
       "    return\n"
       "  }\n"
-      "  func.func @mymain() {\n"
+      "  func.func private @baz() {\n"
+      "    return\n"
+      "  }\n"
+      "  func.func @main() {\n"
       "    call @foo() : () -> ()\n"
+      "    call @bar() : () -> ()\n"
+      "    return\n"
+      "  }\n"
+      "  func.func private @foo() {\n"
+      "    call @bar() : () -> ()\n"
       "    return\n"
       "  }\n"
       "}",
       &context);
-  EXPECT_THAT(walkCalls(
-                  localModule.get(),
-                  [&](func::CallOp callOp) { return WalkResult::advance(); },
-                  /*preOrder=*/true)
-                  ->getName(),
-              "mymain");
+  std::vector<std::string> iteratedFuncs;
+  iterateFuncs(localModule.get(), [&](func::FuncOp funcOp) {
+    iteratedFuncs.push_back(funcOp.getName().str());
+  });
+  EXPECT_THAT(iteratedFuncs, ElementsAre("baz", "bar", "foo", "main"));
 }
 
 TEST_F(UtilsTest, GetShardableValue_AsyncStartOp) {
