[CodeGen] Port ShadowStackGCLowering to new pass manager (#75324)

IIUC the new pass system was designed with parallelism. This pass needs
to add some global variables into the current module, this is not
allowed by

[WritingAnLLVMPass](https://llvm.org/docs/WritingAnLLVMPass.html#the-functionpass-class),
so convert it to module pass, see FIXME in `GetFrameMap`.

Therefore, this will trigger assertion in `CodeGenPassBuilder`:
https://github.com/llvm/llvm-project/blob/effd47ed45e3badd756103346a7c3b9e1e939e5e/llvm/include/llvm/CodeGen/CodeGenPassBuilder.h#L200-L207
Will fix it in future.

Author: paperchalice

llvm/include/llvm/CodeGen/CodeGenPassBuilder.h

@@ -38,6 +38,7 @@
 #include "llvm/CodeGen/ReplaceWithVeclib.h"
 #include "llvm/CodeGen/SafeStack.h"
 #include "llvm/CodeGen/SelectOptimize.h"
+#include "llvm/CodeGen/ShadowStackGCLowering.h"
 #include "llvm/CodeGen/SjLjEHPrepare.h"
 #include "llvm/CodeGen/UnreachableBlockElim.h"
 #include "llvm/CodeGen/WasmEHPrepare.h"
@@ -642,6 +643,9 @@ void CodeGenPassBuilder<Derived>::addIRPasses(AddIRPass &addPass) const {
   // Run GC lowering passes for builtin collectors
   // TODO: add a pass insertion point here
   addPass(GCLoweringPass());
+  // FIXME: `ShadowStackGCLoweringPass` now is a
+  // module pass, so it will trigger assertion.
+  // See comment of `AddingFunctionPasses`
   addPass(ShadowStackGCLoweringPass());
   addPass(LowerConstantIntrinsicsPass());
 

llvm/include/llvm/CodeGen/MachinePassRegistry.def

@@ -26,6 +26,7 @@ MODULE_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis, (PIC))
 MODULE_PASS("pre-isel-intrinsic-lowering", PreISelIntrinsicLoweringPass, ())
 MODULE_PASS("jmc-instrumenter", JMCInstrumenterPass, ())
 MODULE_PASS("lower-emutls", LowerEmuTLSPass, ())
+MODULE_PASS("shadow-stack-gc-lowering", ShadowStackGCLoweringPass, ())
 #undef MODULE_PASS
 
 #ifndef FUNCTION_ANALYSIS
@@ -133,7 +134,6 @@ MACHINE_FUNCTION_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis,
 DUMMY_FUNCTION_PASS("atomic-expand", AtomicExpandPass, ())
 DUMMY_FUNCTION_PASS("codegenprepare", CodeGenPreparePass, ())
 DUMMY_FUNCTION_PASS("gc-lowering", GCLoweringPass, ())
-DUMMY_FUNCTION_PASS("shadow-stack-gc-lowering", ShadowStackGCLoweringPass, ())
 DUMMY_FUNCTION_PASS("stack-protector", StackProtectorPass, ())
 #undef DUMMY_FUNCTION_PASS
 

llvm/include/llvm/CodeGen/ShadowStackGCLowering.h

@@ -0,0 +1,24 @@
+//===- llvm/CodeGen/ShadowStackGCLowering.h ---------------------*- 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_CODEGEN_SHADOWSTACKGCLOWERING_H
+#define LLVM_CODEGEN_SHADOWSTACKGCLOWERING_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class ShadowStackGCLoweringPass
+    : public PassInfoMixin<ShadowStackGCLoweringPass> {
+public:
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_SHADOWSTACKGCLOWERING_H

llvm/lib/CodeGen/ShadowStackGCLowering.cpp

@@ -15,9 +15,11 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/CodeGen/ShadowStackGCLowering.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Analysis/DomTreeUpdater.h"
+#include "llvm/CodeGen/GCMetadata.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constant.h"
@@ -50,7 +52,7 @@ using namespace llvm;
 
 namespace {
 
-class ShadowStackGCLowering : public FunctionPass {
+class ShadowStackGCLoweringImpl {
   /// RootChain - This is the global linked-list that contains the chain of GC
   /// roots.
   GlobalVariable *Head = nullptr;
@@ -64,13 +66,10 @@ class ShadowStackGCLowering : public FunctionPass {
   std::vector<std::pair<CallInst *, AllocaInst *>> Roots;
 
 public:
-  static char ID;
-
-  ShadowStackGCLowering();
+  ShadowStackGCLoweringImpl() = default;
 
-  bool doInitialization(Module &M) override;
-  void getAnalysisUsage(AnalysisUsage &AU) const override;
-  bool runOnFunction(Function &F) override;
+  bool doInitialization(Module &M);
+  bool runOnFunction(Function &F, DomTreeUpdater *DTU);
 
 private:
   bool IsNullValue(Value *V);
@@ -86,8 +85,51 @@ class ShadowStackGCLowering : public FunctionPass {
                                       const char *Name);
 };
 
+class ShadowStackGCLowering : public FunctionPass {
+  ShadowStackGCLoweringImpl Impl;
+
+public:
+  static char ID;
+
+  ShadowStackGCLowering();
+
+  bool doInitialization(Module &M) override { return Impl.doInitialization(M); }
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addPreserved<DominatorTreeWrapperPass>();
+  }
+  bool runOnFunction(Function &F) override {
+    std::optional<DomTreeUpdater> DTU;
+    if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
+      DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy);
+    return Impl.runOnFunction(F, DTU ? &*DTU : nullptr);
+  }
+};
+
 } // end anonymous namespace
 
+PreservedAnalyses ShadowStackGCLoweringPass::run(Module &M,
+                                                 ModuleAnalysisManager &MAM) {
+  auto &Map = MAM.getResult<CollectorMetadataAnalysis>(M);
+  if (Map.StrategyMap.contains("shadow-stack"))
+    return PreservedAnalyses::all();
+
+  ShadowStackGCLoweringImpl Impl;
+  bool Changed = Impl.doInitialization(M);
+  for (auto &F : M) {
+    auto &FAM =
+        MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+    auto *DT = FAM.getCachedResult<DominatorTreeAnalysis>(F);
+    DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
+    Changed |= Impl.runOnFunction(F, DT ? &DTU : nullptr);
+  }
+
+  if (!Changed)
+    return PreservedAnalyses::all();
+  PreservedAnalyses PA;
+  PA.preserve<DominatorTreeAnalysis>();
+  return PA;
+}
+
 char ShadowStackGCLowering::ID = 0;
 char &llvm::ShadowStackGCLoweringID = ShadowStackGCLowering::ID;
 
@@ -104,7 +146,7 @@ ShadowStackGCLowering::ShadowStackGCLowering() : FunctionPass(ID) {
   initializeShadowStackGCLoweringPass(*PassRegistry::getPassRegistry());
 }
 
-Constant *ShadowStackGCLowering::GetFrameMap(Function &F) {
+Constant *ShadowStackGCLoweringImpl::GetFrameMap(Function &F) {
   // doInitialization creates the abstract type of this value.
   Type *VoidPtr = PointerType::getUnqual(F.getContext());
 
@@ -158,7 +200,7 @@ Constant *ShadowStackGCLowering::GetFrameMap(Function &F) {
   return ConstantExpr::getGetElementPtr(FrameMap->getType(), GV, GEPIndices);
 }
 
-Type *ShadowStackGCLowering::GetConcreteStackEntryType(Function &F) {
+Type *ShadowStackGCLoweringImpl::GetConcreteStackEntryType(Function &F) {
   // doInitialization creates the generic version of this type.
   std::vector<Type *> EltTys;
   EltTys.push_back(StackEntryTy);
@@ -170,7 +212,7 @@ Type *ShadowStackGCLowering::GetConcreteStackEntryType(Function &F) {
 
 /// doInitialization - If this module uses the GC intrinsics, find them now. If
 /// not, exit fast.
-bool ShadowStackGCLowering::doInitialization(Module &M) {
+bool ShadowStackGCLoweringImpl::doInitialization(Module &M) {
   bool Active = false;
   for (Function &F : M) {
     if (F.hasGC() && F.getGC() == "shadow-stack") {
@@ -224,13 +266,13 @@ bool ShadowStackGCLowering::doInitialization(Module &M) {
   return true;
 }
 
-bool ShadowStackGCLowering::IsNullValue(Value *V) {
+bool ShadowStackGCLoweringImpl::IsNullValue(Value *V) {
   if (Constant *C = dyn_cast<Constant>(V))
     return C->isNullValue();
   return false;
 }
 
-void ShadowStackGCLowering::CollectRoots(Function &F) {
+void ShadowStackGCLoweringImpl::CollectRoots(Function &F) {
   // FIXME: Account for original alignment. Could fragment the root array.
   //   Approach 1: Null initialize empty slots at runtime. Yuck.
   //   Approach 2: Emit a map of the array instead of just a count.
@@ -258,11 +300,10 @@ void ShadowStackGCLowering::CollectRoots(Function &F) {
   Roots.insert(Roots.begin(), MetaRoots.begin(), MetaRoots.end());
 }
 
-GetElementPtrInst *ShadowStackGCLowering::CreateGEP(LLVMContext &Context,
-                                                    IRBuilder<> &B, Type *Ty,
-                                                    Value *BasePtr, int Idx,
-                                                    int Idx2,
-                                                    const char *Name) {
+GetElementPtrInst *
+ShadowStackGCLoweringImpl::CreateGEP(LLVMContext &Context, IRBuilder<> &B,
+                                     Type *Ty, Value *BasePtr, int Idx,
+                                     int Idx2, const char *Name) {
   Value *Indices[] = {ConstantInt::get(Type::getInt32Ty(Context), 0),
                       ConstantInt::get(Type::getInt32Ty(Context), Idx),
                       ConstantInt::get(Type::getInt32Ty(Context), Idx2)};
@@ -273,9 +314,11 @@ GetElementPtrInst *ShadowStackGCLowering::CreateGEP(LLVMContext &Context,
   return dyn_cast<GetElementPtrInst>(Val);
 }
 
-GetElementPtrInst *ShadowStackGCLowering::CreateGEP(LLVMContext &Context,
-                                            IRBuilder<> &B, Type *Ty, Value *BasePtr,
-                                            int Idx, const char *Name) {
+GetElementPtrInst *ShadowStackGCLoweringImpl::CreateGEP(LLVMContext &Context,
+                                                        IRBuilder<> &B,
+                                                        Type *Ty,
+                                                        Value *BasePtr, int Idx,
+                                                        const char *Name) {
   Value *Indices[] = {ConstantInt::get(Type::getInt32Ty(Context), 0),
                       ConstantInt::get(Type::getInt32Ty(Context), Idx)};
   Value *Val = B.CreateGEP(Ty, BasePtr, Indices, Name);
@@ -285,12 +328,9 @@ GetElementPtrInst *ShadowStackGCLowering::CreateGEP(LLVMContext &Context,
   return dyn_cast<GetElementPtrInst>(Val);
 }
 
-void ShadowStackGCLowering::getAnalysisUsage(AnalysisUsage &AU) const {
-  AU.addPreserved<DominatorTreeWrapperPass>();
-}
-
 /// runOnFunction - Insert code to maintain the shadow stack.
-bool ShadowStackGCLowering::runOnFunction(Function &F) {
+bool ShadowStackGCLoweringImpl::runOnFunction(Function &F,
+                                              DomTreeUpdater *DTU) {
   // Quick exit for functions that do not use the shadow stack GC.
   if (!F.hasGC() || F.getGC() != "shadow-stack")
     return false;
@@ -305,10 +345,6 @@ bool ShadowStackGCLowering::runOnFunction(Function &F) {
   if (Roots.empty())
     return false;
 
-  std::optional<DomTreeUpdater> DTU;
-  if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
-    DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy);
-
   // Build the constant map and figure the type of the shadow stack entry.
   Value *FrameMap = GetFrameMap(F);
   Type *ConcreteStackEntryTy = GetConcreteStackEntryType(F);
@@ -359,8 +395,7 @@ bool ShadowStackGCLowering::runOnFunction(Function &F) {
   AtEntry.CreateStore(NewHeadVal, Head);
 
   // For each instruction that escapes...
-  EscapeEnumerator EE(F, "gc_cleanup", /*HandleExceptions=*/true,
-                      DTU ? &*DTU : nullptr);
+  EscapeEnumerator EE(F, "gc_cleanup", /*HandleExceptions=*/true, DTU);
   while (IRBuilder<> *AtExit = EE.Next()) {
     // Pop the entry from the shadow stack. Don't reuse CurrentHead from
     // AtEntry, since that would make the value live for the entire function.

llvm/lib/Passes/PassBuilder.cpp

@@ -86,6 +86,7 @@
 #include "llvm/CodeGen/LowerEmuTLS.h"
 #include "llvm/CodeGen/SafeStack.h"
 #include "llvm/CodeGen/SelectOptimize.h"
+#include "llvm/CodeGen/ShadowStackGCLowering.h"
 #include "llvm/CodeGen/SjLjEHPrepare.h"
 #include "llvm/CodeGen/TypePromotion.h"
 #include "llvm/CodeGen/WasmEHPrepare.h"

llvm/lib/Passes/PassRegistry.def

@@ -128,6 +128,7 @@ MODULE_PASS("sanmd-module", SanitizerBinaryMetadataPass())
 MODULE_PASS("scc-oz-module-inliner",
             buildInlinerPipeline(OptimizationLevel::Oz,
                                  ThinOrFullLTOPhase::None))
+MODULE_PASS("shadow-stack-gc-lowering", ShadowStackGCLoweringPass())
 MODULE_PASS("strip", StripSymbolsPass())
 MODULE_PASS("strip-dead-debug-info", StripDeadDebugInfoPass())
 MODULE_PASS("strip-dead-prototypes", StripDeadPrototypesPass())