[LLVM][XTHeadVector] Fix RISCVInsertVSETVLI: handle hasAVLImm and constrain temporary register to be GPRNoX0 (llvm#121)

* [LLVM][XTHeadVector] Fix `Info.hasAVLImm()` branch

and constrain the temporary register to be `GPRNoX0`

* [LLVM][XTHeadVector] Update corresponding tests

Author: imkiva

llvm/lib/Target/RISCV/RISCVInsertVSETVLI.cpp

@@ -172,7 +172,7 @@ static bool hasUndefinedMergeOp(const MachineInstr &MI,
     return true;
 
   // If the tied operand is an IMPLICIT_DEF (or a REG_SEQUENCE whose operands
-  // are solely IMPLICIT_DEFS), the pass through lanes are undefined.  
+  // are solely IMPLICIT_DEFS), the pass through lanes are undefined.
   const MachineOperand &UseMO = MI.getOperand(UseOpIdx);
   if (MachineInstr *UseMI = MRI.getVRegDef(UseMO.getReg())) {
     if (UseMI->isImplicitDef())
@@ -755,6 +755,10 @@ class RISCVInsertVSETVLI : public MachineFunctionPass {
   void doLocalPostpass(MachineBasicBlock &MBB);
   void doPRE(MachineBasicBlock &MBB);
   void insertReadVL(MachineBasicBlock &MBB);
+  void emulateXTHeadVectorVSETIVLI(MachineBasicBlock &MBB,
+                                   MachineBasicBlock::iterator &InsertPt,
+                                   const DebugLoc &DL, const VSETVLIInfo &Info,
+                                   uint64_t UImm);
 };
 
 } // end anonymous namespace
@@ -916,11 +920,14 @@ void RISCVInsertVSETVLI::insertVSETVLI(MachineBasicBlock &MBB,
   }
 
   if (Info.hasAVLImm()) {
-    assert(!HasVendorXTHeadV && "XTHeadV extension does not support AVLImm");
-    BuildMI(MBB, InsertPt, DL, TII->get(RISCV::PseudoVSETIVLI))
-        .addReg(RISCV::X0, RegState::Define | RegState::Dead)
-        .addImm(Info.getAVLImm())
-        .addImm(Info.encodeVTYPE());
+    if (HasVendorXTHeadV) {
+      emulateXTHeadVectorVSETIVLI(MBB, InsertPt, DL, Info, Info.getAVLImm());
+    } else {
+      BuildMI(MBB, InsertPt, DL, TII->get(RISCV::PseudoVSETIVLI))
+          .addReg(RISCV::X0, RegState::Define | RegState::Dead)
+          .addImm(Info.getAVLImm())
+          .addImm(Info.encodeVTYPE());
+    }
     return;
   }
 
@@ -943,15 +950,7 @@ void RISCVInsertVSETVLI::insertVSETVLI(MachineBasicBlock &MBB,
     }
     // Otherwise use an AVL of 0 to avoid depending on previous vl.
     if (HasVendorXTHeadV) {
-      // Generate the equivalent of `vsetivli rd, uimm, vtypei` in RVV 0.7
-      auto UImmR = MRI->createVirtualRegister(&RISCV::GPRRegClass);
-      BuildMI(MBB, InsertPt, DL, TII->get(RISCV::PseudoLI))
-          .addReg(UImmR, RegState::Define)
-          .addImm(0);
-      BuildMI(MBB, InsertPt, DL, TII->get(RISCV::PseudoTH_VSETVLI))
-          .addReg(RISCV::X0, RegState::Define | RegState::Dead)
-          .addReg(UImmR, RegState::Kill)
-          .addImm(Info.encodeXTHeadVTYPE());
+      emulateXTHeadVectorVSETIVLI(MBB, InsertPt, DL, Info, 0);
     } else {
       BuildMI(MBB, InsertPt, DL, TII->get(RISCV::PseudoVSETIVLI))
           .addReg(RISCV::X0, RegState::Define | RegState::Dead)
@@ -980,6 +979,30 @@ void RISCVInsertVSETVLI::insertVSETVLI(MachineBasicBlock &MBB,
       .addImm(TypeI);
 }
 
+void RISCVInsertVSETVLI::emulateXTHeadVectorVSETIVLI(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator &InsertPt,
+    const DebugLoc &DL, const VSETVLIInfo &Info, uint64_t UImm) {
+  // RVV 1.0:
+  // vsetvli rd, rs1, vtypei # rd = new vl, rs1 = AVL, vtypei = vtype
+  // vsetivli rd, uimm, vtypei # rd = new vl, uimm = AVL, vtypei = vtype
+
+  // XTHeadVector:
+  // vsetvli rd, rs1, vtypei # rd = new vl, rs1 = AVL, vtypei = new vtype
+
+  // Let's load the AVLImm to a temporary register and use it as AVL in vsetvli
+  auto UImmR = MRI->createVirtualRegister(&RISCV::GPRRegClass);
+  // PseudoTH_VSETVLI requires the rs1 to be a GPRNoX0 register.
+  MRI->constrainRegClass(UImmR, &RISCV::GPRNoX0RegClass);
+
+  BuildMI(MBB, InsertPt, DL, TII->get(RISCV::PseudoLI))
+      .addReg(UImmR, RegState::Define)
+      .addImm(UImm);
+  BuildMI(MBB, InsertPt, DL, TII->get(RISCV::PseudoTH_VSETVLI))
+      .addReg(RISCV::X0, RegState::Define | RegState::Dead)
+      .addReg(UImmR, RegState::Kill)
+      .addImm(Info.encodeXTHeadVTYPE());
+}
+
 void RISCVInsertVSETVLI::insertVSETVLIForCOPY(MachineBasicBlock &MBB) {
   auto findRegSequence = [] (const MachineBasicBlock& BB,
                              MachineBasicBlock::iterator I,
@@ -1709,7 +1732,7 @@ bool RISCVInsertVSETVLI::runOnMachineFunction(MachineFunction &MF) {
     insertReadVL(MBB);
 
   if (HasVendorXTHeadV)
-    for (MachineBasicBlock &MBB : MF)        
+    for (MachineBasicBlock &MBB : MF)
           insertVSETVLIForCOPY(MBB);
 
   BlockInfo.clear();

llvm/test/CodeGen/RISCV/rvv0p71/vmfirst.ll

@@ -22,6 +22,21 @@ entry:
   ret iXLen %a
 }
 
+define iXLen @intrinsic_vmfirst_m_nxv1i1_zero(<vscale x 1 x i1> %0) nounwind {
+; CHECK-LABEL: intrinsic_vmfirst_m_nxv1i1_zero:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    li a0, 0
+; CHECK-NEXT:    th.vsetvli zero, a0, e64, m1, d1
+; CHECK-NEXT:    th.vmfirst.m a0, v0
+; CHECK-NEXT:    ret
+entry:
+  %a = call iXLen @llvm.riscv.th.vmfirst.iXLen.nxv1i1(
+    <vscale x 1 x i1> %0,
+    iXLen 0)
+
+  ret iXLen %a
+}
+
 declare iXLen @llvm.riscv.th.vmfirst.mask.iXLen.nxv1i1(
   <vscale x 1 x i1>,
   <vscale x 1 x i1>,
@@ -52,6 +67,32 @@ entry:
   ret iXLen %a
 }
 
+define iXLen @intrinsic_vmfirst_mask_m_nxv1i1_zero(<vscale x 1 x i1> %0, <vscale x 1 x i1> %1) nounwind {
+; CHECK-LABEL: intrinsic_vmfirst_mask_m_nxv1i1_zero:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    csrr a0, vl
+; CHECK-NEXT:    csrr a1, vtype
+; CHECK-NEXT:    th.vsetvli zero, zero, e8, m1, d1
+; CHECK-NEXT:    th.vmv.v.v v9, v0
+; CHECK-NEXT:    th.vsetvl zero, a0, a1
+; CHECK-NEXT:    csrr a0, vl
+; CHECK-NEXT:    csrr a1, vtype
+; CHECK-NEXT:    th.vsetvli zero, zero, e8, m1, d1
+; CHECK-NEXT:    th.vmv.v.v v0, v8
+; CHECK-NEXT:    th.vsetvl zero, a0, a1
+; CHECK-NEXT:    li a0, 0
+; CHECK-NEXT:    th.vsetvli zero, a0, e64, m1, d1
+; CHECK-NEXT:    th.vmfirst.m a0, v9, v0.t
+; CHECK-NEXT:    ret
+entry:
+  %a = call iXLen @llvm.riscv.th.vmfirst.mask.iXLen.nxv1i1(
+    <vscale x 1 x i1> %0,
+    <vscale x 1 x i1> %1,
+    iXLen 0)
+
+  ret iXLen %a
+}
+
 declare iXLen @llvm.riscv.th.vmfirst.iXLen.nxv2i1(
   <vscale x 2 x i1>,
   iXLen);

llvm/test/CodeGen/RISCV/rvv0p71/vmpopc.ll

@@ -22,6 +22,21 @@ entry:
   ret iXLen %a
 }
 
+define iXLen @intrinsic_vmpopc_m_nxv1i1_zero(<vscale x 1 x i1> %0) nounwind {
+; CHECK-LABEL: intrinsic_vmpopc_m_nxv1i1_zero:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    li a0, 0
+; CHECK-NEXT:    th.vsetvli zero, a0, e64, m1, d1
+; CHECK-NEXT:    th.vmpopc.m a0, v0
+; CHECK-NEXT:    ret
+entry:
+  %a = call iXLen @llvm.riscv.th.vmpopc.iXLen.nxv1i1(
+    <vscale x 1 x i1> %0,
+    iXLen 0)
+
+  ret iXLen %a
+}
+
 declare iXLen @llvm.riscv.th.vmpopc.mask.iXLen.nxv1i1(
   <vscale x 1 x i1>,
   <vscale x 1 x i1>,
@@ -52,6 +67,32 @@ entry:
   ret iXLen %a
 }
 
+define iXLen @intrinsic_vmpopc_mask_m_nxv1i1_zero(<vscale x 1 x i1> %0, <vscale x 1 x i1> %1) nounwind {
+; CHECK-LABEL: intrinsic_vmpopc_mask_m_nxv1i1_zero:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    csrr a0, vl
+; CHECK-NEXT:    csrr a1, vtype
+; CHECK-NEXT:    th.vsetvli zero, zero, e8, m1, d1
+; CHECK-NEXT:    th.vmv.v.v v9, v0
+; CHECK-NEXT:    th.vsetvl zero, a0, a1
+; CHECK-NEXT:    csrr a0, vl
+; CHECK-NEXT:    csrr a1, vtype
+; CHECK-NEXT:    th.vsetvli zero, zero, e8, m1, d1
+; CHECK-NEXT:    th.vmv.v.v v0, v8
+; CHECK-NEXT:    th.vsetvl zero, a0, a1
+; CHECK-NEXT:    li a0, 0
+; CHECK-NEXT:    th.vsetvli zero, a0, e64, m1, d1
+; CHECK-NEXT:    th.vmpopc.m a0, v9, v0.t
+; CHECK-NEXT:    ret
+entry:
+  %a = call iXLen @llvm.riscv.th.vmpopc.mask.iXLen.nxv1i1(
+    <vscale x 1 x i1> %0,
+    <vscale x 1 x i1> %1,
+    iXLen 0)
+
+  ret iXLen %a
+}
+
 declare iXLen @llvm.riscv.th.vmpopc.iXLen.nxv2i1(
   <vscale x 2 x i1>,
   iXLen);