[SimplifyCFG] Allow speculating block containing assume()

[nikic]

Summary

SpeculativelyExecuteBB(), which converts a branch + phi structure into a select, currently bails out if the block contains an assume (because it is not speculatable).

Adjust the fold to ignore ephemeral values (i.e. assumes and values only used in assumes) for cost modelling purposes, and drop them when performing the fold.

Theoretically, we could try to preserve the assume information by generating a assume(br_cond || assume_cond) style assume, but this is very unlikely to to be useful (because we don't do anything useful with assumes of this form) and it would make things substantially more complicated once we take operand bundle assumes into account (which don't really support a || operation). I'd prefer not to do that without good motivation.

[nikic]

llvm/test/Transforms/SimplifyCFG/two-entry-phi-fold-crash.ll

17

A bit non-obvious why this is affected, as there are no assumes here: The load result is unused, so this is effectively an ephemeral value as well.

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

 public:
   bool track(const Instruction *I) {
     if (isEphemeral(I)) {
       EphValues.insert(I);
       return true;
     }
     return false;
   }
+
+  bool contains(const Instruction *I) const { return EphValues.contains(I); }
 };
 } // namespace
 
 /// Determine if we can hoist sink a sole store instruction out of a
 /// conditional block.
 ///
 /// We are looking for code like the following:
 ///   BrBB:
