[MF][MBB]: Add the ability to register callbacks for removal and insertion of MachineInstrs

[aditya_nandakumar]

Summary

This patch adds the ability to watch for insertions/deletions of MachineInstructions. Currently in GISel, we rely on the MachineIRBuilder::recordInsertion to further legalize the intermediate instructions - this won't work if the target builds instructions with BuildMI interface. Similarly, this makes it easy to watch for MI::eraseFromParent to keep various data structures up to date in the legalizer/combiner etc. I suspect there would be non GISel use cases as well.
While this patch adds the ability to register callbacks, nothing makes use of it as I didn't want to put in too much effort before getting feedback. The GISel legalizer would be one place that would be able to use of this.

[qcolombet]

Hi Aditya,

Could you use a delegate pattern like what we do for MachineRegisterInfio and LiveRangeEdit?

Other than that, the direction looks reasonable.

Cheers,
-Quentin

[MatzeB]

I've seen callbacks lead to unintuitive/hard to understand code. Can you motivate the use case some more?

[MatzeB]

Some additional coding style comments

include/llvm/CodeGen/MachineFunction.h

355–356

Use InsertCallbackFn = nullptr; here instead of setting it in init().

401–409

You should try to make these private with some friend declaration for the ilist traits.

lib/CodeGen/MachineFunction.cpp

249–250

I think this function is mostly about freeing memory. Setting the pointer should not be necessary.

[aditya_nandakumar]

In D52127#1237084, @MatzeB wrote:

I've seen callbacks lead to unintuitive/hard to understand code. Can you motivate the use case some more?

One of the motivations was the above - guaranteeing that the legalizer will always legalize intermediate instructions - right now it only guarantees legalization when the intermediate instructions are created with the MachineIRBuilder that's passed in - it would be nice to not rely on that mechanism. Besides insertions, it would be nice to keep track of instruction deletions -for eg in the combiner - it's useful to know which instructions were erased and needs to be removed the various worklists rather than conservatively assuming that the current instruction is always considered erased even if it isn't. Additionally if a combine deletes any other instruction besides the current root, we will currently not know about this and will result in calling combine(MI) on a deleted instruction resulting in UB (same in the legalizer).
Lastly I think this mechanism will allow a CSEing builder to also keep track of new instructions(again without relying on usage of MIRBuilder) as well as deletions (currently very hard).

[aditya_nandakumar]

In D52127#1236991, @qcolombet wrote:

Hi Aditya,

Could you use a delegate pattern like what we do for MachineRegisterInfio and LiveRangeEdit?

Other than that, the direction looks reasonable.

I will update it shortly with that approach. Thanks

Cheers,
-Quentin

[aditya_nandakumar]

Updated based on feedback from Quentin and Matthias.

[MatzeB]

LGTM with nitpicks addressed.

include/llvm/CodeGen/MachineFunction.h

372–373

