[EarlyCSE] Mark the condition of assume intrinsic as true

[mkazantsev]

Summary

EarlyCSE should not just ignore assumes. It should use the fact that its condition is true for all dominated instructions.

[sanjoy]

lgtm

lib/Transforms/Scalar/EarlyCSE.cpp

638 ↗

(On Diff #96525)

I think you can do:

Instruction *CondI;
if (match(Inst, m_Intrinsic<Intrinsic::assume>(m_Instruction(CondI))) &&
    SimpleValue::canHandle(CondI))
  AvailableValues.insert(CondI, ConstantInt::getTrue(BB->getContext()));

[• dberlin]

lib/Transforms/Scalar/EarlyCSE.cpp

638 ↗

(On Diff #96525)

Can you please also just add a test that llvm.assume(false/true) is skipped?

645 ↗

(On Diff #96525)

You aren't skipping it, you should say that :)

[mkazantsev]

Added some tests, addressed comments.

lib/Transforms/Scalar/EarlyCSE.cpp

638 ↗

(On Diff #96525)

@sanjoy , no, because we need to skip assume even if its arg is not an instruction. Added test18 on that.
@dberlin , added tests 16 and 17 on that.

[reames]

LGTM

I think we should constant fold the condition as well as infer it, but that can be a separate patch.

llvm/trunk/lib/Transforms/Scalar/EarlyCSE.cpp

if (isInstructionTriviallyDead(Inst, &TLI)) {
       Inst->eraseFromParent();
       Changed = true;
       ++NumSimplify;
       continue;
     }
 
     // Skip assume intrinsics, they don't really have side effects (although
     // they're marked as such to ensure preservation of control dependencies),
-    // and this pass will not disturb any of the assumption's control
-    // dependencies.
+    // and this pass will not bother with its removal. However, we should mark
+    // its condition as true for all dominated blocks.
     if (match(Inst, m_Intrinsic<Intrinsic::assume>())) {
+      auto *CondI =
+          dyn_cast<Instruction>(cast<CallInst>(Inst)->getArgOperand(0));
+      if (CondI && SimpleValue::canHandle(CondI)) {
+        DEBUG(dbgs() << "EarlyCSE considering assumption: " << *Inst << '\n');
+        AvailableValues.insert(CondI, ConstantInt::getTrue(BB->getContext()));
+      } else
         DEBUG(dbgs() << "EarlyCSE skipping assumption: " << *Inst << '\n');
       continue;
     }
 
     // Skip invariant.start intrinsics since they only read memory, and we can
     // forward values across it. Also, we dont need to consume the last store
     // since the semantics of invariant.start allow us to perform DSE of the
     // last store, if there was a store following invariant.start. Consider:
     //

llvm/trunk/test/Transforms/EarlyCSE/guards.ll

 ; RUN: opt -S -early-cse < %s | FileCheck %s
 ; RUN: opt < %s -S -basicaa -early-cse-memssa | FileCheck %s
 
 declare void @llvm.experimental.guard(i1,...)
 
+declare void @llvm.assume(i1)
+
 define i32 @test0(i32* %ptr, i1 %cond) {
 ; We can do store to load forwarding over a guard, since it does not
 ; clobber memory
 
 ; CHECK-LABEL: @test0(
 ; CHECK-NEXT:  store i32 40, i32* %ptr
 ; CHECK-NEXT:  call void (i1, ...) @llvm.experimental.guard(i1 %cond) [ "deopt"() ]
 ; CHECK-NEXT:  ret i32 40

 
 if.false:
   call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
   br label %merge
 
 merge:
   ret void
 }
+
+define void @test12(i32 %a, i32 %b) {
+; Check that the assume marks its condition as being true (and thus allows to
+; eliminate the dominated guards).
+
+; CHECK-LABEL: @test12(
+; CHECK-NEXT:  %cmp = icmp eq i32 %a, %b
+; CHECK-NEXT:  call void @llvm.assume(i1 %cmp)
+; CHECK-NEXT:  ret void
+
+  %cmp = icmp eq i32 %a, %b
+  call void @llvm.assume(i1 %cmp)
+  call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+  call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+  call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+  ret void
+}
+
+define void @test13(i32 %a, i32 %b, i32* %ptr) {
+; Check that we deal correctly with stores when removing guards due to assume.
+
+; CHECK-LABEL: @test13(
+; CHECK-NEXT:  %cmp = icmp eq i32 %a, %b
+; CHECK-NEXT:  call void @llvm.assume(i1 %cmp)
+; CHECK-NEXT:  store i32 400, i32* %ptr
+; CHECK-NEXT:  ret void
+
+  %cmp = icmp eq i32 %a, %b
+  call void @llvm.assume(i1 %cmp)
+  store i32 100, i32* %ptr
+  call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+  store i32 200, i32* %ptr
+  call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+  store i32 300, i32* %ptr
+  call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+  store i32 400, i32* %ptr
+  ret void
+}
+
+define void @test14(i32 %a, i32 %b, i1 %c, i32* %ptr) {
+; Similar to test13, but with more control flow.
+; TODO: Can we get rid of the store in the end of entry given that it is
+; post-dominated by other stores?
+
+; CHECK-LABEL: @test14(
+; CHECK:       entry:
+; CHECK-NEXT:    %cmp = icmp eq i32 %a, %b
+; CHECK-NEXT:    call void @llvm.assume(i1 %cmp)
+; CHECK-NEXT:    store i32 400, i32* %ptr
+; CHECK-NEXT:    br i1 %c, label %if.true, label %if.false
+; CHECK:       if.true:
+; CHECK-NEXT:    store i32 500, i32* %ptr
+; CHECK-NEXT:    br label %merge
+; CHECK:       if.false:
+; CHECK-NEXT:    store i32 600, i32* %ptr
+; CHECK-NEXT:    br label %merge
+; CHECK:       merge:
+; CHECK-NEXT:    ret void
+
+entry:
+  %cmp = icmp eq i32 %a, %b
+  call void @llvm.assume(i1 %cmp)
+  store i32 100, i32* %ptr
+  call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+  store i32 200, i32* %ptr
+  call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+  store i32 300, i32* %ptr
+  call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+  store i32 400, i32* %ptr
+  br i1 %c, label %if.true, label %if.false
+
+if.true:
+  call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+  store i32 500, i32* %ptr
+  br label %merge
+
+if.false:
+  call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+  store i32 600, i32* %ptr
+  br label %merge
+
+merge:
+  ret void
+}
+
+define void @test15(i32 %a, i32 %b, i1 %c, i32* %ptr) {
+; Make sure that non-dominating assumes do not cause guards removal.
+
+; CHECK-LABEL: @test15(
+; CHECK:       entry:
+; CHECK-NEXT:    %cmp = icmp eq i32 %a, %b
+; CHECK-NEXT:    br i1 %c, label %if.true, label %if.false
+; CHECK:       if.true:
+; CHECK-NEXT:    call void @llvm.assume(i1 %cmp)
+; CHECK-NEXT:    store i32 100, i32* %ptr
+; CHECK-NEXT:    br label %merge
+; CHECK:       if.false:
+; CHECK-NEXT:    store i32 200, i32* %ptr
+; CHECK-NEXT:    br label %merge
+; CHECK:       merge:
+; CHECK-NEXT:    store i32 300, i32* %ptr
+; CHECK-NEXT:    call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+; CHECK-NEXT:    store i32 400, i32* %ptr
+; CHECK-NEXT:    ret void
+
+entry:
+  %cmp = icmp eq i32 %a, %b
+  br i1 %c, label %if.true, label %if.false
+
+if.true:
+  call void @llvm.assume(i1 %cmp)
+  store i32 100, i32* %ptr
+  call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+  br label %merge
+
+if.false:
+  store i32 200, i32* %ptr
+  br label %merge
+
+merge:
+  store i32 300, i32* %ptr
+  call void (i1, ...) @llvm.experimental.guard(i1 %cmp) [ "deopt"() ]
+  store i32 400, i32* %ptr
+  ret void
+}
+
+define void @test16(i32 %a, i32 %b) {
+; Check that we don't bother to do anything with assumes even if we know the
+; condition being true.
+
+; CHECK-LABEL: @test16(
+; CHECK-NEXT:    %cmp = icmp eq i32 %a, %b
+; CHECK-NEXT:    call void @llvm.assume(i1 %cmp)
+; CHECK-NEXT:    call void @llvm.assume(i1 %cmp)
+; CHECK-NEXT:    ret void
+
+  %cmp = icmp eq i32 %a, %b
+  call void @llvm.assume(i1 %cmp)
+  call void @llvm.assume(i1 %cmp)
+  ret void
+}
+
+define void @test17(i32 %a, i32 %b, i1 %c, i32* %ptr) {
+; Check that we don't bother to do anything with assumes even if we know the
+; condition being true or false (includes come control flow).
+
+; CHECK-LABEL: @test17(
+; CHECK:       entry:
+; CHECK-NEXT:    %cmp = icmp eq i32 %a, %b
+; CHECK-NEXT:    br i1 %c, label %if.true, label %if.false
+; CHECK:       if.true:
+; CHECK-NEXT:    call void @llvm.assume(i1 %cmp)
+; CHECK-NEXT:    br label %merge
+; CHECK:       if.false:
+; CHECK-NEXT:    call void @llvm.assume(i1 %cmp)
+; CHECK-NEXT:    br label %merge
+; CHECK:       merge:
+; CHECK-NEXT:    ret void
+
+entry:
+  %cmp = icmp eq i32 %a, %b
+  br i1 %c, label %if.true, label %if.false
+
+if.true:
+  call void @llvm.assume(i1 %cmp)
+  br label %merge
+
+if.false:
+  call void @llvm.assume(i1 %cmp)
+  br label %merge
+
+merge:
+  ret void
+}
+
+define void @test18(i1 %c) {
+; Check that we don't bother to do anything with assumes even if we know the
+; condition being true and not being an instruction.
+
+; CHECK-LABEL: @test18(
+; CHECK-NEXT:    call void @llvm.assume(i1 %c)
+; CHECK-NEXT:    call void @llvm.assume(i1 %c)
+; CHECK-NEXT:    ret void
+
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  ret void
+}