[Tail duplication] Handle source registers with subregisters

[kparzysz]

Summary

When a block is tail-duplicated, the PHI nodes from that block are replaced with appropriate COPY instructions. When those PHI nodes contained use operands with subregisters, the subregisters were dropped from the COPY instructions, resulting in incorrect code.

Keep track of the subregister information and use this information when remapping instructions from the duplicated block.

[qcolombet]

Hi Krzysztof,

This looks mostly good to me. The change on how we add register class constrains seems broken to me and the test case could benefit some check lines and a run of opt -instnamer.

Cheers,
-Quentin

lib/CodeGen/TailDuplicator.cpp

411

Why do you check for the subreg here?
If you believe the adding constraints will not work, you should check the result of constrainRegClass and act accordingly.

921

Use \p in front of variable names.

test/CodeGen/Hexagon/tail-dup-subreg-map.ll

2

FileCheck?

[kparzysz]

lib/CodeGen/TailDuplicator.cpp

411

This is because we are doing a replacement

Reg -> VI.Reg:VI.SubReg

If SubReg is present, then VI.Reg belongs to a class where its sub-registers are compatible with Reg, but VI.Reg itself does not have to be. Actually, I think that the existing code is broken.

[kparzysz]

Expanded handling of the register class constraints when remapping TailBB into its predecessor.

[kparzysz]

lib/CodeGen/TailDuplicator.cpp

921

Sorry, missed this one.

[kparzysz]

Fixed a comment.

[qcolombet]

lib/CodeGen/TailDuplicator.cpp

373

MRI->setRegClass.

389

Instead of using OrigRC, we should use the regclass that is required by MO.
MachineInstr::getRegClassConstraint

392

We need to read from Reg:subReg if Reg has a sub reg.

[kparzysz]

lib/CodeGen/TailDuplicator.cpp

392

Actually, we need to leave the sub-reg alone. NewReg will be inserted into the LocalVRMap as a replacement for Reg. Since there is no Reg anymore (it should be the value taken from VI.Reg:VI:SubReg), the COPY should use the VI.* values instead of Reg.

[kparzysz]

Update the instruction remapping code.

[qcolombet]

Hi Krzysztof,

LGTM.

Thanks,
-Quentin

include/llvm/CodeGen/TailDuplicator.h

public:
   bool tailDuplicateBlocks(MachineFunction &MF);
   static bool isSimpleBB(MachineBasicBlock *TailBB);
   bool shouldTailDuplicate(const MachineFunction &MF, bool IsSimple,
                            MachineBasicBlock &TailBB);
   bool tailDuplicateAndUpdate(MachineFunction &MF, bool IsSimple,
                               MachineBasicBlock *MBB);
 
 private:
+  typedef TargetInstrInfo::RegSubRegPair RegSubRegPair;
+
   void addSSAUpdateEntry(unsigned OrigReg, unsigned NewReg,
                          MachineBasicBlock *BB);
   void processPHI(MachineInstr *MI, MachineBasicBlock *TailBB,
                   MachineBasicBlock *PredBB,
-                  DenseMap<unsigned, unsigned> &LocalVRMap,
-                  SmallVectorImpl<std::pair<unsigned, unsigned>> &Copies,
+                  DenseMap<unsigned, RegSubRegPair> &LocalVRMap,
+                  SmallVectorImpl<std::pair<unsigned, RegSubRegPair>> &Copies,
                   const DenseSet<unsigned> &UsedByPhi, bool Remove);
   void duplicateInstruction(MachineInstr *MI, MachineBasicBlock *TailBB,
                             MachineBasicBlock *PredBB, MachineFunction &MF,
-                            DenseMap<unsigned, unsigned> &LocalVRMap,
+                            DenseMap<unsigned, RegSubRegPair> &LocalVRMap,
                             const DenseSet<unsigned> &UsedByPhi);
   void updateSuccessorsPHIs(MachineBasicBlock *FromBB, bool isDead,
                             SmallVectorImpl<MachineBasicBlock *> &TDBBs,
                             SmallSetVector<MachineBasicBlock *, 8> &Succs);
   bool canCompletelyDuplicateBB(MachineBasicBlock &BB);
   bool duplicateSimpleBB(MachineBasicBlock *TailBB,
                          SmallVectorImpl<MachineBasicBlock *> &TDBBs,
                          const DenseSet<unsigned> &RegsUsedByPhi,
                          SmallVectorImpl<MachineInstr *> &Copies);
   bool tailDuplicate(MachineFunction &MF, bool IsSimple,
                      MachineBasicBlock *TailBB,
                      SmallVectorImpl<MachineBasicBlock *> &TDBBs,
                      SmallVectorImpl<MachineInstr *> &Copies);
