[GlobalISel] Change MachineIRBuilder's SrcOp to contain subregister info

[aemerson]

Summary

r356304 changed the buildCopy() call to allow passing in a subregister index as well as a source reg. However if the CSE builder learns to do CSE on copies then it can cause incorrect behaviour as the CSE lookups for Unknown option "sub-reg" in: {src-reg, sub-reg} pairs will only check for the src-reg, as the sub-reg is added after the buildInstr() call.

This change folds an optional subreg index into SrcOp, and teaches CSEInfo how deal with it.

[arsenm]

I would think we would want to avoid any subregister indexes until as late as possible. Does the verifier reject them for G_* instructions?

[aemerson]

In D59444#1431750, @arsenm wrote:

I would think we would want to avoid any subregister indexes until as late as possible. Does the verifier reject them for G_* instructions?

In practice yes we’d only really use this during selection, but if someone did want to use MachineIRBuilder for emitting copies (as we do in arm64), it should be safe w.r.t CSE as well.

[aditya_nandakumar]

In general I'm not too fond of building very target specific instructions involving subregs with the MachineIRBuilder (As these are only likely to be used in a few places in the selector). Specifically for this case of building a subreg, you can say

Builder.buildInstr(COPY, {Dst},{})
  .addReg(Src, 0, SubReg);

