[llvm] [AMDGPU] Preserve resource requirements through call-graph cycles (PR #221625)

[Matt Arsenault via llvm-commits]

================
@@ -247,33 +252,167 @@ bool AMDGPUMCExpr::evaluateInstPrefSize(MCValue &Res,
 
 bool AMDGPUMCExpr::isSymbolUsedInExpression(const MCSymbol *Sym,
                                             const MCExpr *E) {
-  switch (E->getKind()) {
-  case MCExpr::Constant:
-    return false;
-  case MCExpr::Unary:
-    return isSymbolUsedInExpression(
-        Sym, static_cast<const MCUnaryExpr *>(E)->getSubExpr());
-  case MCExpr::Binary: {
-    const MCBinaryExpr *BE = static_cast<const MCBinaryExpr *>(E);
-    return isSymbolUsedInExpression(Sym, BE->getLHS()) ||
-           isSymbolUsedInExpression(Sym, BE->getRHS());
-  }
-  case MCExpr::SymbolRef: {
-    const MCSymbol &S = static_cast<const MCSymbolRefExpr *>(E)->getSymbol();
-    if (S.isVariable())
-      return isSymbolUsedInExpression(Sym, S.getVariableValue());
-    return &S == Sym;
-  }
-  case MCExpr::Specifier:
-  case MCExpr::Target: {
-    auto *TE = static_cast<const AMDGPUMCExpr *>(E);
-    for (const MCExpr *E : TE->getArgs())
-      if (isSymbolUsedInExpression(Sym, E))
+  SmallVector<const MCExpr *, 16> WorkList{E};
+  SmallPtrSet<const MCExpr *, 16> Seen;
+  while (!WorkList.empty()) {
+    const MCExpr *Expr = WorkList.pop_back_val();
+    if (!Seen.insert(Expr).second)
+      continue;
+    switch (Expr->getKind()) {
+    case MCExpr::Constant:
+      break;
+    case MCExpr::Unary:
+      WorkList.push_back(cast<MCUnaryExpr>(Expr)->getSubExpr());
+      break;
+    case MCExpr::Binary: {
+      const auto *BE = cast<MCBinaryExpr>(Expr);
+      WorkList.push_back(BE->getRHS());
+      WorkList.push_back(BE->getLHS());
+      break;
+    }
+    case MCExpr::SymbolRef: {
+      const MCSymbol &S = cast<MCSymbolRefExpr>(Expr)->getSymbol();
+      if (&S == Sym)
         return true;
-    return false;
+      if (S.isVariable())
+        WorkList.push_back(S.getVariableValue());
+      break;
+    }
+    case MCExpr::Specifier:
+      WorkList.push_back(cast<MCSpecifierExpr>(Expr)->getSubExpr());
+      break;
+    case MCExpr::Target:
+      append_range(WorkList, reverse(cast<AMDGPUMCExpr>(Expr)->getArgs()));
+      break;
+    }
   }
+  return false;
+}
+
+// Resource expressions form a DAG of maxima, boolean unions, and frame-size
+// additions. Evaluate shared subexpressions once per query, including additions
+// between maxima, instead of recursively expanding every path through the DAG.
+static bool evaluateResourceExpr(const AMDGPUMCExpr *Root, MCValue &Res,
+                                 const MCAssembler *Asm) {
+  enum class Phase { Visit, Complete, CheckAbsolute };
+  struct WorkItem {
+    const MCExpr *Expr;
+    Phase Step = Phase::Visit;
+    MCSymbol *ResolvingSymbol = nullptr;
+  };
+  SmallVector<WorkItem, 16> WorkList{{Root}};
+  DenseMap<const MCExpr *, MCValue> Values;
+  SmallPtrSet<const MCExpr *, 16> Active;
+  SmallVector<MCSymbol *, 8> ResolvingSymbols;
+  // Match MCExpr's resolution guard, including on failure. Cache values only
+  // for this query: symbol definitions and assembler layout can change later.
+  auto ClearResolving = scope_exit([&] {
+    for (MCSymbol *Sym : ResolvingSymbols)
+      Sym->setIsResolving(false);
+  });
+  auto EvaluateLeaf = [&](const MCExpr *Expr) {
+    MCValue Value;
+    if (!Expr->evaluateAsRelocatable(Value, Asm))
+      return false;
+    Values.try_emplace(Expr, Value);
+    return true;
+  };
+
+  while (!WorkList.empty()) {
+    WorkItem Item = WorkList.pop_back_val();
+    const MCExpr *Expr = Item.Expr;
+    if (Item.Step == Phase::CheckAbsolute) {
+      // Match normal max/OR evaluation: stop before the next operand if this
+      // operand is not absolute.
+      if (!Values.lookup(Expr).isAbsolute())
+        return false;
+      continue;
+    }
+    if (Item.Step == Phase::Complete) {
+      Active.erase(Expr);
+      if (MCSymbol *Sym = Item.ResolvingSymbol) {
+        MCValue Value = Values.lookup(Sym->getVariableValue());
+        Sym->setIsResolving(false);
+        if (Value.isAbsolute())
+          Values.try_emplace(Expr, Value);
+        // The normal resolver preserves the identity of relocatable aliases.
+        // Clear our guard first so it does not diagnose a spurious cycle.
+        else if (!EvaluateLeaf(Expr))
+          return false;
+        continue;
+      }
----------------
arsenm wrote:

```suggestion
      }
      
```

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