[llvm] [DeadStoreElimination] Optimize tautological assignments (PR #75744)

[Shreyansh Chouhan via llvm-commits]

https://github.com/BK1603 updated https://github.com/llvm/llvm-project/pull/75744

>From c2c68353501887a7ef63f8ef692ecd135b2c5a3b Mon Sep 17 00:00:00 2001
From: Shreyansh Chouhan <chouhan.shreyansh2702 at gmail.com>
Date: Sun, 17 Dec 2023 22:48:51 +0530
Subject: [PATCH] [DeadStoreElimination] Optimize tautological assignments

If a store is immediately dominated by a condition that ensures that the value
being stored in a memory location is already present at that memory location,
consider the store a noop.

Fixes #63419
---
 .../Scalar/DeadStoreElimination.cpp           | 49 +++++++++++
 .../DeadStoreElimination/noop-stores.ll       | 82 +++++++++++++++++++
 2 files changed, 131 insertions(+)

diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index 008dcc53fd44fc..e33a2c3ba717a9 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -1901,6 +1901,52 @@ struct DSEState {
     return true;
   }
 
+  // Check if there is a dominating condition, that implies that the value
+  // being stored in a ptr is already present in the ptr.
+  bool dominatingConditionImpliesValue(MemoryDef *Def) {
+    StoreInst *StoreI = dyn_cast<StoreInst>(Def->getMemoryInst());
+    BasicBlock *StoreBB = StoreI->getParent();
+    Value *StorePtr = StoreI->getPointerOperand();
+    Value *StoreVal = StoreI->getValueOperand();
+
+    DomTreeNode *IDom = DT.getNode(StoreBB)->getIDom();
+    if (!IDom)
+      return false;
+
+    auto *BI = dyn_cast<BranchInst>(IDom->getBlock()->getTerminator());
+    if (!BI || !BI->isConditional())
+      return false;
+
+    // In case both blocks are the same, we cannot optimize the FalseBB when the
+    // condition is true and vice versa. Return false in this case.
+    if (BI->getSuccessor(0) == BI->getSuccessor(1))
+      return false;
+
+    Instruction *ICmpL;
+    ICmpInst::Predicate Pred;
+    if (!match(BI->getCondition(),
+               m_c_ICmp(Pred,
+                        m_CombineAnd(m_Load(m_Specific(StorePtr)),
+                                     m_Instruction(ICmpL)),
+                        m_Specific(StoreVal))) ||
+        !ICmpInst::isEquality(Pred))
+      return false;
+
+    if (Pred == ICmpInst::ICMP_EQ && StoreBB != BI->getSuccessor(0))
+      return false;
+
+    if (Pred == ICmpInst::ICMP_NE && StoreBB != BI->getSuccessor(1))
+      return false;
+
+    MemoryAccess *LoadAcc = MSSA.getMemoryAccess(ICmpL);
+    MemoryAccess *ClobAcc =
+        MSSA.getSkipSelfWalker()->getClobberingMemoryAccess(Def, BatchAA);
+    if (!MSSA.dominates(ClobAcc, LoadAcc))
+      return false;
+
+    return true;
+  }
+
   /// \returns true if \p Def is a no-op store, either because it
   /// directly stores back a loaded value or stores zero to a calloced object.
   bool storeIsNoop(MemoryDef *Def, const Value *DefUO) {
@@ -1931,6 +1977,9 @@ struct DSEState {
     if (!Store)
       return false;
 
+    if (dominatingConditionImpliesValue(Def))
+      return true;
+
     if (auto *LoadI = dyn_cast<LoadInst>(Store->getOperand(0))) {
       if (LoadI->getPointerOperand() == Store->getOperand(1)) {
         // Get the defining access for the load.
diff --git a/llvm/test/Transforms/DeadStoreElimination/noop-stores.ll b/llvm/test/Transforms/DeadStoreElimination/noop-stores.ll
index 3703b8d039ead0..9ba86ee2d5ee39 100644
--- a/llvm/test/Transforms/DeadStoreElimination/noop-stores.ll
+++ b/llvm/test/Transforms/DeadStoreElimination/noop-stores.ll
@@ -795,3 +795,85 @@ join:
   store i8 %v, ptr %q, align 1
   ret void
 }
+
+; Dominating condition implies value already exists, optimize store
+define void @remove_tautological_store(ptr %x, ptr %y, ptr %z) {
+; CHECK-LABEL: @remove_tautological_store(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[VALX:%.*]] = load i32, ptr [[X:%.*]], align 4
+; CHECK-NEXT:    [[CMPX:%.*]] = icmp eq i32 [[VALX]], 4
+; CHECK-NEXT:    br i1 [[CMPX]], label [[IF_X:%.*]], label [[END:%.*]]
+; CHECK:       if.x:
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %valx = load i32, ptr %x, align 4
+  %cmpx = icmp eq i32 %valx, 4
+  br i1 %cmpx, label %if.x, label %end
+
+if.x:
+  store i32 4, ptr %x, align 4
+  br label %end
+
+end:
+  ret void
+}
+
+; Dominating condition implies value already exists, optimize store
+; Should not optimize since value being stored is different from cond check
+define void @remove_tautological_store_fail_diff_value(ptr %x) {
+; CHECK-LABEL: @remove_tautological_store_fail_diff_value(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[VALX:%.*]] = load i32, ptr [[X:%.*]], align 4
+; CHECK-NEXT:    [[CMPX:%.*]] = icmp eq i32 [[VALX]], 4
+; CHECK-NEXT:    br i1 [[CMPX]], label [[IF_X:%.*]], label [[END:%.*]]
+; CHECK:       if.x:
+; CHECK-NEXT:    store i32 5, ptr [[X]], align 4
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %valx = load i32, ptr %x, align 4
+  %cmpx = icmp eq i32 %valx, 4
+  br i1 %cmpx, label %if.x, label %end
+
+if.x:
+  store i32 5, ptr %x, align 4
+  br label %end
+
+end:
+  ret void
+}
+
+; Dominating condition implies value already exists, optimize store
+; Should not optimize since there is a clobbering acc after load
+define void @remove_tautological_store_fail_clobber(ptr %x) {
+; CHECK-LABEL: @remove_tautological_store_fail_clobber(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[VALX:%.*]] = load i32, ptr [[X:%.*]], align 4
+; CHECK-NEXT:    store i32 5, ptr [[X]], align 4
+; CHECK-NEXT:    [[CMPX:%.*]] = icmp eq i32 [[VALX]], 4
+; CHECK-NEXT:    br i1 [[CMPX]], label [[IF_X:%.*]], label [[END:%.*]]
+; CHECK:       if.x:
+; CHECK-NEXT:    store i32 4, ptr [[X]], align 4
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %valx = load i32, ptr %x, align 4
+  store i32 5, ptr %x, align 4
+  %cmpx = icmp eq i32 %valx, 4
+  br i1 %cmpx, label %if.x, label %end
+
+if.x:
+  store i32 4, ptr %x, align 4
+  br label %end
+
+end:
+  ret void
+}
+