[BPF][NewPM] Make BPFTargetMachine properly adjust NPM optimizer pipeline

[aeubanks]

Summary

This involves porting BPFAbstractMemberAccess and BPFPreserveDIType to
NPM, then adding them BPFTargetMachine::registerPassBuilderCallbacks
(the NPM equivalent of adjustPassManager()).

[asbirlea]

I am not familiar with this pass, but every bit of the porting and the callbacks looks on-par to me.

[yonghong-song]

@aeubanks I assume this is related to new pass manager and is not on by default? What is the option I could turn npm on and give a try? Looks like Hexagon is another backend which has implemented npm support. We may plan to implement new (early or module) passes so we need to modify both adjustPassManager and registerPassBuilderCallbacks. Understanding how it really functions will help our future work. The code looks correct, esp. when comparing to what Hexagon has done, but I would like to experiment a little bit before accepting the patch. Thanks!

[yonghong-song]

I think I got how to enable npm in opt. It is -passes option in opt -passes='default<O2>' %s. It looks good. Will stamp it.

I have another patch I plan to land https://reviews.llvm.org/D85570 which inserts a pass for EP_ModuleOptimizerEarly. Does it correspond to NPM registerPipelineStartEPCallback() callback function?

[yonghong-song]

BTW, can you share what is the plan for NPM? Any expected time to turn it on as the default?

[asbirlea]

I think this thread can give some answers: http://lists.llvm.org/pipermail/llvm-dev/2020-July/143666.html
The target deadline is the end of this year, with a coming checkpoint in a couple of weeks to evaluate the status.

[aeubanks]

In D88855#2313492, @yonghong-song wrote:

I think I got how to enable npm in opt. It is -passes option in opt -passes='default<O2>' %s. It looks good. Will stamp it.

Yup that's it.

I have another patch I plan to land https://reviews.llvm.org/D85570 which inserts a pass for EP_ModuleOptimizerEarly. Does it correspond to NPM registerPipelineStartEPCallback() callback function?

Yes, just like was done here.

llvm/lib/Target/BPF/BPF.h

 //===-- BPF.h - Top-level interface for BPF representation ------*- 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_BPF_H
 #define LLVM_LIB_TARGET_BPF_BPF_H
 
 #include "MCTargetDesc/BPFMCTargetDesc.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/Target/TargetMachine.h"
 
 namespace llvm {
 class BPFTargetMachine;
 
 ModulePass *createBPFCheckAndAdjustIR();
 
 FunctionPass *createBPFAbstractMemberAccess(BPFTargetMachine *TM);
 FunctionPass *createBPFPreserveDIType();
 FunctionPass *createBPFISelDag(BPFTargetMachine &TM);
 FunctionPass *createBPFMISimplifyPatchablePass();
 FunctionPass *createBPFMIPeepholePass();
 FunctionPass *createBPFMIPeepholeTruncElimPass();
 FunctionPass *createBPFMIPreEmitPeepholePass();
 FunctionPass *createBPFMIPreEmitCheckingPass();
 
 void initializeBPFCheckAndAdjustIRPass(PassRegistry&);
 
-void initializeBPFAbstractMemberAccessPass(PassRegistry&);
+void initializeBPFAbstractMemberAccessLegacyPassPass(PassRegistry &);
 void initializeBPFPreserveDITypePass(PassRegistry&);
 void initializeBPFMISimplifyPatchablePass(PassRegistry&);
 void initializeBPFMIPeepholePass(PassRegistry&);
 void initializeBPFMIPeepholeTruncElimPass(PassRegistry&);
 void initializeBPFMIPreEmitPeepholePass(PassRegistry&);
 void initializeBPFMIPreEmitCheckingPass(PassRegistry&);
-}
+
+class BPFAbstractMemberAccessPass
+    : public PassInfoMixin<BPFAbstractMemberAccessPass> {
+  BPFTargetMachine *TM;
+
+public:
+  BPFAbstractMemberAccessPass(BPFTargetMachine *TM) : TM(TM) {}
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
+class BPFPreserveDITypePass : public PassInfoMixin<BPFPreserveDITypePass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+} // namespace llvm
 
 #endif

