From 5bb38e84d3d0420e7159ffc34e3dd33ab5088b8e Mon Sep 17 00:00:00 2001 From: Juneyoung Lee Date: Mon, 8 Mar 2021 01:05:31 +0900 Subject: [PATCH] [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 --- llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp | 56 ++++---- .../SimpleLoopUnswitch/nontrivial-unswitch.ll | 156 +++++++++++++++++++++ .../SimpleLoopUnswitch/trivial-unswitch.ll | 150 ++++++++++++++++++++ 3 files changed, 335 insertions(+), 27 deletions(-) diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp index ff6de10..0e5a787 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 Invariants; + bool IsRootAnd = match(&Root, m_LogicalAnd()); + bool IsRootOr = match(&Root, m_LogicalOr()); + // Build a worklist and recurse through operators collecting invariants. SmallVector Worklist; SmallPtrSet 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(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(BI.getCondition())->getOpcode() != Instruction::Or) + if (!match(BI.getCondition(), m_LogicalOr())) return false; } else { - if (cast(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(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(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(BI->getCondition())->getOpcode() != Instruction::Or) { - assert(cast(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(BI->getCondition()); - if (CondI.getOpcode() != Instruction::And && - CondI.getOpcode() != Instruction::Or) + if (!match(&CondI, m_CombineOr(m_LogicalAnd(), m_LogicalOr()))) continue; TinyPtrVector 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(TI); - if (cast(BI.getCondition())->getOpcode() == - Instruction::And) { + if (match(BI.getCondition(), m_LogicalAnd())) { if (SuccBB == BI.getSuccessor(1)) continue; } else { - assert(cast(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 6fff85a..bfbe3e6 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 49a6d93..0870661 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: -- 2.7.4