CodeGen: Move function to get subregister indexes to cover a LaneMask

[arsenm]

Summary

Return the best covering index, and additional needed to complete the
mask. This logically belongs in TargetRegisterInfo, although I ended
up not needing it for why I originally split this out.

[foad]

Seems like an obvious refactoring to me. I think the API could be simplified though. How about returning all indices via the SmallVector, and returning a bool failure indication? Or say that if the vector is empty, it means failure?

llvm/include/llvm/CodeGen/TargetRegisterInfo.h

361–362

Maybe "the additional ones will be returned in" would be clearer?

[arsenm]

In D93394#2457781, @foad wrote:

Seems like an obvious refactoring to me. I think the API could be simplified though. How about returning all indices via the SmallVector, and returning a bool failure indication? Or say that if the vector is empty, it means failure?

I was conflicted about this, but thought returning the single "best" one was useful information

[arsenm]

In D93394#2458010, @arsenm wrote:

In D93394#2457781, @foad wrote:

Seems like an obvious refactoring to me. I think the API could be simplified though. How about returning all indices via the SmallVector, and returning a bool failure indication? Or say that if the vector is empty, it means failure?

I was conflicted about this, but thought returning the single "best" one was useful information

I think the weirder part of the current API is the "best" one isn't included in the complete vector

[arsenm]

ping @qcolombet for API opinion

[qcolombet]

llvm/include/llvm/CodeGen/TargetRegisterInfo.h

366

The API looks fine but I think the comment could be clearer.

E.g., what is the best matching index when we need more than one indexes?

Maybe we could return a bool: a covering exists or not, and just always fill Indexes?

[arsenm]

d696f43bd760c2623bd6da5e5d90289b1639f3bf

llvm/include/llvm/CodeGen/TargetRegisterInfo.h

public:
   /// SubIdx \see LaneBitmask.
   ///
   /// SubIdx == 0 is allowed, it has the lane mask ~0u.
   LaneBitmask getSubRegIndexLaneMask(unsigned SubIdx) const {
     assert(SubIdx < getNumSubRegIndices() && "This is not a subregister index");
     return SubRegIndexLaneMasks[SubIdx];
   }
 
+  /// Try to find a set of subregister indexes to cover \p LaneMask.  Returns
+  /// the best matching index, or NoSubRegister if this is impossible. If
+  /// multiple subregister indexes are required to exactly cover, they will be
+  /// returned in \p Indexes.
+  unsigned getCoveringSubRegIndexes(MachineRegisterInfo &MRI,
+                                    const TargetRegisterClass *RC,
+                                    LaneBitmask LaneMask,
+                                    SmallVectorImpl<unsigned> &Indexes) const;
+
   /// The lane masks returned by getSubRegIndexLaneMask() above can only be
   /// used to determine if sub-registers overlap - they can't be used to
   /// determine if a set of sub-registers completely cover another
   /// sub-register.
   ///
   /// The X86 general purpose registers have two lanes corresponding to the
   /// sub_8bit and sub_8bit_hi sub-registers. Both sub_32bit and sub_16bit have
   /// lane masks '3', but the sub_16bit sub-register doesn't fully cover the

llvm/lib/CodeGen/SplitKit.cpp

SlotIndex SplitEditor::buildCopy(Register FromReg, Register ToReg,
 
   // 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.
   LiveInterval &DestLI = LIS.getInterval(Edit->get(RegIdx));
 
   // 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 = SubRegMask.getNumLanes();
-    PossibleIndexes.push_back(Idx);
-    if (PopCount > BestCover) {
-      BestCover = PopCount;
-      BestIdx = Idx;
-    }
-  }
+  SmallVector<unsigned, 8> Indexes;
+  unsigned BestIdx = TRI.getCoveringSubRegIndexes(MRI, RC, LaneMask, Indexes);
 
   // Abort if we cannot possibly implement the COPY with the given indexes.
   if (BestIdx == 0)
     report_fatal_error("Impossible to implement partial COPY");
 
   SlotIndex Def = buildSingleSubRegCopy(FromReg, ToReg, MBB, InsertBefore,
                                         BestIdx, DestLI, Late, SlotIndex());
 
