[BDCE] clear poison generators after turning a value into zero (PR33695, PR34037)

[spatel]

Summary

Since we have a known miscompile danger here, I'll propose a patch for people to look at.
This is the most basic (yet still incomplete) solution that I could come up with for the miscompiles in:
https://bugs.llvm.org/show_bug.cgi?id=33695
https://bugs.llvm.org/show_bug.cgi?id=34037

As discussed in PR33695, we need to do more than this, and there are probably better implementations, but this is a safe first step to close the hole?

[majnemer]

lib/Transforms/Scalar/BDCE.cpp

56

I recommend using a worklist so that this doesn't blow up the stack.

[hfinkel]

lib/Transforms/Scalar/BDCE.cpp

50

Just to make this a little better: if any of the users are a store, call, invoke, atomicrmw, atomiccmpxchg, or return - anything that will make the value externally observable, we can stop recursing to clear the flags.

54

I can't think of any current metadata that would be invalidated, but assumes might indeed be.

[davide]

I concur with David that this should use a worklist.

[spatel]

Patch updated:

1. Use a loop (worklist) rather than recursion to avoid blowing up the stack.
2. Add an exclusion list based on opcodes to stop the search.
3. Add a test to show that the exclusion list works for a call inst.
4. Add a TODO comment with Nuno's suggestion from the bug report to make this better.

[hfinkel]

lib/Transforms/Scalar/BDCE.cpp

46

I think this is better than the opcode list I suggested. Why don't we just do this instead? Isn't it just something like this?

auto isExternallyVisible = [&](Instruction *I) {
  // Any value we're trivializing should be an integer value, and moreover, any conversion between an integer value and a non-integer value should demand all of the bits, which will cause us to stop looking down the use/def chain, so we should only see integer-typed instructions here.
  assert(I->getType()->isIntegerTy() && "Trivializing a non-integer value?");
  return DB.getDemandedBits(I).isAllOnesValue();
}

80

Range metadata can't be invalid here because it applies only to memory accesses (which demand all bits). As a result, I'd like to not see it in a TODO, although a comment explaining why it's not an issue would be good.

80

Please just include the assume check if it is not complicated. I think this will work:

if (auto *II = dyn_cast<IntrinsicInst>(J)) {
  if (II->getIntrinsicID() == Intrinsic::assume)
   // We can delete the assume here because it can't appear in the worklist multiple times (because it only has one operand).
   II->eraseFromParent();
 }

[spatel]

Patch updated - mostly just cutting and pasting the code Hal posted:

1. Make isExternallyVisible() check the calculated demanded bits rather than a list of opcodes.
2. Add a check for and delete llvm.assume.

I've been trying to find a test for the llvm.assume part, but I haven't gotten anywhere, so I'm probably not understanding that case. An assume takes an i1 param, so it always demands all the bits of that param. And therefore, the assume will never be added to the worklist?

[nlopes]

LGTM, with minor comments: missing the CHECK comment above, and missing a test for the assume removal.

test/Transforms/BDCE/invalidate-assumptions.ll

72

Missing CHECK comment here.

[spatel]

Patch updated - sorry, the last diff was not the one I intended to upload. This is NFC from the last rev, but:

1. Includes Hal's comment about range metadata.
2. Has CHECK lines for the poison_on_call_user_is_ok() test.

I'm still not sure how to test the llvm.assume case - any ideas?

[hfinkel]

The comments about assume, otherwise, LGTM too.

lib/Transforms/Scalar/BDCE.cpp

71

I've changed my mind about assume. I believe, now that you're checking the demanded bits, you can't hit this case. The issue is that the assume demands its operand, and it must do this. Essentially, the value of the condition is an observable of the program (i.e. trivializing values can't change it). If there's a compare feeding the assume, for example, it in turn demands whatever bits are necessary in order to ensure its result. As a result, the compare shouldn't imply anything about its operands that trivializing changes (if it did, then the truth of the condition would not have have universally implied that thing in the first place). Thus, I think you can just remove the assume handling.

72

You need a continue after this. You can't then dereference J.

[spatel]

Patch updated:
Replace the assume deletion with a comment about why it can't happen. :)

[hfinkel]

LGTM

lib/Transforms/Scalar/BDCE.cpp

 // This file implements the Bit-Tracking Dead Code Elimination pass. Some
 // instructions (shifts, some ands, ors, etc.) kill some of their input bits.
 // We track these dead bits and remove instructions that compute only these
 // dead bits.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar/BDCE.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/DemandedBits.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "bdce"
 
 STATISTIC(NumRemoved, "Number of instructions removed (unused)");
 STATISTIC(NumSimplified, "Number of instructions trivialized (dead bits)");
 