+  void appendCopies(MachineBasicBlock *MBB,
+                 SmallVectorImpl<std::pair<unsigned,RegSubRegPair>> &CopyInfos,
+                 SmallVectorImpl<MachineInstr *> &Copies);
 
   void removeDeadBlock(MachineBasicBlock *MBB);
 };
 
 } // End llvm namespace
 
 #endif

lib/CodeGen/TailDuplicator.cpp

else {
     SSAUpdateVRs.push_back(OrigReg);
   }
 }
 
 /// Process PHI node in TailBB by turning it into a copy in PredBB. Remember the
 /// source register that's contributed by PredBB and update SSA update map.
 void TailDuplicator::processPHI(
     MachineInstr *MI, MachineBasicBlock *TailBB, MachineBasicBlock *PredBB,
-    DenseMap<unsigned, unsigned> &LocalVRMap,
-    SmallVectorImpl<std::pair<unsigned, unsigned>> &Copies,
+    DenseMap<unsigned, RegSubRegPair> &LocalVRMap,
+    SmallVectorImpl<std::pair<unsigned, RegSubRegPair>> &Copies,
     const DenseSet<unsigned> &RegsUsedByPhi, bool Remove) {
   unsigned DefReg = MI->getOperand(0).getReg();
   unsigned SrcOpIdx = getPHISrcRegOpIdx(MI, PredBB);
   assert(SrcOpIdx && "Unable to find matching PHI source?");
   unsigned SrcReg = MI->getOperand(SrcOpIdx).getReg();
+  unsigned SrcSubReg = MI->getOperand(SrcOpIdx).getSubReg();
   const TargetRegisterClass *RC = MRI->getRegClass(DefReg);
-  LocalVRMap.insert(std::make_pair(DefReg, SrcReg));
+  LocalVRMap.insert(std::make_pair(DefReg, RegSubRegPair(SrcReg, SrcSubReg)));
 
   // Insert a copy from source to the end of the block. The def register is the
   // available value liveout of the block.
   unsigned NewDef = MRI->createVirtualRegister(RC);
-  Copies.push_back(std::make_pair(NewDef, SrcReg));
+  Copies.push_back(std::make_pair(NewDef, RegSubRegPair(SrcReg, SrcSubReg)));
   if (isDefLiveOut(DefReg, TailBB, MRI) || RegsUsedByPhi.count(DefReg))
     addSSAUpdateEntry(DefReg, NewDef, PredBB);
 
   if (!Remove)
     return;
 
   // Remove PredBB from the PHI node.
   MI->RemoveOperand(SrcOpIdx + 1);
   MI->RemoveOperand(SrcOpIdx);
   if (MI->getNumOperands() == 1)
     MI->eraseFromParent();
 }
 
 /// Duplicate a TailBB instruction to PredBB and update
 /// the source operands due to earlier PHI translation.
 void TailDuplicator::duplicateInstruction(
     MachineInstr *MI, MachineBasicBlock *TailBB, MachineBasicBlock *PredBB,
-    MachineFunction &MF, DenseMap<unsigned, unsigned> &LocalVRMap,
+    MachineFunction &MF,
+    DenseMap<unsigned, RegSubRegPair> &LocalVRMap,
     const DenseSet<unsigned> &UsedByPhi) {
   MachineInstr *NewMI = TII->duplicate(MI, MF);
   if (PreRegAlloc) {
     for (unsigned i = 0, e = NewMI->getNumOperands(); i != e; ++i) {
       MachineOperand &MO = NewMI->getOperand(i);
       if (!MO.isReg())
         continue;
       unsigned Reg = MO.getReg();
       if (!TargetRegisterInfo::isVirtualRegister(Reg))
         continue;
       if (MO.isDef()) {
         const TargetRegisterClass *RC = MRI->getRegClass(Reg);
         unsigned NewReg = MRI->createVirtualRegister(RC);
         MO.setReg(NewReg);
-        LocalVRMap.insert(std::make_pair(Reg, NewReg));
+        LocalVRMap.insert(std::make_pair(Reg, RegSubRegPair(NewReg, 0)));
         if (isDefLiveOut(Reg, TailBB, MRI) || UsedByPhi.count(Reg))
           addSSAUpdateEntry(Reg, NewReg, PredBB);
       } else {
-        DenseMap<unsigned, unsigned>::iterator VI = LocalVRMap.find(Reg);
+        auto VI = LocalVRMap.find(Reg);
         if (VI != LocalVRMap.end()) {
-          MO.setReg(VI->second);
+          // Need to make sure that the register class of the mapped register
+          // will satisfy the constraints of the class of the register being
+          // replaced.
+          auto *OrigRC = MRI->getRegClass(Reg);
+          auto *MappedRC = MRI->getRegClass(VI->second.Reg);
+          const TargetRegisterClass *ConstrRC;
+          if (VI->second.SubReg != 0) {
+            ConstrRC = TRI->getMatchingSuperRegClass(MappedRC, OrigRC,
+                                                     VI->second.SubReg);
+            if (ConstrRC) {
+              // The actual constraining (as in "find appropriate new class")
+              // is done by getMatchingSuperRegClass, so now we only need to
+              // change the class of the mapped register.
+              MRI->setRegClass(VI->second.Reg, ConstrRC);
+            }
+          } else {
+            // For mapped registers that do not have sub-registers, simply
+            // restrict their class to match the original one.
+            ConstrRC = MRI->constrainRegClass(VI->second.Reg, OrigRC);
+          }
+
+          if (ConstrRC) {
+            // If the class constraining succeeded, we can simply replace
+            // the old register with the mapped one.
+            MO.setReg(VI->second.Reg);
+            // We have Reg -> VI.Reg:VI.SubReg, so if Reg is used with a
+            // sub-register, we need to compose the sub-register indices.
+            MO.setSubReg(TRI->composeSubRegIndices(MO.getSubReg(),
+                                                   VI->second.SubReg));
+          } else {
+            // The direct replacement is not possible, due to failing register
+            // class constraints. An explicit COPY is necessary. Create one
+            // that can be reused
+            auto *NewRC = MI->getRegClassConstraint(i, TII, TRI);
+            if (NewRC == nullptr)
+              NewRC = OrigRC;
+            unsigned NewReg = MRI->createVirtualRegister(NewRC);
+            BuildMI(*PredBB, MI, MI->getDebugLoc(),
+                    TII->get(TargetOpcode::COPY), NewReg)
+                .addReg(VI->second.Reg, 0, VI->second.SubReg);
+            LocalVRMap.erase(VI);
+            LocalVRMap.insert(std::make_pair(Reg, RegSubRegPair(NewReg, 0)));
+            MO.setReg(NewReg);
+            // The composed VI.Reg:VI.SubReg is replaced with NewReg, which
+            // is equivalent to the whole register Reg. Hence, Reg:subreg
+            // is same as NewReg:subreg, so keep the sub-register index
+            // unchanged.
+          }
           // Clear any kill flags from this operand.  The new register could
           // have uses after this one, so kills are not valid here.
           MO.setIsKill(false);
-          MRI->constrainRegClass(VI->second, MRI->getRegClass(Reg));
         }
       }
     }
   }
   PredBB->insert(PredBB->instr_end(), NewMI);
 }
 
 /// After FromBB is tail duplicated into its predecessor blocks, the successors
 /// have gained new predecessors. Update the PHI instructions in them
