BPFCheckAndAdjustIR.cpp

Go to the documentation of this file.

//===------------ BPFCheckAndAdjustIR.cpp - Check and Adjust IR -----------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Check IR and adjust IR for verifier friendly codes.
// The following are done for IR checking:
//   - no relocation globals in PHI node.
// The following are done for IR adjustment:
//   - remove __builtin_bpf_passthrough builtins. Target independent IR
//     optimizations are done and those builtins can be removed.
//   - remove llvm.bpf.getelementptr.and.load builtins.
//   - remove llvm.bpf.getelementptr.and.store builtins.
//   - for loads and stores with base addresses from non-zero address space
//     cast base address to zero address space (support for BPF address spaces).
//
//===----------------------------------------------------------------------===//
 
#include "BPF.h"
#include "BPFCORE.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/IR/Analysis.h"
#include "llvm/IR/GlobalVariable.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/IntrinsicsBPF.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PassManager.h"
#include "llvm/IR/ProfDataUtils.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/Value.h"
#include "llvm/Pass.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
 
#define DEBUG_TYPE "bpf-check-and-opt-ir"
 
using namespace llvm;
 
namespace {
 
class BPFCheckAndAdjustIRLegacy final : public ModulePass {
  bool runOnModule(Module &F) override;
 
public:
  static char ID;
  BPFCheckAndAdjustIRLegacy() : ModulePass(ID) {}
  void getAnalysisUsage(AnalysisUsage &AU) const override;
};
} // End anonymous namespace
 
char BPFCheckAndAdjustIRLegacy::ID = 0;
INITIALIZE_PASS(BPFCheckAndAdjustIRLegacy, DEBUG_TYPE,
                "BPF Check And Adjust IR", false, false)
 
ModulePass *llvm::createBPFCheckAndAdjustIRLegacyPass() {
  return new BPFCheckAndAdjustIRLegacy();
}
 
static void checkIR(Module &M) {
  // Ensure relocation global won't appear in PHI node
  // This may happen if the compiler generated the following code:
  //   B1:
  //      g1 = @llvm.skb_buff:0:1...
  //      ...
  //      goto B_COMMON
  //   B2:
  //      g2 = @llvm.skb_buff:0:2...
  //      ...
  //      goto B_COMMON
  //   B_COMMON:
  //      g = PHI(g1, g2)
  //      x = load g
  //      ...
  // If anything likes the above "g = PHI(g1, g2)", issue a fatal error.
  for (Function &F : M)
    for (auto &BB : F)
      for (auto &I : BB) {
        PHINode *PN = dyn_cast<PHINode>(&I);
        if (!PN || PN->use_empty())
          continue;
        for (int i = 0, e = PN->getNumIncomingValues(); i < e; ++i) {
          auto *GV = dyn_cast<GlobalVariable>(PN->getIncomingValue(i));
          if (!GV)
            continue;
          if (GV->hasAttribute(BPFCoreSharedInfo::AmaAttr) ||
              GV->hasAttribute(BPFCoreSharedInfo::TypeIdAttr))
            report_fatal_error("relocation global in PHI node");
        }
      }
}
 
static bool removePassThroughBuiltin(Module &M) {
  // Remove __builtin_bpf_passthrough()'s which are used to prevent
  // certain IR optimizations. Now major IR optimizations are done,
  // remove them.
  bool Changed = false;
  CallInst *ToBeDeleted = nullptr;
  for (Function &F : M)
    for (auto &BB : F)
      for (auto &I : BB) {
        if (ToBeDeleted) {
          ToBeDeleted->eraseFromParent();
          ToBeDeleted = nullptr;
        }
 
        auto *Call = dyn_cast<CallInst>(&I);
        if (!Call)
          continue;
        auto *GV = dyn_cast<GlobalValue>(Call->getCalledOperand());
        if (!GV)
          continue;
        if (!GV->getName().starts_with("llvm.bpf.passthrough"))
          continue;
        Changed = true;
        Value *Arg = Call->getArgOperand(1);
        Call->replaceAllUsesWith(Arg);
        ToBeDeleted = Call;
      }
  return Changed;
}
 
