[JumpThreading] Fix a bug affecting guards and assumes

[reames]

Summary

This patch fixes an issue in jump threading where we'd sometimes use circular logic to fold the condition assumed or guarded to true. The code had an implicit assumption that a fact true at the end of a block must be true for all uses within the block. This is blatantly false. The net effect for guards was that we could remove a guard on a condition which wasn't actual true at runtime and miscompile.

I noticed this while working on another optimization patch. I'm a little terrified that we never noticed this previously.

[reames]

add files missed in initial diff

[sanjoy]

Almost LGTM. Marking as needing changes to make sure we settle the replaceDominatedUsesWith issue before checking in.

lib/Transforms/Scalar/JumpThreading.cpp

834

Is this clang-formatted?

lib/Transforms/Utils/Local.cpp

1759

(In a separate change) What do you think about refactoring out a replaceUsesWithIf that takes a predicate to decide if a certain use needs to be replaced and change replaceNonLocalUsesWith and the two replaceDominatedUsesWith s to use it?

1764

Why not auto instead of Value::use_iterator?

1777

Stray newline?

1809

Please double check, but I don't think this is behavior equivalent. Today if we have

  br <cond>, left, merge

left:
  x = def()
  br merge

merge:
  t = phi(x, 42)

then replaceDominatedUsesWith(x, 100, DT, left) will not replace the use in the phi node. But with your change, we will get all of the uses in of x, including the one in the phi node.

The real problem here is that replaceDominatedUsesWith does not do what it says it does. Above the use of x in the phi node is, in fact, dominated by left, but since replaceDominatedUsesWith "decays" uses to instructions before checking dominance, that fact is lost.

In any case, if you want to make sure replaceNonLocalUsesWith is well tested, I'd nudge you towards adding a test or two in unittests/Transforms/Utils/Local.cpp.

test/Transforms/JumpThreading/assume.ll

89

I'd also add a test with an assume followed by a guard.

[reames]

Anna landed a patch which subsumes this one.

include/llvm/Transforms/Utils/Local.h

 /// the given edge.  Returns the number of replacements made.
 unsigned replaceDominatedUsesWith(Value *From, Value *To, DominatorTree &DT,
                                   const BasicBlockEdge &Edge);
 /// Replace each use of 'From' with 'To' if that use is dominated by
 /// the end of the given BasicBlock. Returns the number of replacements made.
 unsigned replaceDominatedUsesWith(Value *From, Value *To, DominatorTree &DT,
                                   const BasicBlock *BB);
 
+/// Replace each use of 'From' with 'To' if that use is outside the defining
+/// block of 'From'.  Returns the number of replacements made.
+unsigned replaceNonLocalUsesWith(Instruction *From, Value *To);
 
 /// Return true if the CallSite CS calls a gc leaf function.
 ///
 /// A leaf function is a function that does not safepoint the thread during its
 /// execution.  During a call or invoke to such a function, the callers stack
 /// does not have to be made parseable.
 ///
 /// Most passes can and should ignore this information, and it is only used

lib/Transforms/Scalar/JumpThreading.cpp

if (CondBr && CondConst && CondBr->isConditional()) {
         LVI->getPredicateAt(CondCmp->getPredicate(), CondCmp->getOperand(0),
                             CondConst, CondBr);
       if (Ret != LazyValueInfo::Unknown) {
         unsigned ToRemove = Ret == LazyValueInfo::True ? 1 : 0;
         unsigned ToKeep = Ret == LazyValueInfo::True ? 0 : 1;
         CondBr->getSuccessor(ToRemove)->removePredecessor(BB, true);
         BranchInst::Create(CondBr->getSuccessor(ToKeep), CondBr);
         CondBr->eraseFromParent();
-        if (CondCmp->use_empty())
-          CondCmp->eraseFromParent();
-        else if (CondCmp->getParent() == BB) {
-          // If the fact we just learned is true for all uses of the
-          // condition, replace it with a constant value
+
+        // Now, see if we can eagerly update any other uses. WARNING: we used a
+        // context instruction at the end of the current basic block, so the
+        // facts LVI returned may not be true before that point.  (i.e. we may
+        // have a guard or assume part way through the block.)  Thus, we can
+        // not reason about any use which is not dominated by the context
+        // instruction.  Given we don't have dom-tree here, we're restricted to
+        // only handling cases where dominance can be trivially established
+        // from the definition of SSA.
+        if (!CondCmp->use_empty() &&
+            CondCmp->getParent() == BB) {
           auto *CI = Ret == LazyValueInfo::True ?
             ConstantInt::getTrue(CondCmp->getType()) :
             ConstantInt::getFalse(CondCmp->getType());
-          CondCmp->replaceAllUsesWith(CI);
-          CondCmp->eraseFromParent();
+          replaceNonLocalUsesWith(CondCmp, CI);
         }
+        if (CondCmp->use_empty())
+          CondCmp->eraseFromParent();
         return true;
       }
     }
 
     if (CondBr && CondConst && TryToUnfoldSelect(CondCmp, BB))
       return true;
   }
 

