ControlFlowUtils.cpp

Go to the documentation of this file.

//===- ControlFlowUtils.cpp - Control Flow Utilities -----------------------==//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// Utilities to manipulate the CFG and restore SSA for the new control flow.
//
//===----------------------------------------------------------------------===//
 
#include "llvm/Transforms/Utils/ControlFlowUtils.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/ProfDataUtils.h"
#include "llvm/IR/ValueHandle.h"
#include "llvm/Transforms/Utils/Local.h"
 
#define DEBUG_TYPE "control-flow-hub"
 
using namespace llvm;
 
using BBPredicates = DenseMap<BasicBlock *, Instruction *>;
using EdgeDescriptor = ControlFlowHub::BranchDescriptor;
 
// Redirects the terminator of the incoming block to the first guard block in
// the hub. Returns the branch condition from `BB` if it exits.
// - If only one of Succ0 or Succ1 is not null, the corresponding branch
//   successor is redirected to the FirstGuardBlock.
// - Else both are not null, and branch is replaced with an unconditional
//   branch to the FirstGuardBlock.
static Value *redirectToHub(BasicBlock *BB, BasicBlock *Succ0,
                            BasicBlock *Succ1, BasicBlock *FirstGuardBlock) {
  if (auto *Branch = dyn_cast<UncondBrInst>(BB->getTerminator())) {
    assert(Succ0 == Branch->getSuccessor(0));
    assert(!Succ1);
    Branch->setSuccessor(FirstGuardBlock);
    return nullptr;
  }
 
  auto *Branch = cast<CondBrInst>(BB->getTerminator());
  auto *Condition = Branch->getCondition();
 
  assert(Succ0 || Succ1);
  assert(!Succ1 || Succ1 == Branch->getSuccessor(1));
  if (Succ0 && !Succ1) {
    Branch->setSuccessor(0, FirstGuardBlock);
  } else if (Succ1 && !Succ0) {
    Branch->setSuccessor(1, FirstGuardBlock);
  } else {
    Branch->eraseFromParent();
    UncondBrInst::Create(FirstGuardBlock, BB);
  }
 
  return Condition;
}
 
// Setup the branch instructions for guard blocks.
//
// Each guard block terminates in a conditional branch that transfers
// control to the corresponding outgoing block or the next guard
// block. The last guard block has two outgoing blocks as successors.
static void setupBranchForGuard(ArrayRef<BasicBlock *> GuardBlocks,
                                ArrayRef<BasicBlock *> Outgoing,
                                BBPredicates &GuardPredicates) {
  assert(Outgoing.size() > 1);
  assert(GuardBlocks.size() == Outgoing.size() - 1);
  // The guard predicates are PHIs merging values from all the incoming blocks,
  // so the probability of each guard branch cannot be derived locally.
  int I = 0;
  for (int E = GuardBlocks.size() - 1; I != E; ++I) {
    BasicBlock *Out = Outgoing[I];
    auto *Br = CondBrInst::Create(GuardPredicates[Out], Out, GuardBlocks[I + 1],
                                  GuardBlocks[I]);
    setExplicitlyUnknownBranchWeightsIfProfiled(*Br, DEBUG_TYPE);
  }
  BasicBlock *Out = Outgoing[I];
  auto *Br = CondBrInst::Create(GuardPredicates[Out], Out, Outgoing[I + 1],
                                GuardBlocks[I]);
  setExplicitlyUnknownBranchWeightsIfProfiled(*Br, DEBUG_TYPE);
}
 
