[SimpleLoopUnswitch] partial unswitch needs to be careful when replacing invariants with constants

[fedor.sergeev]

Summary

When partial unswitch operates on multiple conditions at once, .e.g:

if (Cond1 || Cond2 || NonInv) ...

it should infer (and replace) values for individual conditions only on one
side of unswitch and not another.

More precisely only these derivations hold true:

(Cond1 || Cond2) == false  =>  Cond1 == Cond2 == false
(Cond1 && Cond2) == true   =>  Cond1 == Cond2 == true

By the way we organize unswitching it means only replacing on "continue" blocks
and never on "unswitched" ones. Since trivial unswitch does not have "unswitched"
blocks it does not have this problem.

Fixes PR 39568.

[asbirlea]

Great catch, thank you for fixing this and for the additional tests.

[chandlerc]

Could you not introduce all the separate test files?

If the main test (which you did enhance) still doesn't cover the cases, just add more variants there?

[chandlerc]

In D54211#1290435, @chandlerc wrote:

Could you not introduce all the separate test files?

I was editting this comemnt and end up editting it into something that reads passive aggressively and not very politely. Really sorry about that, it was not at all my intention.

I was trying to much more genuinely ask "is there a problem with doing this that I'm not seeing?" not being rhetorical or hyperbolic.

Anyways, just felt bad on re-reading my comment and seeing that it ended up very different from my intent / mental process.

[fedor.sergeev]

In D54211#1290438, @chandlerc wrote:

In D54211#1290435, @chandlerc wrote:

Could you not introduce all the separate test files?

I was trying to much more genuinely ask "is there a problem with doing this that I'm not seeing?" not being rhetorical or hyperbolic.

No prob :) I tend to be rather direct in such cases myself, sometimes getting into the trouble...

Honestly, I first did my tests and then figured out that original nontrivial tests do catch 'AND' case.
So yes, with my addition of 'OR' case there in test27 I can freely dump new nontrivial-unswitch test.
And as I have just discovered trivial-unswitch version is covered by existing trivial-unswitch.ll test.

So I'll keep the change in existing tests and dump the new ones.

[fedor.sergeev]

removing the unnecessary tests

[chandlerc]

