[llvm] [CanonicalizeFreezeInLoops] fix duplicate removal (PR #74716)

[Wei Tao via llvm-commits]

https://github.com/Friedrich20 updated https://github.com/llvm/llvm-project/pull/74716

>From 557f8316c19cf9984a3eb51ae38ded8f0e2cb80b Mon Sep 17 00:00:00 2001
From: Wei Tao <friedrich.taow at gmail.com>
Date: Thu, 7 Dec 2023 21:33:40 +0800
Subject: [PATCH] [LLVM][CanonicalizeFreezeInLoops] fix duplicate removal

---
 .../Utils/CanonicalizeFreezeInLoops.cpp       | 58 +++++++++++++------
 .../duplicate_remove.ll                       | 16 +++++
 2 files changed, 57 insertions(+), 17 deletions(-)
 create mode 100644 llvm/test/Transforms/CanonicalizeFreezeInLoops/duplicate_remove.ll

diff --git a/llvm/lib/Transforms/Utils/CanonicalizeFreezeInLoops.cpp b/llvm/lib/Transforms/Utils/CanonicalizeFreezeInLoops.cpp
index fb4d8288537725..cd8c69546617f3 100644
--- a/llvm/lib/Transforms/Utils/CanonicalizeFreezeInLoops.cpp
+++ b/llvm/lib/Transforms/Utils/CanonicalizeFreezeInLoops.cpp
@@ -29,9 +29,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/CanonicalizeFreezeInLoops.h"
+#include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallSet.h"
-#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/Analysis/IVDescriptors.h"
 #include "llvm/Analysis/LoopAnalysisManager.h"
 #include "llvm/Analysis/LoopInfo.h"
@@ -66,19 +66,6 @@ class CanonicalizeFreezeInLoopsImpl {
   ScalarEvolution &SE;
   DominatorTree &DT;
 
-  struct FrozenIndPHIInfo {
-    // A freeze instruction that uses an induction phi
-    FreezeInst *FI = nullptr;
-    // The induction phi, step instruction, the operand idx of StepInst which is
-    // a step value
-    PHINode *PHI;
-    BinaryOperator *StepInst;
-    unsigned StepValIdx = 0;
-
-    FrozenIndPHIInfo(PHINode *PHI, BinaryOperator *StepInst)
-        : PHI(PHI), StepInst(StepInst) {}
-  };
-
   // Can freeze instruction be pushed into operands of I?
   // In order to do this, I should not create a poison after I's flags are
   // stripped.
@@ -99,6 +86,43 @@ class CanonicalizeFreezeInLoopsImpl {
 
 } // anonymous namespace
 
+namespace llvm {
+
+struct FrozenIndPHIInfo {
+  // A freeze instruction that uses an induction phi
+  FreezeInst *FI = nullptr;
+  // The induction phi, step instruction, the operand idx of StepInst which is
+  // a step value
+  PHINode *PHI;
+  BinaryOperator *StepInst;
+  unsigned StepValIdx = 0;
+
+  FrozenIndPHIInfo(PHINode *PHI, BinaryOperator *StepInst)
+      : PHI(PHI), StepInst(StepInst) {}
+};
+
+template <> struct DenseMapInfo<FrozenIndPHIInfo> {
+  static inline FrozenIndPHIInfo getEmptyKey() {
+    return FrozenIndPHIInfo(DenseMapInfo<PHINode *>::getEmptyKey(),
+                            DenseMapInfo<BinaryOperator *>::getEmptyKey());
+  }
+
+  static inline FrozenIndPHIInfo getTombstoneKey() {
+    return FrozenIndPHIInfo(DenseMapInfo<PHINode *>::getTombstoneKey(),
+                            DenseMapInfo<BinaryOperator *>::getTombstoneKey());
+  }
+
+  static unsigned getHashValue(const FrozenIndPHIInfo &Val) {
+    return DenseMapInfo<FreezeInst *>::getHashValue(Val.FI);
+  };
+  static bool isEqual(const FrozenIndPHIInfo &LHS,
+                      const FrozenIndPHIInfo &RHS) {
+    return LHS.FI == RHS.FI;
+  };
+};
+
+} // end namespace llvm
+
 // Given U = (value, user), replace value with freeze(value), and let
 // SCEV forget user. The inserted freeze is placed in the preheader.
 void CanonicalizeFreezeInLoopsImpl::InsertFreezeAndForgetFromSCEV(Use &U) {
@@ -126,7 +150,7 @@ bool CanonicalizeFreezeInLoopsImpl::run() {
   if (!L->isLoopSimplifyForm())
     return false;
 
-  SmallVector<FrozenIndPHIInfo, 4> Candidates;
+  SmallSetVector<FrozenIndPHIInfo, 4> Candidates;
 
   for (auto &PHI : L->getHeader()->phis()) {
     InductionDescriptor ID;
@@ -155,7 +179,7 @@ bool CanonicalizeFreezeInLoopsImpl::run() {
       if (auto *FI = dyn_cast<FreezeInst>(U)) {
         LLVM_DEBUG(dbgs() << "canonfr: found: " << *FI << "\n");
         Info.FI = FI;
-        Candidates.push_back(Info);
+        Candidates.insert(Info);
       }
     };
     for_each(PHI.users(), Visit);
diff --git a/llvm/test/Transforms/CanonicalizeFreezeInLoops/duplicate_remove.ll b/llvm/test/Transforms/CanonicalizeFreezeInLoops/duplicate_remove.ll
new file mode 100644
index 00000000000000..515dd589297527
--- /dev/null
+++ b/llvm/test/Transforms/CanonicalizeFreezeInLoops/duplicate_remove.ll
@@ -0,0 +1,16 @@
+; RUN: opt < %s -canon-freeze -S | FileCheck %s
+define void @check_duplicate_removal(i32 %n) {
+entry:
+br label %loop
+
+loop:
+%t1 = phi i32 [ 0, %entry], [%t3, %loop ]
+%t2 = phi i32 [ 0, %entry], [%t3, %loop ]
+%t3 = add i32 %n, %t2
+%.fr = freeze i32 %t3
+%cond = icmp eq i32 %t2, 0
+br i1 %cond, label %loop, label %exit
+
+exit:
+ret void
+}
\ No newline at end of file