// Assign an index to each outgoing block. At the corresponding guard
// block, compute the branch condition by comparing this index.
static void calcPredicateUsingInteger(ArrayRef<EdgeDescriptor> Branches,
                                      ArrayRef<BasicBlock *> Outgoing,
                                      ArrayRef<BasicBlock *> GuardBlocks,
                                      BBPredicates &GuardPredicates) {
  LLVMContext &Context = GuardBlocks.front()->getContext();
  BasicBlock *FirstGuardBlock = GuardBlocks.front();
  Type *Int32Ty = Type::getInt32Ty(Context);
 
  auto *Phi = PHINode::Create(Int32Ty, Branches.size(), "merged.bb.idx",
                              FirstGuardBlock);
 
  for (auto [BB, Succ0, Succ1] : Branches) {
    Value *Condition = redirectToHub(BB, Succ0, Succ1, FirstGuardBlock);
    Value *IncomingId = nullptr;
    if (Succ0 && Succ1 && Succ0 != Succ1) {
      auto Succ0Iter = find(Outgoing, Succ0);
      auto Succ1Iter = find(Outgoing, Succ1);
      Value *Id0 =
          ConstantInt::get(Int32Ty, std::distance(Outgoing.begin(), Succ0Iter));
      Value *Id1 =
          ConstantInt::get(Int32Ty, std::distance(Outgoing.begin(), Succ1Iter));
      IncomingId = SelectInst::Create(Condition, Id0, Id1, "target.bb.idx",
                                      BB->getTerminator()->getIterator());
    } else {
      // Get the index of the non-null successor, or when both successors
      // are the same block, use that block's index directly.
      auto SuccIter = Succ0 ? find(Outgoing, Succ0) : find(Outgoing, Succ1);
      IncomingId =
          ConstantInt::get(Int32Ty, std::distance(Outgoing.begin(), SuccIter));
    }
    Phi->addIncoming(IncomingId, BB);
  }
 
  for (int I = 0, E = Outgoing.size() - 1; I != E; ++I) {
    BasicBlock *Out = Outgoing[I];
    LLVM_DEBUG(dbgs() << "Creating integer guard for " << Out->getName()
                      << "\n");
    auto *Cmp = ICmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_EQ, Phi,
                                 ConstantInt::get(Int32Ty, I),
                                 Out->getName() + ".predicate", GuardBlocks[I]);
    GuardPredicates[Out] = Cmp;
  }
}
 
// Determine the branch condition to be used at each guard block from the
// original boolean values.
static void calcPredicateUsingBooleans(
    ArrayRef<EdgeDescriptor> Branches, ArrayRef<BasicBlock *> Outgoing,
    SmallVectorImpl<BasicBlock *> &GuardBlocks, BBPredicates &GuardPredicates,
    SmallVectorImpl<WeakVH> &DeletionCandidates) {
  LLVMContext &Context = GuardBlocks.front()->getContext();
  auto *BoolTrue = ConstantInt::getTrue(Context);
  auto *BoolFalse = ConstantInt::getFalse(Context);
  BasicBlock *FirstGuardBlock = GuardBlocks.front();
 
  // The predicate for the last outgoing is trivially true, and so we
  // process only the first N-1 successors.
  for (int I = 0, E = Outgoing.size() - 1; I != E; ++I) {
    BasicBlock *Out = Outgoing[I];
    LLVM_DEBUG(dbgs() << "Creating boolean guard for " << Out->getName()
                      << "\n");
 
    auto *Phi =
        PHINode::Create(Type::getInt1Ty(Context), Branches.size(),
                        StringRef("Guard.") + Out->getName(), FirstGuardBlock);
    GuardPredicates[Out] = Phi;
  }
 
  for (auto [BB, Succ0, Succ1] : Branches) {
    Value *Condition = redirectToHub(BB, Succ0, Succ1, FirstGuardBlock);
 
    // Optimization: Consider an incoming block A with both successors
    // Succ0 and Succ1 in the set of outgoing blocks. The predicates
    // for Succ0 and Succ1 complement each other. If Succ0 is visited
    // first in the loop below, control will branch to Succ0 using the
    // corresponding predicate. But if that branch is not taken, then
    // control must reach Succ1, which means that the incoming value of
    // the predicate from `BB` is true for Succ1.
    bool OneSuccessorDone = false;
    for (int I = 0, E = Outgoing.size() - 1; I != E; ++I) {
      BasicBlock *Out = Outgoing[I];
      PHINode *Phi = cast<PHINode>(GuardPredicates[Out]);
      if (Out != Succ0 && Out != Succ1) {
        Phi->addIncoming(BoolFalse, BB);
      } else if (!Succ0 || !Succ1 || Succ0 == Succ1 || OneSuccessorDone) {
        // Optimization: When only one successor is an outgoing block,
        // or both successors are the same block, the incoming predicate
        // from `BB` is always true.
        Phi->addIncoming(BoolTrue, BB);
      } else {
        assert(Succ0 && Succ1);
        if (Out == Succ0) {
          Phi->addIncoming(Condition, BB);
        } else {
          Value *Inverted = invertCondition(Condition);
          DeletionCandidates.push_back(Condition);
          Phi->addIncoming(Inverted, BB);
        }
        OneSuccessorDone = true;
      }
    }
  }
}
 