bool SimplifyCFGOpt::SpeculativelyExecuteBB(BranchInst *BI, BasicBlock *ThenBB,
   // - All of their uses are in ThenBB.
   SmallDenseMap<Instruction *, unsigned, 4> SinkCandidateUseCounts;
 
   SmallVector<Instruction *, 4> SpeculatedDbgIntrinsics;
 
   unsigned SpeculatedInstructions = 0;
   Value *SpeculatedStoreValue = nullptr;
   StoreInst *SpeculatedStore = nullptr;
-  for (Instruction &I : drop_end(*ThenBB)) {
+  EphemeralValueTracker EphTracker;
+  for (Instruction &I : reverse(drop_end(*ThenBB))) {
     // Skip debug info.
     if (isa<DbgInfoIntrinsic>(I)) {
       SpeculatedDbgIntrinsics.push_back(&I);
       continue;
     }
 
     // Skip pseudo probes. The consequence is we lose track of the branch
     // probability for ThenBB, which is fine since the optimization here takes
     // place regardless of the branch probability.
     if (isa<PseudoProbeInst>(I)) {
       // The probe should be deleted so that it will not be over-counted when
       // the samples collected on the non-conditional path are counted towards
       // the conditional path. We leave it for the counts inference algorithm to
       // figure out a proper count for an unknown probe.
       SpeculatedDbgIntrinsics.push_back(&I);
       continue;
     }
 
+    // Ignore ephemeral values, they will be dropped by the transform.
+    if (EphTracker.track(&I))
+      continue;
+
     // Only speculatively execute a single instruction (not counting the
     // terminator) for now.
     ++SpeculatedInstructions;
     if (SpeculatedInstructions > 1)
       return false;
 
     // Don't hoist the instruction if it's unsafe or expensive.
     if (!isSafeToSpeculativelyExecute(&I) &&
bool SimplifyCFGOpt::SpeculativelyExecuteBB(BranchInst *BI, BasicBlock *ThenBB,
   }
 
   // Metadata can be dependent on the condition we are hoisting above.
   // Conservatively strip all metadata on the instruction. Drop the debug loc
   // to avoid making it appear as if the condition is a constant, which would
   // be misleading while debugging.
   // Similarly strip attributes that maybe dependent on condition we are
   // hoisting above.
-  for (auto &I : *ThenBB) {
+  for (auto &I : make_early_inc_range(*ThenBB)) {
     if (!SpeculatedStoreValue || &I != SpeculatedStore)
       I.setDebugLoc(DebugLoc());
     I.dropUndefImplyingAttrsAndUnknownMetadata();
+
+    // Drop ephemeral values.
+    if (EphTracker.contains(&I)) {
+      I.replaceAllUsesWith(PoisonValue::get(I.getType()));
+      I.eraseFromParent();
+    }
   }
 
   // Hoist the instructions.
   BB->getInstList().splice(BI->getIterator(), ThenBB->getInstList(),
                            ThenBB->begin(), std::prev(ThenBB->end()));
 
   // Insert selects and rewrite the PHI operands.
   IRBuilder<NoFolder> Builder(BI);

llvm/test/CodeGen/ARM/2010-05-18-PostIndexBug.ll

-; RUN: llc < %s -mtriple=armv7-apple-darwin   | FileCheck %s -check-prefix=ARM
-; RUN: llc < %s -mtriple=thumbv7-apple-darwin | FileCheck %s -check-prefix=THUMB
+; RUN: llc < %s -mtriple=armv7-apple-darwin -arm-atomic-cfg-tidy=0 | FileCheck %s -check-prefix=ARM
+; RUN: llc < %s -mtriple=thumbv7-apple-darwin -arm-atomic-cfg-tidy=0 | FileCheck %s -check-prefix=THUMB
 ; rdar://7998649
 
 %struct.foo = type { i64, i64 }
 
 define zeroext i8 @t(%struct.foo* %this, i1 %tst) noreturn optsize {
 entry:
 ; ARM-LABEL:       t:
 ; ARM-DAG:       mov r[[ADDR:[0-9]+]], #8

llvm/test/Transforms/SimplifyCFG/assume.ll

;
   call void @llvm.assume(i1 undef)
   ret void
 
 }
 
 define i32 @speculate_block_with_assume_basic(i1 %c, i32 %x) {
 ; CHECK-LABEL: @speculate_block_with_assume_basic(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF:%.*]], label [[JOIN:%.*]]
-; CHECK:       if:
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[X:%.*]], 0
-; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
-; CHECK-NEXT:    br label [[JOIN]]
-; CHECK:       join:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ 1, [[IF]] ]
-; CHECK-NEXT:    ret i32 [[PHI]]
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[C:%.*]], i32 1, i32 0
+; CHECK-NEXT:    ret i32 [[SPEC_SELECT]]
 ;
 entry:
   br i1 %c, label %if, label %join
 
 if:
   %cmp = icmp ne i32 %x, 0
   call void @llvm.assume(i1 %cmp)
   br label %join
 
 join:
   %phi = phi i32 [ 0, %entry ], [ 1, %if ]
   ret i32 %phi
 }
 
 define i32 @speculate_block_with_assume_extra_instr(i1 %c, i32 %x) {
 ; CHECK-LABEL: @speculate_block_with_assume_extra_instr(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF:%.*]], label [[JOIN:%.*]]
-; CHECK:       if:
 ; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[X:%.*]], 1
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[ADD]], 0
-; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
-; CHECK-NEXT:    br label [[JOIN]]
-; CHECK:       join:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[ADD]], [[IF]] ]
-; CHECK-NEXT:    ret i32 [[PHI]]
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[C:%.*]], i32 [[ADD]], i32 0
+; CHECK-NEXT:    ret i32 [[SPEC_SELECT]]
 ;
 entry:
   br i1 %c, label %if, label %join
 
 if:
   %add = add i32 %x, 1
   %cmp = icmp ne i32 %add, 0
   call void @llvm.assume(i1 %cmp)
   br label %join
 
 join:
   %phi = phi i32 [ 0, %entry ], [ %add, %if ]
   ret i32 %phi
 }
 
+; We only allow speculating one instruction. Here %add and %add2 are used by
+; the assume, but not ephemeral, because they are also used by %phi.
 define i32 @speculate_block_with_assume_extra_instrs_too_many(i1 %c, i32 %x) {
 ; CHECK-LABEL: @speculate_block_with_assume_extra_instrs_too_many(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF:%.*]], label [[JOIN:%.*]]
 ; CHECK:       if:
 ; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[X:%.*]], 1
 ; CHECK-NEXT:    [[ADD2:%.*]] = add i32 [[ADD]], 1
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[ADD2]], 0

 join:
   %phi = phi i32 [ 0, %entry ], [ %add2, %if ]
   ret i32 %phi
 }
 
 define i32 @speculate_block_with_assume_extra_instrs_okay(i1 %c, i32 %x) {
 ; CHECK-LABEL: @speculate_block_with_assume_extra_instrs_okay(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF:%.*]], label [[JOIN:%.*]]
-; CHECK:       if:
 ; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[X:%.*]], 1
