SplitKit: Correctly implement partial copies

[MatzeB]

Summary

• This fixes a bug where subregisters incompatible with the given vreg were choosen.
• Implement the case where multiple copies are necessary to cover a given lanemask.

rdar://30388692

[rampitec]

Thank you!

[rampitec]

defFromParent supposed to return a single def, how does it work if there is more than one? I suppose it does work as you should have tests passed...

[kparzysz]

lib/CodeGen/SplitKit.cpp

551

Is there any guarantee that this loop will actually terminate? In other words, do we know that the set PossibleIndexes covers the lanes completely?

572

Would it make sense to add || BestCover <= 0?

[MatzeB]

Redo a lot of this patch:

• Add more tests so that the code runs in more situations and screwed up liveranges are detected.
• Use instruction bundles during splitting/regalloc and expand them into a series of instruction at VirtRegMap time. I spend days trying to get this to work by directly creating a sequence of instructions in splitkit but it turned out to be extremely hard to stop liverange splitting between the copies in an infinite loop or getting a spill/reload for every one of the copies in a series; the bundle based approach only needed nearly no changes outside of splitkit itself except for the expansion code and fixing some virtregmap bugs with bundles.

[MatzeB]

In D30438#688098, @rampitec wrote:

defFromParent supposed to return a single def, how does it work if there is more than one? I suppose it does work as you should have tests passed...

Indeed this turned out to screw up the liveness and I just wasn't able to catch it with the small amount of tests we had before. Hence this new version of the patch.

[MatzeB]

Upload the correct patch this time!

Redo a lot of this patch:

• Add more tests so that the code runs in more situations and screwed up liveranges are detected.
• Use instruction bundles during splitting/regalloc and expand them into a series of instruction at VirtRegMap time. (I spend days trying to get this to work by directly creating a sequence of instructions in splitkit but it turned out to be extremely hard to stop liverange splitting between the copies in an infinite loop or getting a spill/reload for every one of the copies in a series; the bundle based approach only needed nearly no changes outside of splitkit itself except for the expansion code and fixing some virtregmap bugs with bundles.)

[MatzeB]

lib/CodeGen/SplitKit.cpp

551

This specific loop only checks each subregister index once and will therefore terminate.

The later loop is repeated as we are able to remove bits from the lanemask. If we fail to find a subregister index that removes any of the bits we will end up in the report_fatal_error() cases.

572