// Capture the existing control flow as guard predicates, and redirect
// control flow from \p Incoming block through the \p GuardBlocks to the
// \p Outgoing blocks.
//
// There is one guard predicate for each outgoing block OutBB. The
// predicate represents whether the hub should transfer control flow
// to OutBB. These predicates are NOT ORTHOGONAL. The Hub evaluates
// them in the same order as the Outgoing set-vector, and control
// branches to the first outgoing block whose predicate evaluates to true.
//
// The last guard block has two outgoing blocks as successors since the
// condition for the final outgoing block is trivially true. So we create one
// less block (including the first guard block) than the number of outgoing
// blocks.
static void convertToGuardPredicates(
    ArrayRef<EdgeDescriptor> Branches, ArrayRef<BasicBlock *> Outgoing,
    SmallVectorImpl<BasicBlock *> &GuardBlocks,
    SmallVectorImpl<WeakVH> &DeletionCandidates, const StringRef Prefix,
    std::optional<unsigned> MaxControlFlowBooleans) {
  BBPredicates GuardPredicates;
  Function *F = Outgoing.front()->getParent();
 
  for (int I = 0, E = Outgoing.size() - 1; I != E; ++I)
    GuardBlocks.push_back(
        BasicBlock::Create(F->getContext(), Prefix + ".guard", F));
 
  // When we are using an integer to record which target block to jump to, we
  // are creating less live values, actually we are using one single integer to
  // store the index of the target block. When we are using booleans to store
  // the branching information, we need (N-1) boolean values, where N is the
  // number of outgoing block.
  if (!MaxControlFlowBooleans || Outgoing.size() <= *MaxControlFlowBooleans)
    calcPredicateUsingBooleans(Branches, Outgoing, GuardBlocks, GuardPredicates,
                               DeletionCandidates);
  else
    calcPredicateUsingInteger(Branches, Outgoing, GuardBlocks, GuardPredicates);
 
  setupBranchForGuard(GuardBlocks, Outgoing, GuardPredicates);
}
 
// After creating a control flow hub, the operands of PHINodes in an outgoing
// block Out no longer match the predecessors of that block. The edges to Out
// that are redirected through the hub are now replaced by just one edge from
// the hub; other edges to Out, even from incoming blocks of the hub, are kept.
// To match this new control flow, the values for the redirected edges are
// moved from each PHINode to a new PHINode in the first guard block of the hub.
//
// If Out is itself an incoming block of the hub (a self-loop that is
// redirected through the hub), the new PHINode receives the value that the
// original PHINode received along the self-loop. If that value is the original
// PHINode itself, it becomes the new PHINode when the original is replaced.
static void reconnectPhis(BasicBlock *Out, BasicBlock *GuardBlock,
                          ArrayRef<EdgeDescriptor> Incoming,
                          BasicBlock *FirstGuardBlock) {
  auto I = Out->begin();
  while (I != Out->end() && isa<PHINode>(I)) {
    auto *Phi = cast<PHINode>(I);
    auto *NewPhi =
        PHINode::Create(Phi->getType(), Incoming.size(),
                        Phi->getName() + ".moved", FirstGuardBlock->begin());
    bool AllUndef = true;
    bool AllPoison = true;
    for (auto [BB, Succ0, Succ1] : Incoming) {
      Value *V = PoisonValue::get(Phi->getType());
      // Only edges from BB to Out that go through the hub move their value.
      // BB may have another edge to Out that is not redirected.
      if ((Succ0 == Out || Succ1 == Out) && Phi->getBasicBlockIndex(BB) != -1) {
        V = Phi->removeIncomingValue(BB, false);
        // When both successors are the same (Succ0 == Succ1), there are two
        // edges from BB to Out, so we need to remove the second PHI entry too.
        if (Succ0 && Succ1 && Succ0 == Succ1 &&
            Phi->getBasicBlockIndex(BB) != -1)
          Phi->removeIncomingValue(BB, false);
        AllUndef &= isa<UndefValue>(V);
        AllPoison &= isa<PoisonValue>(V);
      }
 
      NewPhi->addIncoming(V, BB);
    }
    assert(NewPhi->getNumIncomingValues() == Incoming.size());
    Value *NewV = NewPhi;
    // If all moved values are undef or poison, use that instead of a PHI. Only
    // use poison if they are all poison: replacing undef with poison would make
    // the program less defined.
    if (AllUndef) {
      NewPhi->eraseFromParent();
      NewV = AllPoison ? PoisonValue::get(Phi->getType())
                       : UndefValue::get(Phi->getType());
    }
    if (Phi->getNumOperands() == 0) {
      Phi->replaceAllUsesWith(NewV);
      I = Phi->eraseFromParent();
      continue;
    }
    Phi->addIncoming(NewV, GuardBlock);
    ++I;
  }
}
 