LGTM too (but always feel free to submit w/ Alina's review)

lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

static void unswitchNontrivialInvariants(
   if (BI) {
     // If we unswitched a branch which collapses the condition to a known
     // constant we want to replace all the uses of the invariants within both
     // the original and cloned blocks. We do this here so that we can use the
     // now updated dominator tree to identify which side the users are on.
     assert(UnswitchedSuccBBs.size() == 1 &&
            "Only one possible unswitched block for a branch!");
     BasicBlock *ClonedPH = ClonedPHs.begin()->second;
+
+    // When considering multiple partially-unswitched invariants
+    // we cant just go replace them with constants in both branches.
+    //
+    // For 'AND' we infer that true branch ("continue") means true
+    // for each invariant operand.
+    // For 'OR' we can infer that false branch ("continue") means false
+    // for each invariant operand.
+    // So it happens that for multiple-partial case we dont replace
+    // in the unswitched branch.
+    bool ReplaceUnswitched = FullUnswitch || (Invariants.size() == 1);
+
     ConstantInt *UnswitchedReplacement =
         Direction ? ConstantInt::getTrue(BI->getContext())
                   : ConstantInt::getFalse(BI->getContext());
     ConstantInt *ContinueReplacement =
         Direction ? ConstantInt::getFalse(BI->getContext())
                   : ConstantInt::getTrue(BI->getContext());
     for (Value *Invariant : Invariants)
       for (auto UI = Invariant->use_begin(), UE = Invariant->use_end();
            UI != UE;) {
         // Grab the use and walk past it so we can clobber it in the use list.
         Use *U = &*UI++;
         Instruction *UserI = dyn_cast<Instruction>(U->getUser());
         if (!UserI)
           continue;
 
         // Replace it with the 'continue' side if in the main loop body, and the
         // unswitched if in the cloned blocks.
         if (DT.dominates(LoopPH, UserI->getParent()))
           U->set(ContinueReplacement);
-        else if (DT.dominates(ClonedPH, UserI->getParent()))
+        else if (ReplaceUnswitched &&
+                 DT.dominates(ClonedPH, UserI->getParent()))
           U->set(UnswitchedReplacement);
       }
   }
 
   // We can change which blocks are exit blocks of all the cloned sibling
   // loops, the current loop, and any parent loops which shared exit blocks
   // with the current loop. As a consequence, we need to re-form LCSSA for
   // them. But we shouldn't need to re-form LCSSA for any child loops.

test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch.ll

 ; The 'loop_b' unswitched loop.
 ;
 ; CHECK:       entry.split.us:
 ; CHECK-NEXT:    br label %loop_begin.us
 ;
 ; CHECK:       loop_begin.us:
 ; CHECK-NEXT:    %[[V1_US:.*]] = load i1, i1* %ptr1
 ; CHECK-NEXT:    %[[V2_US:.*]] = load i1, i1* %ptr2
-; CHECK-NEXT:    %[[AND1_US:.*]] = and i1 %[[V1_US]], false
+; CHECK-NEXT:    %[[AND1_US:.*]] = and i1 %[[V1_US]], %cond1
 ; CHECK-NEXT:    %[[OR1_US:.*]] = or i1 %[[V2_US]], %cond2
 ; CHECK-NEXT:    %[[AND2_US:.*]] = and i1 %[[AND1_US]], %[[OR1_US]]
-; CHECK-NEXT:    %[[AND3_US:.*]] = and i1 %[[AND2_US]], false
+; CHECK-NEXT:    %[[AND3_US:.*]] = and i1 %[[AND2_US]], %cond3
 ; CHECK-NEXT:    br label %loop_b.us
 ;
 ; CHECK:       loop_b.us:
 ; CHECK-NEXT:    call i32 @b()
 ; CHECK-NEXT:    br label %latch.us
 ;
 ; CHECK:       latch.us:
 ; CHECK-NEXT:    %[[V3_US:.*]] = load i1, i1* %ptr3
loop_exit:
   ret i32 0
 ; CHECK:       loop_exit.split:
 ; CHECK-NEXT:    br label %loop_exit
 ;
 ; CHECK:       loop_exit:
 ; CHECK-NEXT:    ret
 }
 
-; Non-trivial unswitching of a switch.
-define i32 @test27(i1* %ptr, i32 %cond) {
+; Non-trivial partial loop unswitching of multiple invariant inputs to an `or`
+; chain. Basically an inverted version of corresponding `and` test (test26).
+define i32 @test27(i1* %ptr1, i1* %ptr2, i1* %ptr3, i1 %cond1, i1 %cond2, i1 %cond3) {
 ; CHECK-LABEL: @test27(
 entry:
   br label %loop_begin
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    %[[INV_OR:.*]] = or i1 %cond3, %cond1
+; CHECK-NEXT:    br i1 %[[INV_OR]], label %entry.split.us, label %entry.split
+
+loop_begin:
+  %v1 = load i1, i1* %ptr1
+  %v2 = load i1, i1* %ptr2
+  %cond_or1 = or i1 %v1, %cond1
+  %cond_and1 = and i1 %v2, %cond2
+  %cond_or2 = or i1 %cond_or1, %cond_and1
+  %cond_or3 = or i1 %cond_or2, %cond3
+  br i1 %cond_or3, label %loop_b, label %loop_a
+; The 'loop_b' unswitched loop.
+;
+; CHECK:       entry.split.us:
+; CHECK-NEXT:    br label %loop_begin.us
+;
+; CHECK:       loop_begin.us:
+; CHECK-NEXT:    %[[V1_US:.*]] = load i1, i1* %ptr1
+; CHECK-NEXT:    %[[V2_US:.*]] = load i1, i1* %ptr2
+; CHECK-NEXT:    %[[OR1_US:.*]] = or i1 %[[V1_US]], %cond1
+; CHECK-NEXT:    %[[AND1_US:.*]] = and i1 %[[V2_US]], %cond2
+; CHECK-NEXT:    %[[OR2_US:.*]] = or i1 %[[OR1_US]], %[[AND1_US]]
+; CHECK-NEXT:    %[[OR3_US:.*]] = or i1 %[[OR2_US]], %cond3
+; CHECK-NEXT:    br label %loop_b.us
+;
+; CHECK:       loop_b.us:
+; CHECK-NEXT:    call i32 @b()
+; CHECK-NEXT:    br label %latch.us
+;
+; CHECK:       latch.us:
+; CHECK-NEXT:    %[[V3_US:.*]] = load i1, i1* %ptr3
+; CHECK-NEXT:    br i1 %[[V3_US]], label %loop_begin.us, label %loop_exit.split.us
+;
+; CHECK:       loop_exit.split.us:
+; CHECK-NEXT:    br label %loop_exit
+
+; The original loop.
+;
+; CHECK:       entry.split:
+; CHECK-NEXT:    br label %loop_begin
+;
+; CHECK:       loop_begin:
+; CHECK-NEXT:    %[[V1:.*]] = load i1, i1* %ptr1
+; CHECK-NEXT:    %[[V2:.*]] = load i1, i1* %ptr2
+; CHECK-NEXT:    %[[OR1:.*]] = or i1 %[[V1]], false
+; CHECK-NEXT:    %[[AND1:.*]] = and i1 %[[V2]], %cond2
+; CHECK-NEXT:    %[[OR2:.*]] = or i1 %[[OR1]], %[[AND1]]
+; CHECK-NEXT:    %[[OR3:.*]] = or i1 %[[OR2]], false
+; CHECK-NEXT:    br i1 %[[OR3]], label %loop_b, label %loop_a
+
+loop_a:
+  call i32 @a()
+  br label %latch
+; CHECK:       loop_a:
+; CHECK-NEXT:    call i32 @a()
+; CHECK-NEXT:    br label %latch
+
+loop_b:
+  call i32 @b()
+  br label %latch
+; CHECK:       loop_b:
+; CHECK-NEXT:    call i32 @b()
+; CHECK-NEXT:    br label %latch
+
+latch:
+  %v3 = load i1, i1* %ptr3
+  br i1 %v3, label %loop_begin, label %loop_exit
+; CHECK:       latch:
+; CHECK-NEXT:    %[[V3:.*]] = load i1, i1* %ptr3
+; CHECK-NEXT:    br i1 %[[V3]], label %loop_begin, label %loop_exit.split
+
+loop_exit:
+  ret i32 0
+; CHECK:       loop_exit.split:
+; CHECK-NEXT:    br label %loop_exit
+;
+; CHECK:       loop_exit:
+; CHECK-NEXT:    ret
+}
+
+; Non-trivial unswitching of a switch.
+define i32 @test28(i1* %ptr, i32 %cond) {
+; CHECK-LABEL: @test28(
+entry:
+  br label %loop_begin
+; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    switch i32 %cond, label %[[ENTRY_SPLIT_LATCH:.*]] [
 ; CHECK-NEXT:      i32 0, label %[[ENTRY_SPLIT_A:.*]]
 ; CHECK-NEXT:      i32 1, label %[[ENTRY_SPLIT_B:.*]]
 ; CHECK-NEXT:      i32 2, label %[[ENTRY_SPLIT_C:.*]]
 ; CHECK-NEXT:    ]
 
 loop_begin:
   switch i32 %cond, label %latch [

 ; CHECK:       loop_exit:
 ; CHECK-NEXT:    ret i32 0
 }
 
 ; A test case designed to exercise unusual properties of switches: they
 ; can introduce multiple edges to successors. These need lots of special case
 ; handling as they get collapsed in many cases (domtree, the unswitch itself)
 ; but not in all cases (the PHI node operands).
-define i32 @test28(i32 %arg) {
-; CHECK-LABEL: @test28(
+define i32 @test29(i32 %arg) {
+; CHECK-LABEL: @test29(
 entry:
   br label %header
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    switch i32 %arg, label %[[ENTRY_SPLIT_C:.*]] [
 ; CHECK-NEXT:      i32 0, label %[[ENTRY_SPLIT_A:.*]]
 ; CHECK-NEXT:      i32 1, label %[[ENTRY_SPLIT_A]]
 ; CHECK-NEXT:      i32 2, label %[[ENTRY_SPLIT_B:.*]]
 ; CHECK-NEXT:      i32 3, label %[[ENTRY_SPLIT_C]]

 ; CHECK-NEXT:    %[[EXIT_PHI1:.*]] = phi i32 [ %[[LCSSA_C]], %[[LOOP_EXIT_C]] ], [ %[[LCSSA_B]], %[[LOOP_EXIT_B]] ]
 ; CHECK-NEXT:    br label %exit
 
 ; CHECK:       exit:
 ; CHECK-NEXT:    %[[EXIT_PHI2:.*]] = phi i32 [ %[[EXIT_PHI1]], %[[EXIT_SPLIT]] ], [ %[[LCSSA_A]], %[[LOOP_EXIT_A]] ]
 ; CHECK-NEXT:    ret i32 %[[EXIT_PHI2]]
 }
 
-; Similar to @test28 but designed to have one of the duplicate edges be
+; Similar to @test29 but designed to have one of the duplicate edges be
 ; a loop exit edge as those can in some cases be special. Among other things,
 ; this includes an LCSSA phi with multiple entries despite being a dedicated
 ; exit block.
-define i32 @test29(i32 %arg) {
-; CHECK-LABEL: define i32 @test29(
+define i32 @test30(i32 %arg) {
+; CHECK-LABEL: define i32 @test30(
 entry:
   br label %header
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    switch i32 %arg, label %[[ENTRY_SPLIT_EXIT:.*]] [
 ; CHECK-NEXT:      i32 -1, label %[[ENTRY_SPLIT_EXIT]]
 ; CHECK-NEXT:      i32 0, label %[[ENTRY_SPLIT_A:.*]]
 ; CHECK-NEXT:      i32 1, label %[[ENTRY_SPLIT_B:.*]]
 ; CHECK-NEXT:      i32 2, label %[[ENTRY_SPLIT_B]]

 ; rest of the switch cannot be unswitched because its cost is too high.
 ; However, the unswitching of the trivial edge creates a new switch in the
 ; outer loop. *This* switch isn't trivial, but has a low cost to unswitch. When
 ; we unswitch this switch from the outer loop, we will remove it completely and
 ; create a clone of the inner loop on one side. This clone will then again be
 ; viable for unswitching the inner-most loop. This lets us check that the
 ; unswitching doesn't end up cycling infinitely even when the cycle is
 ; indirect and due to revisiting a loop after cloning.
-define void @test30(i32 %arg) {
-; CHECK-LABEL: define void @test30(
+define void @test31(i32 %arg) {
+; CHECK-LABEL: define void @test31(
 entry:
   br label %outer.header
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    switch i32 %arg, label %[[ENTRY_SPLIT:.*]] [
 ; CHECK-NEXT:      i32 1, label %[[ENTRY_SPLIT_US:.*]]
 ; CHECK-NEXT:      i32 2, label %[[ENTRY_SPLIT_US]]
 ; CHECK-NEXT:    ]
 ;