[SelectionDAG] Propagate PCSections through SDNodes

[melver]

Summary

Add a new entry to SDNodeExtraInfo to propagate PCSections through
SelectionDAG.

Depends on D130881

[melver]

Rebase.

[melver]

Update RAUW unit test.

[melver]

Rebase.

llvm/include/llvm/CodeGen/SelectionDAG.h

class SelectionDAG {
   SDDbgInfo *DbgInfo;
 
   using CallSiteInfo = MachineFunction::CallSiteInfo;
   using CallSiteInfoImpl = MachineFunction::CallSiteInfoImpl;
 
   struct NodeExtraInfo {
     CallSiteInfo CSInfo;
     MDNode *HeapAllocSite = nullptr;
+    MDNode *PCSections = nullptr;
     bool NoMerge = false;
   };
   /// Out-of-line extra information for SDNodes.
   DenseMap<const SDNode *, NodeExtraInfo> SDEI;
 
   /// PersistentId counter to be used when inserting the next
   /// SDNode to this SelectionDAG. We do not place that under
   /// `#if LLVM_ENABLE_ABI_BREAKING_CHECKS` intentionally because
DAGNodeDeletedListener(SelectionDAG &DAG,
         : DAGUpdateListener(DAG), Callback(std::move(Callback)) {}
 
     void NodeDeleted(SDNode *N, SDNode *E) override { Callback(N, E); }
 
    private:
     virtual void anchor();
   };
 
+  struct DAGNodeInsertedListener : public DAGUpdateListener {
+    std::function<void(SDNode *)> Callback;
+
+    DAGNodeInsertedListener(SelectionDAG &DAG,
+                            std::function<void(SDNode *)> Callback)
+        : DAGUpdateListener(DAG), Callback(std::move(Callback)) {}
+
+    void NodeInserted(SDNode *N) override { Callback(N); }
+
+  private:
+    virtual void anchor();
+  };
+
   /// Help to insert SDNodeFlags automatically in transforming. Use
   /// RAII to save and resume flags in current scope.
   class FlagInserter {
     SelectionDAG &DAG;
     SDNodeFlags Flags;
     FlagInserter *LastInserter;
 
   public:
public:
   void addHeapAllocSite(const SDNode *Node, MDNode *MD) {
     SDEI[Node].HeapAllocSite = MD;
   }
   /// Return HeapAllocSite associated with Node, or nullptr if none exists.
   MDNode *getHeapAllocSite(const SDNode *Node) const {
     auto I = SDEI.find(Node);
     return I != SDEI.end() ? I->second.HeapAllocSite : nullptr;
   }
+  /// Set PCSections to be associated with Node.
+  void addPCSections(const SDNode *Node, MDNode *MD) {
+    SDEI[Node].PCSections = MD;
+  }
+  /// Return PCSections associated with Node, or nullptr if none exists.
+  MDNode *getPCSections(const SDNode *Node) const {
+    auto It = SDEI.find(Node);
+    return It != SDEI.end() ? It->second.PCSections : nullptr;
+  }
   /// Set NoMergeSiteInfo to be associated with Node if NoMerge is true.
   void addNoMergeSiteInfo(const SDNode *Node, bool NoMerge) {
     if (NoMerge)
       SDEI[Node].NoMerge = NoMerge;
   }
   /// Return NoMerge info associated with Node.
   bool getNoMergeSiteInfo(const SDNode *Node) const {
     auto I = SDEI.find(Node);

llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp

auto EmitNode =
     if (MI->isCandidateForCallSiteEntry() &&
         DAG->getTarget().Options.EmitCallSiteInfo)
       MF.addCallArgsForwardingRegs(MI, DAG->getCallSiteInfo(Node));
 
     if (DAG->getNoMergeSiteInfo(Node)) {
       MI->setFlag(MachineInstr::MIFlag::NoMerge);
     }
 
+    if (MDNode *MD = DAG->getPCSections(Node))
+      MI->setPCSections(MF, MD);
+
     return MI;
   };
 