static bool removeCompareBuiltin(Module &M) {
  // Remove __builtin_bpf_compare()'s which are used to prevent
  // certain IR optimizations. Now major IR optimizations are done,
  // remove them.
  bool Changed = false;
  CallInst *ToBeDeleted = nullptr;
  for (Function &F : M)
    for (auto &BB : F)
      for (auto &I : BB) {
        if (ToBeDeleted) {
          ToBeDeleted->eraseFromParent();
          ToBeDeleted = nullptr;
        }
 
        auto *Call = dyn_cast<CallInst>(&I);
        if (!Call)
          continue;
        auto *GV = dyn_cast<GlobalValue>(Call->getCalledOperand());
        if (!GV)
          continue;
        if (!GV->getName().starts_with("llvm.bpf.compare"))
          continue;
 
        Changed = true;
        Value *Arg0 = Call->getArgOperand(0);
        Value *Arg1 = Call->getArgOperand(1);
        Value *Arg2 = Call->getArgOperand(2);
 
        auto OpVal = cast<ConstantInt>(Arg0)->getValue().getZExtValue();
        CmpInst::Predicate Opcode = (CmpInst::Predicate)OpVal;
 
        auto *ICmp = new ICmpInst(Opcode, Arg1, Arg2);
        ICmp->insertBefore(Call->getIterator());
 
        Call->replaceAllUsesWith(ICmp);
        ToBeDeleted = Call;
      }
  return Changed;
}
 
struct MinMaxSinkInfo {
  ICmpInst *ICmp;
  Value *Other;
  ICmpInst::Predicate Predicate;
  CallInst *MinMax;
  ZExtInst *ZExt;
  SExtInst *SExt;
 
  MinMaxSinkInfo(ICmpInst *ICmp, Value *Other, ICmpInst::Predicate Predicate)
      : ICmp(ICmp), Other(Other), Predicate(Predicate), MinMax(nullptr),
        ZExt(nullptr), SExt(nullptr) {}
};
 
