[llvm] [Transforms] Speed up SSAUpdater::FindExistingPHI (PR #100281)

[William Junda Huang via llvm-commits]

https://github.com/huangjd created https://github.com/llvm/llvm-project/pull/100281

In SSAUpdater::FindExistingPHI, the cleanup function is inefficient for large function with many blocks because it clears the Phi value reference for every block if not matched for every phi value, even if most blocks are not modified by CheckIfPHIMatches. This behavior is particularly slow for large functions because the complexity is Θ(# PHI * # BBs).


Updated the behavior to only clear modified blocks, which in practice has a much less complexity because of early exit on PHI mismatch. 

>From 3b2ffc540c6c5a6794d6a49d3c6e08f154dc864a Mon Sep 17 00:00:00 2001
From: William Huang <williamjhuang at google.com>
Date: Tue, 23 Jul 2024 19:21:41 -0400
Subject: [PATCH] [Transforms] Speed up SSAUpdater::FindExistingPHI

In SSAUpdater::FindExistingPHI, the cleanup function is inefficient for
large function with many blocks because it clears the Phi value
reference for every block if not matched for every phi value, even if
most blocks are not modified by CheckIfPHIMatches. Updated the behavior
to only clear modified blocks.
---
 .../llvm/Transforms/Utils/SSAUpdaterImpl.h    | 24 ++++++++++++++-----
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/llvm/include/llvm/Transforms/Utils/SSAUpdaterImpl.h b/llvm/include/llvm/Transforms/Utils/SSAUpdaterImpl.h
index 28ff6c4c7927d..24cbeaf72f82b 100644
--- a/llvm/include/llvm/Transforms/Utils/SSAUpdaterImpl.h
+++ b/llvm/include/llvm/Transforms/Utils/SSAUpdaterImpl.h
@@ -418,10 +418,6 @@ class SSAUpdaterImpl {
         RecordMatchingPHIs(BlockList);
         break;
       }
-      // Match failed: clear all the PHITag values.
-      for (typename BlockListTy::iterator I = BlockList->begin(),
-             E = BlockList->end(); I != E; ++I)
-        (*I)->PHITag = nullptr;
     }
   }
 
@@ -429,10 +425,21 @@ class SSAUpdaterImpl {
   /// in the BBMap.
   bool CheckIfPHIMatches(PhiT *PHI) {
     SmallVector<PhiT *, 20> WorkList;
+    SmallVector<BBInfo *, 20> TaggedBlocks;
     WorkList.push_back(PHI);
 
+    // Match failed: clear all the PHITag values. Only need to clear visited
+    // blocks.
+    auto OnFalseCleanup = [&]() {
+      for (BBInfo *TaggedBlock : TaggedBlocks) {
+        TaggedBlock->PHITag = nullptr;
+      }
+    };
+
     // Mark that the block containing this PHI has been visited.
-    BBMap[PHI->getParent()]->PHITag = PHI;
+    BBInfo *PHIBlock = BBMap[PHI->getParent()];
+    PHIBlock->PHITag = PHI;
+    TaggedBlocks.push_back(PHIBlock);
 
     while (!WorkList.empty()) {
       PHI = WorkList.pop_back_val();
@@ -450,21 +457,26 @@ class SSAUpdaterImpl {
         if (PredInfo->AvailableVal) {
           if (IncomingVal == PredInfo->AvailableVal)
             continue;
+          OnFalseCleanup();
           return false;
         }
 
         // Check if the value is a PHI in the correct block.
         PhiT *IncomingPHIVal = Traits::ValueIsPHI(IncomingVal, Updater);
-        if (!IncomingPHIVal || IncomingPHIVal->getParent() != PredInfo->BB)
+        if (!IncomingPHIVal || IncomingPHIVal->getParent() != PredInfo->BB) {
+          OnFalseCleanup();
           return false;
+        }
 
         // If this block has already been visited, check if this PHI matches.
         if (PredInfo->PHITag) {
           if (IncomingPHIVal == PredInfo->PHITag)
             continue;
+          OnFalseCleanup();
           return false;
         }
         PredInfo->PHITag = IncomingPHIVal;
+        TaggedBlocks.push_back(PredInfo);
 
         WorkList.push_back(IncomingPHIVal);
       }