[RISCV] Keep AVLReg define instr inside VSETVLInfo

Currently, the vsetvli pass track the define instruction through MRI->getVRegDef due to the SSA form.

This patch keeps the AVLReg DefMI within VSETVLInfo during construction. And replace MRI->getVRegDef(AVLReg) with getAVLRegDefMI().

This information is useful when vsetvli pass live in post-ra situation.

The testcases don't change because the VReg always has a unique def in SSA.

Author: BeMg

llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp

@@ -153,7 +153,7 @@ static std::optional<unsigned> getEEWForLoadStore(const MachineInstr &MI) {
   }
 }
 
-static bool isNonZeroLoadImmediate(MachineInstr &MI) {
+static bool isNonZeroLoadImmediate(const MachineInstr &MI) {
   return MI.getOpcode() == RISCV::ADDI &&
     MI.getOperand(1).isReg() && MI.getOperand(2).isImm() &&
     MI.getOperand(1).getReg() == RISCV::X0 &&
@@ -438,6 +438,8 @@ class VSETVLIInfo {
     unsigned AVLImm;
   };
 
+  const MachineInstr *AVLDefMI;
+
   enum : uint8_t {
     Uninitialized,
     AVLIsReg,
@@ -454,7 +456,7 @@ class VSETVLIInfo {
 
 public:
   VSETVLIInfo()
-      : AVLImm(0), TailAgnostic(false), MaskAgnostic(false),
+      : AVLImm(0), AVLDefMI(nullptr), TailAgnostic(false), MaskAgnostic(false),
         SEWLMULRatioOnly(false) {}
 
   static VSETVLIInfo getUnknown() {
@@ -478,6 +480,8 @@ class VSETVLIInfo {
     State = AVLIsImm;
   }
 
+  void setAVLDefMI(const MachineInstr *DefMI) { AVLDefMI = DefMI; }
+
   bool hasAVLImm() const { return State == AVLIsImm; }
   bool hasAVLReg() const { return State == AVLIsReg; }
   Register getAVLReg() const {
@@ -489,13 +493,16 @@ class VSETVLIInfo {
     return AVLImm;
   }
 
+  const MachineInstr *getAVLDefMI() const { return AVLDefMI; }
+
   void setAVL(VSETVLIInfo Info) {
     assert(Info.isValid());
     if (Info.isUnknown())
       setUnknown();
-    else if (Info.hasAVLReg())
+    else if (Info.hasAVLReg()) {
       setAVLReg(Info.getAVLReg());
-    else {
+      setAVLDefMI(Info.getAVLDefMI());
+    } else {
       assert(Info.hasAVLImm());
       setAVLImm(Info.getAVLImm());
     }
@@ -512,8 +519,7 @@ class VSETVLIInfo {
     if (hasAVLReg()) {
       if (getAVLReg() == RISCV::X0)
         return true;
-      if (MachineInstr *MI = MRI.getVRegDef(getAVLReg());
-          MI && isNonZeroLoadImmediate(*MI))
+      if (getAVLDefMI() && isNonZeroLoadImmediate(*getAVLDefMI()))
         return true;
       return false;
     }
@@ -792,7 +798,8 @@ INITIALIZE_PASS(RISCVInsertVSETVLI, DEBUG_TYPE, RISCV_INSERT_VSETVLI_NAME,
 
 // Return a VSETVLIInfo representing the changes made by this VSETVLI or
 // VSETIVLI instruction.
-static VSETVLIInfo getInfoForVSETVLI(const MachineInstr &MI) {
+static VSETVLIInfo getInfoForVSETVLI(const MachineInstr &MI,
+                                     const MachineRegisterInfo &MRI) {
   VSETVLIInfo NewInfo;
   if (MI.getOpcode() == RISCV::PseudoVSETIVLI) {
     NewInfo.setAVLImm(MI.getOperand(1).getImm());
@@ -803,6 +810,8 @@ static VSETVLIInfo getInfoForVSETVLI(const MachineInstr &MI) {
     assert((AVLReg != RISCV::X0 || MI.getOperand(0).getReg() != RISCV::X0) &&
            "Can't handle X0, X0 vsetvli yet");
     NewInfo.setAVLReg(AVLReg);
+    if (AVLReg.isVirtual())
+      NewInfo.setAVLDefMI(MRI.getVRegDef(AVLReg));
   }
   NewInfo.setVTYPE(MI.getOperand(2).getImm());
 
@@ -875,6 +884,8 @@ static VSETVLIInfo computeInfoForInstr(const MachineInstr &MI, uint64_t TSFlags,
         InstrInfo.setAVLImm(Imm);
     } else {
       InstrInfo.setAVLReg(VLOp.getReg());
+      if (VLOp.getReg().isVirtual())
+        InstrInfo.setAVLDefMI(MRI->getVRegDef(VLOp.getReg()));
     }
   } else {
     assert(isScalarExtractInstr(MI));
@@ -892,9 +903,10 @@ static VSETVLIInfo computeInfoForInstr(const MachineInstr &MI, uint64_t TSFlags,
   // register AVLs to avoid extending live ranges without being sure we can
   // kill the original source reg entirely.
   if (InstrInfo.hasAVLReg() && InstrInfo.getAVLReg().isVirtual()) {
-    MachineInstr *DefMI = MRI->getVRegDef(InstrInfo.getAVLReg());
-    if (DefMI && isVectorConfigInstr(*DefMI)) {
-      VSETVLIInfo DefInstrInfo = getInfoForVSETVLI(*DefMI);
+    if (InstrInfo.getAVLDefMI() &&
+        isVectorConfigInstr(*InstrInfo.getAVLDefMI())) {
+      VSETVLIInfo DefInstrInfo =
+          getInfoForVSETVLI(*InstrInfo.getAVLDefMI(), *MRI);
       if (DefInstrInfo.hasSameVLMAX(InstrInfo) &&
           (DefInstrInfo.hasAVLImm() || DefInstrInfo.getAVLReg() == RISCV::X0)) {
         InstrInfo.setAVL(DefInstrInfo);
@@ -934,9 +946,9 @@ void RISCVInsertVSETVLI::insertVSETVLI(MachineBasicBlock &MBB,
     // same, we can use the X0, X0 form.
     if (Info.hasSameVLMAX(PrevInfo) && Info.hasAVLReg() &&
         Info.getAVLReg().isVirtual()) {
-      if (MachineInstr *DefMI = MRI->getVRegDef(Info.getAVLReg())) {
-        if (isVectorConfigInstr(*DefMI)) {
-          VSETVLIInfo DefInfo = getInfoForVSETVLI(*DefMI);
+      if (Info.getAVLDefMI()) {
+        if (isVectorConfigInstr(*Info.getAVLDefMI())) {
+          VSETVLIInfo DefInfo = getInfoForVSETVLI(*Info.getAVLDefMI(), *MRI);
           if (DefInfo.hasSameAVL(PrevInfo) && DefInfo.hasSameVLMAX(PrevInfo)) {
             BuildMI(MBB, InsertPt, DL, TII->get(RISCV::PseudoVSETVLIX0))
                 .addReg(RISCV::X0, RegState::Define | RegState::Dead)
@@ -1056,9 +1068,9 @@ bool RISCVInsertVSETVLI::needVSETVLI(const MachineInstr &MI,
   // VSETVLI here.
   if (Require.hasAVLReg() && Require.getAVLReg().isVirtual() &&
       CurInfo.hasCompatibleVTYPE(Used, Require)) {
-    if (MachineInstr *DefMI = MRI->getVRegDef(Require.getAVLReg())) {
-      if (isVectorConfigInstr(*DefMI)) {
-        VSETVLIInfo DefInfo = getInfoForVSETVLI(*DefMI);
+    if (Require.getAVLDefMI()) {
+      if (isVectorConfigInstr(*Require.getAVLDefMI())) {
+        VSETVLIInfo DefInfo = getInfoForVSETVLI(*Require.getAVLDefMI(), *MRI);
         if (DefInfo.hasSameAVL(CurInfo) && DefInfo.hasSameVLMAX(CurInfo))
           return false;
       }
@@ -1145,13 +1157,15 @@ void RISCVInsertVSETVLI::transferBefore(VSETVLIInfo &Info,
 void RISCVInsertVSETVLI::transferAfter(VSETVLIInfo &Info,
                                        const MachineInstr &MI) const {
   if (isVectorConfigInstr(MI)) {
-    Info = getInfoForVSETVLI(MI);
+    Info = getInfoForVSETVLI(MI, *MRI);
     return;
   }
 
   if (RISCV::isFaultFirstLoad(MI)) {
     // Update AVL to vl-output of the fault first load.
     Info.setAVLReg(MI.getOperand(1).getReg());
+    if (MI.getOperand(1).getReg().isVirtual())
+      Info.setAVLDefMI(MRI->getVRegDef(MI.getOperand(1).getReg()));
     return;
   }
 
@@ -1270,7 +1284,7 @@ bool RISCVInsertVSETVLI::needVSETVLIPHI(const VSETVLIInfo &Require,
 
     // We found a VSET(I)VLI make sure it matches the output of the
     // predecessor block.
-    VSETVLIInfo DefInfo = getInfoForVSETVLI(*DefMI);
+    VSETVLIInfo DefInfo = getInfoForVSETVLI(*DefMI, *MRI);
     if (!DefInfo.hasSameAVL(PBBInfo.Exit) ||
         !DefInfo.hasSameVTYPE(PBBInfo.Exit))
       return true;
@@ -1418,7 +1432,7 @@ void RISCVInsertVSETVLI::doPRE(MachineBasicBlock &MBB) {
   // we need to prove the value is available at the point we're going
   // to insert the vsetvli at.
   if (AvailableInfo.hasAVLReg() && RISCV::X0 != AvailableInfo.getAVLReg()) {
-    MachineInstr *AVLDefMI = MRI->getVRegDef(AvailableInfo.getAVLReg());
+    const MachineInstr *AVLDefMI = AvailableInfo.getAVLDefMI();
     if (!AVLDefMI)
       return;
     // This is an inline dominance check which covers the case of
@@ -1504,8 +1518,8 @@ static bool canMutatePriorConfig(const MachineInstr &PrevMI,
     if (Used.VLZeroness) {
       if (isVLPreservingConfig(PrevMI))
         return false;
-      if (!getInfoForVSETVLI(PrevMI).hasEquallyZeroAVL(getInfoForVSETVLI(MI),
-                                                       MRI))
+      if (!getInfoForVSETVLI(PrevMI, MRI)
+               .hasEquallyZeroAVL(getInfoForVSETVLI(MI, MRI), MRI))
         return false;
     }