[SimplifyCFG] Improve store speculation check

[nikic]

Summary

isSafeToSpeculateStore() looks for a preceding store to the same location to make sure that introducing a new store of the same value is safe. It currently bails on intervening mayHaveSideEffect() instructions. However, I believe just checking mayWriteToMemory() is sufficient there -- we just need to make sure that we know which value was stored, we don't care if we can unwind in the meantime.

While looking into this, I started having some doubts about the correctness of the transform with regard to thread safety. While we don't try to hoist non-simple stores, I believe we also need to make sure that the preceding store is simple as well. Otherwise we could introduce a spurious non-atomic write after an atomic write -- I believe under our memory model this would result in a subsequent undef atomic read, even if the second write stores the same value as the first.

Example: https://alive2.llvm.org/ce/z/q_3YAL

[jdoerfert]

looks reasonable to me.

[lebedev.ri]

I think this should be fine.
The check does seem sufficient, at least for readonly/readnone calls: https://alive2.llvm.org/ce/z/q_3YAL

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

static Value *isSafeToSpeculateStore(Instruction *I, BasicBlock *BrBB,
   // Skip pseudo probe intrinsic calls which are not really killing any memory
   // accesses.
   for (Instruction &CurI : reverse(BrBB->instructionsWithoutDebug(true))) {
     if (!MaxNumInstToLookAt)
       break;
     --MaxNumInstToLookAt;
 
     // Could be calling an instruction that affects memory like free().
-    if (CurI.mayHaveSideEffects() && !isa<StoreInst>(CurI))
+    if (CurI.mayWriteToMemory() && !isa<StoreInst>(CurI))
       return nullptr;
 
     if (auto *SI = dyn_cast<StoreInst>(&CurI)) {
-      // Found the previous store make sure it stores to the same location.
+      // Found the previous store to same location and type. Make sure it is
+      // simple, to avoid introducing a spurious non-atomic write after an
+      // atomic write.
       if (SI->getPointerOperand() == StorePtr &&
-          SI->getValueOperand()->getType() == StoreTy)
+          SI->getValueOperand()->getType() == StoreTy && SI->isSimple())
         // Found the previous store, return its value operand.
         return SI->getValueOperand();
       return nullptr; // Unknown store.
     }
   }
 
   return nullptr;
 }

llvm/test/Transforms/SimplifyCFG/speculate-store.ll

if.then:
   br label %ret.end
 
 ret.end:
   ret void
 }
 
 define void @readonly_call(ptr %ptr, i1 %cmp) {
 ; CHECK-LABEL: @readonly_call(
+; CHECK-NEXT:  ret.end:
 ; CHECK-NEXT:    store i32 0, ptr [[PTR:%.*]], align 4
 ; CHECK-NEXT:    call void @unknown_fun() #[[ATTR0:[0-9]+]]
-; CHECK-NEXT:    br i1 [[CMP:%.*]], label [[IF_THEN:%.*]], label [[RET_END:%.*]]
-; CHECK:       if.then:
-; CHECK-NEXT:    store i32 1, ptr [[PTR]], align 4
-; CHECK-NEXT:    br label [[RET_END]]
-; CHECK:       ret.end:
+; CHECK-NEXT:    [[SPEC_STORE_SELECT:%.*]] = select i1 [[CMP:%.*]], i32 1, i32 0
+; CHECK-NEXT:    store i32 [[SPEC_STORE_SELECT]], ptr [[PTR]], align 4
 ; CHECK-NEXT:    ret void
 ;
   store i32 0, ptr %ptr
   call void @unknown_fun() readonly
   br i1 %cmp, label %if.then, label %ret.end
 
 if.then:
   store i32 1, ptr %ptr
   br label %ret.end
 
 ret.end:
   ret void
 }
 
 define void @atomic_and_simple(ptr %ptr, i1 %cmp) {
 ; CHECK-LABEL: @atomic_and_simple(
-; CHECK-NEXT:  ret.end:
 ; CHECK-NEXT:    store atomic i32 0, ptr [[PTR:%.*]] seq_cst, align 4
-; CHECK-NEXT:    [[SPEC_STORE_SELECT:%.*]] = select i1 [[CMP:%.*]], i32 1, i32 0
-; CHECK-NEXT:    store i32 [[SPEC_STORE_SELECT]], ptr [[PTR]], align 4
+; CHECK-NEXT:    br i1 [[CMP:%.*]], label [[IF_THEN:%.*]], label [[RET_END:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    store i32 1, ptr [[PTR]], align 4
+; CHECK-NEXT:    br label [[RET_END]]
+; CHECK:       ret.end:
 ; CHECK-NEXT:    ret void
 ;
   store atomic i32 0, ptr %ptr seq_cst, align 4
   br i1 %cmp, label %if.then, label %ret.end
 
 if.then:
   store i32 1, ptr %ptr
   br label %ret.end
 
 ret.end:
   ret void
 }
 
 ; CHECK: !0 = !{!"branch_weights", i32 3, i32 5}
 !0 = !{!"branch_weights", i32 3, i32 5}