-  // Greedy heuristic: Keep iterating keeping the best covering subreg index
-  // each time.
-  LaneBitmask LanesLeft = LaneMask & ~(TRI.getSubRegIndexLaneMask(BestIdx));
-  while (LanesLeft.any()) {
-    unsigned BestIdx = 0;
-    int BestCover = std::numeric_limits<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 = (SubRegMask & LanesLeft).getNumLanes()
-                - (SubRegMask & ~LanesLeft).getNumLanes();
-      if (Cover > BestCover) {
-        BestCover = Cover;
-        BestIdx = Idx;
-      }
-    }
-
-    if (BestIdx == 0)
-      report_fatal_error("Impossible to implement partial COPY");
-
+  for (unsigned BestIdx : Indexes) {
     buildSingleSubRegCopy(FromReg, ToReg, MBB, InsertBefore, BestIdx,
                           DestLI, Late, Def);
-    LanesLeft &= ~TRI.getSubRegIndexLaneMask(BestIdx);
   }
 
   return Def;
 }
 
 VNInfo *SplitEditor::defFromParent(unsigned RegIdx,
                                    VNInfo *ParentVNI,
                                    SlotIndex UseIdx,

llvm/lib/CodeGen/TargetRegisterInfo.cpp

if (RegSize)
       return RegSize;
     // Since Reg is not a generic register, it must have a register class.
     RC = MRI.getRegClass(Reg);
   }
   assert(RC && "Unable to deduce the register class");
   return getRegSizeInBits(*RC);
 }
 
+unsigned TargetRegisterInfo::getCoveringSubRegIndexes(
+    MachineRegisterInfo &MRI, const TargetRegisterClass *RC,
+    LaneBitmask LaneMask, SmallVectorImpl<unsigned> &NeededIndexes) const {
+  SmallVector<unsigned, 8> PossibleIndexes;
+  unsigned BestIdx = 0;
+  unsigned BestCover = 0;
+  // assert(RC == MRI.getRegClass(ToReg) && "Should have same reg class");
+  for (unsigned Idx = 1, E = getNumSubRegIndices(); Idx < E; ++Idx) {
+    // Is this index even compatible with the given class?
+    if (getSubClassWithSubReg(RC, Idx) != RC)
+      continue;
+    LaneBitmask SubRegMask = 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 = SubRegMask.getNumLanes();
+    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)
+    return 0;
+
+  // Greedy heuristic: Keep iterating keeping the best covering subreg index
+  // each time.
+  LaneBitmask LanesLeft = LaneMask & ~getSubRegIndexLaneMask(BestIdx);
+  while (LanesLeft.any()) {
+    unsigned BestIdx = 0;
+    int BestCover = std::numeric_limits<int>::min();
+    for (unsigned Idx : PossibleIndexes) {
+      LaneBitmask SubRegMask = 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 = (SubRegMask & LanesLeft).getNumLanes() -
+                  (SubRegMask & ~LanesLeft).getNumLanes();
+      if (Cover > BestCover) {
+        BestCover = Cover;
+        BestIdx = Idx;
+      }
+    }
+
+    if (BestIdx == 0)
+      return 0; // Impossible to handle
+
+    NeededIndexes.push_back(BestIdx);
+
+    LanesLeft &= ~getSubRegIndexLaneMask(BestIdx);
+  }
+
+  return BestIdx;
+}
+
 Register
 TargetRegisterInfo::lookThruCopyLike(Register SrcReg,
                                      const MachineRegisterInfo *MRI) const {
   while (true) {
     const MachineInstr *MI = MRI->getVRegDef(SrcReg);
     if (!MI->isCopyLike())
       return SrcReg;