which is not too much larger than (and doesn't require making every SrcOp bigger).

Builder.buildInstr(COPY, {Dst}, {{Src, SubReg}});

While adding CSE support for COPYs with Subregs is great, I'm not sure how often it'll trigger and be useful (If it does occur then great). Also should we start adding support for building other kinds of target instructions - Adding immediate operands/target index nodes etc?
My approach for building non generic instructions have been with using the .addImm(..).addReg(...)... pattern - but if others strongly feel the need to add the ability to pass in all of the operands at once to the builder interfaces, and enabling CSE for them, then we can go ahead with this change.

[aemerson]

In D59444#1433415, @aditya_nandakumar wrote:

In general I'm not too fond of building very target specific instructions involving subregs with the MachineIRBuilder (As these are only likely to be used in a few places in the selector). Specifically for this case of building a subreg, you can say

Builder.buildInstr(COPY, {Dst},{})
  .addReg(Src, 0, SubReg);

which is not too much larger than (and doesn't require making every SrcOp bigger).

Builder.buildInstr(COPY, {Dst}, {{Src, SubReg}});

While adding CSE support for COPYs with Subregs is great, I'm not sure how often it'll trigger and be useful (If it does occur then great). Also should we start adding support for building other kinds of target instructions - Adding immediate operands/target index nodes etc?
My approach for building non generic instructions have been with using the .addImm(..).addReg(...)... pattern - but if others strongly feel the need to add the ability to pass in all of the operands at once to the builder interfaces, and enabling CSE for them, then we can go ahead with this change.

In that case I'll revert r356304 keeping in mind that we should take care not to enable COPY CSE in the selector in future, as it won't be subreg safe.

llvm/include/llvm/CodeGen/GlobalISel/CSEInfo.h

 // TODO: Find a better place for this.
 // Commonly used for O0 config.
 class CSEConfigConstantOnly : public CSEConfig {
 public:
   virtual ~CSEConfigConstantOnly() = default;
   virtual bool shouldCSEOpc(unsigned Opc) override;
 };
 
+class CSEConfigCopyOnly : public CSEConfig {
+public:
+  virtual ~CSEConfigCopyOnly() = default;
+  virtual bool shouldCSEOpc(unsigned Opc) override;
+};
+
 /// The CSE Analysis object.
 /// This installs itself as a delegate to the MachineFunction to track
 /// new instructions as well as deletions. It however will not be able to
 /// track instruction mutations. In such cases, recordNewInstruction should be
 /// called (for eg inside MachineIRBuilder::recordInsertion).
 /// Also because of how just the instruction can be inserted without adding any
 /// operands to the instruction, instructions are uniqued and inserted lazily.
 /// CSEInfo should assert when trying to enter an incomplete instruction into
class GISelInstProfileBuilder {
   const MachineRegisterInfo &MRI;
 
 public:
   GISelInstProfileBuilder(FoldingSetNodeID &ID, const MachineRegisterInfo &MRI)
       : ID(ID), MRI(MRI) {}
   // Profiling methods.
   const GISelInstProfileBuilder &addNodeIDOpcode(unsigned Opc) const;
   const GISelInstProfileBuilder &addNodeIDRegType(const LLT &Ty) const;
-  const GISelInstProfileBuilder &addNodeIDRegType(const unsigned) const;
+  const GISelInstProfileBuilder &addNodeIDRegType(const unsigned,
+                                                  unsigned) const;
 
   const GISelInstProfileBuilder &
   addNodeIDRegType(const TargetRegisterClass *RC) const;
   const GISelInstProfileBuilder &addNodeIDRegType(const RegisterBank *RB) const;
 
   const GISelInstProfileBuilder &addNodeIDRegNum(unsigned Reg) const;
 
   const GISelInstProfileBuilder &addNodeIDImmediate(int64_t Imm) const;

llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

public:
 
   DstType getDstOpKind() const { return Ty; }
 
 private:
   DstType Ty;
 };
 
 class SrcOp {
+  struct RegSubRegTy {
+    unsigned Reg;
+    unsigned SubReg;
+  };
   union {
     MachineInstrBuilder SrcMIB;
-    unsigned Reg;
+    struct RegSubRegTy RegPair;
     CmpInst::Predicate Pred;
   };
 
 public:
   enum class SrcType { Ty_Reg, Ty_MIB, Ty_Predicate };
-  SrcOp(unsigned R) : Reg(R), Ty(SrcType::Ty_Reg) {}
-  SrcOp(const MachineOperand &Op) : Reg(Op.getReg()), Ty(SrcType::Ty_Reg) {}
+  SrcOp(unsigned R, unsigned S = 0) : RegPair({R, S}), Ty(SrcType::Ty_Reg) {}
+  SrcOp(const MachineOperand &Op)
+      : RegPair({Op.getReg(), 0}), Ty(SrcType::Ty_Reg) {}
   SrcOp(const MachineInstrBuilder &MIB) : SrcMIB(MIB), Ty(SrcType::Ty_MIB) {}
   SrcOp(const CmpInst::Predicate P) : Pred(P), Ty(SrcType::Ty_Predicate) {}
 
   void addSrcToMIB(MachineInstrBuilder &MIB) const {
     switch (Ty) {
     case SrcType::Ty_Predicate:
       MIB.addPredicate(Pred);
       break;
     case SrcType::Ty_Reg:
-      MIB.addUse(Reg);
+      MIB.addUse(RegPair.Reg, 0, RegPair.SubReg);
       break;
     case SrcType::Ty_MIB:
       MIB.addUse(SrcMIB->getOperand(0).getReg());
       break;
     }
   }
 
   LLT getLLTTy(const MachineRegisterInfo &MRI) const {
     switch (Ty) {
     case SrcType::Ty_Predicate:
       llvm_unreachable("Not a register operand");
     case SrcType::Ty_Reg:
-      return MRI.getType(Reg);
+      return MRI.getType(RegPair.Reg);
     case SrcType::Ty_MIB:
       return MRI.getType(SrcMIB->getOperand(0).getReg());
     }
     llvm_unreachable("Unrecognised SrcOp::SrcType enum");
   }
 
   unsigned getReg() const {
     switch (Ty) {
     case SrcType::Ty_Predicate:
       llvm_unreachable("Not a register operand");
     case SrcType::Ty_Reg:
-      return Reg;
+      return RegPair.Reg;
     case SrcType::Ty_MIB:
       return SrcMIB->getOperand(0).getReg();
     }
     llvm_unreachable("Unrecognised SrcOp::SrcType enum");
   }
 
+  unsigned getSubReg() const { return RegPair.SubReg; }
+
   CmpInst::Predicate getPredicate() const {
     switch (Ty) {
     case SrcType::Ty_Predicate:
       return Pred;
     default:
       llvm_unreachable("Not a register operand");
     }
   }
public:
 
   /// Build and insert \p Res = COPY Op
   ///
   /// Register-to-register COPY sets \p Res to \p Op.
   ///
   /// \pre setBasicBlock or setMI must have been called.
   ///
   /// \return a MachineInstrBuilder for the newly created instruction.
-  MachineInstrBuilder buildCopy(const DstOp &Res, const SrcOp &Op,
-                                unsigned Subreg = 0);
+  MachineInstrBuilder buildCopy(const DstOp &Res, const SrcOp &Op);
 
   /// Build and insert `Res = G_LOAD Addr, MMO`.
   ///
   /// Loads the value stored at \p Addr. Puts the result in \p Res.
   ///
   /// \pre setBasicBlock or setMI must have been called.
   /// \pre \p Res must be a generic virtual register.
   /// \pre \p Addr must be a generic virtual register with pointer type.

llvm/lib/CodeGen/GlobalISel/CSEInfo.cpp

case TargetOpcode::G_TRUNC:
     return true;
   }
   return false;
 }
 
 bool CSEConfigConstantOnly::shouldCSEOpc(unsigned Opc) {
   return Opc == TargetOpcode::G_CONSTANT;
 }
+
+bool CSEConfigCopyOnly::shouldCSEOpc(unsigned Opc) {
+  return Opc == TargetOpcode::COPY;
+}
+
 /// -----------------------------------------
 
 /// -------- GISelCSEInfo -------------//
 void GISelCSEInfo::setMF(MachineFunction &MF) {
   this->MF = &MF;
   this->MRI = &MF.getRegInfo();
 }
 
void GISelCSEInfo::handleRecordedInsts() {
   while (!TemporaryInsts.empty()) {
     auto *MI = TemporaryInsts.pop_back_val();
     handleRecordedInst(MI);
   }
 }
 
 bool GISelCSEInfo::shouldCSE(unsigned Opc) const {
   // Only GISel opcodes are CSEable
-  if (!isPreISelGenericOpcode(Opc))
+  if (!isPreISelGenericOpcode(Opc) && Opc != TargetOpcode::COPY)
     return false;
   assert(CSEOpt.get() && "CSEConfig not set");
   return CSEOpt->shouldCSEOpc(Opc);
 }
 
 void GISelCSEInfo::erasingInstr(MachineInstr &MI) { handleRemoveInst(&MI); }
 void GISelCSEInfo::createdInstr(MachineInstr &MI) { recordNewInstruction(&MI); }
 void GISelCSEInfo::changingInstr(MachineInstr &MI) {

 
 const GISelInstProfileBuilder &
 GISelInstProfileBuilder::addNodeIDRegNum(unsigned Reg) const {
   ID.AddInteger(Reg);
   return *this;
 }
 
 const GISelInstProfileBuilder &
-GISelInstProfileBuilder::addNodeIDRegType(const unsigned Reg) const {
-  addNodeIDMachineOperand(MachineOperand::CreateReg(Reg, false));
+GISelInstProfileBuilder::addNodeIDRegType(const unsigned Reg,
+                                          unsigned SubReg) const {
+  addNodeIDMachineOperand(MachineOperand::CreateReg(
+      Reg, false, false, false, false, false, false, SubReg));
   return *this;
 }
 
 const GISelInstProfileBuilder &
 GISelInstProfileBuilder::addNodeIDMBB(const MachineBasicBlock *MBB) const {
   ID.AddPointer(MBB);
   return *this;
 }
 
 const GISelInstProfileBuilder &
 GISelInstProfileBuilder::addNodeIDFlag(unsigned Flag) const {
   if (Flag)
     ID.AddInteger(Flag);
   return *this;
 }
 
 const GISelInstProfileBuilder &GISelInstProfileBuilder::addNodeIDMachineOperand(
     const MachineOperand &MO) const {
   if (MO.isReg()) {
     unsigned Reg = MO.getReg();
     if (!MO.isDef())
       addNodeIDRegNum(Reg);
+    const auto *TRI = MRI.getTargetRegisterInfo();
+    if (TRI->isPhysicalRegister(Reg))
+      return *this;
     LLT Ty = MRI.getType(Reg);
     if (Ty.isValid())
       addNodeIDRegType(Ty);
     auto *RB = MRI.getRegBankOrNull(Reg);
     if (RB)
       addNodeIDRegType(RB);
     auto *RC = MRI.getRegClassOrNull(Reg);
     if (RC)

llvm/lib/CodeGen/GlobalISel/CSEMIRBuilder.cpp

 }
 
 void CSEMIRBuilder::profileSrcOp(const SrcOp &Op,
                                  GISelInstProfileBuilder &B) const {
   switch (Op.getSrcOpKind()) {
   case SrcOp::SrcType::Ty_Predicate:
     B.addNodeIDImmediate(static_cast<int64_t>(Op.getPredicate()));
     break;
+  case SrcOp::SrcType::Ty_MIB:
+    B.addNodeIDRegType(Op.getReg(), 0);
+    break;
   default:
-    B.addNodeIDRegType(Op.getReg());
+    B.addNodeIDRegType(Op.getReg(), Op.getSubReg());
     break;
   }
 }
 
 void CSEMIRBuilder::profileMBBOpcode(GISelInstProfileBuilder &B,
                                      unsigned Opc) const {
   // First add the MBB (Local CSE).
   B.addNodeIDMBB(&getMBB());

llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp

 }
 
 MachineInstrBuilder MachineIRBuilder::buildBrIndirect(unsigned Tgt) {
   assert(getMRI()->getType(Tgt).isPointer() && "invalid branch destination");
   return buildInstr(TargetOpcode::G_BRINDIRECT).addUse(Tgt);
 }
 
 MachineInstrBuilder MachineIRBuilder::buildCopy(const DstOp &Res,
-                                                const SrcOp &Op,
-                                                unsigned Subreg) {
-  auto Copy = buildInstr(TargetOpcode::COPY, Res, Op);
-  Copy->getOperand(1).setSubReg(Subreg);
-  return Copy;
+                                                const SrcOp &Op) {
+  return buildInstr(TargetOpcode::COPY, Res, Op);
 }
 
 MachineInstrBuilder MachineIRBuilder::buildConstant(const DstOp &Res,
                                                     const ConstantInt &Val) {
   LLT Ty = Res.getLLTTy(*getMRI());
   LLT EltTy = Ty.getScalarType();
   assert(EltTy.getScalarSizeInBits() == Val.getBitWidth() &&
          "creating constant with the wrong size");

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

 /// Dst = COPY SubRegCopy
 static bool selectSubregisterCopy(MachineInstr &I, MachineRegisterInfo &MRI,
                                   const RegisterBankInfo &RBI, unsigned SrcReg,
                                   const TargetRegisterClass *From,
                                   const TargetRegisterClass *To,
                                   unsigned SubReg) {
   MachineIRBuilder MIB(I);
   auto Copy = MIB.buildCopy({From}, {SrcReg});
-  auto SubRegCopy = MIB.buildCopy({To}, {Copy}, SubReg);
+  auto SubRegCopy = MIB.buildCopy({To}, {Copy.getReg(0), SubReg});
   MachineOperand &RegOp = I.getOperand(1);
   RegOp.setReg(SubRegCopy.getReg(0));
 
   // It's possible that the destination register won't be constrained. Make
   // sure that happens.
   if (!TargetRegisterInfo::isPhysicalRegister(I.getOperand(0).getReg()))
     RBI.constrainGenericRegister(I.getOperand(0).getReg(), *To, MRI);
 
case TargetOpcode::G_EXTRACT: {
     if (SrcSize < 64) {
       assert(SrcSize == 32 && DstTy.getSizeInBits() == 16 &&
              "unexpected G_EXTRACT types");
       return constrainSelectedInstRegOperands(I, TII, TRI, RBI);
     }
 
     unsigned DstReg = MRI.createGenericVirtualRegister(LLT::scalar(64));
     MIB.setInsertPt(MIB.getMBB(), std::next(I.getIterator()));
-    MIB.buildCopy({I.getOperand(0).getReg()}, {DstReg}, AArch64::sub_32);
+    MIB.buildCopy({I.getOperand(0).getReg()}, {DstReg, AArch64::sub_32});
     RBI.constrainGenericRegister(I.getOperand(0).getReg(),
                                  AArch64::GPR32RegClass, MRI);
     I.getOperand(0).setReg(DstReg);
 
     return constrainSelectedInstRegOperands(I, TII, TRI, RBI);
   }
 
   case TargetOpcode::G_INSERT: {
MachineInstr *AArch64InstructionSelector::emitExtractVectorElt(
   }
 
   // The register that we're going to copy into.
   unsigned InsertReg = VecReg;
   if (!DstReg)
     DstReg = MRI.createVirtualRegister(DstRC);
   // If the lane index is 0, we just use a subregister COPY.
   if (LaneIdx == 0) {
-    auto Copy = MIRBuilder.buildCopy({*DstReg}, {VecReg}, ExtractSubReg);
+    auto Copy = MIRBuilder.buildCopy({*DstReg}, {VecReg, ExtractSubReg});
     RBI.constrainGenericRegister(*DstReg, *DstRC, MRI);
     return &*Copy;
   }
 
   // Lane copies require 128-bit wide registers. If we're dealing with an
   // unpacked vector, then we need to move up to that width. Insert an implicit
   // def and a subregister insert to get us there.
   if (VecTy.getSizeInBits() != 128) {
if (NarrowTy.getSizeInBits() * NumElts == 128) {
     }
   }
 
   // Now that we've created any necessary subregister inserts, we can
   // create the copies.
   //
   // Perform the first copy separately as a subregister copy.
   unsigned CopyTo = I.getOperand(0).getReg();
-  auto FirstCopy = MIB.buildCopy({CopyTo}, {InsertRegs[0]}, ExtractSubReg);
+  auto FirstCopy = MIB.buildCopy({CopyTo}, {InsertRegs[0], ExtractSubReg});
   constrainSelectedInstRegOperands(*FirstCopy, TII, TRI, RBI);
 
   // Now, perform the remaining copies as vector lane copies.
   unsigned LaneIdx = 1;
   for (unsigned InsReg : InsertRegs) {
     unsigned CopyTo = I.getOperand(LaneIdx).getReg();
     MachineInstr &CopyInst =
         *BuildMI(MBB, I, I.getDebugLoc(), TII.get(CopyOpc), CopyTo)
IndexLoad =
         emitScalarToVector(64, &AArch64::FPR128RegClass,
                            IndexLoad->getOperand(0).getReg(), MIRBuilder);
 
     auto TBL1 = MIRBuilder.buildInstr(
         AArch64::TBLv16i8One, {&AArch64::FPR128RegClass},
         {Concat->getOperand(0).getReg(), IndexLoad->getOperand(0).getReg()});
     constrainSelectedInstRegOperands(*TBL1, TII, TRI, RBI);
 
-    auto Copy =
-        MIRBuilder.buildCopy({I.getOperand(0).getReg()}, {TBL1}, AArch64::dsub);
+    auto Copy = MIRBuilder.buildCopy({I.getOperand(0).getReg()},
+                                     {TBL1.getReg(0), AArch64::dsub});
     RBI.constrainGenericRegister(Copy.getReg(0), AArch64::FPR64RegClass, MRI);
     I.eraseFromParent();
     return true;
   }
 
   // For TBL2 we need to emit a REG_SEQUENCE to tie together two consecutive
   // Q registers for regalloc.
   auto RegSeq = MIRBuilder
if (SubReg != AArch64::ssub && SubReg != AArch64::dsub) {
       LLVM_DEBUG(dbgs() << "Unsupported destination size! (" << DstSize
                         << "\n");
       return false;
     }
 
     unsigned Reg = MRI.createVirtualRegister(RC);
     unsigned DstReg = I.getOperand(0).getReg();
 
-    MIRBuilder.buildCopy({DstReg}, {DstVec}, SubReg);
+    MIRBuilder.buildCopy({DstReg}, {DstVec, SubReg});
     MachineOperand &RegOp = I.getOperand(1);
     RegOp.setReg(Reg);
     RBI.constrainGenericRegister(DstReg, *RC, MRI);
   } else {
     // We don't need a subregister copy. Save a copy by re-using the
     // destination register on the final insert.
     assert(PrevMI && "PrevMI was null?");
     PrevMI->getOperand(0).setReg(I.getOperand(0).getReg());

llvm/unittests/CodeGen/GlobalISel/CSETest.cpp

auto MIBAdd1 =
       CSEB.buildInstr(TargetOpcode::G_ADD, {s16}, {MIBInput, MIBInput});
   // We should CSE constants only. Adds should not be CSEd.
   EXPECT_TRUE(MIBAdd1->getOpcode() != TargetOpcode::COPY);
   EXPECT_TRUE(&*MIBAdd1 != &*MIBAdd);
   // We should CSE constant.
   auto MIBZeroTmp = CSEB.buildConstant(s16, 0);
   EXPECT_TRUE(&*MIBZero == &*MIBZeroTmp);
 }
+
+TEST_F(GISelMITest, TestCSESubRegOp) {
+  if (!TM)
+    return;
+
+  LLT s32{LLT::scalar(32)};
+  GISelCSEInfo CSEInfo;
+  CSEInfo.setCSEConfig(make_unique<CSEConfigCopyOnly>());
+  CSEInfo.analyze(*MF);
+  B.setCSEInfo(&CSEInfo);
+  CSEMIRBuilder CSEB(B.getState());
+
+  CSEB.setInsertPt(*EntryMBB, EntryMBB->end());
+  auto Copy1 = CSEB.buildInstr(TargetOpcode::COPY, {s32}, {Copies[0]});
+  auto MIBCopy = CSEB.buildInstr(TargetOpcode::COPY, {s32}, {Copies[0]});
+  EXPECT_EQ(&*MIBCopy, &*Copy1);
+
+  // Check copy of subregs
+  auto SubRegCopy1 =
+      CSEB.buildInstr(TargetOpcode::COPY, {s32}, {{Copies[0], 2}});
+  auto SubRegCopy2 =
+      CSEB.buildInstr(TargetOpcode::COPY, {s32}, {{Copies[0], 2}});
+  EXPECT_EQ(&*SubRegCopy1, &*SubRegCopy2);
+}
 } // namespace