[CodeGen] Finish untangling LRE::scanRemattable [nfc]

[preames]

This is an attempt to simplify the rematerialization logic in InlineSpiller and SplitKit. I'd earlier done the same for RegisterCoalescer in 57b673.

The basic idea of this change is that we don't need to check whether an instruction is rematerializable early. Instead, we can defer the check to the point where we're actually trying to materialize something. We also don't need to indirect that query through a VNI key, and can instead just check the instruction directly at the use site.

I would appreciate careful review on this one. I have a bad feeling that I might be missing something subtle here. In the same area of code, I recently tried to strengthen asserts and hit a very subtle interaction with removing dead remats. That case is properly addressed and commented in this patch, but there might be something else. One worry is compile time; I don't directly see how checking the remat callback on each def's use (instead of once per def) is going to be expensive compared to code which is already walking all of those uses, but maybe I'm missing a case?

[llvmbot]

@llvm/pr-subscribers-llvm-regalloc

Author: Philip Reames (preames)

Changes

This is an attempt to simplify the rematerialization logic in InlineSpiller and SplitKit. I'd earlier done the same for RegisterCoalescer in 57b673.

The basic idea of this change is that we don't need to check whether an instruction is rematerializable early. Instead, we can defer the check to the point where we're actually trying to materialize something. We also don't need to indirect that query through a VNI key, and can instead just check the instruction directly at the use site.

I would appreciate careful review on this one. I have a bad feeling that I might be missing something subtle here. In the same area of code, I recently tried to strengthen asserts and hit a very subtle interaction with removing dead remats. That case is properly addressed and commented in this patch, but there might be something else. One worry is compile time; I don't directly see how checking the remat callback on each def's use (instead of once per def) is going to be expensive compared to code which is already walking all of those uses, but maybe I'm missing a case?

---

Full diff: https://github.com/llvm/llvm-project/pull/161963.diff

4 Files Affected:

• (modified) llvm/include/llvm/CodeGen/LiveRangeEdit.h (+2-17)
• (modified) llvm/lib/CodeGen/InlineSpiller.cpp (+14-6)
• (modified) llvm/lib/CodeGen/LiveRangeEdit.cpp (+3-32)
• (modified) llvm/lib/CodeGen/SplitKit.cpp (+1-3)

