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

[Weibo He 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;
+    });
----------------
NewSigma wrote:

This two-lookup approach seems a bit clunky. Could we leverage something like `InstVisitor` to streamline it?

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