lib/Transforms/Utils/Local.cpp

unsigned KnownIDs[] = {
       LLVMContext::MD_noalias,         LLVMContext::MD_range,
       LLVMContext::MD_invariant_load,  LLVMContext::MD_nonnull,
       LLVMContext::MD_invariant_group, LLVMContext::MD_align,
       LLVMContext::MD_dereferenceable,
       LLVMContext::MD_dereferenceable_or_null};
   combineMetadata(K, J, KnownIDs);
 }
 
+unsigned llvm::replaceNonLocalUsesWith(Instruction *From, Value *To) {
+  assert(From->getType() == To->getType());
+
+  BasicBlock *DefBB = From->getParent();
+  
+  unsigned Count = 0;
+  for (Value::use_iterator UI = From->use_begin(), UE = From->use_end();
+       UI != UE; ) {
+    Use &U = *UI++;
+    auto *I = cast<Instruction>(U.getUser());
+    if (DefBB != I->getParent()) {
+      U.set(To);
+      DEBUG(dbgs() << "Replace non-local use of '"
+            << From->getName() << "' as "
+            << *To << " in " << *U << "\n");
+      ++Count;
+    }
+  }
+  return Count;
+
+}
+
+
 unsigned llvm::replaceDominatedUsesWith(Value *From, Value *To,
                                         DominatorTree &DT,
                                         const BasicBlockEdge &Root) {
   assert(From->getType() == To->getType());
   
   unsigned Count = 0;
   for (Value::use_iterator UI = From->use_begin(), UE = From->use_end();
        UI != UE; ) {
unsigned llvm::replaceDominatedUsesWith(Value *From, Value *To,
   return Count;
 }
 
 unsigned llvm::replaceDominatedUsesWith(Value *From, Value *To,
                                         DominatorTree &DT,
                                         const BasicBlock *BB) {
   assert(From->getType() == To->getType());
 
+  // This "micro optimization" is here mostly to ensure this code path gets
+  // test coverage.
+  if (auto *FromI = dyn_cast<Instruction>(From))
+    if (FromI->getParent() == BB)
+      return replaceNonLocalUsesWith(FromI, To);
+
   unsigned Count = 0;
   for (Value::use_iterator UI = From->use_begin(), UE = From->use_end();
        UI != UE;) {
     Use &U = *UI++;
     auto *I = cast<Instruction>(U.getUser());
     if (DT.properlyDominates(BB, I->getParent())) {
       U.set(To);
       DEBUG(dbgs() << "Replace dominated use of '" << From->getName() << "' as "

test/Transforms/JumpThreading/assume.ll

if.then:                                          ; preds = %entry
   tail call void (...) @bar() #1
   br label %return
 
 return:                                           ; preds = %entry, %if.then
   %retval.0 = phi i32 [ 1, %if.then ], [ 0, %entry ]
   ret i32 %retval.0
 }
 
+@g = external global i32
+
+; Check that we do prove a fact using an assume within the block and
+; that we *don't* use circular logic to remove the assume.
+; CHECK-LABEL: @local_assume
+; CHECK: %notnull = icmp ne i32* %array, null
+; CHECK-NEXT: call void @llvm.assume(i1 %notnull)
+; CHECK-NEXT: ret void
+define void @local_assume(i32* %array) {
+  %notnull = icmp ne i32* %array, null
+  call void @llvm.assume(i1 %notnull)
+  br i1 %notnull, label %normal, label %error
+
+normal:
+  ret void
+
+error:
+  store atomic i32 0, i32* @g unordered, align 4
+  ret void
+}
+
+; Check that we do prove a fact using an assume within the block and
+; that we *don't* use circular logic to remove the guard.
+; CHECK-LABEL: @local_guard
+; CHECK: %notnull = icmp ne i32* %array, null
+; CHECK-NEXT: call void (i1, ...) @llvm.experimental.guard(i1 %notnull) [ "deopt"(i32 0) ]
+; CHECK-NEXT: ret void
+
+define void @local_guard(i32* %array) {
+  %notnull = icmp ne i32* %array, null
+  call void (i1, ...) @llvm.experimental.guard(i1 %notnull) ["deopt" (i32 0)]
+  br i1 %notnull, label %normal, label %error
+
+normal:
+  ret void
+
+error:
+  store atomic i32 0, i32* @g unordered, align 4
+  ret void
+}
+
 ; Function Attrs: nounwind
 declare void @llvm.assume(i1) #1
+declare void @llvm.experimental.guard(i1 %cnd, ...)
 
 declare void @bar(...)
 
 declare void @car(...)
 
 attributes #0 = { nounwind uwtable }
 attributes #1 = { nounwind }