diff --git a/llvm/include/llvm/CodeGen/LiveRangeEdit.h b/llvm/include/llvm/CodeGen/LiveRangeEdit.h
index 6473138a801f7..d0ed3ff660d9b 100644
--- a/llvm/include/llvm/CodeGen/LiveRangeEdit.h
+++ b/llvm/include/llvm/CodeGen/LiveRangeEdit.h
@@ -75,24 +75,14 @@ class LiveRangeEdit : private MachineRegisterInfo::Delegate {
   /// FirstNew - Index of the first register added to NewRegs.
   const unsigned FirstNew;
 
-  /// ScannedRemattable - true when remattable values have been identified.
-  bool ScannedRemattable = false;
-
   /// DeadRemats - The saved instructions which have already been dead after
   /// rematerialization but not deleted yet -- to be done in postOptimization.
   SmallPtrSet<MachineInstr *, 32> *DeadRemats;
 
-  /// Remattable - Values defined by remattable instructions as identified by
-  /// tii.isTriviallyReMaterializable().
-  SmallPtrSet<const VNInfo *, 4> Remattable;
-
   /// Rematted - Values that were actually rematted, and so need to have their
   /// live range trimmed or entirely removed.
   SmallPtrSet<const VNInfo *, 4> Rematted;
 
-  /// scanRemattable - Identify the Parent values that may rematerialize.
-  void scanRemattable();
-
   /// foldAsLoad - If LI has a single use and a single def that can be folded as
   /// a load, eliminate the register by folding the def into the use.
   bool foldAsLoad(LiveInterval *LI, SmallVectorImpl<MachineInstr *> &Dead);
@@ -175,11 +165,6 @@ class LiveRangeEdit : private MachineRegisterInfo::Delegate {
 
   Register create() { return createFrom(getReg()); }
 
-  /// anyRematerializable - Return true if any parent values may be
-  /// rematerializable.  This function must be called before
-  /// canRematerializeAt is called..
-  bool anyRematerializable();
-
   /// Remat - Information needed to rematerialize at a specific location.
   struct Remat {
     const VNInfo *const ParentVNI;  // parent_'s value at the remat location.
@@ -189,9 +174,9 @@ class LiveRangeEdit : private MachineRegisterInfo::Delegate {
     explicit Remat(const VNInfo *ParentVNI) : ParentVNI(ParentVNI) {}
   };
 
-  /// canRematerializeAt - Determine if ParentVNI can be rematerialized at
+  /// canRematerializeAt - Determine if RM.Orig can be rematerialized at
   /// UseIdx. It is assumed that parent_.getVNINfoAt(UseIdx) == ParentVNI.
-  bool canRematerializeAt(Remat &RM, VNInfo *OrigVNI, SlotIndex UseIdx);
+  bool canRematerializeAt(Remat &RM, SlotIndex UseIdx);
 
   /// rematerializeAt - Rematerialize RM.ParentVNI into DestReg by inserting an
   /// instruction into MBB before MI. The new instruction is mapped, but
diff --git a/llvm/lib/CodeGen/InlineSpiller.cpp b/llvm/lib/CodeGen/InlineSpiller.cpp
index 0c2b74c907d2a..8867e0c0352b4 100644
--- a/llvm/lib/CodeGen/InlineSpiller.cpp
+++ b/llvm/lib/CodeGen/InlineSpiller.cpp
@@ -671,10 +671,21 @@ bool InlineSpiller::reMaterializeFor(LiveInterval &VirtReg, MachineInstr &MI) {
 
   LiveInterval &OrigLI = LIS.getInterval(Original);
   VNInfo *OrigVNI = OrigLI.getVNInfoAt(UseIdx);
-  LiveRangeEdit::Remat RM(ParentVNI);
-  RM.OrigMI = LIS.getInstructionFromIndex(OrigVNI->def);
+  assert(OrigVNI && "corrupted sub-interval");
+  MachineInstr *DefMI = LIS.getInstructionFromIndex(OrigVNI->def);
+  // This can happen if for two reasons: 1) This could be a phi valno,
+  // or 2) the remat def has already been removed from the original
+  // live interval; this happens if we rematted to all uses, and
+  // then further split one of those live ranges.
+  if (!DefMI) {
+    markValueUsed(&VirtReg, ParentVNI);
+    LLVM_DEBUG(dbgs() << "\tcannot remat missing def for " << UseIdx << '\t' << MI);
+    return false;
+  }
 
-  if (!Edit->canRematerializeAt(RM, OrigVNI, UseIdx)) {
+  LiveRangeEdit::Remat RM(ParentVNI);
+  RM.OrigMI = DefMI;
+  if (!Edit->canRematerializeAt(RM, UseIdx)) {
     markValueUsed(&VirtReg, ParentVNI);
     LLVM_DEBUG(dbgs() << "\tcannot remat for " << UseIdx << '\t' << MI);
     return false;
@@ -739,9 +750,6 @@ bool InlineSpiller::reMaterializeFor(LiveInterval &VirtReg, MachineInstr &MI) {
 /// reMaterializeAll - Try to rematerialize as many uses as possible,
 /// and trim the live ranges after.
 void InlineSpiller::reMaterializeAll() {
-  if (!Edit->anyRematerializable())
-    return;
-
   UsedValues.clear();
 
   // Try to remat before all uses of snippets.
diff --git a/llvm/lib/CodeGen/LiveRangeEdit.cpp b/llvm/lib/CodeGen/LiveRangeEdit.cpp
index 59bc82dc267b5..209e95fcdd47a 100644
--- a/llvm/lib/CodeGen/LiveRangeEdit.cpp
+++ b/llvm/lib/CodeGen/LiveRangeEdit.cpp
@@ -68,41 +68,12 @@ Register LiveRangeEdit::createFrom(Register OldReg) {
   return VReg;
 }
 
-void LiveRangeEdit::scanRemattable() {
-  for (VNInfo *VNI : getParent().valnos) {
-    if (VNI->isUnused())
-      continue;
-    Register Original = VRM->getOriginal(getReg());
-    LiveInterval &OrigLI = LIS.getInterval(Original);
-    VNInfo *OrigVNI = OrigLI.getVNInfoAt(VNI->def);
-    if (!OrigVNI)
-      continue;
-    MachineInstr *DefMI = LIS.getInstructionFromIndex(OrigVNI->def);
-    if (!DefMI)
-      continue;
-    if (TII.isReMaterializable(*DefMI))
-      Remattable.insert(OrigVNI);
-  }
-  ScannedRemattable = true;
-}
-
-bool LiveRangeEdit::anyRematerializable() {
-  if (!ScannedRemattable)
-    scanRemattable();
-  return !Remattable.empty();
-}
-
-bool LiveRangeEdit::canRematerializeAt(Remat &RM, VNInfo *OrigVNI,
-                                       SlotIndex UseIdx) {
-  assert(ScannedRemattable && "Call anyRematerializable first");
+bool LiveRangeEdit::canRematerializeAt(Remat &RM, SlotIndex UseIdx) {
+  assert(RM.OrigMI && "No defining instruction for remattable value");
 
-  // Use scanRemattable info.
-  if (!Remattable.count(OrigVNI))
+  if (!TII.isReMaterializable(*RM.OrigMI))
     return false;
 
-  // No defining instruction provided.
-  assert(RM.OrigMI && "No defining instruction for remattable value");
-
   // Verify that all used registers are available with the same values.
   if (!VirtRegAuxInfo::allUsesAvailableAt(RM.OrigMI, UseIdx, LIS, MRI, TII))
     return false;
diff --git a/llvm/lib/CodeGen/SplitKit.cpp b/llvm/lib/CodeGen/SplitKit.cpp
index f118ee521925f..f9ecb2c97b2e0 100644
--- a/llvm/lib/CodeGen/SplitKit.cpp
+++ b/llvm/lib/CodeGen/SplitKit.cpp
@@ -376,8 +376,6 @@ void SplitEditor::reset(LiveRangeEdit &LRE, ComplementSpillMode SM) {
   if (SpillMode)
     LICalc[1].reset(&VRM.getMachineFunction(), LIS.getSlotIndexes(), &MDT,
                     &LIS.getVNInfoAllocator());
-
-  Edit->anyRematerializable();
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
@@ -638,7 +636,7 @@ VNInfo *SplitEditor::defFromParent(unsigned RegIdx, const VNInfo *ParentVNI,
     LiveRangeEdit::Remat RM(ParentVNI);
     RM.OrigMI = LIS.getInstructionFromIndex(OrigVNI->def);
     if (RM.OrigMI && TII.isAsCheapAsAMove(*RM.OrigMI) &&
-        Edit->canRematerializeAt(RM, OrigVNI, UseIdx)) {
+        Edit->canRematerializeAt(RM, UseIdx)) {
       if (!rematWillIncreaseRestriction(RM.OrigMI, MBB, UseIdx)) {
         SlotIndex Def = Edit->rematerializeAt(MBB, I, Reg, RM, TRI, Late);
         ++NumRemats;

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[preames]

Doesn't look like much compile time impact if any: https://llvm-compile-time-tracker.com/compare.php?from=25126117b5781e96453e5c5b1a9a6a6f8aa3989c&to=fe927580c21af14d706275f2859fd3ba6f9718cc&stat=instructions:u