-; CHECK-NEXT:    [[ADD2:%.*]] = add i32 [[ADD]], 1
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[ADD2]], 0
-; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
-; CHECK-NEXT:    br label [[JOIN]]
-; CHECK:       join:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[ADD]], [[IF]] ]
-; CHECK-NEXT:    ret i32 [[PHI]]
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[C:%.*]], i32 [[ADD]], i32 0
+; CHECK-NEXT:    ret i32 [[SPEC_SELECT]]
 ;
 entry:
   br i1 %c, label %if, label %join
 
 if:
   %add = add i32 %x, 1
   %add2 = add i32 %add, 1
   %cmp = icmp ne i32 %add2, 0
   call void @llvm.assume(i1 %cmp)
   br label %join
 
 join:
   %phi = phi i32 [ 0, %entry ], [ %add, %if ]
   ret i32 %phi
 }
 
 define i32 @speculate_block_with_assume_operand_bundle(i1 %c, ptr %p) {
 ; CHECK-LABEL: @speculate_block_with_assume_operand_bundle(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF:%.*]], label [[JOIN:%.*]]
-; CHECK:       if:
-; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr [[P:%.*]]) ]
-; CHECK-NEXT:    br label [[JOIN]]
-; CHECK:       join:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ 1, [[IF]] ]
-; CHECK-NEXT:    ret i32 [[PHI]]
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[C:%.*]], i32 1, i32 0
+; CHECK-NEXT:    ret i32 [[SPEC_SELECT]]
 ;
 entry:
   br i1 %c, label %if, label %join
 
 if:
   call void @llvm.assume(i1 true) ["nonnull"(ptr %p)]
   br label %join
 
 join:
   %phi = phi i32 [ 0, %entry ], [ 1, %if ]
   ret i32 %phi
 }
 
 declare void @llvm.assume(i1) nounwind
 

llvm/test/Transforms/SimplifyCFG/two-entry-phi-fold-crash.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -S -simplifycfg -simplifycfg-require-and-preserve-domtree=1 2>&1 | FileCheck %s
 
 declare i32 @blam(i8*, i32)
 
 define i32 @wibble(i8* %arg, i8** %arg1) {
 ; CHECK-LABEL: @wibble(
 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    br label [[BB2:%.*]]
 ; CHECK:       bb2:
-; CHECK-NEXT:    [[BORG:%.*]] = phi i32 [ 0, [[BB:%.*]] ], [ [[BORG]], [[BB8:%.*]] ]
-; CHECK-NEXT:    [[BORG3:%.*]] = phi i32 [ 8, [[BB]] ], [ [[BORG10:%.*]], [[BB8]] ]
+; CHECK-NEXT:    [[BORG:%.*]] = phi i32 [ 0, [[BB:%.*]] ], [ [[BORG]], [[BB2]] ]
+; CHECK-NEXT:    [[BORG3:%.*]] = phi i32 [ 8, [[BB]] ], [ [[SPEC_SELECT:%.*]], [[BB2]] ]
 ; CHECK-NEXT:    [[BORG4:%.*]] = tail call i32 @blam(i8* [[ARG:%.*]], i32 [[BORG]])
 ; CHECK-NEXT:    [[BORG5:%.*]] = icmp eq i32 [[BORG4]], 0
-; CHECK-NEXT:    br i1 [[BORG5]], label [[BB8]], label [[BB6:%.*]]
-; CHECK:       bb6:
-; CHECK-NEXT:    [[BORG7:%.*]] = load i8*, i8** [[ARG1:%.*]], align 4
-; CHECK-NEXT:    br label [[BB8]]
-; CHECK:       bb8:
-; CHECK-NEXT:    [[BORG10]] = phi i32 [ [[BORG4]], [[BB6]] ], [ [[BORG3]], [[BB2]] ]
+; CHECK-NEXT:    [[SPEC_SELECT]] = select i1 [[BORG5]], i32 [[BORG3]], i32 [[BORG4]]
 ; CHECK-NEXT:    [[BORG11:%.*]] = icmp ult i32 [[BORG]], 2
 ; CHECK-NEXT:    br i1 [[BORG11]], label [[BB2]], label [[BB12:%.*]]
 ; CHECK:       bb12:
 ; CHECK-NEXT:    ret i32 1
 ;
 bb:
   br label %bb2