[GlobalISel]: Fix some non determinism exposed in CSE due to not notifying observers about mutations + add verification for CSE

https://reviews.llvm.org/D67133

While investigating some non determinism (CSE doesn't produce wrong
code, it just doesn't CSE some times) in GISel CSE on an out of tree
target, I realized that the core issue was that there were lots of code
that mutates (setReg, setRegClass etc), but doesn't notify observers
(CSE in this case but this could be any other observer). In order to
make the Observer be available in various parts of code and to avoid
having to thread it through various API, the MachineFunction now has the
observer as field. This allows it to be easily used in helper functions
such as constrainOperandRegClass.
Also added some invariant verification method in CSEInfo which can
catch these issues (when CSE is enabled).

llvm/include/llvm/CodeGen/GlobalISel/CSEInfo.h

@@ -120,6 +120,8 @@ public:
 
   void setMF(MachineFunction &MF);
 
+  Error verify();
+
   /// Records a newly created inst in a list and lazily insert it to the CSEMap.
   /// Sometimes, this method might be called with a partially constructed
   /// MachineInstr,

llvm/include/llvm/CodeGen/GlobalISel/GISelChangeObserver.h

@@ -101,7 +101,7 @@ public:
   void MF_HandleRemoval(MachineInstr &MI) override { erasingInstr(MI); }
 };
 
-/// A simple RAII based CSEInfo installer.
+/// A simple RAII based Delegate installer.
 /// Use this in a scope to install a delegate to the MachineFunction and reset
 /// it at the end of the scope.
 class RAIIDelegateInstaller {
@@ -113,5 +113,27 @@ public:
   ~RAIIDelegateInstaller();
 };
 
+/// A simple RAII based Observer installer.
+/// Use this in a scope to install the Observer to the MachineFunction and reset
+/// it at the end of the scope.
+class RAIIMFObserverInstaller {
+  MachineFunction &MF;
+
+public:
+  RAIIMFObserverInstaller(MachineFunction &MF, GISelChangeObserver &Observer);
+  ~RAIIMFObserverInstaller();
+};
+
+/// Class to install both of the above.
+class RAIIMFObsDelInstaller {
+  RAIIDelegateInstaller DelI;
+  RAIIMFObserverInstaller ObsI;
+
+public:
+  RAIIMFObsDelInstaller(MachineFunction &MF, GISelObserverWrapper &Wrapper)
+      : DelI(MF, &Wrapper), ObsI(MF, Wrapper) {}
+  ~RAIIMFObsDelInstaller() = default;
+};
+
 } // namespace llvm
 #endif

llvm/include/llvm/CodeGen/MachineFunction.h

@@ -72,6 +72,7 @@ class TargetRegisterClass;
 class TargetSubtargetInfo;
 struct WasmEHFuncInfo;
 struct WinEHFuncInfo;
+class GISelChangeObserver;
 
 template <> struct ilist_alloc_traits<MachineBasicBlock> {
   void deleteNode(MachineBasicBlock *MBB);
@@ -396,6 +397,7 @@ public:
 
 private:
   Delegate *TheDelegate = nullptr;
+  GISelChangeObserver *Observer = nullptr;
 
   using CallSiteInfoMap = DenseMap<const MachineInstr *, CallSiteInfo>;
   /// Map a call instruction to call site arguments forwarding info.
@@ -444,6 +446,10 @@ public:
     TheDelegate = delegate;
   }
 
+  void setObserver(GISelChangeObserver *O) { Observer = O; }
+
+  GISelChangeObserver *getObserver() const { return Observer; }
+
   MachineModuleInfo &getMMI() const { return MMI; }
   MCContext &getContext() const { return Ctx; }
 

llvm/lib/CodeGen/GlobalISel/CSEInfo.cpp

@@ -261,6 +261,39 @@ void GISelCSEInfo::releaseMemory() {
 #endif
 }
 
