[Attributor] Preserve BBs and instructions added in AA manifests

[ggeorgakoudis]

Summary

Manifesting AbstractAttributes may add new BBs in the IR. This patch provides an interface to register those BBs in the Attributor so that those BBs and containing instructions are not deleted as dead.

[jdoerfert]

LG, Please add registerManifestAddedBasicBlock into the AAKernelInfo when we create the new state machine.

[jdoerfert]

And also add a commit message explaining this.

[ggeorgakoudis]

Fix for comments

llvm/include/llvm/Transforms/IPO/Attributor.h

struct Attributor {
   /// Record that \p I is deleted after information was manifested. This also
   /// triggers deletion of trivially dead istructions.
   void deleteAfterManifest(Instruction &I) { ToBeDeletedInsts.insert(&I); }
 
   /// Record that \p BB is deleted after information was manifested. This also
   /// triggers deletion of trivially dead istructions.
   void deleteAfterManifest(BasicBlock &BB) { ToBeDeletedBlocks.insert(&BB); }
 
+  // Record that \p BB is added during the manifest of an AA. Added basic blocks
+  // are preserved in the IR.
+  void registerManifestAddedBasicBlock(BasicBlock &BB) {
+    ManifestAddedBlocks.insert(&BB);
+  }
+
   /// Record that \p F is deleted after information was manifested.
   void deleteAfterManifest(Function &F) {
     if (DeleteFns)
       ToBeDeletedFunctions.insert(&F);
   }
 
   /// If \p V is assumed to be a constant, return it, if it is unclear yet,
   /// return None, otherwise return `nullptr`.
private:
   /// The current initialization chain length. Tracked to avoid stack overflows.
   unsigned InitializationChainLength = 0;
 
   /// Functions, blocks, and instructions we delete after manifest is done.
   ///
   ///{
   SmallPtrSet<Function *, 8> ToBeDeletedFunctions;
   SmallPtrSet<BasicBlock *, 8> ToBeDeletedBlocks;
+  SmallPtrSet<BasicBlock *, 8> ManifestAddedBlocks;
   SmallDenseSet<WeakVH, 8> ToBeDeletedInsts;
   ///}
 
   /// Callback to get an OptimizationRemarkEmitter from a Function *.
   Optional<OptimizationRemarkGetter> OREGetter;
 
   /// The name of the pass to emit remarks for.
   const char *PassName = "";

llvm/lib/Transforms/IPO/Attributor.cpp

 bool Attributor::isAssumedDead(const Instruction &I,
                                const AbstractAttribute *QueryingAA,
                                const AAIsDead *FnLivenessAA,
                                bool &UsedAssumedInformation,
                                bool CheckBBLivenessOnly, DepClassTy DepClass) {
   const IRPosition::CallBaseContext *CBCtx =
       QueryingAA ? QueryingAA->getCallBaseContext() : nullptr;
 
+  if (ManifestAddedBlocks.contains(I.getParent()))
+    return false;
+
   if (!FnLivenessAA)
     FnLivenessAA =
         lookupAAFor<AAIsDead>(IRPosition::function(*I.getFunction(), CBCtx),
                               QueryingAA, DepClassTy::NONE);
 
   // If we have a context instruction and a liveness AA we use it.
   if (FnLivenessAA &&
       FnLivenessAA->getIRPosition().getAnchorScope() == I.getFunction() &&

 ChangeStatus Attributor::cleanupIR() {
   TimeTraceScope TimeScope("Attributor::cleanupIR");
   // Delete stuff at the end to avoid invalid references and a nice order.
   LLVM_DEBUG(dbgs() << "\n[Attributor] Delete/replace at least "
                     << ToBeDeletedFunctions.size() << " functions and "
                     << ToBeDeletedBlocks.size() << " blocks and "
                     << ToBeDeletedInsts.size() << " instructions and "
                     << ToBeChangedValues.size() << " values and "
-                    << ToBeChangedUses.size() << " uses\n");
+                    << ToBeChangedUses.size() << " uses. "
+                    << "Preserve manifest added " << ManifestAddedBlocks.size()
+                    << " blocks\n");
 
   SmallVector<WeakTrackingVH, 32> DeadInsts;
   SmallVector<Instruction *, 32> TerminatorsToFold;
 
   auto ReplaceUse = [&](Use *U, Value *NewV) {
     Value *OldV = U->get();
 
     // If we plan to replace NewV we need to update it at this point.
ChangeStatus Attributor::cleanupIR() {
 
   if (unsigned NumDeadBlocks = ToBeDeletedBlocks.size()) {
     SmallVector<BasicBlock *, 8> ToBeDeletedBBs;
     ToBeDeletedBBs.reserve(NumDeadBlocks);
     for (BasicBlock *BB : ToBeDeletedBlocks) {
       assert(isRunOn(*BB->getParent()) &&
              "Cannot delete a block outside the current SCC!");
       CGModifiedFunctions.insert(BB->getParent());
+      // Do not delete BBs added during manifests of AAs.
+      if (ManifestAddedBlocks.contains(BB))
+        continue;
       ToBeDeletedBBs.push_back(BB);
     }
     // Actually we do not delete the blocks but squash them into a single
     // unreachable but untangling branches that jump here is something we need
     // to do in a more generic way.
     DetatchDeadBlocks(ToBeDeletedBBs, nullptr);
   }
 

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

ChangeStatus buildCustomStateMachine(Attributor &A) {
     BasicBlock *StateMachineIfCascadeCurrentBB =
         BasicBlock::Create(Ctx, "worker_state_machine.parallel_region.check",
                            Kernel, UserCodeEntryBB);
     BasicBlock *StateMachineEndParallelBB =
         BasicBlock::Create(Ctx, "worker_state_machine.parallel_region.end",
                            Kernel, UserCodeEntryBB);
     BasicBlock *StateMachineDoneBarrierBB = BasicBlock::Create(
         Ctx, "worker_state_machine.done.barrier", Kernel, UserCodeEntryBB);
+    A.registerManifestAddedBasicBlock(*InitBB);
+    A.registerManifestAddedBasicBlock(*UserCodeEntryBB);
+    A.registerManifestAddedBasicBlock(*StateMachineBeginBB);
+    A.registerManifestAddedBasicBlock(*StateMachineFinishedBB);
+    A.registerManifestAddedBasicBlock(*StateMachineIsActiveCheckBB);
+    A.registerManifestAddedBasicBlock(*StateMachineIfCascadeCurrentBB);
+    A.registerManifestAddedBasicBlock(*StateMachineEndParallelBB);
+    A.registerManifestAddedBasicBlock(*StateMachineDoneBarrierBB);
 
     const DebugLoc &DLoc = KernelInitCB->getDebugLoc();
     ReturnInst::Create(Ctx, StateMachineFinishedBB)->setDebugLoc(DLoc);
 
     InitBB->getTerminator()->eraseFromParent();
     Instruction *IsWorker =
         ICmpInst::Create(ICmpInst::ICmp, llvm::CmpInst::ICMP_NE, KernelInitCB,
                          ConstantInt::get(KernelInitCB->getType(), -1),