BestCover <= 0 can happen and is fine (it just means the best we could find is an index that covers more unnecessary lanes that we already covered anyway, than lanes we haven't covered yet).

[rampitec]

lib/CodeGen/VirtRegMap.cpp

384 ↗

(On Diff #91662)

I think it can break something which was bundled not by SplitKit. I wish we could mark these somehow... It will not happen now because we do not use bundles, but potentially a problem. Maybe insert some sort of a marker into the bundle or slot, or create a new MIFlag? Ideally of course it shall not use bundles just to keep LIS correct :(

[rampitec]

The other question, what would happen if target implements its own addOptimizedRegAlloc and skips VirtRegRewriter? I guess in this case copies will remain bundled. Looks like NVPTX does not use it.

[MatzeB]

In D30438#700956, @rampitec wrote:

The other question, what would happen if target implements its own addOptimizedRegAlloc and skips VirtRegRewriter? I guess in this case copies will remain bundled. Looks like NVPTX does not use it.

If they have their own regalloc they probably cannot use SplitKit either. If they do use SplitKit, well they have to implement their own expansion mechanism (though it seems SplitKit has references to VirtRegMap so I don't think you realistically will use SplitKit without the VirtRegRewriter).

lib/CodeGen/VirtRegMap.cpp

384 ↗

(On Diff #91662)

I was hoping that targets would always use a target specific bundle header instructions, so this would be the only situation in which we have a headerless bundle. Admittedly this is somewhat weak, but I hoped we can delay this discussion until we actually have targets using bundles pre-regalloc more.

[rampitec]

lib/CodeGen/VirtRegMap.cpp

384 ↗

(On Diff #91662)

Its almost the moment I'm thinking of adding it ;) It will not be all copy though.

[MatzeB]

lib/CodeGen/VirtRegMap.cpp

384 ↗

(On Diff #91662)

But do you plan to use bundles without a "header" (i.e. the first instruction is not some target-specific instruction that describes the type of the bundle?)

[rampitec]

lib/CodeGen/VirtRegMap.cpp

384 ↗

(On Diff #91662)

OK, agree, I do not.

[kparzysz]

lib/CodeGen/SplitKit.cpp

572

Well, wasn't that the original problem, i.e. that we copied too much (which then ended up clobbering lanes marked as unused)?

[MatzeB]

lib/CodeGen/SplitKit.cpp

572

No, the PossibleIndexes vector only contains indexes that cover bits that we are allowed to cover. The BestCover is solely about minimizing the amount of lanes we copy twice. We will not copy lanes that we are not supposed to copy.

[kparzysz]

lib/CodeGen/SplitKit.cpp

572

Ok, I see it now.

597

This doesn't need to be addressed in this patch, but on a general note---this would be a case where the compiler corners itself into such a situation. It's not a symptom of simply a bug, but rather a design that can end up in a dead end. Ideally, we should avoid getting here early enough to have other options.

test/CodeGen/AMDGPU/splitkit.mir

35 ↗

(On Diff #91662)

This is the only CHECK like in this function. Did it use to crash?

[MatzeB]

lib/CodeGen/SplitKit.cpp

597

I think this cannot happen. As for liveness to split up into a given lanemask we need an operand in the function that uses a subregister index with such a lanemask. But that would mean we can use the same index to implement the COPY.

I'm not a 100% sure that I am not missing something here so I left this check in to better abort than silently infinite loop if I did :).

[MatzeB]

test/CodeGen/AMDGPU/splitkit.mir

35 ↗

(On Diff #91662)

This should indeed have more CHECKs. I'll add them.

[qcolombet]

LGTM with more checks in the test.

test/CodeGen/AMDGPU/splitkit.mir

35 ↗

(On Diff #91662)

Please add more checks indeed.

lib/CodeGen/SplitKit.h

private:
   void extendPHIKillRanges();
 
   /// rewriteAssigned - Rewrite all uses of Edit.getReg() to assigned registers.
   void rewriteAssigned(bool ExtendRanges);
 
   /// deleteRematVictims - Delete defs that are dead after rematerializing.
   void deleteRematVictims();
 
+  /// Add a copy instruction copying \p FromReg to \p ToReg before
+  /// \p InsertBefore. This can be invoked with a \p LaneMask which may make it
+  /// necessary to construct a sequence of copies to cover it exactly.
+  SlotIndex buildCopy(unsigned FromReg, unsigned ToReg, LaneBitmask LaneMask,
+                      MachineBasicBlock &MBB,
+                      MachineBasicBlock::iterator InsertBefore, bool Late);
+
 public:
   /// Create a new SplitEditor for editing the LiveInterval analyzed by SA.
   /// Newly created intervals will be appended to newIntervals.
   SplitEditor(SplitAnalysis &SA, AliasAnalysis &AA, LiveIntervals&,
               VirtRegMap&, MachineDominatorTree&,
               MachineBlockFrequencyInfo &);
 
   /// reset - Prepare for a new split.

lib/CodeGen/SplitKit.cpp

 #include "llvm/CodeGen/LiveRangeEdit.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/VirtRegMap.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetInstrInfo.h"
 #include "llvm/Target/TargetMachine.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "regalloc"
 
void SplitEditor::forceRecompute(unsigned RegIdx, const VNInfo *ParentVNI) {
   // This was previously a single mapping. Make sure the old def is represented
   // by a trivial live range.
   addDeadDef(LIS.getInterval(Edit->get(RegIdx)), VNI, false);
 
   // Mark as complex mapped, forced.
   VFP = ValueForcePair(nullptr, true);
 }
 
+SlotIndex SplitEditor::buildCopy(unsigned FromReg, unsigned ToReg,
+                                 LaneBitmask LaneMask, MachineBasicBlock &MBB,
+                                 MachineBasicBlock::iterator InsertBefore,
+                                 bool Late) {
+  const MCInstrDesc &Desc = TII.get(TargetOpcode::COPY);
+  SlotIndexes &Indexes = *LIS.getSlotIndexes();
+  SlotIndex Def;
+  if (LaneMask.all() || LaneMask == MRI.getMaxLaneMaskForVReg(FromReg)) {
+    // The full vreg is copied.
+    MachineInstr *CopyMI =
+        BuildMI(MBB, InsertBefore, DebugLoc(), Desc, ToReg).addReg(FromReg);
+    Def = Indexes.insertMachineInstrInMaps(*CopyMI, Late).getRegSlot();
+  } else {
+    // Only a subset of lanes needs to be copied. The following is a simple
+    // heuristic to construct a sequence of COPYs. We could add a target
+    // specific callback if this turns out to be suboptimal.
+
+    // First pass: Try to find a perfectly matching subregister index. If none
+    // exists find the one covering the most lanemask bits.
+    SmallVector<unsigned, 8> PossibleIndexes;
+    unsigned BestIdx = 0;
+    unsigned BestCover = 0;
+    const TargetRegisterClass *RC = MRI.getRegClass(FromReg);
+    assert(RC == MRI.getRegClass(ToReg) && "Should have same reg class");
+    for (unsigned Idx = 1, E = TRI.getNumSubRegIndices(); Idx < E; ++Idx) {
+      // Is this index even compatible with the given class?
+      if (TRI.getSubClassWithSubReg(RC, Idx) != RC)
+        continue;
+      LaneBitmask SubRegMask = TRI.getSubRegIndexLaneMask(Idx);
+      // Early exit if we found a perfect match.
+      if (SubRegMask == LaneMask) {
+        BestIdx = Idx;
+        break;
+      }
+
+      // The index must not cover any lanes outside \p LaneMask.
+      if ((SubRegMask & ~LaneMask).any())
+        continue;
+
+      unsigned PopCount = countPopulation(SubRegMask.getAsInteger());
+      PossibleIndexes.push_back(Idx);
+      if (PopCount > BestCover) {
+        BestCover = PopCount;
+        BestIdx = Idx;
+      }
+    }
+
+    // Abort if we cannot possibly implement the COPY with the given indexes.
+    if (BestIdx == 0)
+      report_fatal_error("Impossible to implement partial COPY");
+
+    MachineInstr *CopyMI = BuildMI(MBB, InsertBefore, DebugLoc(), Desc)
+        .addReg(ToReg, RegState::Define | RegState::Undef, BestIdx)
+        .addReg(FromReg, 0, BestIdx);
+    Def = Indexes.insertMachineInstrInMaps(*CopyMI, Late).getRegSlot();
+
+    // Greedy heuristic: Keep iterating keeping the best covering subreg index
+    // each time.
+    LaneBitmask LanesLeft =
+        LaneMask & ~(TRI.getSubRegIndexLaneMask(BestCover));
+    while (LanesLeft.any()) {
+      unsigned BestIdx = 0;
+      int BestCover = INT_MIN;
+      for (unsigned Idx : PossibleIndexes) {
+        LaneBitmask SubRegMask = TRI.getSubRegIndexLaneMask(Idx);
+        // Early exit if we found a perfect match.
+        if (SubRegMask == LanesLeft) {
+          BestIdx = Idx;
+          break;
+        }
+
+        // Try to cover as much of the remaining lanes as possible but
+        // as few of the already covered lanes as possible.
+        int Cover = countPopulation((SubRegMask & LanesLeft).getAsInteger())
+                  - countPopulation((SubRegMask & ~LanesLeft).getAsInteger());
+        if (Cover > BestCover) {
+          BestCover = Cover;
+          BestIdx = Idx;
+        }
+      }
+
+      if (BestIdx == 0)
+        report_fatal_error("Impossible to implement partial COPY");
+
+      MachineInstr *CopyMI = BuildMI(MBB, InsertBefore, DebugLoc(), Desc)
+          .addReg(ToReg, RegState::Define, BestIdx)
+          .addReg(FromReg, 0, BestIdx);
+      Def = Indexes.insertMachineInstrInMaps(*CopyMI, Late).getRegSlot();
+      LanesLeft &= ~TRI.getSubRegIndexLaneMask(BestIdx);
+    }
+  }
+
+  return Def;
+}
+
 VNInfo *SplitEditor::defFromParent(unsigned RegIdx,
                                    VNInfo *ParentVNI,
                                    SlotIndex UseIdx,
                                    MachineBasicBlock &MBB,
                                    MachineBasicBlock::iterator I) {
-  MachineInstr *CopyMI = nullptr;
   SlotIndex Def;
   LiveInterval *LI = &LIS.getInterval(Edit->get(RegIdx));
 
   // We may be trying to avoid interference that ends at a deleted instruction,
   // so always begin RegIdx 0 early and all others late.
   bool Late = RegIdx != 0;
 
   // Attempt cheap-as-a-copy rematerialization.
   unsigned Original = VRM.getOriginal(Edit->get(RegIdx));
   LiveInterval &OrigLI = LIS.getInterval(Original);
   VNInfo *OrigVNI = OrigLI.getVNInfoAt(UseIdx);
 
   bool DidRemat = false;
+  unsigned Reg = LI->reg;
   if (OrigVNI) {
     LiveRangeEdit::Remat RM(ParentVNI);
     RM.OrigMI = LIS.getInstructionFromIndex(OrigVNI->def);
     if (Edit->canRematerializeAt(RM, OrigVNI, UseIdx, true)) {
-      Def = Edit->rematerializeAt(MBB, I, LI->reg, RM, TRI, Late);
+      Def = Edit->rematerializeAt(MBB, I, Reg, RM, TRI, Late);
       ++NumRemats;
       DidRemat = true;
     }
   }
   if (!DidRemat) {
-    // Can't remat, just insert a copy from parent.
-    CopyMI = BuildMI(MBB, I, DebugLoc(), TII.get(TargetOpcode::COPY), LI->reg)
-               .addReg(Edit->getReg());
-    Def = LIS.getSlotIndexes()
-              ->insertMachineInstrInMaps(*CopyMI, Late)
-              .getRegSlot();
+    LaneBitmask LaneMask;
     if (LI->hasSubRanges()) {
-      LaneBitmask LM = LaneBitmask::getNone();
+      LaneMask = LaneBitmask::getNone();
       for (LiveInterval::SubRange &S : LI->subranges())
-        LM |= S.LaneMask;
-
-      if (MRI.getMaxLaneMaskForVReg(LI->reg) != LM) {
-        // Find subreg for the lane mask.
-        unsigned SubIdx = 0;
-        for (unsigned I = 1, E = TRI.getNumSubRegIndices(); I < E; ++I) {
-          if (TRI.getSubRegIndexLaneMask(I) == LM) {
-            SubIdx = I;
-            break;
-          }
-        }
-        if (SubIdx == 0)
-          report_fatal_error("Cannot find subreg index to cover all alive lanes");
-        CopyMI->getOperand(0).setSubReg(SubIdx);
-        CopyMI->getOperand(1).setSubReg(SubIdx);
-        CopyMI->getOperand(0).setIsUndef(true);
-      }
+        LaneMask |= S.LaneMask;
+    } else {
+      LaneMask = LaneBitmask::getAll();
     }
+
+    Def = buildCopy(Edit->getReg(), Reg, LaneMask, MBB, I, Late);
     ++NumCopies;
   }
 
   // Define the value in Reg.
   return defValue(RegIdx, ParentVNI, Def, false);
 }
 
 /// Create a new virtual register and live interval.