From f933abff2bfb15f69fa35836d075e91aae5acc18 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333@gmail.com>
Date: Thu, 7 Nov 2024 20:48:33 +0800
Subject: [PATCH] [DSE] Add `tryToShortenBegin` support for memcpy

---
 .../Scalar/DeadStoreElimination.cpp           | 29 +++++++++++++++++--
 .../DeadStoreElimination/libcalls-chk.ll      |  4 ++-
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index b619248c59de0..a221ce75d772c 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -203,9 +203,23 @@ static bool isShortenableAtTheEnd(Instruction *I) {
 /// Returns true if the beginning of this instruction can be safely shortened
 /// in length.
 static bool isShortenableAtTheBeginning(Instruction *I) {
-  // FIXME: Handle only memset for now. Supporting memcpy/memmove should be
-  // easily done by offsetting the source address.
-  return isa<AnyMemSetInst>(I);
+  if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
+    switch (II->getIntrinsicID()) {
+    default:
+      return false;
+    case Intrinsic::memset:
+    case Intrinsic::memcpy:
+    case Intrinsic::memset_inline:
+    case Intrinsic::memcpy_inline:
+    case Intrinsic::memcpy_element_unordered_atomic:
+    case Intrinsic::memset_element_unordered_atomic:
+      // Do shorten memory intrinsics.
+      // FIXME: Add memmove if it's also safe to transform.
+      return true;
+    }
+  }
+
+  return false;
 }
 
 static std::optional<TypeSize> getPointerSize(const Value *V,
@@ -644,6 +658,15 @@ static bool tryToShorten(Instruction *DeadI, int64_t &DeadStart,
         DeadI->getIterator());
     NewDestGEP->setDebugLoc(DeadIntrinsic->getDebugLoc());
     DeadIntrinsic->setDest(NewDestGEP);
+
+    if (auto *MTI = dyn_cast<AnyMemTransferInst>(DeadIntrinsic)) {
+      Instruction *NewSrcGEP = GetElementPtrInst::CreateInBounds(
+          Type::getInt8Ty(DeadIntrinsic->getContext()), MTI->getRawSource(),
+          Indices, "", DeadI->getIterator());
+      NewSrcGEP->setDebugLoc(DeadIntrinsic->getDebugLoc());
+      MTI->setSource(NewSrcGEP);
+      MTI->setSourceAlignment(PrefAlign);
+    }
   }
 
   // Update attached dbg.assign intrinsics. Assume 8-bit byte.
diff --git a/llvm/test/Transforms/DeadStoreElimination/libcalls-chk.ll b/llvm/test/Transforms/DeadStoreElimination/libcalls-chk.ll
index 737f99f539523..b2cea8566c594 100644
--- a/llvm/test/Transforms/DeadStoreElimination/libcalls-chk.ll
+++ b/llvm/test/Transforms/DeadStoreElimination/libcalls-chk.ll
@@ -130,7 +130,9 @@ define void @dse_strncpy_memcpy_chk_test2(ptr noalias %out, ptr noalias %in, i64
 
 define void @test_memcpy_intrinsic_and_memcpy_chk(ptr %A, ptr %B, ptr noalias %C) {
 ; CHECK-LABEL: @test_memcpy_intrinsic_and_memcpy_chk(
-; CHECK-NEXT:    tail call void @llvm.memcpy.p0.p0.i64(ptr [[A:%.*]], ptr [[B:%.*]], i64 48, i1 false)
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i8, ptr [[A:%.*]], i64 1
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i8, ptr [[B:%.*]], i64 1
+; CHECK-NEXT:    tail call void @llvm.memcpy.p0.p0.i64(ptr align 1 [[TMP1]], ptr align 1 [[TMP2]], i64 47, i1 false)
 ; CHECK-NEXT:    [[CALL:%.*]] = call ptr @__memcpy_chk(ptr [[A]], ptr [[C:%.*]], i64 1, i64 10)
 ; CHECK-NEXT:    ret void
 ;