+/// If an instruction is trivialized (dead), then the chain of users of that
+/// instruction may need to be cleared of assumptions that can no longer be
+/// guaranteed correct.
+static void clearAssumptionsOfUsers(Instruction *I) {
+  // If a value may be visible outside of this def/use chain, then we don't need
+  // to continue processing its users. Ie, all bits of this value should be
+  // demanded.
+  // TODO: "Any value in the def/use chain that has all bits set as non-dead
+  //        (by the demanded bit analysis) cannot have its value changed: if all
+  //        bits are needed, all bits must keep the same value before/after
+  //        optimization. That means that our 'taint' propagation (of removing
+  //        attributes/assumptions) can end there." - Nuno Lopes (PR33695)
+  auto isExternallyVisible = [](Instruction *I) {
+    switch (I->getOpcode()) {
+      case Instruction::Store:
+      case Instruction::Call:
+      case Instruction::Invoke:
+      case Instruction::AtomicRMW:
+      case Instruction::AtomicCmpXchg:
+      case Instruction::Ret:
+        return true;
+      default:
+        return false;
+    }
+  };
+
+  // Initialize the worklist with eligible direct users.
+  SmallVector<Instruction *, 16> WorkList;
+  for (User *JU : I->users()) {
+    auto *J = dyn_cast<Instruction>(JU);
+    if (J && !isExternallyVisible(J))
+      WorkList.push_back(J);
+  }
+
+  // DFS through subsequent users while tracking visits to avoid cycles.
+  SmallPtrSet<Instruction *, 16> Visited;
+  while (!WorkList.empty()) {
+    Instruction *J = WorkList.pop_back_val();
+
+    // NSW, NUW, and exact are based on operands that might have changed.
+    J->dropPoisonGeneratingFlags();
+    // FIXME: llvm.assume (and range metadata?) may also be invalid now.
+
+    Visited.insert(J);
+
+    for (User *KU : J->users()) {
+      auto *K = dyn_cast<Instruction>(KU);
+      if (K && !Visited.count(K) && !isExternallyVisible(K))
+        WorkList.push_back(K);
+    }
+  }
+}
+
 static bool bitTrackingDCE(Function &F, DemandedBits &DB) {
   SmallVector<Instruction*, 128> Worklist;
   bool Changed = false;
   for (Instruction &I : instructions(F)) {
     // If the instruction has side effects and no non-dbg uses,
     // skip it. This way we avoid computing known bits on an instruction
     // that will not help us.
     if (I.mayHaveSideEffects() && I.use_empty())
       continue;
 
     if (I.getType()->isIntegerTy() &&
         !DB.getDemandedBits(&I).getBoolValue()) {
       // For live instructions that have all dead bits, first make them dead by
       // replacing all uses with something else. Then, if they don't need to
       // remain live (because they have side effects, etc.) we can remove them.
       DEBUG(dbgs() << "BDCE: Trivializing: " << I << " (all bits dead)\n");
+
+      clearAssumptionsOfUsers(&I);
+
       // FIXME: In theory we could substitute undef here instead of zero.
       // This should be reconsidered once we settle on the semantics of
       // undef, poison, etc.
       Value *Zero = ConstantInt::get(I.getType(), 0);
       ++NumSimplified;
       I.replaceNonMetadataUsesWith(Zero);
       Changed = true;
     }

test/Transforms/BDCE/invalidate-assumptions.ll

 ; RUN: opt -bdce %s -S | FileCheck %s
 
-; FIXME: The 'nuw' on the subtract allows us to deduce that %setbit is not demanded.
+; The 'nuw' on the subtract allows us to deduce that %setbit is not demanded.
 ; But if we change that value to '0', then the 'nuw' is no longer valid. If we don't
 ; remove the 'nuw', another pass (-instcombine) may make a transform based on an
 ; that incorrect assumption and we can miscompile:
 ; https://bugs.llvm.org/show_bug.cgi?id=33695
 
 define i1 @PR33695(i1 %b, i8 %x) {
 ; CHECK-LABEL: @PR33695(
 ; CHECK-NEXT:    [[SETBIT:%.*]] = or i8 %x, 64
 ; CHECK-NEXT:    [[LITTLE_NUMBER:%.*]] = zext i1 %b to i8
 ; CHECK-NEXT:    [[BIG_NUMBER:%.*]] = shl i8 0, 1
-; CHECK-NEXT:    [[SUB:%.*]] = sub nuw i8 [[BIG_NUMBER]], [[LITTLE_NUMBER]]
+; CHECK-NEXT:    [[SUB:%.*]] = sub i8 [[BIG_NUMBER]], [[LITTLE_NUMBER]]
 ; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i8 [[SUB]] to i1
 ; CHECK-NEXT:    ret i1 [[TRUNC]]
 ;
   %setbit = or i8 %x, 64
   %little_number = zext i1 %b to i8
   %big_number = shl i8 %setbit, 1
   %sub = sub nuw i8 %big_number, %little_number
   %trunc = trunc i8 %sub to i1
   ret i1 %trunc
 }
 
-; FIXME: Similar to above, but now with more no-wrap.
+; Similar to above, but now with more no-wrap.
 ; https://bugs.llvm.org/show_bug.cgi?id=34037
 
 define i64 @PR34037(i64 %m, i32 %r, i64 %j, i1 %b, i32 %k, i64 %p) {
 ; CHECK-LABEL: @PR34037(
 ; CHECK-NEXT:    [[CONV:%.*]] = zext i32 %r to i64
 ; CHECK-NEXT:    [[AND:%.*]] = and i64 %m, 0
 ; CHECK-NEXT:    [[NEG:%.*]] = xor i64 0, 34359738367
 ; CHECK-NEXT:    [[OR:%.*]] = or i64 %j, 0
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i64 0, 29
 ; CHECK-NEXT:    [[CONV1:%.*]] = select i1 %b, i64 7, i64 0
-; CHECK-NEXT:    [[SUB:%.*]] = sub nuw nsw i64 [[SHL]], [[CONV1]]
+; CHECK-NEXT:    [[SUB:%.*]] = sub i64 [[SHL]], [[CONV1]]
 ; CHECK-NEXT:    [[CONV2:%.*]] = zext i32 %k to i64
-; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i64 [[SUB]], [[CONV2]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul i64 [[SUB]], [[CONV2]]
 ; CHECK-NEXT:    [[CONV4:%.*]] = and i64 %p, 65535
 ; CHECK-NEXT:    [[AND5:%.*]] = and i64 [[MUL]], [[CONV4]]
 ; CHECK-NEXT:    ret i64 [[AND5]]
 ;
   %conv = zext i32 %r to i64
   %and = and i64 %m, %conv
   %neg = xor i64 %and, 34359738367
   %or = or i64 %j, %neg
   %shl = shl i64 %or, 29
   %conv1 = select i1 %b, i64 7, i64 0
   %sub = sub nuw nsw i64 %shl, %conv1
   %conv2 = zext i32 %k to i64
   %mul = mul nsw i64 %sub, %conv2
   %conv4 = and i64 %p, 65535
   %and5 = and i64 %mul, %conv4
   ret i64 %and5
 }
 
+
+; This is a manufactured example based on the 1st test to prove that the
+; assumption-killing algorithm stops at the call. Ie, it does not remove
+; nsw/nuw from the 'add' because a call demands all bits of its argument.
+
+declare i1 @foo(i1)
+define i1 @poison_on_call_user_is_ok(i1 %b, i8 %x) {
+; CHECK-LABEL: @poison_on_call_user_is_ok(
+; CHECK-NEXT:    [[SETBIT:%.*]] = or i8 %x, 64
+; CHECK-NEXT:    [[LITTLE_NUMBER:%.*]] = zext i1 %b to i8
+; CHECK-NEXT:    [[BIG_NUMBER:%.*]] = shl i8 0, 1
+; CHECK-NEXT:    [[SUB:%.*]] = sub i8 [[BIG_NUMBER]], [[LITTLE_NUMBER]]
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i8 [[SUB]] to i1
+; CHECK-NEXT:    [[CALL_RESULT:%.*]] = call i1 @foo(i1 [[TRUNC]])
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw i1 [[CALL_RESULT]], true
+; CHECK-NEXT:    [[MUL:%.*]] = mul i1 [[TRUNC]], [[ADD]]
+; CHECK-NEXT:    ret i1 [[MUL]]
+;
+  %setbit = or i8 %x, 64
+  %little_number = zext i1 %b to i8
+  %big_number = shl i8 %setbit, 1
+  %sub = sub nuw i8 %big_number, %little_number
+  %trunc = trunc i8 %sub to i1
+  %call_result = call i1 @foo(i1 %trunc)
+  %add = add nsw nuw i1 %call_result, 1
+  %mul = mul i1 %trunc, %add
+  ret i1 %mul
+}
+