[BPF] Prevent load and store sinking in BPF backend

llvm/include/llvm/IR/IntrinsicsBPF.td

def int_bpf_btf_type_id : ClangBuiltin<"__builtin_bpf_btf_type_id">,
               [IntrNoMem]>;
   def int_bpf_preserve_type_info : ClangBuiltin<"__builtin_bpf_preserve_type_info">,
               Intrinsic<[llvm_i32_ty], [llvm_i32_ty, llvm_i64_ty],
               [IntrNoMem]>;
   def int_bpf_preserve_enum_value : ClangBuiltin<"__builtin_bpf_preserve_enum_value">,
               Intrinsic<[llvm_i64_ty], [llvm_i32_ty, llvm_ptr_ty, llvm_i64_ty],
               [IntrNoMem]>;
   def int_bpf_passthrough : ClangBuiltin<"__builtin_bpf_passthrough">,
-              Intrinsic<[llvm_any_ty], [llvm_i32_ty, llvm_any_ty], [IntrNoMem]>;
+              Intrinsic<[llvm_any_ty], [llvm_i32_ty, llvm_any_ty], [IntrNoMem]> {
+    bit DisableDefaultAttributes = false;
+  }
   def int_bpf_compare : ClangBuiltin<"__builtin_bpf_compare">,
               Intrinsic<[llvm_i1_ty], [llvm_i32_ty, llvm_anyint_ty, llvm_anyint_ty],
               [IntrNoMem]>;
+  def int_bpf_dummy_use : ClangBuiltin<"__builtin_bpf_dummy_use">,
+              Intrinsic<[], [llvm_ptr_ty], [IntrArgMemOnly, NoCapture<ArgIndex<0>>]> {
+    bit DisableDefaultAttributes = false;
+  }
 }

llvm/lib/Target/BPF/BPF.h

public:
 
   static bool isRequired() { return true; }
 };
 
 class BPFAdjustOptPass : public PassInfoMixin<BPFAdjustOptPass> {
 public:
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };
+
+class BPFPreventLoadSinking : public PassInfoMixin<BPFPreventLoadSinking> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+
+  static bool isRequired() { return true; }
+};
+
 } // namespace llvm
 
 #endif

