[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 ↗

(On Diff #82996)

(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 ↗

(On Diff #82996)

Why not auto instead of Value::use_iterator?

1777 ↗

(On Diff #82996)

Stray newline?

1809 ↗

(On Diff #82996)

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.

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;
   }
 

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 }