[llvm] [AMDGPU] Promote allocas used by phi/select (PR #221870)

[via llvm-commits]

================
@@ -375,48 +398,111 @@ bool AMDGPUPromoteAllocaImpl::run(Function &F, bool PromoteToLDS) {
                                   : (MaxVGPRs * 32)) /
       VGPRBudgetRatio;
 
-  std::vector<AllocaAnalysis> Allocas;
+  SmallMapVector<AllocaInst *, AllocaAnalysis, 4> AllocaAnalysisMap;
   for (Instruction &I : F.getEntryBlock()) {
-    if (AllocaInst *AI = dyn_cast<AllocaInst>(&I)) {
-      // Array allocations are probably not worth handling, since an allocation
-      // of the array type is the canonical form.
-      if (!AI->isStaticAlloca() || AI->isArrayAllocation())
-        continue;
+    auto *AI = dyn_cast<AllocaInst>(&I);
+    // Array allocations are probably not worth handling, since an allocation
+    // of the array type is the canonical form.
+    if (!AI || !AI->isStaticAlloca() || AI->isArrayAllocation())
+      continue;
+    LLVM_DEBUG(dbgs() << "Analyzing: " << *AI << '\n');
+    AllocaAnalysis AA{AI};
+    if (!collectAllocaUses(AA))
+      continue;
+    AllocaAnalysisMap.insert({AI, std::move(AA)});
+  }
 
-      LLVM_DEBUG(dbgs() << "Analyzing: " << *AI << '\n');
+  std::vector<AllocaAnalysis> Allocas;
+  for (auto &AA : AllocaAnalysisMap.values()) {
+    if (AA.Uses.empty())
+      continue;
 
-      AllocaAnalysis AA{AI};
-      if (collectAllocaUses(AA)) {
-        analyzePromoteToVector(AA);
-        if (PromoteToLDS)
-          analyzePromoteToLDS(AA);
-        if (AA.Vector.Ty || AA.LDS.Enable) {
-          scoreAlloca(AA);
-          Allocas.push_back(std::move(AA));
-        }
+    bool Promotable = true;
+    SmallVector<AllocaInst *> Worklist;
+    LLVM_DEBUG(dbgs() << "Process leader alloca " << *AA.getLeaderAlloca()
+                      << '\n');
+    Worklist.append(AA.Links);
+    // Pull the information from links into the leader alloca.
+    while (!Worklist.empty()) {
+      AllocaInst *CurLink = Worklist.pop_back_val();
+      auto LinkIter = AllocaAnalysisMap.find(CurLink);
+      if (LinkIter != AllocaAnalysisMap.end()) {
+        AllocaAnalysis &LinkAA = LinkIter->second;
+        // Skip if the linked alloca was done or points to the leader alloca
+        // itself.
+        if (LinkAA.Uses.empty() || &LinkAA == &AA)
+          continue;
+
+        LLVM_DEBUG({
+          dbgs() << "  Process link: " << *LinkAA.getLeaderAlloca() << '\n';
+          for (auto *X : LinkAA.Uses)
+            dbgs() << "    Add User " << *X->getUser() << '\n';
+        });
+
+        AA.Uses.insert_range(LinkAA.Uses);
+        LinkAA.Uses.clear();
+        AA.HaveUnpromotableMerge |= LinkAA.HaveUnpromotableMerge;
+        AA.Members.push_back(LinkAA.getLeaderAlloca());
+
+        // Add indirect links to worklist, so that their info are properly
+        // propagated to the leader alloca.
+        Worklist.append(LinkAA.Links);
+        LinkAA.Links.clear();
+      } else {
+        LLVM_DEBUG(dbgs() << "  The alloca is not promotable\n");
+        Promotable = false;
+        break;
       }
----------------
ruiling wrote:

done

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