static bool sinkMinMaxInBB(BasicBlock &BB,
                           const std::function<bool(Instruction *)> &Filter) {
  // Check if V is:
  //   (fn %a %b) or (ext (fn %a %b))
  // Where:
  //   ext := sext | zext
  //   fn  := smin | umin | smax | umax
  auto IsMinMaxCall = [=](Value *V, MinMaxSinkInfo &Info) {
    if (auto *ZExt = dyn_cast<ZExtInst>(V)) {
      V = ZExt->getOperand(0);
      Info.ZExt = ZExt;
    } else if (auto *SExt = dyn_cast<SExtInst>(V)) {
      V = SExt->getOperand(0);
      Info.SExt = SExt;
    }
 
    auto *Call = dyn_cast<CallInst>(V);
    if (!Call)
      return false;
 
    auto *Called = dyn_cast<Function>(Call->getCalledOperand());
    if (!Called)
      return false;
 
    switch (Called->getIntrinsicID()) {
    case Intrinsic::smin:
    case Intrinsic::umin:
    case Intrinsic::smax:
    case Intrinsic::umax:
      break;
    default:
      return false;
    }
 
    if (!Filter(Call))
      return false;
 
    Info.MinMax = Call;
 
    return true;
  };
 
  auto ZeroOrSignExtend = [](IRBuilder<> &Builder, Value *V,
                             MinMaxSinkInfo &Info) {
    if (Info.SExt) {
      if (Info.SExt->getType() == V->getType())
        return V;
      return Builder.CreateSExt(V, Info.SExt->getType());
    }
    if (Info.ZExt) {
      if (Info.ZExt->getType() == V->getType())
        return V;
      return Builder.CreateZExt(V, Info.ZExt->getType());
    }
    return V;
  };
 
  bool Changed = false;
  SmallVector<MinMaxSinkInfo, 2> SinkList;
 
  // Check BB for instructions like:
  //   insn := (icmp %a (fn ...)) | (icmp (fn ...)  %a)
  //
  // Where:
  //   fn := min | max | (sext (min ...)) | (sext (max ...))
  //
  // Put such instructions to SinkList.
  for (Instruction &I : BB) {
    ICmpInst *ICmp = dyn_cast<ICmpInst>(&I);
    if (!ICmp)
      continue;
    if (!ICmp->isRelational())
      continue;
    MinMaxSinkInfo First(ICmp, ICmp->getOperand(1),
                         ICmpInst::getSwappedPredicate(ICmp->getPredicate()));
    MinMaxSinkInfo Second(ICmp, ICmp->getOperand(0), ICmp->getPredicate());
    bool FirstMinMax = IsMinMaxCall(ICmp->getOperand(0), First);
    bool SecondMinMax = IsMinMaxCall(ICmp->getOperand(1), Second);
    if (!(FirstMinMax ^ SecondMinMax))
      continue;
    SinkList.push_back(FirstMinMax ? First : Second);
  }
 
  // Iterate SinkList and replace each (icmp ...) with corresponding
  // `x < a && x < b` or similar expression.
  for (auto &Info : SinkList) {
    ICmpInst *ICmp = Info.ICmp;
    CallInst *MinMax = Info.MinMax;
    Intrinsic::ID IID = MinMax->getCalledFunction()->getIntrinsicID();
    ICmpInst::Predicate P = Info.Predicate;
    if (ICmpInst::isSigned(P) && IID != Intrinsic::smin &&
        IID != Intrinsic::smax)
      continue;
 
    IRBuilder<> Builder(ICmp);
    Value *X = Info.Other;
    Value *A = ZeroOrSignExtend(Builder, MinMax->getArgOperand(0), Info);
    Value *B = ZeroOrSignExtend(Builder, MinMax->getArgOperand(1), Info);
    bool IsMin = IID == Intrinsic::smin || IID == Intrinsic::umin;
    bool IsMax = IID == Intrinsic::smax || IID == Intrinsic::umax;
    bool IsLess = ICmpInst::isLE(P) || ICmpInst::isLT(P);
    bool IsGreater = ICmpInst::isGE(P) || ICmpInst::isGT(P);
    assert(IsMin ^ IsMax);
    assert(IsLess ^ IsGreater);
 
    Value *Replacement;
    Value *LHS = Builder.CreateICmp(P, X, A);
    Value *RHS = Builder.CreateICmp(P, X, B);
    if ((IsLess && IsMin) || (IsGreater && IsMax))
      // x < min(a, b) -> x < a && x < b
      // x > max(a, b) -> x > a && x > b
      Replacement = Builder.CreateLogicalAnd(LHS, RHS);
    else
      // x > min(a, b) -> x > a || x > b
      // x < max(a, b) -> x < a || x < b
      Replacement = Builder.CreateLogicalOr(LHS, RHS);
 
    if (SelectInst *ReplacementInst = dyn_cast<SelectInst>(Replacement))
      setExplicitlyUnknownBranchWeightsIfProfiled(*ReplacementInst, DEBUG_TYPE);
 
    ICmp->replaceAllUsesWith(Replacement);
 
    Instruction *ToRemove[] = {ICmp, Info.ZExt, Info.SExt, MinMax};
    for (Instruction *I : ToRemove)
      if (I && I->use_empty())
        I->eraseFromParent();
 
    Changed = true;
  }
 
  return Changed;
}
 
