[llvm] [RegisterCoalescer] Materialize undef subrange PHI inputs (PR #204867)

[Lukas Sommer via llvm-commits]

================
@@ -3692,6 +3833,682 @@ bool RegisterCoalescer::isHighCostLiveInterval(LiveInterval &LI) {
   return true;
 }
 
+static LaneBitmask
+getDefMaskForUndefPHIRepairSubReg(const TargetRegisterInfo &TRI,
+                                  LaneBitmask MaxMask, unsigned SubReg) {
+  // If no subregister is used, the mask is the full mask of the register class
+  // as given by MaxMask
+  if (SubReg == 0)
+    return MaxMask;
+
+  // Otherwise ask the target for the lane mask for the given subregister and
+  // make sure it is a subset of the maximum mask for the register class
+  LaneBitmask DefMask = TRI.getSubRegIndexLaneMask(SubReg);
+  return DefMask & MaxMask;
+}
+
+static bool getUndefPHIRepairSubRegIndexes(const TargetRegisterInfo &TRI,
+                                           const TargetRegisterClass *RC,
+                                           LaneBitmask MissingMask,
+                                           LaneBitmask AllowedMask,
+                                           SmallVectorImpl<unsigned> &Indexes) {
+  assert((MissingMask & ~AllowedMask).none() &&
+         "missing lanes must be legal to repair");
+
+  // Find the best subregister indexes to cover the missing lanes where each
+  // index covers as many missing lanes as possible, and as few extra lanes as
+  // possible. When legal single-lane indexes exist, this can fall back to one
+  // index per missing lane.
+  LaneBitmask LanesLeft = MissingMask;
+  while (LanesLeft.any()) {
+    unsigned BestIdx = 0;
+    unsigned BestCover = 0;
+    unsigned BestExtra = std::numeric_limits<unsigned>::max();
+
+    for (unsigned Idx = 1, E = TRI.getNumSubRegIndices(); Idx < E; ++Idx) {
+      // Skip subregister indexes that don't belong to the register class
+      if (TRI.getSubClassWithSubReg(RC, Idx) != RC)
+        continue;
+
+      LaneBitmask SubRegMask = TRI.getSubRegIndexLaneMask(Idx);
+      LaneBitmask Covered = SubRegMask & LanesLeft;
+
+      // The subregister doesn't cover any of the lanes we're still missing
+      if (Covered.none())
+        continue;
+
+      // The chosen subregister would clobber lanes that are not allowed to be
+      // defined by the repair
+      if ((SubRegMask & ~AllowedMask).any())
+        continue;
+
+      // Skip subregisters that overlap with missing lanes we have already
+      // covered with other subregisters
+      if ((SubRegMask & (MissingMask & ~LanesLeft)).any())
+        continue;
+
+      unsigned Cover = Covered.getNumLanes();
+      unsigned Extra = (SubRegMask & ~LanesLeft).getNumLanes();
+      if (Cover > BestCover || (Cover == BestCover && Extra < BestExtra)) {
+        BestIdx = Idx;
+        BestCover = Cover;
+        BestExtra = Extra;
+      }
+    }
+
+    if (BestIdx == 0)
+      return false;
+
+    Indexes.push_back(BestIdx);
+    LanesLeft &= ~TRI.getSubRegIndexLaneMask(BestIdx);
+  }
+
+  return true;
+}
+
+static void copyLiveInterval(LiveInterval &Dst, const LiveInterval &Src,
+                             BumpPtrAllocator &Allocator) {
+  Dst.assign(Src, Allocator);
+  for (const LiveInterval::SubRange &SR : Src.subranges())
+    Dst.createSubRangeFrom(Allocator, SR.LaneMask, SR);
+}
+
+static void addUndefPHIRepair(SmallVectorImpl<UndefPHIRepair> &Repairs,
+                              MachineBasicBlock *PredMBB,
+                              MachineBasicBlock *SuccMBB,
+                              LaneBitmask LaneMask) {
+  // A predecessor can be missing several lanes for the same successor PHI.
+  // Combining them lets us emit the fewest IMPLICIT_DEFs that cover the
+  // register class instead of repairing each subrange independently.
+  for (UndefPHIRepair &R : Repairs) {
+    if (R.PredMBB == PredMBB && R.SuccMBB == SuccMBB) {
+      R.LaneMask |= LaneMask;
+      return;
+    }
+  }
+  Repairs.push_back({PredMBB, SuccMBB, LaneMask});
+}
+
+static void collectUndefPHIRepairs(
+    LiveInterval &LI, const MachineRegisterInfo &MRI,
+    const SlotIndexes &Indexes, SmallVectorImpl<UndefPHIRepair> &Repairs,
+    LaneBitmask TreatLiveOutAsMissing = LaneBitmask::getNone()) {
+  LaneBitmask VRegMask = MRI.getMaxLaneMaskForVReg(LI.reg());
+
+  for (LiveInterval::SubRange &SR : LI.subranges()) {
+    if (SR.empty())
+      continue;
+
+    // Check if the lanes of this subrange are marked as undefined by a
+    // subregister read-undef definition. We can use that to determine whether
+    // a repair will really be needed further down
+    SmallVector<SlotIndex, 4> Undefs;
+    LaneBitmask UndefMask = SR.LaneMask & VRegMask;
+    if (UndefMask.any())
+      LI.computeSubRangeUndefs(Undefs, UndefMask, MRI, Indexes);
+
+    for (const VNInfo *VNI : SR.valnos) {
+      if (VNI->isUnused() || !VNI->isPHIDef())
+        continue;
+
+      MachineBasicBlock *MBB = Indexes.getMBBFromIndex(VNI->def);
+      for (MachineBasicBlock *PredMBB : MBB->predecessors()) {
+        SlotIndex PredEnd = Indexes.getMBBEndIdx(PredMBB);
+
+        // For the pre-coalescing check, treat any lane of this subrange that
+        // overlaps with TreatLiveOutAsMissing as missing
+        bool ForceMissing =
+            (SR.LaneMask & TreatLiveOutAsMissing & VRegMask).any();
+
+        // PHI subranges are live-in to the block only if the same lane is live
+        // out of each predecessor. A missing live-out lane means this edge has
+        // an undef PHI input that must be made explicit by an IMPLICIT_DEF
+        if (!ForceMissing && SR.liveAt(PredEnd.getPrevSlot()))
+          continue;
+
+        // No need to repair anything if the predecessor is dominated by all the
+        // blocks corresponding to the SlotIndexes in Undefs
+        if (!ForceMissing &&
+            LiveRangeCalc::isJointlyDominated(PredMBB, Undefs, Indexes))
+          continue;
+
+        LLVM_DEBUG({
+          if (ForceMissing)
+            dbgs() << "\t\tTreating shrink-marked PHI lanes as missing on "
+                   << printMBBReference(*PredMBB) << " -> "
+                   << printMBBReference(*MBB) << ": "
+                   << PrintLaneMask(SR.LaneMask & TreatLiveOutAsMissing &
+                                    VRegMask)
+                   << '\n';
+        });
+
+        addUndefPHIRepair(Repairs, PredMBB, MBB, SR.LaneMask);
+      }
+    }
+  }
+}
+
+static bool repairNeedsEdgeSplit(const LiveInterval &LI,
+                                 const SlotIndexes &Indexes,
+                                 const UndefPHIRepair &Repair) {
+  // No need to split if the edge recorded by Repair is the only outgoing edge
+  // from the predecessor
+  if (Repair.PredMBB->succ_size() <= 1)
+    return false;
+
+  // The predecessor has multiple successors. Inserting an IMPLICIT_DEF would
+  // affect all outgoing edges, not just the one edge whose PHI input is undef.
+  // If the main interval is live at the end of the predecessor, there is a real
+  // live out value for that vreg on some outgoing path. Require an edge split
+  // to avoid clobbering that value.
+  SlotIndex LastPredSlot = Indexes.getMBBEndIdx(Repair.PredMBB).getPrevSlot();
+  return LI.liveAt(LastPredSlot);
+}
+
+static bool containsPHIDef(const LiveRange &LR) {
+  for (const VNInfo *VNI : LR.valnos)
+    if (!VNI->isUnused() && VNI->isPHIDef())
+      return true;
+  return false;
----------------
sommerlukas wrote:

Nit: You could use `llvm::any_of` here.

https://github.com/llvm/llvm-project/pull/204867