llvm/lib/Target/BPF/BPFAbstractMemberAccess.cpp

 #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/PassManager.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"
 #include <stack>
 
 #define DEBUG_TYPE "bpf-abstract-member-access"
Instruction *BPFCoreSharedInfo::insertPassThrough(Module *M, BasicBlock *BB,
   BB->getInstList().insert(Before->getIterator(), NewInst);
   return NewInst;
 }
 } // namespace llvm
 
 using namespace llvm;
 
 namespace {
-
-class BPFAbstractMemberAccess final : public FunctionPass {
-  bool runOnFunction(Function &F) override;
-
+class BPFAbstractMemberAccess final {
 public:
-  static char ID;
-  TargetMachine *TM;
-  // Add optional BPFTargetMachine parameter so that BPF backend can add the phase
-  // with target machine to find out the endianness. The default constructor (without
-  // parameters) is used by the pass manager for managing purposes.
-  BPFAbstractMemberAccess(BPFTargetMachine *TM = nullptr)
-      : FunctionPass(ID), TM(TM) {}
+  BPFAbstractMemberAccess(BPFTargetMachine *TM) : TM(TM) {}
+
+  bool run(Function &F);
 
   struct CallInfo {
     uint32_t Kind;
     uint32_t AccessIndex;
     Align RecordAlignment;
     MDNode *Metadata;
     Value *Base;
   };
   typedef std::stack<std::pair<CallInst *, CallInfo>> CallInfoStack;
 
 private:
   enum : uint32_t {
     BPFPreserveArrayAI = 1,
     BPFPreserveUnionAI = 2,
     BPFPreserveStructAI = 3,
     BPFPreserveFieldInfoAI = 4,
   };
 
+  TargetMachine *TM;
   const DataLayout *DL = nullptr;
   Module *M = nullptr;
 
   std::map<std::string, GlobalVariable *> GEPGlobals;
   // A map to link preserve_*_access_index instrinsic calls.
   std::map<CallInst *, std::pair<CallInst *, CallInfo>> AIChain;
   // A map to hold all the base preserve_*_access_index instrinsic calls.
   // The base call is not an input of any other preserve_*
private:
 
   Value *computeBaseAndAccessKey(CallInst *Call, CallInfo &CInfo,
                                  std::string &AccessKey, MDNode *&BaseMeta);
   MDNode *computeAccessKey(CallInst *Call, CallInfo &CInfo,
                            std::string &AccessKey, bool &IsInt32Ret);
   uint64_t getConstant(const Value *IndexValue);
   bool transformGEPChain(CallInst *Call, CallInfo &CInfo);
 };
+
+class BPFAbstractMemberAccessLegacyPass final : public FunctionPass {
+  BPFTargetMachine *TM;
+
+  bool runOnFunction(Function &F) override {
+    return BPFAbstractMemberAccess(TM).run(F);
+  }
+
+public:
+  static char ID;
+
+  // Add optional BPFTargetMachine parameter so that BPF backend can add the
+  // phase with target machine to find out the endianness. The default
+  // constructor (without parameters) is used by the pass manager for managing
+  // purposes.
+  BPFAbstractMemberAccessLegacyPass(BPFTargetMachine *TM = nullptr)
+      : FunctionPass(ID), TM(TM) {}
+};
+
 } // End anonymous namespace
 
-char BPFAbstractMemberAccess::ID = 0;
-INITIALIZE_PASS(BPFAbstractMemberAccess, DEBUG_TYPE,
+char BPFAbstractMemberAccessLegacyPass::ID = 0;
+INITIALIZE_PASS(BPFAbstractMemberAccessLegacyPass, DEBUG_TYPE,
                 "BPF Abstract Member Access", false, false)
 
 FunctionPass *llvm::createBPFAbstractMemberAccess(BPFTargetMachine *TM) {
-  return new BPFAbstractMemberAccess(TM);
+  return new BPFAbstractMemberAccessLegacyPass(TM);
 }
 
-bool BPFAbstractMemberAccess::runOnFunction(Function &F) {
+bool BPFAbstractMemberAccess::run(Function &F) {
   LLVM_DEBUG(dbgs() << "********** Abstract Member Accesses **********\n");
 
   M = F.getParent();
   if (!M)
     return false;
 
   // Bail out if no debug info.
   if (M->debug_compile_units().empty())
bool BPFAbstractMemberAccess::doTransformation(Function &F) {
   // patterns similar to GEP.
   collectAICallChains(F);
 
   for (auto &C : BaseAICalls)
     Transformed = transformGEPChain(C.first, C.second) || Transformed;
 
   return removePreserveAccessIndexIntrinsic(F) || Transformed;
 }
+
+PreservedAnalyses
+BPFAbstractMemberAccessPass::run(Function &F, FunctionAnalysisManager &AM) {
+  return BPFAbstractMemberAccess(TM).run(F) ? PreservedAnalyses::none()
+                                            : PreservedAnalyses::all();
+}

llvm/lib/Target/BPF/BPFPreserveDIType.cpp

 
 #include "BPF.h"
 #include "BPFCORE.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.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-preserve-di-type"
 
 namespace llvm {
 constexpr StringRef BPFCoreSharedInfo::TypeIdAttr;
 } // namespace llvm
 
 using namespace llvm;
 
 namespace {
 
-class BPFPreserveDIType final : public FunctionPass {
-  bool runOnFunction(Function &F) override;
-
-public:
-  static char ID;
-  BPFPreserveDIType() : FunctionPass(ID) {}
-
-private:
-  Module *M = nullptr;
-
-  bool doTransformation(Function &F);
-};
-} // End anonymous namespace
-
-char BPFPreserveDIType::ID = 0;
-INITIALIZE_PASS(BPFPreserveDIType, DEBUG_TYPE, "BPF Preserve Debuginfo Type",
-                false, false)
-
-FunctionPass *llvm::createBPFPreserveDIType() { return new BPFPreserveDIType(); }
-
-bool BPFPreserveDIType::runOnFunction(Function &F) {
+static bool BPFPreserveDITypeImpl(Function &F) {
   LLVM_DEBUG(dbgs() << "********** preserve debuginfo type **********\n");
 
-  M = F.getParent();
-  if (!M)
-    return false;
+  Module *M = F.getParent();
 
   // Bail out if no debug info.
   if (M->debug_compile_units().empty())
     return false;
 
-  return doTransformation(F);
-}
-
-bool BPFPreserveDIType::doTransformation(Function &F) {
   std::vector<CallInst *> PreserveDITypeCalls;
 
   for (auto &BB : F) {
     for (auto &I : BB) {
       auto *Call = dyn_cast<CallInst>(&I);
       if (!Call)
         continue;
 
Instruction *PassThroughInst =
         BPFCoreSharedInfo::insertPassThrough(M, BB, LDInst, Call);
     Call->replaceAllUsesWith(PassThroughInst);
     Call->eraseFromParent();
     Count++;
   }
 
   return true;
 }
+
+class BPFPreserveDIType final : public FunctionPass {
+  bool runOnFunction(Function &F) override;
+
+public:
+  static char ID;
+  BPFPreserveDIType() : FunctionPass(ID) {}
+};
+} // End anonymous namespace
+
+char BPFPreserveDIType::ID = 0;
+INITIALIZE_PASS(BPFPreserveDIType, DEBUG_TYPE, "BPF Preserve Debuginfo Type",
+                false, false)
+
+FunctionPass *llvm::createBPFPreserveDIType() {
+  return new BPFPreserveDIType();
+}
+
+bool BPFPreserveDIType::runOnFunction(Function &F) {
+  return BPFPreserveDITypeImpl(F);
+}
+
+PreservedAnalyses BPFPreserveDITypePass::run(Function &F,
+                                             FunctionAnalysisManager &AM) {
+  return BPFPreserveDITypeImpl(F) ? PreservedAnalyses::none()
+                                  : PreservedAnalyses::all();
+}

llvm/lib/Target/BPF/BPFTargetMachine.h

public:
 
   TargetPassConfig *createPassConfig(PassManagerBase &PM) override;
 
   TargetLoweringObjectFile *getObjFileLowering() const override {
     return TLOF.get();
   }
 
   void adjustPassManager(PassManagerBuilder &) override;
+  void registerPassBuilderCallbacks(PassBuilder &PB,
+                                    bool DebugPassManager) override;
 };
 }
 
 #endif

llvm/lib/Target/BPF/BPFTargetMachine.cpp

 #include "BPFTargetMachine.h"
 #include "BPF.h"
 #include "MCTargetDesc/BPFMCAsmInfo.h"
 #include "TargetInfo/BPFTargetInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetLoweringObjectFileImpl.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/LegacyPassManager.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Passes/PassBuilder.h"
 #include "llvm/Support/FormattedStream.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
 #include "llvm/Transforms/Scalar.h"
+#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"));
 
 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();
-  initializeBPFAbstractMemberAccessPass(PR);
+  initializeBPFAbstractMemberAccessLegacyPassPass(PR);
   initializeBPFPreserveDITypePass(PR);
   initializeBPFCheckAndAdjustIRPass(PR);
   initializeBPFMIPeepholePass(PR);
   initializeBPFMIPeepholeTruncElimPass(PR);
 }
 
 // DataLayout: little or big endian
 static std::string computeDataLayout(const Triple &TT) {
Builder.addExtension(
   Builder.addExtension(
       PassManagerBuilder::EP_Peephole,
       [&](const PassManagerBuilder &, legacy::PassManagerBase &PM) {
         PM.add(createCFGSimplificationPass(
             SimplifyCFGOptions().hoistCommonInsts(true)));
       });
 }
 
+void BPFTargetMachine::registerPassBuilderCallbacks(PassBuilder &PB,
+                                                    bool DebugPassManager) {
+  PB.registerPipelineStartEPCallback([=](ModulePassManager &MPM) {
+    FunctionPassManager FPM(DebugPassManager);
+    FPM.addPass(BPFAbstractMemberAccessPass(this));
+    FPM.addPass(BPFPreserveDITypePass());
+    MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
+  });
+  PB.registerPeepholeEPCallback([=](FunctionPassManager &FPM,
+                                    PassBuilder::OptimizationLevel Level) {
+    FPM.addPass(SimplifyCFGPass(SimplifyCFGOptions().hoistCommonInsts(true)));
+  });
+}
+
 void BPFPassConfig::addIRPasses() {
   addPass(createBPFCheckAndAdjustIR());
   TargetPassConfig::addIRPasses();
 }
 
 // Install an instruction selector pass using
 // the ISelDag to gen BPF code.
 bool BPFPassConfig::addInstSelector() {

llvm/test/CodeGen/BPF/CORE/store-addr.ll

 ; RUN: opt -O2 %s | llvm-dis > %t1
 ; RUN: llc -filetype=asm -o - %t1 | FileCheck %s
 ; RUN: llc -mattr=+alu32 -filetype=asm -o - %t1 | FileCheck %s
+; RUN: opt -passes='default<O2>' %s | llvm-dis > %t1
+; RUN: llc -filetype=asm -o - %t1 | FileCheck %s
+; RUN: llc -mattr=+alu32 -filetype=asm -o - %t1 | FileCheck %s
 ; Source code:
 ;   struct t {
 ;     int a;
 ;   } __attribute__((preserve_access_index));
 ;   int foo(void *);
 ;   int test(struct t *arg) {
 ;       long param[1];
 ;       param[0] = (long)&arg->a;