llvm/lib/Target/BPF/BPFCheckAndAdjustIR.cpp - rust-lang/llvm-project - Git at Google

 //===------------ 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/Analysis/LoopInfo.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicsBPF.h"
 #include "llvm/IR/Module.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 BPFCheckAndAdjustIR final : public ModulePass {
   bool runOnModule(Module &F) override;

 public:
   static char ID;
   BPFCheckAndAdjustIR() : ModulePass(ID) {}
   virtual void getAnalysisUsage(AnalysisUsage &AU) const override;

 private:
   void checkIR(Module &M);
   bool adjustIR(Module &M);
   bool removePassThroughBuiltin(Module &M);
   bool removeCompareBuiltin(Module &M);
   bool sinkMinMax(Module &M);
   bool removeGEPBuiltins(Module &M);
   bool insertASpaceCasts(Module &M);
 };
 } // End anonymous namespace

 char BPFCheckAndAdjustIR::ID = 0;
 INITIALIZE_PASS(BPFCheckAndAdjustIR, DEBUG_TYPE, "BPF Check And Adjust IR",
                 false, false)

 ModulePass *llvm::createBPFCheckAndAdjustIR() {
   return new BPFCheckAndAdjustIR();
 }

 void BPFCheckAndAdjustIR::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");
         }
       }
 }

 bool BPFCheckAndAdjustIR::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;
 }

 bool BPFCheckAndAdjustIR::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);

     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/
 bool BPFCheckAndAdjustIR::sinkMinMax(Module &M) {
   bool Changed = false;

   for (Function &F : M) {
     if (F.isDeclaration())
       continue;

     LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>(F).getLoopInfo();
     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 BPFCheckAndAdjustIR::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 ...)).
 bool BPFCheckAndAdjustIR::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->getContext());
   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();
   }
 }

 // Support for BPF address spaces:
 // - for each function in the module M, update pointer operand of
 //   each memory access instruction (load/store/cmpxchg/atomicrmw)
 //   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
 bool BPFCheckAndAdjustIR::insertASpaceCasts(Module &M) {
   bool Changed = false;
   for (Function &F : M) {
     DenseMap<Value *, Value *> CastsCache;
     for (BasicBlock &BB : F) {
       for (Instruction &I : BB) {
         unsigned PtrOpNum;

         if (auto *LD = dyn_cast<LoadInst>(&I))
           PtrOpNum = LD->getPointerOperandIndex();
         else if (auto *ST = dyn_cast<StoreInst>(&I))
           PtrOpNum = ST->getPointerOperandIndex();
         else if (auto *CmpXchg = dyn_cast<AtomicCmpXchgInst>(&I))
           PtrOpNum = CmpXchg->getPointerOperandIndex();
         else if (auto *RMW = dyn_cast<AtomicRMWInst>(&I))
           PtrOpNum = RMW->getPointerOperandIndex();
         else
           continue;

         aspaceWrapOperand(CastsCache, &I, PtrOpNum);
       }
     }
     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;
 }

 bool BPFCheckAndAdjustIR::adjustIR(Module &M) {
   bool Changed = removePassThroughBuiltin(M);
   Changed = removeCompareBuiltin(M) || Changed;
   Changed = sinkMinMax(M) || Changed;
   Changed = removeGEPBuiltins(M) || Changed;
   Changed = insertASpaceCasts(M) || Changed;
   return Changed;
 }

 bool BPFCheckAndAdjustIR::runOnModule(Module &M) {
   checkIR(M);
   return adjustIR(M);
 }
	//===------------ 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/Analysis/LoopInfo.h"
	#include "llvm/IR/GlobalVariable.h"
	#include "llvm/IR/IRBuilder.h"
	#include "llvm/IR/Instruction.h"
	#include "llvm/IR/Instructions.h"
	#include "llvm/IR/IntrinsicsBPF.h"
	#include "llvm/IR/Module.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 BPFCheckAndAdjustIR final : public ModulePass {
	bool runOnModule(Module &F) override;

	public:
	static char ID;
	BPFCheckAndAdjustIR() : ModulePass(ID) {}
	virtual void getAnalysisUsage(AnalysisUsage &AU) const override;

	private:
	void checkIR(Module &M);
	bool adjustIR(Module &M);
	bool removePassThroughBuiltin(Module &M);
	bool removeCompareBuiltin(Module &M);
	bool sinkMinMax(Module &M);
	bool removeGEPBuiltins(Module &M);
	bool insertASpaceCasts(Module &M);
	};
	} // End anonymous namespace

	char BPFCheckAndAdjustIR::ID = 0;
	INITIALIZE_PASS(BPFCheckAndAdjustIR, DEBUG_TYPE, "BPF Check And Adjust IR",
	false, false)

	ModulePass *llvm::createBPFCheckAndAdjustIR() {
	return new BPFCheckAndAdjustIR();
	}

	void BPFCheckAndAdjustIR::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");
	}
	}
	}

	bool BPFCheckAndAdjustIR::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;
	}

	bool BPFCheckAndAdjustIR::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);

	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/
	bool BPFCheckAndAdjustIR::sinkMinMax(Module &M) {
	bool Changed = false;

	for (Function &F : M) {
	if (F.isDeclaration())
	continue;

	LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>(F).getLoopInfo();
	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 BPFCheckAndAdjustIR::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 ...)).
	bool BPFCheckAndAdjustIR::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->getContext());
	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();
	}
	}

	// Support for BPF address spaces:
	// - for each function in the module M, update pointer operand of
	// each memory access instruction (load/store/cmpxchg/atomicrmw)
	// 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
	bool BPFCheckAndAdjustIR::insertASpaceCasts(Module &M) {
	bool Changed = false;
	for (Function &F : M) {
	DenseMap<Value , Value > CastsCache;
	for (BasicBlock &BB : F) {
	for (Instruction &I : BB) {
	unsigned PtrOpNum;

	if (auto *LD = dyn_cast<LoadInst>(&I))
	PtrOpNum = LD->getPointerOperandIndex();
	else if (auto *ST = dyn_cast<StoreInst>(&I))
	PtrOpNum = ST->getPointerOperandIndex();
	else if (auto *CmpXchg = dyn_cast<AtomicCmpXchgInst>(&I))
	PtrOpNum = CmpXchg->getPointerOperandIndex();
	else if (auto *RMW = dyn_cast<AtomicRMWInst>(&I))
	PtrOpNum = RMW->getPointerOperandIndex();
	else
	continue;

	aspaceWrapOperand(CastsCache, &I, PtrOpNum);
	}
	}
	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;
	}

	bool BPFCheckAndAdjustIR::adjustIR(Module &M) {
	bool Changed = removePassThroughBuiltin(M);
	Changed = removeCompareBuiltin(M) \|\| Changed;
	Changed = sinkMinMax(M) \|\| Changed;
	Changed = removeGEPBuiltins(M) \|\| Changed;
	Changed = insertASpaceCasts(M) \|\| Changed;
	return Changed;
	}

	bool BPFCheckAndAdjustIR::runOnModule(Module &M) {
	checkIR(M);
	return adjustIR(M);
	}