// Do the following transformation:
//
//   x < min(a, b) -> x < a && x < b
//   x > min(a, b) -> x > a || x > b
//   x < max(a, b) -> x < a || x < b
//   x > max(a, b) -> x > a && x > b
//
// Such patterns are introduced by LICM.cpp:hoistMinMax()
// transformation and might lead to BPF verification failures for
// older kernels.
//
// To minimize "collateral" changes only do it for icmp + min/max
// calls when icmp is inside a loop and min/max is outside of that
// loop.
//
// Verification failure happens when:
// - RHS operand of some `icmp LHS, RHS` is replaced by some RHS1;
// - verifier can recognize RHS as a constant scalar in some context;
// - verifier can't recognize RHS1 as a constant scalar in the same
//   context;
//
// The "constant scalar" is not a compile time constant, but a register
// that holds a scalar value known to verifier at some point in time
// during abstract interpretation.
//
// See also:
//   https://lore.kernel.org/bpf/20230406164505.1046801-1-yhs@fb.com/
static bool sinkMinMax(Module &M,
                       function_ref<LoopInfo &(Function &)> GetLoopInfo) {
  bool Changed = false;
 
  for (Function &F : M) {
    if (F.isDeclaration())
      continue;
 
    LoopInfo &LI = GetLoopInfo(F);
    for (Loop *L : LI)
      for (BasicBlock *BB : L->blocks()) {
        // Filter out instructions coming from the same loop
        Loop *BBLoop = LI.getLoopFor(BB);
        auto OtherLoopFilter = [&](Instruction *I) {
          return LI.getLoopFor(I->getParent()) != BBLoop;
        };
        Changed |= sinkMinMaxInBB(*BB, OtherLoopFilter);
      }
  }
 
  return Changed;
}
 
void BPFCheckAndAdjustIRLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
  AU.addRequired<LoopInfoWrapperPass>();
}
 
static void unrollGEPLoad(CallInst *Call) {
  auto [GEP, Load] = BPFPreserveStaticOffsetPass::reconstructLoad(Call);
  GEP->insertBefore(Call->getIterator());
  Load->insertBefore(Call->getIterator());
  Call->replaceAllUsesWith(Load);
  Call->eraseFromParent();
}
 
static void unrollGEPStore(CallInst *Call) {
  auto [GEP, Store] = BPFPreserveStaticOffsetPass::reconstructStore(Call);
  GEP->insertBefore(Call->getIterator());
  Store->insertBefore(Call->getIterator());
  Call->eraseFromParent();
}
 
static bool removeGEPBuiltinsInFunc(Function &F) {
  SmallVector<CallInst *> GEPLoads;
  SmallVector<CallInst *> GEPStores;
  for (auto &BB : F)
    for (auto &Insn : BB)
      if (auto *Call = dyn_cast<CallInst>(&Insn))
        if (auto *Called = Call->getCalledFunction())
          switch (Called->getIntrinsicID()) {
          case Intrinsic::bpf_getelementptr_and_load:
            GEPLoads.push_back(Call);
            break;
          case Intrinsic::bpf_getelementptr_and_store:
            GEPStores.push_back(Call);
            break;
          }
 
  if (GEPLoads.empty() && GEPStores.empty())
    return false;
 
  for_each(GEPLoads, unrollGEPLoad);
  for_each(GEPStores, unrollGEPStore);
 
  return true;
}
 
// Rewrites the following builtins:
// - llvm.bpf.getelementptr.and.load
// - llvm.bpf.getelementptr.and.store
// As (load (getelementptr ...)) or (store (getelementptr ...)).
static bool removeGEPBuiltins(Module &M) {
  bool Changed = false;
  for (auto &F : M)
    Changed = removeGEPBuiltinsInFunc(F) || Changed;
  return Changed;
}
 
// Wrap ToWrap with cast to address space zero:
// - if ToWrap is a getelementptr,
//   wrap it's base pointer instead and return a copy;
// - if ToWrap is Instruction, insert address space cast
//   immediately after ToWrap;
// - if ToWrap is not an Instruction (function parameter
//   or a global value), insert address space cast at the
//   beginning of the Function F;
// - use Cache to avoid inserting too many casts;
static Value *aspaceWrapValue(DenseMap<Value *, Value *> &Cache, Function *F,
                              Value *ToWrap) {
  auto It = Cache.find(ToWrap);
  if (It != Cache.end())
    return It->getSecond();
 
  if (auto *GEP = dyn_cast<GetElementPtrInst>(ToWrap)) {
    Value *Ptr = GEP->getPointerOperand();
    Value *WrappedPtr = aspaceWrapValue(Cache, F, Ptr);
    auto *GEPTy = cast<PointerType>(GEP->getType());
    auto *NewGEP = GEP->clone();
    NewGEP->insertAfter(GEP->getIterator());
    NewGEP->mutateType(PointerType::getUnqual(GEPTy->getContext()));
    NewGEP->setOperand(GEP->getPointerOperandIndex(), WrappedPtr);
    NewGEP->setName(GEP->getName());
    Cache[ToWrap] = NewGEP;
    return NewGEP;
  }
 
  IRBuilder IB(*F->getParent());
  if (Instruction *InsnPtr = dyn_cast<Instruction>(ToWrap))
    IB.SetInsertPoint(*InsnPtr->getInsertionPointAfterDef());
  else
    IB.SetInsertPoint(F->getEntryBlock().getFirstInsertionPt());
  auto *ASZeroPtrTy = IB.getPtrTy(0);
  auto *ACast = IB.CreateAddrSpaceCast(ToWrap, ASZeroPtrTy, ToWrap->getName());
  Cache[ToWrap] = ACast;
  return ACast;
}
 