+Error GISelCSEInfo::verify() {
+#ifndef NDEBUG
+  handleRecordedInsts();
+  // For each instruction in map from MI -> UMI,
+  // Profile(MI) and make sure UMI is found for that profile.
+  for (auto &It : InstrMapping) {
+    FoldingSetNodeID TmpID;
+    GISelInstProfileBuilder(TmpID, *MRI).addNodeID(It.first);
+    void *InsertPos;
+    UniqueMachineInstr *FoundNode =
+        CSEMap.FindNodeOrInsertPos(TmpID, InsertPos);
+    if (FoundNode != It.second)
+      return createStringError(std::errc::not_supported,
+                               "CSEMap mismatch, InstrMapping has MIs without "
+                               "corresponding Nodes in CSEMap");
+  }
+
+  // For every node in the CSEMap, make sure that the InstrMapping
+  // points to it.
+  for (auto It = CSEMap.begin(), End = CSEMap.end(); It != End; ++It) {
+    const UniqueMachineInstr &UMI = *It;
+    if (!InstrMapping.count(UMI.MI))
+      return createStringError(std::errc::not_supported,
+                               "Node in CSE without InstrMapping", UMI.MI);
+
+    if (InstrMapping[UMI.MI] != &UMI)
+      return createStringError(std::make_error_code(std::errc::not_supported),
+                               "Mismatch in CSE mapping");
+  }
+#endif
+  return Error::success();
+}
+
 void GISelCSEInfo::print() {
   LLVM_DEBUG(for (auto &It
                   : OpcodeHitTable) {
@@ -370,6 +403,7 @@ GISelCSEInfo &
 GISelCSEAnalysisWrapper::get(std::unique_ptr<CSEConfigBase> CSEOpt,
                              bool Recompute) {
   if (!AlreadyComputed || Recompute) {
+    Info.releaseMemory();
     Info.setCSEConfig(std::move(CSEOpt));
     Info.analyze(*MF);
     AlreadyComputed = true;

llvm/lib/CodeGen/GlobalISel/GISelChangeObserver.cpp

@@ -38,3 +38,11 @@ RAIIDelegateInstaller::RAIIDelegateInstaller(MachineFunction &MF,
 }
 
 RAIIDelegateInstaller::~RAIIDelegateInstaller() { MF.resetDelegate(Delegate); }
+
+RAIIMFObserverInstaller::RAIIMFObserverInstaller(MachineFunction &MF,
+                                                 GISelChangeObserver &Observer)
+    : MF(MF) {
+  MF.setObserver(&Observer);
+}
+
+RAIIMFObserverInstaller::~RAIIMFObserverInstaller() { MF.setObserver(nullptr); }

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

@@ -2381,6 +2381,7 @@ bool IRTranslator::runOnMachineFunction(MachineFunction &CurMF) {
     WrapperObserver.addObserver(&Verifier);
 #endif // ifndef NDEBUG
     RAIIDelegateInstaller DelInstall(*MF, &WrapperObserver);
+    RAIIMFObserverInstaller ObsInstall(*MF, WrapperObserver);
     for (const BasicBlock *BB : RPOT) {
       MachineBasicBlock &MBB = getMBB(*BB);
       // Set the insertion point of all the following translations to

llvm/lib/CodeGen/GlobalISel/Legalizer.cpp

@@ -28,6 +28,7 @@
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Target/TargetMachine.h"
 
 #include <iterator>
@@ -180,7 +181,7 @@ Legalizer::legalizeMachineFunction(MachineFunction &MF, const LegalizerInfo &LI,
 
   // Now install the observer as the delegate to MF.
   // This will keep all the observers notified about new insertions/deletions.
-  RAIIDelegateInstaller DelInstall(MF, &WrapperObserver);
+  RAIIMFObsDelInstaller Installer(MF, WrapperObserver);
   LegalizerHelper Helper(MF, LI, WrapperObserver, MIRBuilder);
   LegalizationArtifactCombiner ArtCombiner(MIRBuilder, MRI, LI);
   auto RemoveDeadInstFromLists = [&WrapperObserver](MachineInstr *DeadMI) {
@@ -305,6 +306,7 @@ bool Legalizer::runOnMachineFunction(MachineFunction &MF) {
     // We want CSEInfo in addition to WorkListObserver to observe all changes.
     AuxObservers.push_back(CSEInfo);
   }
+  assert(!CSEInfo || !errorToBool(CSEInfo->verify()));
 
   const LegalizerInfo &LI = *MF.getSubtarget().getLegalizerInfo();
   MFResult Result = legalizeMachineFunction(MF, LI, AuxObservers, *MIRBuilder);
@@ -324,5 +326,11 @@ bool Legalizer::runOnMachineFunction(MachineFunction &MF) {
     reportGISelFailure(MF, TPC, MORE, R);
     return false;
   }
+  // If for some reason CSE was not enabled, make sure that we invalidate the
+  // CSEInfo object (as we currently declare that the analysis is preserved).
+  // The next time get on the wrapper is called, it will force it to recompute
+  // the analysis.
+  if (!EnableCSE)
+    Wrapper.setComputed(false);
   return Result.Changed;
 }

llvm/lib/CodeGen/GlobalISel/Utils.cpp

@@ -12,6 +12,7 @@
 #include "llvm/CodeGen/GlobalISel/Utils.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/Twine.h"
+#include "llvm/CodeGen/GlobalISel/GISelChangeObserver.h"
 #include "llvm/CodeGen/GlobalISel/RegisterBankInfo.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
@@ -62,6 +63,15 @@ Register llvm::constrainOperandRegClass(
               TII.get(TargetOpcode::COPY), Reg)
           .addReg(ConstrainedReg);
     }
+  } else {
+    if (GISelChangeObserver *Observer = MF.getObserver()) {
+      if (!RegMO.isDef()) {
+        MachineInstr *RegDef = MRI.getVRegDef(Reg);
+        Observer->changedInstr(*RegDef);
+      }
+      Observer->changingAllUsesOfReg(MRI, Reg);
+      Observer->finishedChangingAllUsesOfReg();
+    }
   }
   return ConstrainedReg;
 }