[llvm] 5bb38e8 - [LoopUnswitch] unswitch if cond is in select form of and/or as well

Juneyoung Lee via llvm-commits llvm-commits at lists.llvm.org
Sun Mar 7 08:20:09 PST 2021


Author: Juneyoung Lee
Date: 2021-03-08T01:19:43+09:00
New Revision: 5bb38e84d3d0420e7159ffc34e3dd33ab5088b8e

URL: https://github.com/llvm/llvm-project/commit/5bb38e84d3d0420e7159ffc34e3dd33ab5088b8e
DIFF: https://github.com/llvm/llvm-project/commit/5bb38e84d3d0420e7159ffc34e3dd33ab5088b8e.diff

LOG: [LoopUnswitch] unswitch if cond is in select form of and/or as well

Hello all,
I'm trying to fix unsafe propagation of poison values in and/or conditions by using
equivalent select forms (`select i1 A, i1 B, i1 false` and `select i1 A, i1 true, i1 false`)
instead.
D93065 has links to patches for this.

This patch allows unswitch to happen if the condition is in this form as well.
`collectHomogenousInstGraphLoopInvariants` is updated to keep traversal if
Root and the visiting I matches both m_LogicalOr()/m_LogicalAnd().
Other than this, the remaining changes are almost straightforward and simply replaces
Instruction::And/Or check with match(m_LogicalOr()/m_LogicalAnd()).

Reviewed By: nikic

Differential Revision: https://reviews.llvm.org/D97756

Added: 
    

Modified: 
    llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
    llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch.ll
    llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index ff6de10616f7..0e5a787cea70 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -38,6 +38,7 @@
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/Use.h"
 #include "llvm/IR/Value.h"
 #include "llvm/InitializePasses.h"
@@ -63,6 +64,7 @@
 #define DEBUG_TYPE "simple-loop-unswitch"
 
 using namespace llvm;
+using namespace llvm::PatternMatch;
 
 STATISTIC(NumBranches, "Number of branches unswitched");
 STATISTIC(NumSwitches, "Number of switches unswitched");
@@ -116,6 +118,9 @@ collectHomogenousInstGraphLoopInvariants(Loop &L, Instruction &Root,
          "Only need to walk the graph if root itself is not invariant.");
   TinyPtrVector<Value *> Invariants;
 
+  bool IsRootAnd = match(&Root, m_LogicalAnd());
+  bool IsRootOr  = match(&Root, m_LogicalOr());
+
   // Build a worklist and recurse through operators collecting invariants.
   SmallVector<Instruction *, 4> Worklist;
   SmallPtrSet<Instruction *, 8> Visited;
@@ -136,12 +141,13 @@ collectHomogenousInstGraphLoopInvariants(Loop &L, Instruction &Root,
 
       // If not an instruction with the same opcode, nothing we can do.
       Instruction *OpI = dyn_cast<Instruction>(OpV);
-      if (!OpI || OpI->getOpcode() != Root.getOpcode())
-        continue;
 
-      // Visit this operand.
-      if (Visited.insert(OpI).second)
-        Worklist.push_back(OpI);
+      if (OpI && ((IsRootAnd && match(OpI, m_LogicalAnd())) ||
+                  (IsRootOr  && match(OpI, m_LogicalOr())))) {
+        // Visit this operand.
+        if (Visited.insert(OpI).second)
+          Worklist.push_back(OpI);
+      }
     }
   } while (!Worklist.empty());
 