// Wrap a pointer operand OpNum of instruction I
// with cast to address space zero
static void aspaceWrapOperand(DenseMap<Value *, Value *> &Cache, Instruction *I,
                              unsigned OpNum) {
  Value *OldOp = I->getOperand(OpNum);
  if (OldOp->getType()->getPointerAddressSpace() == 0)
    return;
 
  Value *NewOp = aspaceWrapValue(Cache, I->getFunction(), OldOp);
  I->setOperand(OpNum, NewOp);
  // Check if there are any remaining users of old GEP,
  // delete those w/o users
  for (;;) {
    auto *OldGEP = dyn_cast<GetElementPtrInst>(OldOp);
    if (!OldGEP)
      break;
    if (!OldGEP->use_empty())
      break;
    OldOp = OldGEP->getPointerOperand();
    OldGEP->eraseFromParent();
  }
}
 
static Value *wrapPtrIfASNotZero(DenseMap<Value *, Value *> &Cache,
                                 CallInst *CI, Value *P) {
  if (auto *PTy = dyn_cast<PointerType>(P->getType())) {
    if (PTy->getAddressSpace() == 0)
      return P;
  }
  return aspaceWrapValue(Cache, CI->getFunction(), P);
}
 
static Instruction *aspaceMemSet(Intrinsic::ID ID,
                                 DenseMap<Value *, Value *> &Cache,
                                 CallInst *CI) {
  auto *MI = cast<MemIntrinsic>(CI);
  IRBuilder<> B(CI);
 
  Value *OldDst = CI->getArgOperand(0);
  Value *NewDst = wrapPtrIfASNotZero(Cache, CI, OldDst);
  if (OldDst == NewDst)
    return nullptr;
 
  // memset(new_dst, val, len, align, isvolatile, md)
  Value *Val = CI->getArgOperand(1);
  Value *Len = CI->getArgOperand(2);
 
  auto *MS = cast<MemSetInst>(CI);
  MaybeAlign Align = MS->getDestAlign();
  bool IsVolatile = MS->isVolatile();
 
  if (ID == Intrinsic::memset)
    return B.CreateMemSet(NewDst, Val, Len, Align, IsVolatile,
                          MI->getAAMetadata());
  else
    return B.CreateMemSetInline(NewDst, Align, Val, Len, IsVolatile,
                                MI->getAAMetadata());
}
 
static Instruction *aspaceMemCpy(Intrinsic::ID ID,
                                 DenseMap<Value *, Value *> &Cache,
                                 CallInst *CI) {
  auto *MI = cast<MemIntrinsic>(CI);
  IRBuilder<> B(CI);
 
  Value *OldDst = CI->getArgOperand(0);
  Value *OldSrc = CI->getArgOperand(1);
  Value *NewDst = wrapPtrIfASNotZero(Cache, CI, OldDst);
  Value *NewSrc = wrapPtrIfASNotZero(Cache, CI, OldSrc);
  if (OldDst == NewDst && OldSrc == NewSrc)
    return nullptr;
 
  // memcpy(new_dst, dst_align, new_src, src_align, len, isvolatile, md)
  Value *Len = CI->getArgOperand(2);
 
  auto *MT = cast<MemTransferInst>(CI);
  MaybeAlign DstAlign = MT->getDestAlign();
  MaybeAlign SrcAlign = MT->getSourceAlign();
  bool IsVolatile = MT->isVolatile();
 
  return B.CreateMemTransferInst(ID, NewDst, DstAlign, NewSrc, SrcAlign, Len,
                                 IsVolatile, MI->getAAMetadata());
}
 