if (RS && !TailBB->livein_empty()) {
           // If a register is previously livein to the tail but it's not live
           // at the end of predecessor BB, then it should be added to its
           // livein list.
           PredBB->addLiveIn(LI);
       }
     }
 
     // Clone the contents of TailBB into PredBB.
-    DenseMap<unsigned, unsigned> LocalVRMap;
-    SmallVector<std::pair<unsigned, unsigned>, 4> CopyInfos;
+    DenseMap<unsigned, RegSubRegPair> LocalVRMap;
+    SmallVector<std::pair<unsigned, RegSubRegPair>, 4> CopyInfos;
     // Use instr_iterator here to properly handle bundles, e.g.
     // ARM Thumb2 IT block.
     MachineBasicBlock::instr_iterator I = TailBB->instr_begin();
     while (I != TailBB->instr_end()) {
       MachineInstr *MI = &*I;
       ++I;
       if (MI->isPHI()) {
         // Replace the uses of the def of the PHI with the register coming
         // from PredBB.
         processPHI(MI, TailBB, PredBB, LocalVRMap, CopyInfos, UsedByPhi, true);
       } else {
         // Replace def of virtual registers with new registers, and update
         // uses with PHI source register or the new registers.
         duplicateInstruction(MI, TailBB, PredBB, MF, LocalVRMap, UsedByPhi);
       }
     }
