[llvm] [Mem2Reg] Don't use single store optimization for potentially poison value (PR #97711)

[Nikita Popov via llvm-commits]

https://github.com/nikic created https://github.com/llvm/llvm-project/pull/97711

If there is a single store, then loads must either load the stored value or uninitialized memory (undef). If the stored value may be poison, then replacing an uninitialized memory load with it would be incorrect. Fall back to the generic code in that case.

This PR only fixes the case where there is a literal poison store -- the case where the value is non-trivially poison will still get miscompiled by phi simplification later, see #96631.

Fixes https://github.com/llvm/llvm-project/issues/97702.

>From fcbd9469eb59d3d16e8315bdc09180742f331f20 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Thu, 4 Jul 2024 12:50:13 +0200
Subject: [PATCH] [Mem2Reg] Don't use single store optimization for potentially
 poison value

If there is a single store, then loads must either load the stored
value or uninitialized memory (undef). If the stored value may be
poison, then replacing an uninitialized memory load with it would
be incorrect. Fall back to the generic code in that case.

This PR only fixes the case where there is a literal poison store
-- the case where the value is non-trivially poison will be
miscompiled by phi simplification later, see #96631.

Fixes https://github.com/llvm/llvm-project/issues/97702.
---
 llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp | 10 ++++++++--
 llvm/test/Transforms/Mem2Reg/single-store.ll          |  3 +--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp b/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
index 6e021a5e2d05a..cd5ab55c2122f 100644
--- a/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
+++ b/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
@@ -525,7 +525,14 @@ rewriteSingleStoreAlloca(AllocaInst *AI, AllocaInfo &Info, LargeBlockInfo &LBI,
                          SmallSet<DbgAssignIntrinsic *, 8> *DbgAssignsToDelete,
                          SmallSet<DbgVariableRecord *, 8> *DVRAssignsToDelete) {
   StoreInst *OnlyStore = Info.OnlyStore;
-  bool StoringGlobalVal = !isa<Instruction>(OnlyStore->getOperand(0));
+  Value *ReplVal = OnlyStore->getOperand(0);
+  // Loads may either load the stored value or uninitialized memory (undef).
+  // If the stored value may be poison, then replacing an uninitialized memory
+  // load with it would be incorrect.
+  if (!isGuaranteedNotToBePoison(ReplVal))
+    return false;
+
+  bool StoringGlobalVal = !isa<Instruction>(ReplVal);
   BasicBlock *StoreBB = OnlyStore->getParent();
   int StoreIndex = -1;
 
@@ -565,7 +572,6 @@ rewriteSingleStoreAlloca(AllocaInst *AI, AllocaInfo &Info, LargeBlockInfo &LBI,
     }
 
     // Otherwise, we *can* safely rewrite this load.
-    Value *ReplVal = OnlyStore->getOperand(0);
     // If the replacement value is the load, this must occur in unreachable
     // code.
     if (ReplVal == LI)
diff --git a/llvm/test/Transforms/Mem2Reg/single-store.ll b/llvm/test/Transforms/Mem2Reg/single-store.ll
index b82e26158a361..f864227c49145 100644
--- a/llvm/test/Transforms/Mem2Reg/single-store.ll
+++ b/llvm/test/Transforms/Mem2Reg/single-store.ll
@@ -1,7 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
 ; RUN: opt -S -passes=mem2reg < %s | FileCheck %s
 
-; FIXME: This is a miscompile.
 define i8 @single_store_literal_poison(i1 %cond) {
 ; CHECK-LABEL: define i8 @single_store_literal_poison(
 ; CHECK-SAME: i1 [[COND:%.*]]) {
@@ -9,7 +8,7 @@ define i8 @single_store_literal_poison(i1 %cond) {
 ; CHECK:       [[IF]]:
 ; CHECK-NEXT:    br label %[[EXIT]]
 ; CHECK:       [[EXIT]]:
-; CHECK-NEXT:    ret i8 poison
+; CHECK-NEXT:    ret i8 undef
 ;
   %a = alloca i8, align 1
   br i1 %cond, label %if, label %exit