   // If this is the first BB, emit byval parameter dbg_value's.
   if (HasDbg && BB->getParent()->begin() == MachineFunction::iterator(BB)) {
     SDDbgInfo::DbgIterator PDI = DAG->ByvalParmDbgBegin();
     SDDbgInfo::DbgIterator PDE = DAG->ByvalParmDbgEnd();
     for (; PDI != PDE; ++PDI) {

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

 }
 
 // Default null implementations of the callbacks.
 void SelectionDAG::DAGUpdateListener::NodeDeleted(SDNode*, SDNode*) {}
 void SelectionDAG::DAGUpdateListener::NodeUpdated(SDNode*) {}
 void SelectionDAG::DAGUpdateListener::NodeInserted(SDNode *) {}
 
 void SelectionDAG::DAGNodeDeletedListener::anchor() {}
+void SelectionDAG::DAGNodeInsertedListener::anchor() {}
 
 #define DEBUG_TYPE "selectiondag"
 
 static cl::opt<bool> EnableMemCpyDAGOpt("enable-memcpy-dag-opt",
        cl::Hidden, cl::init(true),
        cl::desc("Gang up loads and stores generated by inlining of memcpy"));
 
 static cl::opt<int> MaxLdStGlue("ldstmemcpy-glue-max",

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

void SelectionDAGBuilder::visit(const Instruction &I) {
   }
 
   // Increase the SDNodeOrder if dealing with a non-debug instruction.
   if (!isa<DbgInfoIntrinsic>(I))
     ++SDNodeOrder;
 
   CurInst = &I;
 
+  // Set inserted listener only if required.
+  bool NodeInserted = false;
+  std::unique_ptr<SelectionDAG::DAGNodeInsertedListener> InsertedListener;
+  MDNode *PCSectionsMD = I.getMetadata(LLVMContext::MD_pcsections);
+  if (PCSectionsMD) {
+    InsertedListener = std::make_unique<SelectionDAG::DAGNodeInsertedListener>(
+        DAG, [&](SDNode *) { NodeInserted = true; });
+  }
+
   visit(I.getOpcode(), I);
 
   if (!I.isTerminator() && !HasTailCall &&
       !isa<GCStatepointInst>(I)) // statepoints handle their exports internally
     CopyToExportRegsIfNeeded(&I);
 
+  // Handle metadata.
+  if (PCSectionsMD) {
+    auto It = NodeMap.find(&I);
+    if (It != NodeMap.end()) {
+      DAG.addPCSections(It->second.getNode(), PCSectionsMD);
+    } else if (NodeInserted) {
+      // This should not happen; if it does, don't let it go unnoticed so we can
+      // fix it. Relevant visit*() function is probably missing a setValue().
+      errs() << "warning: loosing !pcsections metadata ["
+             << I.getModule()->getName() << "]\n";
+      LLVM_DEBUG(I.dump());
+      assert(false);
+    }
+  }
+
   CurInst = nullptr;
 }
 
 void SelectionDAGBuilder::visitPHI(const PHINode &) {
   llvm_unreachable("SelectionDAGBuilder shouldn't visit PHI nodes!");
 }
 
 void SelectionDAGBuilder::visit(unsigned Opcode, const User &I) {
if (CB.TrueBB == NextBlock(SwitchBB)) {
     SDValue True = DAG.getConstant(1, dl, Cond.getValueType());
     Cond = DAG.getNode(ISD::XOR, dl, Cond.getValueType(), Cond, True);
   }
 
   SDValue BrCond = DAG.getNode(ISD::BRCOND, dl,
                                MVT::Other, getControlRoot(), Cond,
                                DAG.getBasicBlock(CB.TrueBB));
 
+  setValue(CurInst, BrCond);
+
   // Insert the false branch. Do this even if it's a fall through branch,
   // this makes it easier to do DAG optimizations which require inverting
   // the branch condition.
   BrCond = DAG.getNode(ISD::BR, dl, MVT::Other, BrCond,
                        DAG.getBasicBlock(CB.FalseBB));
 
   DAG.setRoot(BrCond);
 }
for (unsigned i = 0; i != NumValues; ++i, ++ChainI) {
     SDValue St =
         DAG.getStore(Root, dl, Val, Add, MachinePointerInfo(PtrV, Offsets[i]),
                      Alignment, MMOFlags, AAInfo);
     Chains[ChainI] = St;
   }
 
   SDValue StoreNode = DAG.getNode(ISD::TokenFactor, dl, MVT::Other,
                                   makeArrayRef(Chains.data(), ChainI));
+  setValue(&I, StoreNode);
   DAG.setRoot(StoreNode);
 }
 
