[llvm] [Coroutines] Only rematerialize when operands are available after the suspend (PR #209195)

[Christian Ulmann via llvm-commits]

================
@@ -28,6 +32,83 @@ using namespace coro;
 // "coro-frame", which results in leaner debug spew.
 #define DEBUG_TYPE "coro-suspend-crossing"
 
+// Returns true if \p Root can be rematerialized in the resume function without
+// introducing a new spill. This holds when a value is:
+// 1. A constant: genuinely available at no cost, materialized as an immediate
+//    in the resume function.
+// 2. An argument: treated as available heuristically. An argument crossing the
+//    suspend is actually spilled, but treating it as available preserves
+//    beneficial shared-operand rematerialization.
+// 3. A value that already crosses a suspend point for an independent use, i.e.,
+//    it is spilled regardless.
+// 4. A materializable value all of whose operands are themselves available.
+static bool isAvailableAfterSuspend(
+    Value *Root, const std::function<bool(Instruction &)> &Materializable,
+    const SuspendCrossingInfo &Checker, SmallDenseMap<Value *, bool> &Memo) {
+  SmallVector<Value *> Stack;
+  // Materializable nodes whose operands have been scheduled but not yet folded.
+  // Distinguishes a node being visited a second time (operands ready) from the
+  // first visit, and lets an operand reached through a cycle resolve to false.
+  SmallPtrSet<Value *, 16> Opened;
+  Stack.push_back(Root);
+  while (!Stack.empty()) {
+    Value *V = Stack.back();
+    if (Memo.contains(V)) {
+      Stack.pop_back();
+      continue;
+    }
+
+    // Leaves that resolve without inspecting operands. Non-instructions
+    // (constants, arguments) are available; a value that already crosses a
+    // suspend independently is spilled regardless, so referencing its slot is
+    // available at no new cost.
+    auto *I = dyn_cast<Instruction>(V);
+    if (!I || Checker.isDefinitionAcrossSuspend(*I)) {
+      Memo[V] = true;
+      Stack.pop_back();
+      continue;
+    }
+    if (!Materializable(*I)) {
+      Memo[V] = false;
+      Stack.pop_back();
+      continue;
+    }
+
+    // Materializable: available iff all operands are available. On the first
+    // visit, schedule the unresolved operands above V and revisit V once they
+    // are folded.
+    if (Opened.insert(V).second) {
+      for (Use &U : I->operands())
+        if (!Memo.contains(U.get()))
+          Stack.push_back(U.get());
+      continue;
+    }
+
+    // Second visit: operands are resolved, except any reached through a cycle
+    // (only possible via non-materializable PHIs, which never open). An
+    // unresolved operand is treated as not available, conservatively.
+    Stack.pop_back();
+    Memo[V] = all_of(I->operands(), [&](Use &U) {
+      auto It = Memo.find(U.get());
+      return It != Memo.end() && It->second;
+    });
----------------
Dinistro wrote:

I don't fully understand your concern. Are you referring to the whole post-order DFS construct here? The `InstVisitor` cannot be used for this, but graph traits could be used. Unfortunately, setting up the graph traits for this kind of traversal is more code than what is written already.

https://github.com/llvm/llvm-project/pull/209195