[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 bcbe91ee63ef03b3c0551ea47a42501007624f88 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       | 65 +++++++++++++++++++
 2 files changed, 114 insertions(+)

diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index dd0a290252dae3..cecf96c2fb979a 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -1931,6 +1931,54 @@ struct DSEState {
     if (!Store)
       return false;
 
+    // If there is a dominating condition, that ensures that the value
+    // being stored in a memory location is already present at the
+    // memory location, the store is a noop.
+    BasicBlock *StoreBB = DefI->getParent();
+    Value *StorePtr = Store->getPointerOperand();
+    Value *StoreVal = Store->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;
+
+    // Load can be either of the two operands
+    ICmpInst *CI = dyn_cast<ICmpInst>(BI->getCondition());
+    if (!CI || !CI->isEquality())
+      return false;
+    // Check if load is from the operand that we are storing to.
+    LoadInst *ICmpL = nullptr;
+    Value *ICmpV = nullptr;
+    for (Value *Op : CI->operands()) {
+      LLVM_DEBUG(dbgs() << "Operand: " << *Op << "\n");
+      auto *LI = dyn_cast<LoadInst>(Op);
+      if (LI && LI->getPointerOperand() == StorePtr)
+        ICmpL = LI;
+      if (Op == StoreVal)
+        ICmpV = Op;
+    }
+
+    if (!ICmpL || !ICmpV)
+      return false;
+
+    MemoryAccess *LoadAcc, *ClobAcc;
+    LoadAcc = MSSA.getMemoryAccess(ICmpL)->getDefiningAccess();
+    ClobAcc = MSSA.getSkipSelfWalker()->getClobberingMemoryAccess(Def, BatchAA);
+
+    if (!MSSA.dominates(ClobAcc, LoadAcc))
+      return false;
+
+    ICmpInst::Predicate Pred = CI->getPredicate();
+    if (Pred == ICmpInst::ICMP_EQ && StoreBB == BI->getSuccessor(0))
+      return true;
+
+    if (Pred == ICmpInst::ICMP_NE && StoreBB == BI->getSuccessor(1))
+      return true;
+
     if (auto *LoadI = dyn_cast<LoadInst>(Store->getOperand(0))) {
       if (LoadI->getPointerOperand() == Store->getOperand(1)) {
         // Get the defining access for the load.
@@ -2009,6 +2057,7 @@ struct DSEState {
         continue;
 
       Instruction *DefInst = Def->getMemoryInst();
+      LLVM_DEBUG(dbgs() << "DefInst: " << *DefInst << "\n");
       auto MaybeDefLoc = getLocForWrite(DefInst);
       if (!MaybeDefLoc || !isRemovable(DefInst))
         continue;
diff --git a/llvm/test/Transforms/DeadStoreElimination/noop-stores.ll b/llvm/test/Transforms/DeadStoreElimination/noop-stores.ll
index 3703b8d039ead0..c0ee304eca619a 100644
--- a/llvm/test/Transforms/DeadStoreElimination/noop-stores.ll
+++ b/llvm/test/Transforms/DeadStoreElimination/noop-stores.ll
@@ -795,3 +795,68 @@ join:
   store i8 %v, ptr %q, align 1
   ret void
 }
+
+; If the condition ensures that the value is already present in
+; in the variable, treat the store as noop and optimize. Should
+; not optimize stores if a clobber happens after the load that
+; is used in condition check.
+ at x = global i32 0, align 4
+ at y = global i32 0, align 4
+ at z = global i32 0, align 4
+define void @remove_tautological_store() {
+; CHECK-LABEL: @remove_tautological_store(
+; 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:       test.y:
+; CHECK-NEXT:    [[VALY:%.*]] = load i32, ptr @y, align 4
+; CHECK-NEXT:    [[CMPY:%.*]] = icmp eq i32 [[VALY]], 4
+; CHECK-NEXT:    br i1 [[CMPY]], label [[IF_Y:%.*]], label [[END]]
+; CHECK:       test.z:
+; CHECK-NEXT:    [[VALZ:%.*]] = load i32, ptr @z, align 4
+; CHECK-NEXT:    [[CMPZ:%.*]] = icmp eq i32 [[VALZ]], 4
+; CHECK-NEXT:    br i1 [[CMPZ]], label [[IF_Z:%.*]], label [[END]]
+; CHECK:       if.x:
+; CHECK-NEXT:    br label [[TEST_Y:%.*]]
+; CHECK:       if.y:
+; CHECK-NEXT:    store i32 5, ptr @y, align 4
+; CHECK-NEXT:    br label [[TEST_Z:%.*]]
+; CHECK:       if.z:
+; 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
+
+test.y:
+  %valy = load i32, ptr @y, align 4
+  %cmpy = icmp eq i32 %valy, 4
+  br i1 %cmpy, label %if.y, label %end
+
+test.z:
+  %valz = load i32, ptr @z, align 4
+  %cmpz = icmp eq i32 %valz, 4
+  br i1 %cmpz, label %if.z, label %end
+
+
+if.x:
+  store i32 4, ptr @x, align 4
+  br label %test.y
+
+if.y:
+  store i32 5, ptr @y, align 4
+  br label %test.z
+
+if.z:
+  store i32 4, ptr @z, align 4
+  br label %end
+
+end:
+  ret void
+}