std::pair<BasicBlock *, bool> ControlFlowHub::finalize(
    DomTreeUpdater *DTU, SmallVectorImpl<BasicBlock *> &GuardBlocks,
    const StringRef Prefix, std::optional<unsigned> MaxControlFlowBooleans) {
#ifndef NDEBUG
  SmallPtrSet<BasicBlock *, 8> Incoming;
#endif
  SetVector<BasicBlock *> Outgoing;
 
  for (auto [BB, Succ0, Succ1] : Branches) {
#ifndef NDEBUG
    assert(
        (Incoming.insert(BB).second || isa<CallBrInst>(BB->getTerminator())) &&
        "Duplicate entry for incoming block.");
#endif
    if (Succ0)
      Outgoing.insert(Succ0);
    if (Succ1)
      Outgoing.insert(Succ1);
  }
 
  assert(Outgoing.size() && "No outgoing edges");
 
  if (Outgoing.size() < 2)
    return {Outgoing.front(), false};
 
  SmallVector<DominatorTree::UpdateType, 16> Updates;
  if (DTU) {
    for (auto [BB, Succ0, Succ1] : Branches) {
      if (Succ0)
        Updates.push_back({DominatorTree::Delete, BB, Succ0});
      // Only add Succ1 if it's different from Succ0 to avoid duplicate updates
      if (Succ1 && Succ1 != Succ0)
        Updates.push_back({DominatorTree::Delete, BB, Succ1});
    }
  }
 
  SmallVector<WeakVH, 8> DeletionCandidates;
  convertToGuardPredicates(Branches, Outgoing.getArrayRef(), GuardBlocks,
                           DeletionCandidates, Prefix, MaxControlFlowBooleans);
  BasicBlock *FirstGuardBlock = GuardBlocks.front();
 
  // Update the PHINodes in each outgoing block to match the new control flow.
  for (int I = 0, E = GuardBlocks.size(); I != E; ++I)
    reconnectPhis(Outgoing[I], GuardBlocks[I], Branches, FirstGuardBlock);
  // Process the Nth (last) outgoing block with the (N-1)th (last) guard block.
  reconnectPhis(Outgoing.back(), GuardBlocks.back(), Branches, FirstGuardBlock);
 
  if (DTU) {
    int NumGuards = GuardBlocks.size();
 
    for (auto [BB, Succ0, Succ1] : Branches)
      Updates.push_back({DominatorTree::Insert, BB, FirstGuardBlock});
 
    for (int I = 0; I != NumGuards - 1; ++I) {
      Updates.push_back({DominatorTree::Insert, GuardBlocks[I], Outgoing[I]});
      Updates.push_back(
          {DominatorTree::Insert, GuardBlocks[I], GuardBlocks[I + 1]});
    }
    // The second successor of the last guard block is an outgoing block instead
    // of having a "next" guard block.
    Updates.push_back({DominatorTree::Insert, GuardBlocks[NumGuards - 1],
                       Outgoing[NumGuards - 1]});
    Updates.push_back({DominatorTree::Insert, GuardBlocks[NumGuards - 1],
                       Outgoing[NumGuards]});
    DTU->applyUpdates(Updates);
  }
 
  for (auto &I : DeletionCandidates) {
    if (auto *Inst = dyn_cast_if_present<Instruction>(I)) {
      if (Inst->use_empty())
        Inst->eraseFromParent();
    }
  }
 
  return {FirstGuardBlock, true};
}