[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 c8fec5f1530201c98853fbc12676a1b0f208a1b4 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                       | 32 ++++++++++
 2 files changed, 74 insertions(+), 16 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..5564e1a19542d2 100644
--- a/llvm/lib/Transforms/Utils/CanonicalizeFreezeInLoops.cpp
+++ b/llvm/lib/Transforms/Utils/CanonicalizeFreezeInLoops.cpp
@@ -29,9 +29,10 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/CanonicalizeFreezeInLoops.h"
+#include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallSet.h"
-#include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/IVDescriptors.h"
 #include "llvm/Analysis/LoopAnalysisManager.h"
 #include "llvm/Analysis/LoopInfo.h"
@@ -66,19 +67,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 +87,44 @@ 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 +152,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 +181,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..a46bb00ba7fa60
--- /dev/null
+++ b/llvm/test/Transforms/CanonicalizeFreezeInLoops/duplicate_remove.ll
@@ -0,0 +1,32 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s --passes=canon-freeze -S | FileCheck %s
+
+define void @check_duplicate_removal(i32 %n) {
+; CHECK-LABEL: define void @check_duplicate_removal(
+; CHECK-SAME: i32 [[N:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[N_FROZEN:%.*]] = freeze i32 [[N]]
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[T1:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[T3:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[T2:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[T3]], [[LOOP]] ]
+; CHECK-NEXT:    [[T3]] = add i32 [[N_FROZEN]], [[T2]]
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[T2]], 0
+; CHECK-NEXT:    br i1 [[COND]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+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
+}