Use const MachineInstr& here to make it obvious that the parameters cannot be nullptr. (We didn't do that in older APIs but I think we should do it for new APIs we create).

380–381

should also use references.

401

Add doxygen comment to public API.

410

Add a doxygen comment.

[aditya_nandakumar]

Submitted in r342696. Thanks Matthias for the review.

include/llvm/CodeGen/MachineFunction.h

class MachineFunction {
   /// unchanged the references to Fn, Target, MMI, and FunctionNumber.
   void clear();
   /// Allocate and initialize the different members.
   /// In particular, the XXXInfo data structure.
   /// \pre Fn, Target, MMI, and FunctionNumber are properly set.
   void init();
 
 public:
   struct VariableDbgInfo {
     const DILocalVariable *Var;
     const DIExpression *Expr;
     // The Slot can be negative for fixed stack objects.
     int Slot;
     const DILocation *Loc;
 
     VariableDbgInfo(const DILocalVariable *Var, const DIExpression *Expr,
                     int Slot, const DILocation *Loc)
         : Var(Var), Expr(Expr), Slot(Slot), Loc(Loc) {}
   };
+
+  class Delegate {
+    virtual void anchor();
+
+  public:
+    virtual ~Delegate() = default;
+    virtual void MF_HandleInsertion(const MachineInstr *MI) = 0;
+    virtual void MF_HandleRemoval(const MachineInstr *MI) = 0;
+  };
+
+private:
+  Delegate *TheDelegate = nullptr;
+
+  // Callbacks for insertion and removal.
+  void handleInsertion(const MachineInstr *MI);
+  void handleRemoval(const MachineInstr *MI);
+  friend struct ilist_traits<MachineInstr>;
+
+public:
   using VariableDbgInfoMapTy = SmallVector<VariableDbgInfo, 4>;
   VariableDbgInfoMapTy VariableDbgInfos;
 
   MachineFunction(const Function &F, const TargetMachine &Target,
                   const TargetSubtargetInfo &STI, unsigned FunctionNum,
                   MachineModuleInfo &MMI);
   MachineFunction(const MachineFunction &) = delete;
   MachineFunction &operator=(const MachineFunction &) = delete;
   ~MachineFunction();
 
   /// Reset the instance as if it was just created.
   void reset() {
     clear();
     init();
   }
 
+  void resetDelegate(Delegate *delegate) {
+    // Ensure another delegate does not take over unless the current
+    // delegate first unattaches itself. If we ever need to multicast
+    // notifications, we will need to change to using a list.
+    assert(TheDelegate == delegate &&
+           "Only the current delegate can perform reset!");
+    TheDelegate = nullptr;
+  }
+
+  void setDelegate(Delegate *delegate) {
+    assert(delegate && !TheDelegate &&
+           "Attempted to set delegate to null, or to change it without "
+           "first resetting it!");
+
+    TheDelegate = delegate;
+  }
+
   MachineModuleInfo &getMMI() const { return MMI; }
   MCContext &getContext() const { return Ctx; }
 
   PseudoSourceValueManager &getPSVManager() const { return *PSVManager; }
 
   /// Return the DataLayout attached to the Module associated to this MF.
   const DataLayout &getDataLayout() const;
 

lib/CodeGen/MachineBasicBlock.cpp

 void ilist_traits<MachineInstr>::addNodeToList(MachineInstr *N) {
   assert(!N->getParent() && "machine instruction already in a basic block");
   N->setParent(Parent);
 
   // Add the instruction's register operands to their corresponding
   // use/def lists.
   MachineFunction *MF = Parent->getParent();
   N->AddRegOperandsToUseLists(MF->getRegInfo());
+  MF->handleInsertion(N);
 }
 
 /// When we remove an instruction from a basic block list, we update its parent
 /// pointer and remove its operands from reg use/def lists if appropriate.
 void ilist_traits<MachineInstr>::removeNodeFromList(MachineInstr *N) {
   assert(N->getParent() && "machine instruction not in a basic block");
 
   // Remove from the use/def lists.
-  if (MachineFunction *MF = N->getMF())
+  if (MachineFunction *MF = N->getMF()) {
+    MF->handleRemoval(N);
     N->RemoveRegOperandsFromUseLists(MF->getRegInfo());
+  }
 
   N->setParent(nullptr);
 }
 
 /// When moving a range of instructions from one MBB list to another, we need to
 /// update the parent pointers and the use/def lists.
 void ilist_traits<MachineInstr>::transferNodesFromList(ilist_traits &FromList,
                                                        instr_iterator First,

lib/CodeGen/MachineFunction.cpp

static const char *getPropertyName(MachineFunctionProperties::Property Prop) {
   case P::NoVRegs: return "NoVRegs";
   case P::RegBankSelected: return "RegBankSelected";
   case P::Selected: return "Selected";
   case P::TracksLiveness: return "TracksLiveness";
   }
   llvm_unreachable("Invalid machine function property");
 }
 
+// Pin the vtable to this file.
+void MachineFunction::Delegate::anchor() {}
+
 void MachineFunctionProperties::print(raw_ostream &OS) const {
   const char *Separator = "";
   for (BitVector::size_type I = 0; I < Properties.size(); ++I) {
     if (!Properties[I])
       continue;
     OS << Separator << getPropertyName(static_cast<Property>(I));
     Separator = ", ";
   }

 MachineFunction::MachineFunction(const Function &F, const TargetMachine &Target,
                                  const TargetSubtargetInfo &STI,
                                  unsigned FunctionNum, MachineModuleInfo &mmi)
     : F(F), Target(Target), STI(&STI), Ctx(mmi.getContext()), MMI(mmi) {
   FunctionNumber = FunctionNum;
   init();
 }
 
+void MachineFunction::handleInsertion(const MachineInstr *MI) {
+  if (TheDelegate)
+    TheDelegate->MF_HandleInsertion(MI);
+}
+
+void MachineFunction::handleRemoval(const MachineInstr *MI) {
+  if (TheDelegate)
+    TheDelegate->MF_HandleRemoval(MI);
+}
+
 void MachineFunction::init() {
   // Assume the function starts in SSA form with correct liveness.
   Properties.set(MachineFunctionProperties::Property::IsSSA);
   Properties.set(MachineFunctionProperties::Property::TracksLiveness);
   if (STI->getRegisterInfo())
     RegInfo = new (Allocator) MachineRegisterInfo(this);
   else
     RegInfo = nullptr;
if (JumpTableInfo) {
     JumpTableInfo->~MachineJumpTableInfo();
     Allocator.Deallocate(JumpTableInfo);
   }
 
   if (WinEHInfo) {
     WinEHInfo->~WinEHFuncInfo();
     Allocator.Deallocate(WinEHInfo);
   }
 }
 
 const DataLayout &MachineFunction::getDataLayout() const {
   return F.getParent()->getDataLayout();
 }
 
 /// Get the JumpTableInfo for this function.
 /// If it does not already exist, allocate one.
 MachineJumpTableInfo *MachineFunction::
 getOrCreateJumpTableInfo(unsigned EntryKind) {