static Instruction *aspaceMemMove(DenseMap<Value *, Value *> &Cache,
                                  CallInst *CI) {
  auto *MI = cast<MemIntrinsic>(CI);
  IRBuilder<> B(CI);
 
  Value *OldDst = CI->getArgOperand(0);
  Value *OldSrc = CI->getArgOperand(1);
  Value *NewDst = wrapPtrIfASNotZero(Cache, CI, OldDst);
  Value *NewSrc = wrapPtrIfASNotZero(Cache, CI, OldSrc);
  if (OldDst == NewDst && OldSrc == NewSrc)
    return nullptr;
 
  // memmove(new_dst, dst_align, new_src, src_align, len, isvolatile, md)
  Value *Len = CI->getArgOperand(2);
 
  auto *MT = cast<MemTransferInst>(CI);
  MaybeAlign DstAlign = MT->getDestAlign();
  MaybeAlign SrcAlign = MT->getSourceAlign();
  bool IsVolatile = MT->isVolatile();
 
  return B.CreateMemMove(NewDst, DstAlign, NewSrc, SrcAlign, Len, IsVolatile,
                         MI->getAAMetadata());
}
 
// Support for BPF address spaces:
// - for each function in the module M, update pointer operand of
//   each memory access instruction (load/store/cmpxchg/atomicrmw)
//   or intrinsic call insns (memset/memcpy/memmove)
//   by casting it from non-zero address space to zero address space, e.g:
//
//   (load (ptr addrspace (N) %p) ...)
//     -> (load (addrspacecast ptr addrspace (N) %p to ptr))
//
// - assign section with name .addr_space.N for globals defined in
//   non-zero address space N
static bool insertASpaceCasts(Module &M) {
  bool Changed = false;
  for (Function &F : M) {
    DenseMap<Value *, Value *> CastsCache;
    for (BasicBlock &BB : F) {
      for (Instruction &I : llvm::make_early_inc_range(BB)) {
        unsigned PtrOpNum;
 
        if (auto *LD = dyn_cast<LoadInst>(&I)) {
          PtrOpNum = LD->getPointerOperandIndex();
          aspaceWrapOperand(CastsCache, &I, PtrOpNum);
          continue;
        }
        if (auto *ST = dyn_cast<StoreInst>(&I)) {
          PtrOpNum = ST->getPointerOperandIndex();
          aspaceWrapOperand(CastsCache, &I, PtrOpNum);
          continue;
        }
        if (auto *CmpXchg = dyn_cast<AtomicCmpXchgInst>(&I)) {
          PtrOpNum = CmpXchg->getPointerOperandIndex();
          aspaceWrapOperand(CastsCache, &I, PtrOpNum);
          continue;
        }
        if (auto *RMW = dyn_cast<AtomicRMWInst>(&I)) {
          PtrOpNum = RMW->getPointerOperandIndex();
          aspaceWrapOperand(CastsCache, &I, PtrOpNum);
          continue;
        }
 
        auto *CI = dyn_cast<CallInst>(&I);
        if (!CI)
          continue;
 
        Function *Callee = CI->getCalledFunction();
        if (!Callee || !Callee->isIntrinsic())
          continue;
 
        // Check memset/memcpy/memmove
        Intrinsic::ID ID = Callee->getIntrinsicID();
        bool IsSet = ID == Intrinsic::memset || ID == Intrinsic::memset_inline;
        bool IsCpy = ID == Intrinsic::memcpy || ID == Intrinsic::memcpy_inline;
        bool IsMove = ID == Intrinsic::memmove;
        if (!IsSet && !IsCpy && !IsMove)
          continue;
 
        Instruction *New;
        if (IsSet)
          New = aspaceMemSet(ID, CastsCache, CI);
        else if (IsCpy)
          New = aspaceMemCpy(ID, CastsCache, CI);
        else
          New = aspaceMemMove(CastsCache, CI);
 
        if (!New)
          continue;
 
        I.replaceAllUsesWith(New);
        New->takeName(&I);
        I.eraseFromParent();
      }
    }
    Changed |= !CastsCache.empty();
  }
  // Merge all globals within same address space into single
  // .addr_space.<addr space no> section
  for (GlobalVariable &G : M.globals()) {
    if (G.getAddressSpace() == 0 || G.hasSection())
      continue;
    SmallString<16> SecName;
    raw_svector_ostream OS(SecName);
    OS << ".addr_space." << G.getAddressSpace();
    G.setSection(SecName);
    // Prevent having separate section for constants
    G.setConstant(false);
  }
  return Changed;
}
 