-    MachineBasicBlock::iterator Loc = PredBB->getFirstTerminator();
-    for (unsigned i = 0, e = CopyInfos.size(); i != e; ++i) {
-      Copies.push_back(BuildMI(*PredBB, Loc, DebugLoc(),
-                               TII->get(TargetOpcode::COPY), CopyInfos[i].first)
-                           .addReg(CopyInfos[i].second));
-    }
+    appendCopies(PredBB, CopyInfos, Copies);
 
     // Simplify
     TII->AnalyzeBranch(*PredBB, PredTBB, PredFBB, PredCond, true);
 
     NumInstrDups += TailBB->size() - 1; // subtract one for removed branch
 
     // Update the CFG.
     PredBB->removeSuccessor(PredBB->succ_begin());
bool TailDuplicator::tailDuplicate(MachineFunction &MF, bool IsSimple,
   // AnalyzeBranch.
   if (PrevBB->succ_size() == 1 &&
       !TII->AnalyzeBranch(*PrevBB, PriorTBB, PriorFBB, PriorCond, true) &&
       PriorCond.empty() && !PriorTBB && TailBB->pred_size() == 1 &&
       !TailBB->hasAddressTaken()) {
     DEBUG(dbgs() << "\nMerging into block: " << *PrevBB
                  << "From MBB: " << *TailBB);
     if (PreRegAlloc) {
-      DenseMap<unsigned, unsigned> LocalVRMap;
-      SmallVector<std::pair<unsigned, unsigned>, 4> CopyInfos;
+      DenseMap<unsigned, RegSubRegPair> LocalVRMap;
+      SmallVector<std::pair<unsigned, RegSubRegPair>, 4> CopyInfos;
       MachineBasicBlock::iterator I = TailBB->begin();
       // Process PHI instructions first.
       while (I != TailBB->end() && I->isPHI()) {
         // Replace the uses of the def of the PHI with the register coming
         // from PredBB.
         MachineInstr *MI = &*I++;
         processPHI(MI, TailBB, PrevBB, LocalVRMap, CopyInfos, UsedByPhi, true);
       }
 
       // Now copy the non-PHI instructions.
       while (I != TailBB->end()) {
         // Replace def of virtual registers with new registers, and update
         // uses with PHI source register or the new registers.
         MachineInstr *MI = &*I++;
         assert(!MI->isBundle() && "Not expecting bundles before regalloc!");
         duplicateInstruction(MI, TailBB, PrevBB, MF, LocalVRMap, UsedByPhi);
         MI->eraseFromParent();
       }
-      MachineBasicBlock::iterator Loc = PrevBB->getFirstTerminator();
-      for (unsigned i = 0, e = CopyInfos.size(); i != e; ++i) {
-        Copies.push_back(BuildMI(*PrevBB, Loc, DebugLoc(),
-                                 TII->get(TargetOpcode::COPY),
-                                 CopyInfos[i].first)
-                             .addReg(CopyInfos[i].second));
-      }
+      appendCopies(PrevBB, CopyInfos, Copies);
     } else {
       // No PHIs to worry about, just splice the instructions over.
       PrevBB->splice(PrevBB->end(), TailBB, TailBB->begin(), TailBB->end());
     }
     PrevBB->removeSuccessor(PrevBB->succ_begin());
     assert(PrevBB->succ_empty());
     PrevBB->transferSuccessors(TailBB);
     TDBBs.push_back(PrevBB);
for (SmallSetVector<MachineBasicBlock *, 8>::iterator PI = Preds.begin(),
     MachineBasicBlock *PredBB = *PI;
     if (std::find(TDBBs.begin(), TDBBs.end(), PredBB) != TDBBs.end())
       continue;
 
     // EH edges
     if (PredBB->succ_size() != 1)
       continue;
 
-    DenseMap<unsigned, unsigned> LocalVRMap;
-    SmallVector<std::pair<unsigned, unsigned>, 4> CopyInfos;
+    DenseMap<unsigned, RegSubRegPair> LocalVRMap;
+    SmallVector<std::pair<unsigned, RegSubRegPair>, 4> CopyInfos;
     MachineBasicBlock::iterator I = TailBB->begin();
     // Process PHI instructions first.
     while (I != TailBB->end() && I->isPHI()) {
       // Replace the uses of the def of the PHI with the register coming
       // from PredBB.
       MachineInstr *MI = &*I++;
       processPHI(MI, TailBB, PredBB, LocalVRMap, CopyInfos, UsedByPhi, false);
     }
-    MachineBasicBlock::iterator Loc = PredBB->getFirstTerminator();
-    for (unsigned i = 0, e = CopyInfos.size(); i != e; ++i) {
-      Copies.push_back(BuildMI(*PredBB, Loc, DebugLoc(),
-                               TII->get(TargetOpcode::COPY), CopyInfos[i].first)
-                           .addReg(CopyInfos[i].second));
-    }
+    appendCopies(PredBB, CopyInfos, Copies);
   }
 
   return Changed;
 }
 
+/// At the end of the block \p MBB generate COPY instructions between registers
+/// described by \p CopyInfos. Append resulting instructions to \p Copies.
+void TailDuplicator::appendCopies(MachineBasicBlock *MBB,
+      SmallVectorImpl<std::pair<unsigned,RegSubRegPair>> &CopyInfos,
+      SmallVectorImpl<MachineInstr*> &Copies) {
+  MachineBasicBlock::iterator Loc = MBB->getFirstTerminator();
+  const MCInstrDesc &CopyD = TII->get(TargetOpcode::COPY);
+  for (auto &CI : CopyInfos) {
+    auto C = BuildMI(*MBB, Loc, DebugLoc(), CopyD, CI.first)
+                .addReg(CI.second.Reg, 0, CI.second.SubReg);
+    Copies.push_back(C);
+  }
+}
+
 /// Remove the specified dead machine basic block from the function, updating
 /// the CFG.
 void TailDuplicator::removeDeadBlock(MachineBasicBlock *MBB) {
   assert(MBB->pred_empty() && "MBB must be dead!");
   DEBUG(dbgs() << "\nRemoving MBB: " << *MBB);
 
   // Remove all successors.
   while (!MBB->succ_empty())
     MBB->removeSuccessor(MBB->succ_end() - 1);
 
   // Remove the block.
   MBB->eraseFromParent();
 }
 
 } // End llvm namespace