@@ -416,10 +422,10 @@ static bool unswitchTrivialBranch(Loop &L, BranchInst &BI, DominatorTree &DT,
   // and the condition is a graph of `and` operations.
   if (!FullUnswitch) {
     if (ExitDirection) {
-      if (cast<Instruction>(BI.getCondition())->getOpcode() != Instruction::Or)
+      if (!match(BI.getCondition(), m_LogicalOr()))
         return false;
     } else {
-      if (cast<Instruction>(BI.getCondition())->getOpcode() != Instruction::And)
+      if (!match(BI.getCondition(), m_LogicalAnd()))
         return false;
     }
   }
@@ -497,13 +503,13 @@ static bool unswitchTrivialBranch(Loop &L, BranchInst &BI, DominatorTree &DT,
     // Only unswitching a subset of inputs to the condition, so we will need to
     // build a new branch that merges the invariant inputs.
     if (ExitDirection)
-      assert(cast<Instruction>(BI.getCondition())->getOpcode() ==
-                 Instruction::Or &&
-             "Must have an `or` of `i1`s for the condition!");
+      assert(match(BI.getCondition(), m_LogicalOr()) &&
+             "Must have an `or` of `i1`s or `select i1 X, true, Y`s for the "
+             "condition!");
     else
-      assert(cast<Instruction>(BI.getCondition())->getOpcode() ==
-                 Instruction::And &&
-             "Must have an `and` of `i1`s for the condition!");
+      assert(match(BI.getCondition(), m_LogicalAnd()) &&
+             "Must have an `and` of `i1`s or `select i1 X, Y, false`s for the"
+             " condition!");
     buildPartialUnswitchConditionalBranch(*OldPH, Invariants, ExitDirection,
                                           *UnswitchedBB, *NewPH);
   }
@@ -1989,11 +1995,10 @@ static void unswitchNontrivialInvariants(
   bool Direction = true;
   int ClonedSucc = 0;
   if (!FullUnswitch) {
-    if (cast<Instruction>(BI->getCondition())->getOpcode() != Instruction::Or) {
-      assert(cast<Instruction>(BI->getCondition())->getOpcode() ==
-                 Instruction::And &&
-             "Only `or` and `and` instructions can combine invariants being "
-             "unswitched.");
+    if (!match(BI->getCondition(), m_LogicalOr())) {
+      assert(match(BI->getCondition(), m_LogicalAnd()) &&
+             "Only `or`, `and`, an `select` instructions can combine "
+             "invariants being unswitched.");
       Direction = false;
       ClonedSucc = 1;
     }
@@ -2646,8 +2651,7 @@ unswitchBestCondition(Loop &L, DominatorTree &DT, LoopInfo &LI,
     }
 
     Instruction &CondI = *cast<Instruction>(BI->getCondition());
-    if (CondI.getOpcode() != Instruction::And &&
-      CondI.getOpcode() != Instruction::Or)
+    if (!match(&CondI, m_CombineOr(m_LogicalAnd(), m_LogicalOr())))
       continue;
 
     TinyPtrVector<Value *> Invariants =
@@ -2759,18 +2763,16 @@ unswitchBestCondition(Loop &L, DominatorTree &DT, LoopInfo &LI,
         continue;
 
       // If this is a partial unswitch candidate, then it must be a conditional
-      // branch with a condition of either `or` or `and`. In that case, one of
-      // the successors is necessarily duplicated, so don't even try to remove
-      // its cost.
+      // branch with a condition of either `or`, `and`, or their corresponding
+      // select forms. In that case, one of the successors is necessarily
+      // duplicated, so don't even try to remove its cost.
       if (!FullUnswitch) {
         auto &BI = cast<BranchInst>(TI);
-        if (cast<Instruction>(BI.getCondition())->getOpcode() ==
-            Instruction::And) {
+        if (match(BI.getCondition(), m_LogicalAnd())) {
           if (SuccBB == BI.getSuccessor(1))
             continue;
         } else {
-          assert(cast<Instruction>(BI.getCondition())->getOpcode() ==
-                     Instruction::Or &&
+          assert(match(BI.getCondition(), m_LogicalOr()) &&
                  "Only `and` and `or` conditions can result in a partial "
                  "unswitch!");
           if (SuccBB == BI.getSuccessor(0))

diff  --git a/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch.ll b/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch.ll
index 6fff85a2255c..bfbe3e65daa6 100644
--- a/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch.ll
+++ b/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch.ll
@@ -4201,3 +4201,159 @@ exit:
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 }
+
+; Non-trivial partial loop unswitching of multiple invariant inputs to an `and`
+; chain (select version).
+define i32 @test32(i1* %ptr1, i1* %ptr2, i1* %ptr3, i1 %cond1, i1 %cond2) {
+; CHECK-LABEL: @test32(
+entry:
+  br label %loop_begin
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    %[[INV_AND:.*]] = and i1 %cond2, %cond1
+; CHECK-NEXT:    br i1 %[[INV_AND]], label %entry.split, label %entry.split.us
+
+loop_begin:
+  %v1 = load i1, i1* %ptr1
+  %v2 = load i1, i1* %ptr2
+  %cond_and1 = select i1 %v1, i1 %cond1, i1 false
+  %cond_and2 = select i1 %cond_and1, i1 %cond2, i1 false
+  br i1 %cond_and2, label %loop_a, label %loop_b
+; The 'loop_b' unswitched loop.
+;
+; CHECK:       entry.split.us:
+; CHECK-NEXT:    br label %loop_begin.us
+;
+; CHECK:       loop_begin.us:
+; CHECK-NEXT:    %[[V2_US]] = load i1, i1* %ptr2, align 1
+; 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, align 1
+; 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:    %[[AND1:.*]] = select i1 %[[V1]], i1 true, i1 false
+; CHECK-NEXT:    %[[AND2:.*]] = select i1 %[[AND1]], i1 true, i1 false
+; CHECK-NEXT:    br i1 %[[AND2]], label %loop_a, label %loop_b
+
+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, align 1
+; 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 partial loop unswitching of multiple invariant inputs to an `or`
+; chain (select version).
+define i32 @test33(i1* %ptr1, i1* %ptr2, i1* %ptr3, i1 %cond1, i1 %cond2) {
+; CHECK-LABEL: @test33(
+entry:
+  br label %loop_begin
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    %[[INV_OR:.*]] = or i1 %cond2, %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_and1 = select i1 %v1, i1 true, i1 %cond1
+  %cond_and2 = select i1 %cond_and1, i1 true, i1 %cond2
+  br i1 %cond_and2, 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:    %[[V2_US]] = load i1, i1* %ptr2, align 1
+; 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, align 1
+; 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:    %[[AND1:.*]] = select i1 %[[V1]], i1 true, i1 false
+; CHECK-NEXT:    %[[AND2:.*]] = select i1 %[[AND1]], i1 true, i1 false
+; CHECK-NEXT:    br i1 %[[AND2]], 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, align 1
+; 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
+}

diff  --git a/llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch.ll b/llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch.ll
index 49a6d93d0bef..08706612da55 100644
--- a/llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch.ll
+++ b/llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch.ll
@@ -490,6 +490,99 @@ loop_exit:
 ; CHECK-NEXT:    ret
 }
 
+define i32 @test_partial_condition_unswitch_and_select(i32* %var, i1 %cond1, i1 %cond2) {
+; CHECK-LABEL: @test_partial_condition_unswitch_and_select(
+entry:
+  br label %loop_begin
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 %cond1, label %entry.split, label %loop_exit.split
+;
+; CHECK:       entry.split:
+; CHECK-NEXT:    br i1 %cond2, label %entry.split.split, label %loop_exit
+;
+; CHECK:       entry.split.split:
+; CHECK-NEXT:    br label %loop_begin
+
+loop_begin:
+  br i1 %cond1, label %continue, label %loop_exit
+; CHECK:       loop_begin:
+; CHECK-NEXT:    br label %continue
+
+continue:
+  %var_val = load i32, i32* %var
+  %var_cond = trunc i32 %var_val to i1
+  %cond_and = select i1 %var_cond, i1 %cond2, i1 false
+  br i1 %cond_and, label %do_something, label %loop_exit
+; CHECK:       continue:
+; CHECK-NEXT:    %[[VAR:.*]] = load i32
+; CHECK-NEXT:    %[[VAR_COND:.*]] = trunc i32 %[[VAR]] to i1
+; CHECK-NEXT:    %[[COND_AND:.*]] = select i1 %[[VAR_COND]], i1 true, i1 false
+; CHECK-NEXT:    br i1 %[[COND_AND]], label %do_something, label %loop_exit
+
+do_something:
+  call void @some_func() noreturn nounwind
+  br label %loop_begin
+; CHECK:       do_something:
+; CHECK-NEXT:    call
+; CHECK-NEXT:    br label %loop_begin
+
+loop_exit:
+  ret i32 0
+; CHECK:       loop_exit:
+; CHECK-NEXT:    br label %loop_exit.split
+;
+; CHECK:       loop_exit.split:
+; CHECK-NEXT:    ret
+}
+
+define i32 @test_partial_condition_unswitch_or_simple_select(i32* %var, i1 %cond1, i1 %cond2) {
+; CHECK-LABEL: @test_partial_condition_unswitch_or_simple_select(
+entry:
+  br label %loop_begin
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 %cond1, label %entry.split, label %loop_exit.split
+;
+; CHECK:       entry.split:
+; CHECK-NEXT:    br i1 %cond2, label %loop_exit.split1, label %entry.split.split
+;
+; CHECK:       entry.split.split:
+; CHECK-NEXT:    br label %loop_begin
+
+loop_begin:
+  br i1 %cond1, label %continue, label %loop_exit
+; CHECK:       loop_begin:
+; CHECK-NEXT:    br label %continue
+
+continue:
+  %var_val = load i32, i32* %var
+  %var_cond = trunc i32 %var_val to i1
+  %cond_or = select i1 %var_cond, i1 true, i1 %cond2
+  br i1 %cond_or, label %loop_exit, label %do_something
+; CHECK:       continue:
+; CHECK-NEXT:    %[[VAR:.*]] = load i32
+; CHECK-NEXT:    %[[VAR_COND:.*]] = trunc i32 %[[VAR]] to i1
+; CHECK-NEXT:    %[[COND_OR:.*]] = select i1 %[[VAR_COND]], i1 true, i1 false
+; CHECK-NEXT:    br i1 %[[COND_OR]], label %loop_exit, label %do_something
+
+do_something:
+  call void @some_func() noreturn nounwind
+  br label %loop_begin
+; CHECK:       do_something:
+; CHECK-NEXT:    call
+; CHECK-NEXT:    br label %loop_begin
+
+loop_exit:
+  ret i32 0
+; CHECK:       loop_exit:
+; CHECK-NEXT:    br label %loop_exit.split1
+;
+; CHECK:       loop_exit.split1:
+; CHECK-NEXT:    br label %loop_exit.split
+;
+; CHECK:       loop_exit.split:
+; CHECK-NEXT:    ret
+}
+
 define i32 @test_partial_condition_unswitch_or(i32* %var, i1 %cond1, i1 %cond2, i1 %cond3, i1 %cond4, i1 %cond5, i1 %cond6) {
 ; CHECK-LABEL: @test_partial_condition_unswitch_or(
 entry:
@@ -541,6 +634,63 @@ loop_exit:
 ; CHECK-NEXT:    ret
 }
 
+; Check that loop unswitch looks through a combination of or and select instructions.
+; Note that cond6 can be unswitched because `select i1 %cond_or5, i1 true, i1 false` is
+; both logical-or and logical-and.
+define i32 @test_partial_condition_unswitch_or_select(i32* %var, i1 %cond1, i1 %cond2, i1 %cond3, i1 %cond4, i1 %cond5, i1 %cond6) {
+; CHECK-LABEL: @test_partial_condition_unswitch_or_select(
+entry:
+  br label %loop_begin
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    %[[INV_OR1:.*]] = or i1 %cond4, %cond2
+; CHECK-NEXT:    %[[INV_OR2:.*]] = or i1 %[[INV_OR1]], %cond3
+; CHECK-NEXT:    %[[INV_OR3:.*]] = or i1 %[[INV_OR2]], %cond1
+; CHECK-NEXT:    br i1 %[[INV_OR3]], label %loop_exit.split, label %entry.split
+;
+; CHECK:       entry.split:
+; CHECK-NEXT:    br i1 %cond6, label %loop_exit.split1, label %entry.split.split
+;
+; CHECK:       entry.split.split:
+; CHECK-NEXT:    br label %loop_begin
+
+loop_begin:
+  %var_val = load i32, i32* %var
+  %var_cond = trunc i32 %var_val to i1
+  %cond_or1 = or i1 %var_cond, %cond1
+  %cond_or2 = or i1 %cond2, %cond3
+  %cond_or3 = or i1 %cond_or1, %cond_or2
+  %cond_xor1 = xor i1 %cond5, %var_cond
+  %cond_and1 = and i1 %cond6, %var_cond
+  %cond_or4 = or i1 %cond_xor1, %cond_and1
+  %cond_or5 = select i1 %cond_or3, i1 true, i1 %cond_or4
+  %cond_or6 = select i1 %cond_or5, i1 true, i1 %cond4
+  br i1 %cond_or6, label %loop_exit, label %do_something
+; CHECK:       loop_begin:
+; CHECK-NEXT:    %[[VAR:.*]] = load i32
+; CHECK-NEXT:    %[[VAR_COND:.*]] = trunc i32 %[[VAR]] to i1
+; CHECK-NEXT:    %[[COND_OR1:.*]] = or i1 %[[VAR_COND]], false
+; CHECK-NEXT:    %[[COND_OR2:.*]] = or i1 false, false
+; CHECK-NEXT:    %[[COND_OR3:.*]] = or i1 %[[COND_OR1]], %[[COND_OR2]]
+; CHECK-NEXT:    %[[COND_XOR:.*]] = xor i1 %cond5, %[[VAR_COND]]
+; CHECK-NEXT:    %[[COND_AND:.*]] = and i1 false, %[[VAR_COND]]
+; CHECK-NEXT:    %[[COND_OR4:.*]] = or i1 %[[COND_XOR]], %[[COND_AND]]
+; CHECK-NEXT:    %[[COND_OR5:.*]] = select i1 %[[COND_OR3]], i1 true, i1 %[[COND_OR4]]
+; CHECK-NEXT:    %[[COND_OR6:.*]] = select i1 %[[COND_OR5]], i1 true, i1 false
+; CHECK-NEXT:    br i1 %[[COND_OR6]], label %loop_exit, label %do_something
+
+do_something:
+  call void @some_func() noreturn nounwind
+  br label %loop_begin
+; CHECK:       do_something:
+; CHECK-NEXT:    call
+; CHECK-NEXT:    br label %loop_begin
+
+loop_exit:
+  ret i32 0
+; CHECK:       loop_exit.split:
+; CHECK-NEXT:    ret
+}
+
 define i32 @test_partial_condition_unswitch_with_lcssa_phi1(i32* %var, i1 %cond, i32 %x) {
 ; CHECK-LABEL: @test_partial_condition_unswitch_with_lcssa_phi1(
 entry:


        


More information about the llvm-commits mailing list