[DSE] Use precise loc for memset_chk during overwrite checks

memset_chk may not write the number of bytes specified by the third
argument, if it is larger than the destination size (specified as 4th
argument).

Reviewed By: asbirlea

Differential Revision: https://reviews.llvm.org/D115167

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

@@ -879,6 +879,26 @@ struct DSEState {
     CodeMetrics::collectEphemeralValues(&F, &AC, EphValues);
   }
 
+  LocationSize strengthenLocationSize(const Instruction *I,
+                                      LocationSize Size) const {
+    if (auto *CB = dyn_cast<CallBase>(I)) {
+      LibFunc F;
+      if (TLI.getLibFunc(*CB, F) && TLI.has(F) && F == LibFunc_memset_chk) {
+        // Use the precise location size specified by the 3rd argument
+        // for determining KillingI overwrites DeadLoc if it is a memset_chk
+        // instruction. memset_chk will write either the amount specified as 3rd
+        // argument or the function will immediately abort and exit the program.
+        // NOTE: AA may determine NoAlias if it can prove that the access size
+        // is larger than the allocation size due to that being UB. To avoid
+        // returning potentially invalid NoAlias results by AA, limit the use of
+        // the precise location size to isOverwrite.
+        if (const auto *Len = dyn_cast<ConstantInt>(CB->getArgOperand(2)))
+          return LocationSize::precise(Len->getZExtValue());
+      }
+    }
+    return Size;
+  }
+
   /// Return 'OW_Complete' if a store to the 'KillingLoc' location (by \p
   /// KillingI instruction) completely overwrites a store to the 'DeadLoc'
   /// location (by \p DeadI instruction).
@@ -898,23 +918,25 @@ struct DSEState {
     if (!isGuaranteedLoopIndependent(DeadI, KillingI, DeadLoc))
       return OW_Unknown;
 
+    LocationSize KillingLocSize =
+        strengthenLocationSize(KillingI, KillingLoc.Size);
     const Value *DeadPtr = DeadLoc.Ptr->stripPointerCasts();
     const Value *KillingPtr = KillingLoc.Ptr->stripPointerCasts();
     const Value *DeadUndObj = getUnderlyingObject(DeadPtr);
     const Value *KillingUndObj = getUnderlyingObject(KillingPtr);
 
     // Check whether the killing store overwrites the whole object, in which
     // case the size/offset of the dead store does not matter.
-    if (DeadUndObj == KillingUndObj && KillingLoc.Size.isPrecise()) {
+    if (DeadUndObj == KillingUndObj && KillingLocSize.isPrecise()) {
       uint64_t KillingUndObjSize = getPointerSize(KillingUndObj, DL, TLI, &F);
       if (KillingUndObjSize != MemoryLocation::UnknownSize &&
-          KillingUndObjSize == KillingLoc.Size.getValue())
+          KillingUndObjSize == KillingLocSize.getValue())
         return OW_Complete;
     }
 
     // FIXME: Vet that this works for size upper-bounds. Seems unlikely that we'll
     // get imprecise values here, though (except for unknown sizes).
-    if (!KillingLoc.Size.isPrecise() || !DeadLoc.Size.isPrecise()) {
+    if (!KillingLocSize.isPrecise() || !DeadLoc.Size.isPrecise()) {
       // In case no constant size is known, try to an IR values for the number
       // of bytes written and check if they match.
       const auto *KillingMemI = dyn_cast<MemIntrinsic>(KillingI);
@@ -931,7 +953,7 @@ struct DSEState {
       return isMaskedStoreOverwrite(KillingI, DeadI, BatchAA);
     }
 
-    const uint64_t KillingSize = KillingLoc.Size.getValue();
+    const uint64_t KillingSize = KillingLocSize.getValue();
     const uint64_t DeadSize = DeadLoc.Size.getValue();
 
     // Query the alias information

llvm/test/Transforms/DeadStoreElimination/libcalls-chk.ll

@@ -12,8 +12,7 @@ declare void @use(ptr)
 ; strncpy -> __memset_chk, full overwrite
 define void @dse_strncpy_memset_chk_test1(ptr noalias %out, ptr noalias %in, i64 %n) {
 ; CHECK-LABEL: @dse_strncpy_memset_chk_test1(
-; CHECK-NEXT:    [[CALL:%.*]] = tail call ptr @strncpy(ptr [[OUT:%.*]], ptr [[IN:%.*]], i64 100)
-; CHECK-NEXT:    [[CALL_2:%.*]] = tail call ptr @__memset_chk(ptr [[OUT]], i32 42, i64 100, i64 [[N:%.*]])
+; CHECK-NEXT:    [[CALL_2:%.*]] = tail call ptr @__memset_chk(ptr [[OUT:%.*]], i32 42, i64 100, i64 [[N:%.*]])
 ; CHECK-NEXT:    ret void
 ;
   %call = tail call ptr @strncpy(ptr %out, ptr %in, i64 100)
@@ -23,8 +22,7 @@ define void @dse_strncpy_memset_chk_test1(ptr noalias %out, ptr noalias %in, i64
 
 define void @dse_memset_chk_eliminate_store1(ptr %out, i64 %n) {
 ; CHECK-LABEL: @dse_memset_chk_eliminate_store1(
-; CHECK-NEXT:    store i8 10, ptr [[OUT:%.*]], align 1
-; CHECK-NEXT:    [[CALL_2:%.*]] = tail call ptr @__memset_chk(ptr [[OUT]], i32 42, i64 100, i64 [[N:%.*]])
+; CHECK-NEXT:    [[CALL_2:%.*]] = tail call ptr @__memset_chk(ptr [[OUT:%.*]], i32 42, i64 100, i64 [[N:%.*]])
 ; CHECK-NEXT:    ret void
 ;
   store i8 10, ptr %out
@@ -48,7 +46,6 @@ define void @dse_memset_chk_eliminate_store2(ptr %out, i64 %n) {
 define void @dse_memset_chk_eliminates_store_local_object_escapes_after(i64 %n) {
 ; CHECK-LABEL: @dse_memset_chk_eliminates_store_local_object_escapes_after(
 ; CHECK-NEXT:    [[A:%.*]] = alloca [200 x i8], align 1
-; CHECK-NEXT:    store i8 10, ptr [[A]], align 1
 ; CHECK-NEXT:    [[OUT_100:%.*]] = getelementptr i8, ptr [[A]], i64 100
 ; CHECK-NEXT:    store i8 10, ptr [[OUT_100]], align 1
 ; CHECK-NEXT:    [[CALL_2:%.*]] = tail call ptr @__memset_chk(ptr [[A]], i32 42, i64 100, i64 [[N:%.*]])
@@ -68,7 +65,6 @@ define void @dse_memset_chk_eliminates_store_local_object_escapes_before(i64 %n)
 ; CHECK-LABEL: @dse_memset_chk_eliminates_store_local_object_escapes_before(
 ; CHECK-NEXT:    [[A:%.*]] = alloca [200 x i8], align 1
 ; CHECK-NEXT:    call void @use(ptr [[A]])
-; CHECK-NEXT:    store i8 10, ptr [[A]], align 1
 ; CHECK-NEXT:    [[OUT_100:%.*]] = getelementptr i8, ptr [[A]], i64 100
 ; CHECK-NEXT:    store i8 0, ptr [[OUT_100]], align 1
 ; CHECK-NEXT:    [[CALL_2:%.*]] = tail call ptr @__memset_chk(ptr [[A]], i32 42, i64 100, i64 [[N:%.*]])