test/CodeGen/Hexagon/tail-dup-subreg-map.ll

+; RUN: llc -march=hexagon < %s | FileCheck %s
+; REQUIRES: asserts
+
+; When tail-duplicating a block with PHI nodes that use subregisters, the
+; subregisters were dropped by the tail duplicator, resulting in invalid
+; COPY instructions being generated.
+
+; CHECK: = extractu(r{{[0-9]+}}, #15, #17)
+
+target triple = "hexagon"
+
+%struct.0 = type { i64, i16 }
+%struct.1 = type { i64, i64 }
+
+declare hidden fastcc void @foo(%struct.0* noalias nocapture, i8 signext, i8 zeroext, i32, i64, i64) unnamed_addr #0
+
+define void @fred(%struct.0* noalias nocapture sret %agg.result, %struct.1* byval nocapture readonly align 8 %a) #1 {
+entry:
+  %0 = load i64, i64* undef, align 8
+  switch i32 undef, label %if.else [
+    i32 32767, label %if.then
+    i32 0, label %if.then7
+  ]
+
+if.then:                                          ; preds = %entry
+  ret void
+
+if.then7:                                         ; preds = %entry
+  br i1 undef, label %if.then.i, label %if.else16.i
+
+if.then.i:                                        ; preds = %if.then7
+  br i1 undef, label %if.then5.i, label %if.else.i
+
+if.then5.i:                                       ; preds = %if.then.i
+  %shl.i21 = shl i64 %0, 0
+  br label %if.end.i
+
+if.else.i:                                        ; preds = %if.then.i
+  %shl12.i = shl i64 %0, undef
+  br label %if.end.i
+
+if.end.i:                                         ; preds = %if.else.i, %if.then5.i
+  %aSig0.0 = phi i64 [ undef, %if.then5.i ], [ %shl12.i, %if.else.i ]
+  %storemerge43.i = phi i64 [ %shl.i21, %if.then5.i ], [ 0, %if.else.i ]
+  %sub15.i = sub nsw i32 -63, undef
+  br label %if.end13
+
+if.else16.i:                                      ; preds = %if.then7
+  br label %if.end13
+
+if.else:                                          ; preds = %entry
+  %or12 = or i64 undef, 281474976710656
+  br label %if.end13
+
+if.end13:                                         ; preds = %if.else, %if.else16.i, %if.end.i
+  %aSig1.1 = phi i64 [ %0, %if.else ], [ %storemerge43.i, %if.end.i ], [ undef, %if.else16.i ]
+  %aSig0.2 = phi i64 [ %or12, %if.else ], [ %aSig0.0, %if.end.i ], [ undef, %if.else16.i ]
+  %aExp.0 = phi i32 [ undef, %if.else ], [ %sub15.i, %if.end.i ], [ undef, %if.else16.i ]
+  %shl2.i = shl i64 %aSig0.2, 15
+  %shr.i = lshr i64 %aSig1.1, 49
+  %or.i = or i64 %shl2.i, %shr.i
+  tail call fastcc void @foo(%struct.0* noalias %agg.result, i8 signext 80, i8 zeroext undef, i32 %aExp.0, i64 %or.i, i64 undef)
+  unreachable
+}
+
+attributes #0 = { norecurse nounwind }
+attributes #1 = { nounwind }