[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().

386–394

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

lib/CodeGen/MachineFunction.cpp

238–239

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

376–377

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).

384–385

should also use references.

386

Add doxygen comment to public API.

395

Add a doxygen comment.

[aditya_nandakumar]

Submitted in r342696. Thanks Matthias for the review.

include/llvm/CodeGen/MachineFunction.h

class MachineFunction {
   /// objects and get rid of all the XXXInfo data structure, but keep
   /// 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();
 
+  using MICallbackFnTy = std::function<void(const MachineInstr*)>;
+  MICallbackFnTy InsertCallbackFn;
+  MICallbackFnTy RemoveCallbackFn;
+
 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) {}
   };
   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 doInsertCallback(const MachineInstr *MI) {
+    if (InsertCallbackFn)
+      InsertCallbackFn(MI);
+  }
+
+  void doRemoveCallback(const MachineInstr *MI) {
+    if (RemoveCallbackFn)
+      RemoveCallbackFn(MI);
+  }
+
+  void setInsertCallback(MICallbackFnTy Fn) {
+    InsertCallbackFn = std::move(Fn);
+  }
+
+  void setRemoveCallback(MICallbackFnTy Fn) {
+    RemoveCallbackFn = std::move(Fn);
+  }
+
   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->doInsertCallback(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->doInsertCallback(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

void MachineFunction::init() {
 
   assert(Target.isCompatibleDataLayout(getDataLayout()) &&
          "Can't create a MachineFunction using a Module with a "
          "Target-incompatible DataLayout attached\n");
 
   PSVManager =
     llvm::make_unique<PseudoSourceValueManager>(*(getSubtarget().
                                                   getInstrInfo()));
+  InsertCallbackFn = nullptr;
+  RemoveCallbackFn = nullptr;
 }
 
 MachineFunction::~MachineFunction() {
   clear();
 }
 
 void MachineFunction::clear() {
   Properties.reset();
if (JumpTableInfo) {
     JumpTableInfo->~MachineJumpTableInfo();
     Allocator.Deallocate(JumpTableInfo);
   }
 
   if (WinEHInfo) {
     WinEHInfo->~WinEHFuncInfo();
     Allocator.Deallocate(WinEHInfo);
   }
+  InsertCallbackFn = nullptr;
+  RemoveCallbackFn = nullptr;
 }
 
 const DataLayout &MachineFunction::getDataLayout() const {
   return F.getParent()->getDataLayout();
 }
 
 /// Get the JumpTableInfo for this function.
 /// If it does not already exist, allocate one.