// There is no range check in front of a jump table when the default of the
// switch is unreachable. The compiler knows that the value is one of the cases,
// e.g. the discriminant of an enum that was loaded from memory. The verifier
// doesn't and rejects the load from the table.
//
// Anything can be the default of such a switch. Take the destination that most
// cases go to and drop these cases. A switch that becomes compares needs as
// many of them as before.
static bool giveSwitchesDefault(Module &M) {
  bool Changed = false;
 
  for (Function &F : M) {
    for (BasicBlock &BB : F) {
      auto *SI = dyn_cast<SwitchInst>(BB.getTerminator());
 
      if (!SI || !SI->getNumCases() ||
          !isa<UnreachableInst>(*SI->getDefaultDest()->getFirstNonPHIIt()))
        continue;
 
      MapVector<BasicBlock *, unsigned> Cnt;
      for (auto &Case : SI->cases())
        ++Cnt[Case.getCaseSuccessor()];
      BasicBlock *Dest = max_element(Cnt, [](const auto &A, const auto &B) {
                           return A.second < B.second;
                         })->first;
 
      for (PHINode &Phi : Dest->phis())
        Phi.addIncoming(Phi.getIncomingValueForBlock(&BB), &BB);
      SI->getDefaultDest()->removePredecessor(&BB);
      SI->setDefaultDest(Dest);
 
      SwitchInstProfUpdateWrapper SIW(*SI);
      for (auto It = SIW->case_begin(); It != SIW->case_end();) {
        if (It->getCaseSuccessor() != Dest) {
          ++It;
          continue;
        }
        Dest->removePredecessor(&BB);
        It = SIW.removeCase(It);
      }
      Changed = true;
    }
  }
  return Changed;
}
 
static bool adjustIR(Module &M,
                     function_ref<LoopInfo &(Function &)> GetLoopInfo) {
  bool Changed = removePassThroughBuiltin(M);
  Changed = removeCompareBuiltin(M) || Changed;
  Changed = sinkMinMax(M, GetLoopInfo) || Changed;
  Changed = removeGEPBuiltins(M) || Changed;
  Changed = insertASpaceCasts(M) || Changed;
  Changed = giveSwitchesDefault(M) || Changed;
  return Changed;
}
 
bool BPFCheckAndAdjustIRLegacy::runOnModule(Module &M) {
  checkIR(M);
  return adjustIR(M, [&](Function &F) -> LoopInfo & {
    return getAnalysis<LoopInfoWrapperPass>(F).getLoopInfo();
  });
}
 
PreservedAnalyses BPFCheckAndAdjustIRPass::run(Module &M,
                                               ModuleAnalysisManager &MAM) {
  checkIR(M);
  bool Changed = adjustIR(M, [&](Function &F) -> LoopInfo & {
    return MAM.getResult<FunctionAnalysisManagerModuleProxy>(M)
        .getManager()
        .getResult<LoopAnalysis>(F);
  });
  return Changed ? PreservedAnalyses::none().preserveSet<CFGAnalyses>()
                 : PreservedAnalyses::all();
}