 void SelectionDAGBuilder::visitMaskedStore(const CallInst &I,
                                            bool IsCompressing) {
   SDLoc sdl = getCurSDLoc();
 
   auto getMaskedStoreOps = [&](Value *&Ptr, Value *&Mask, Value *&Src0,
void SelectionDAGBuilder::visitFence(const FenceInst &I) {
   SDLoc dl = getCurSDLoc();
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
   SDValue Ops[3];
   Ops[0] = getRoot();
   Ops[1] = DAG.getTargetConstant((unsigned)I.getOrdering(), dl,
                                  TLI.getFenceOperandTy(DAG.getDataLayout()));
   Ops[2] = DAG.getTargetConstant(I.getSyncScopeID(), dl,
                                  TLI.getFenceOperandTy(DAG.getDataLayout()));
-  DAG.setRoot(DAG.getNode(ISD::ATOMIC_FENCE, dl, MVT::Other, Ops));
+  SDValue N = DAG.getNode(ISD::ATOMIC_FENCE, dl, MVT::Other, Ops);
+  setValue(&I, N);
+  DAG.setRoot(N);
 }
 
 void SelectionDAGBuilder::visitAtomicLoad(const LoadInst &I) {
   SDLoc dl = getCurSDLoc();
   AtomicOrdering Order = I.getOrdering();
   SyncScope::ID SSID = I.getSyncScopeID();
 
   SDValue InChain = getRoot();
void SelectionDAGBuilder::visitAtomicStore(const StoreInst &I) {
   if (Val.getValueType() != MemVT)
     Val = DAG.getPtrExtOrTrunc(Val, dl, MemVT);
   SDValue Ptr = getValue(I.getPointerOperand());
 
   if (TLI.lowerAtomicStoreAsStoreSDNode(I)) {
     // TODO: Once this is better exercised by tests, it should be merged with
     // the normal path for stores to prevent future divergence.
     SDValue S = DAG.getStore(InChain, dl, Val, Ptr, MMO);
+    setValue(&I, S);
     DAG.setRoot(S);
     return;
   }
   SDValue OutChain = DAG.getAtomic(ISD::ATOMIC_STORE, dl, MemVT, InChain,
                                    Ptr, Val, MMO);
 
-
+  setValue(&I, OutChain);
   DAG.setRoot(OutChain);
 }
 
 /// visitTargetIntrinsic - Lower a call of a target intrinsic to an INTRINSIC
 /// node.
 void SelectionDAGBuilder::visitTargetIntrinsic(const CallInst &I,
                                                unsigned Intrinsic) {
   // Ignore the callsite's attributes. A specific call site may be marked with

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

public:
 
   /// NodeDeleted - Handle nodes deleted from the graph. If the node being
   /// deleted is the current ISelPosition node, update ISelPosition.
   ///
   void NodeDeleted(SDNode *N, SDNode *E) override {
     if (ISelPosition == SelectionDAG::allnodes_iterator(N))
       ++ISelPosition;
   }
+
+  /// NodeInserted - Handle new nodes inserted into the graph: propagate
+  /// metadata from root nodes that also applies to new nodes, in case the root
+  /// is later deleted.
+  void NodeInserted(SDNode *N) override {
+    SDNode *CurNode = &*ISelPosition;
+    if (MDNode *MD = DAG.getPCSections(CurNode))
+      DAG.addPCSections(N, MD);
+  }
 };
 
 } // end anonymous namespace
 
 // This function is used to enforce the topological node id property
 // leveraged during instruction selection. Before the selection process all
 // nodes are given a non-negative id such that all nodes have a greater id than
 // their operands. As this holds transitively we can prune checks that a node N
// Select target instructions for the DAG.
     // Create a dummy node (which is not added to allnodes), that adds
     // a reference to the root node, preventing it from being deleted,
     // and tracking any changes of the root.
     HandleSDNode Dummy(CurDAG->getRoot());
     SelectionDAG::allnodes_iterator ISelPosition (CurDAG->getRoot().getNode());
     ++ISelPosition;
 
     // Make sure that ISelPosition gets properly updated when nodes are deleted
-    // in calls made from this function.
+    // in calls made from this function. New nodes inherit relevant metadata.
     ISelUpdater ISU(*CurDAG, ISelPosition);
 
     // The AllNodes list is now topological-sorted. Visit the
     // nodes by starting at the end of the list (the root of the
     // graph) and preceding back toward the beginning (the entry
     // node).
     while (ISelPosition != CurDAG->allnodes_begin()) {
       SDNode *Node = &*--ISelPosition;

llvm/unittests/CodeGen/AArch64SelectionDAGTest.cpp

TEST_F(AArch64SelectionDAGTest, ReplaceAllUsesWith) {
   EVT IntVT = EVT::getIntegerVT(Context, 8);
 
   SDValue N0 = DAG->getConstant(0x42, Loc, IntVT);
   SDValue N1 = DAG->getRegister(0, IntVT);
   // Construct node to fill arbitrary ExtraInfo.
   SDValue N2 = DAG->getNode(ISD::SUB, Loc, IntVT, N0, N1);
   EXPECT_FALSE(DAG->getHeapAllocSite(N2.getNode()));
   EXPECT_FALSE(DAG->getNoMergeSiteInfo(N2.getNode()));
+  EXPECT_FALSE(DAG->getPCSections(N2.getNode()));
   MDNode *MD = MDNode::get(Context, None);
   DAG->addHeapAllocSite(N2.getNode(), MD);
   DAG->addNoMergeSiteInfo(N2.getNode(), true);
+  DAG->addPCSections(N2.getNode(), MD);
   EXPECT_EQ(DAG->getHeapAllocSite(N2.getNode()), MD);
   EXPECT_TRUE(DAG->getNoMergeSiteInfo(N2.getNode()));
+  EXPECT_EQ(DAG->getPCSections(N2.getNode()), MD);
 
   SDValue Root = DAG->getNode(ISD::ADD, Loc, IntVT, N2, N2);
   EXPECT_EQ(Root->getOperand(0)->getOpcode(), ISD::SUB);
   // Create new node and check that ExtraInfo is propagated on RAUW.
   SDValue New = DAG->getNode(ISD::ADD, Loc, IntVT, N1, N1);
   EXPECT_FALSE(DAG->getHeapAllocSite(New.getNode()));
   EXPECT_FALSE(DAG->getNoMergeSiteInfo(New.getNode()));
+  EXPECT_FALSE(DAG->getPCSections(New.getNode()));
 
   DAG->ReplaceAllUsesWith(N2, New);
   EXPECT_EQ(Root->getOperand(0), New);
   EXPECT_EQ(DAG->getHeapAllocSite(New.getNode()), MD);
   EXPECT_TRUE(DAG->getNoMergeSiteInfo(New.getNode()));
+  EXPECT_EQ(DAG->getPCSections(New.getNode()), MD);
 }
 
 } // end namespace llvm