llvm/lib/Target/BPF/BPFAbstractMemberAccess.cpp

 
 #define DEBUG_TYPE "bpf-abstract-member-access"
 
 namespace llvm {
 constexpr StringRef BPFCoreSharedInfo::AmaAttr;
 uint32_t BPFCoreSharedInfo::SeqNum;
 
 Instruction *BPFCoreSharedInfo::insertPassThrough(Module *M, BasicBlock *BB,
-                                                  Instruction *Input,
+                                                  Value *Input,
                                                   Instruction *Before) {
   Function *Fn = Intrinsic::getDeclaration(
       M, Intrinsic::bpf_passthrough, {Input->getType(), Input->getType()});
   Constant *SeqNumVal = ConstantInt::get(Type::getInt32Ty(BB->getContext()),
                                          BPFCoreSharedInfo::SeqNum++);
 
   auto *NewInst = CallInst::Create(Fn, {SeqNumVal, Input});
   BB->getInstList().insert(Before->getIterator(), NewInst);

llvm/lib/Target/BPF/BPFCORE.h

 //===- BPFCORE.h - Common info for Compile-Once Run-EveryWhere  -*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TARGET_BPF_BPFCORE_H
 #define LLVM_LIB_TARGET_BPF_BPFCORE_H
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/IR/Value.h"
 
 namespace llvm {
 
 class BasicBlock;
 class Instruction;
 class Module;
 
 class BPFCoreSharedInfo {
public:
   static constexpr StringRef AmaAttr = "btf_ama";
   /// The attribute attached to globals representing a type id
   static constexpr StringRef TypeIdAttr = "btf_type_id";
 
   /// llvm.bpf.passthrough builtin seq number
   static uint32_t SeqNum;
 
   /// Insert a bpf passthrough builtin function.
-  static Instruction *insertPassThrough(Module *M, BasicBlock *BB,
-                                        Instruction *Input,
+  static Instruction *insertPassThrough(Module *M, BasicBlock *BB, Value *Input,
                                         Instruction *Before);
 };
 
 } // namespace llvm
 
 #endif

llvm/lib/Target/BPF/BPFCheckAndAdjustIR.cpp

 
 #include "BPF.h"
 #include "BPFCORE.h"
 #include "BPFTargetMachine.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/GlobalVariable.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/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Pass.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 
 #define DEBUG_TYPE "bpf-check-and-opt-ir"
for (auto &BB : F)
         }
 
         auto *Call = dyn_cast<CallInst>(&I);
         if (!Call)
           continue;
         auto *GV = dyn_cast<GlobalValue>(Call->getCalledOperand());
         if (!GV)
           continue;
-        if (!GV->getName().startswith("llvm.bpf.passthrough"))
-          continue;
+        if (GV->getName().startswith("llvm.bpf.passthrough")) {
           Changed = true;
           Value *Arg = Call->getArgOperand(1);
           Call->replaceAllUsesWith(Arg);
           ToBeDeleted = Call;
+          continue;
+        }
+        if (Call->getIntrinsicID() == Intrinsic::bpf_dummy_use) {
+          Changed = true;
+          ToBeDeleted = Call;
+          continue;
+        }
       }
   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.

llvm/lib/Target/BPF/BPFPreventLoadSinking.cpp

+//===------------ BPFPreventLoadSinking.cpp -------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+//
+//===----------------------------------------------------------------------===//
+
+#include "BPF.h"
+#include "BPFCORE.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/CFG.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/PassManager.h"
+#include "llvm/Support/Casting.h"
+
+#define DEBUG_TYPE "bpf-prevent-load-sinking"
+
+using namespace llvm;
+
+namespace {
+// This is a simplified version of LockstepReverseIterator from SimplifyCFG
+class LockstepReverseIterator {
+  SmallVector<Instruction *, 4> Insts;
+  bool Fail = false;
+
+public:
+  LockstepReverseIterator(ArrayRef<BasicBlock *> Blocks) {
+    for (auto *BB : Blocks) {
+      Instruction *Inst = BB->getTerminator();
+      for (Inst = Inst->getPrevNode(); isa_and_nonnull<DbgInfoIntrinsic>(Inst);)
+        Inst = Inst->getPrevNode();
+      if (!Inst) {
+        // Block wasn't big enough.
+        Fail = true;
+        return;
+      }
+      Insts.push_back(Inst);
+    }
+  }
+
+  void operator--() {
+    if (Fail)
+      return;
+    for (auto *&Inst : Insts) {
+      for (Inst = Inst->getPrevNode(); isa_and_nonnull<DbgInfoIntrinsic>(Inst);)
+        Inst = Inst->getPrevNode();
+      // Already at beginning of block.
+      if (!Inst) {
+        Fail = true;
+        return;
+      }
+    }
+  }
+
+  ArrayRef<Instruction *> operator*() const { return Insts; }
+
+  bool isValid() { return !Fail; }
+};
+} // anonymous namespace
+
+static bool maybeSunk(ArrayRef<Instruction *> Insts) {
+  Instruction *I0 = Insts[0];
+  for (auto *I : Insts)
+    if (!I->isSameOperationAs(I0))
+      return false;
+  return true;
+}
+
+// Scan unconditional basic block predecessors looking for
+// instructions that might get sunk by SimplifyCFG.
+//
+// If load or store could get sunk add a call to llvm.bpf.passthrough
+// or llvm.bpf.dummy.use in order to prevent sinking.
+//
+// It is hard to determine for sure if particular instruction would be
+// sunk w/o repeating 90% of logic from SimplifyCFG, so this
+// transformation has to be conservative.
+static bool maybePatchBB(BasicBlock *BB) {
+  SmallVector<BasicBlock *, 4> UnconditionalPreds;
+  for (auto *PredBB : predecessors(BB)) {
+    auto *PredBr = dyn_cast<BranchInst>(PredBB->getTerminator());
+    if (PredBr && PredBr->isUnconditional())
+      UnconditionalPreds.push_back(PredBB);
+  }
+  if (UnconditionalPreds.size() < 2)
+    return false;
+
+  auto *M = BB->getModule();
+
+  LockstepReverseIterator LRI(UnconditionalPreds);
+  while (LRI.isValid() && maybeSunk(*LRI)) {
+    Instruction *I0 = (*LRI)[0];
+    auto SameAsI0 = [=](const Instruction *I) {
+      for (unsigned OpNum = 0; OpNum < I0->getNumOperands(); ++OpNum)
+        if (I0->getOperand(OpNum) != I->getOperand(OpNum))
+          return false;
+      return true;
+    };
+    if (!all_of(*LRI, SameAsI0)) {
+      if (isa<LoadInst>(I0)) {
+        auto *PassThrough = BPFCoreSharedInfo::insertPassThrough(
+            M, I0->getParent(), I0, I0->getNextNode());
+        I0->replaceUsesWithIf(PassThrough, [=](auto &Use) {
+          return Use.getUser() != PassThrough;
+        });
+        return true;
+      }
+      if (auto *Store = dyn_cast<StoreInst>(I0)) {
+        Function *Fn =
+            Intrinsic::getDeclaration(M, Intrinsic::bpf_dummy_use, {});
+        auto *DummyUse = CallInst::Create(Fn, {Store->getPointerOperand()}, "",
+                                          Store->getNextNode());
+        DummyUse->setAAMetadata(Store->getAAMetadata());
+        return true;
+      }
+    }
+    --LRI;
+  }
+  return false;
+}
+
+static bool preventLoadSinking(Function &F) {
+  bool Changed = false;
+  for (auto &BB : F)
+    Changed |= maybePatchBB(&BB);
+  return Changed;
+}
+
+PreservedAnalyses BPFPreventLoadSinking::run(Function &F,
+                                             FunctionAnalysisManager &AM) {
+  return preventLoadSinking(F) ? PreservedAnalyses::none()
+                               : PreservedAnalyses::all();
+}

llvm/lib/Target/BPF/BPFTargetMachine.cpp

 #include "llvm/Transforms/Scalar/SimplifyCFG.h"
 #include "llvm/Transforms/Utils/SimplifyCFGOptions.h"
 using namespace llvm;
 
 static cl::
 opt<bool> DisableMIPeephole("disable-bpf-peephole", cl::Hidden,
                             cl::desc("Disable machine peepholes for BPF"));
 
+static cl::opt<bool> DisablePreventLoadSinking("disable-prevent-load-sinking",
+                                               cl::Hidden, cl::desc(""));
+
 extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeBPFTarget() {
   // Register the target.
   RegisterTargetMachine<BPFTargetMachine> X(getTheBPFleTarget());
   RegisterTargetMachine<BPFTargetMachine> Y(getTheBPFbeTarget());
   RegisterTargetMachine<BPFTargetMachine> Z(getTheBPFTarget());
 
   PassRegistry &PR = *PassRegistry::getPassRegistry();
   initializeBPFAbstractMemberAccessLegacyPassPass(PR);
PB.registerPipelineStartEPCallback(
         FPM.addPass(BPFPreserveDITypePass());
         FPM.addPass(BPFIRPeepholePass());
         MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
       });
   PB.registerPeepholeEPCallback([=](FunctionPassManager &FPM,
                                     OptimizationLevel Level) {
     FPM.addPass(SimplifyCFGPass(SimplifyCFGOptions().hoistCommonInsts(true)));
   });
+  PB.registerScalarOptimizerLateEPCallback(
+      [=](FunctionPassManager &FPM, OptimizationLevel Level) {
+        if (!DisablePreventLoadSinking)
+          FPM.addPass(BPFPreventLoadSinking());
+      });
   PB.registerPipelineEarlySimplificationEPCallback(
       [=](ModulePassManager &MPM, OptimizationLevel) {
         MPM.addPass(BPFAdjustOptPass());
       });
 }
 
 void BPFPassConfig::addIRPasses() {
   addPass(createBPFCheckAndAdjustIR());

llvm/lib/Target/BPF/CMakeLists.txt

add_llvm_target(BPFCodeGen
   BPFCheckAndAdjustIR.cpp
   BPFFrameLowering.cpp
   BPFInstrInfo.cpp
   BPFIRPeephole.cpp
   BPFISelDAGToDAG.cpp
   BPFISelLowering.cpp
   BPFMCInstLower.cpp
   BPFPreserveDIType.cpp
+  BPFPreventLoadSinking.cpp
   BPFRegisterInfo.cpp
   BPFSelectionDAGInfo.cpp
   BPFSubtarget.cpp
   BPFTargetMachine.cpp
   BPFMIPeephole.cpp
   BPFMIChecking.cpp
   BPFMISimplifyPatchable.cpp
   BTFDebug.cpp