[RISCV] Improve register allocation around vector masks

[frasercrmck]

Summary

With vector mask registers only allocatable to V0 (VMV0Regs) it is
relatively simple to generate code which uses multiple masks and naively
requires spilling.

This patch aims to improve codegen in such cases by telling LLVM it can
use VRRegs to hold masks. This will prevent spilling in many cases by
having LLVM copy to an available VR register.

[frasercrmck]

This has been on my list for a while since @rogfer01 brought it up after D94294. I just encountered it myself so thought I'd nip it in the bud.

I'll include the diff on the test here. I can precommit it if people think it's adequate.

--- a/llvm/test/CodeGen/RISCV/rvv/mask-reg-alloc.mir
+++ b/llvm/test/CodeGen/RISCV/rvv/mask-reg-alloc.mir
@@ -15,24 +15,10 @@ body:             |
     liveins: $v0, $v1, $v2, $v3
     ; CHECK-LABEL: name: mask_reg_alloc
     ; CHECK: liveins: $v0, $v1, $v2, $v3
-    ; CHECK: CFI_INSTRUCTION def_cfa_offset 0
-    ; CHECK: $x10 = PseudoReadVLENB
-    ; CHECK: $x10 = SLLI killed $x10, 1
-    ; CHECK: $x2 = SUB $x2, killed $x10
-    ; CHECK: PseudoVSPILL_M1 $v0, $x2 :: (store unknown-size into %stack.1, align 8)
-    ; CHECK: $x10 = PseudoReadVLENB
-    ; CHECK: $x10 = ADD $x2, killed $x10
-    ; CHECK: PseudoVSPILL_M1 $v1, killed $x10 :: (store unknown-size into %stack.0, align 8)
-    ; CHECK: renamable $v0 = PseudoVRELOAD_M1 $x2 :: (load unknown-size from %stack.1, align 8)
     ; CHECK: renamable $v25 = PseudoVMERGE_VIM_M1 killed renamable $v2, 1, killed renamable $v0, $noreg, -1, implicit $vl, implicit $vtype
-    ; CHECK: $x10 = PseudoReadVLENB
-    ; CHECK: $x10 = ADD $x2, killed $x10
-    ; CHECK: renamable $v0 = PseudoVRELOAD_M1 killed $x10 :: (load unknown-size from %stack.0, align 8)
+    ; CHECK: renamable $v0 = COPY killed renamable $v1
     ; CHECK: renamable $v26 = PseudoVMERGE_VIM_M1 killed renamable $v3, 1, killed renamable $v0, $noreg, -1, implicit $vl, implicit $vtype
     ; CHECK: renamable $v0 = PseudoVADD_VV_M1 killed renamable $v25, killed renamable $v26, $noreg, -1, implicit $vl, implicit $vtype
-    ; CHECK: $x10 = PseudoReadVLENB
-    ; CHECK: $x10 = SLLI killed $x10, 1
-    ; CHECK: $x2 = ADD $x2, killed $x10
     ; CHECK: PseudoRET implicit $v0
     %0:vr = COPY $v0
     %1:vr = COPY $v1

[frasercrmck]

• update test run line

[craig.topper]

LGTM

[frasercrmck]

• rebase on main and pre-committed test case

llvm/lib/Target/RISCV/RISCVRegisterInfo.h

bool requiresFrameIndexScavenging(const MachineFunction &MF) const override {
     return true;
   }
 
   const TargetRegisterClass *
   getPointerRegClass(const MachineFunction &MF,
                      unsigned Kind = 0) const override {
     return &RISCV::GPRRegClass;
   }
+
+  const TargetRegisterClass *
+  getLargestLegalSuperClass(const TargetRegisterClass *RC,
+                            const MachineFunction &) const override;
 };
 }
 
 #endif

llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp

RISCVRegisterInfo::getCallPreservedMask(const MachineFunction & MF,
   case RISCVABI::ABI_ILP32F:
   case RISCVABI::ABI_LP64F:
     return CSR_ILP32F_LP64F_RegMask;
   case RISCVABI::ABI_ILP32D:
   case RISCVABI::ABI_LP64D:
     return CSR_ILP32D_LP64D_RegMask;
   }
 }
+
+const TargetRegisterClass *
+RISCVRegisterInfo::getLargestLegalSuperClass(const TargetRegisterClass *RC,
+                                             const MachineFunction &) const {
+  if (RC == &RISCV::VMV0RegClass)
+    return &RISCV::VRRegClass;
+  return RC;
+}

llvm/test/CodeGen/RISCV/rvv/mask-reg-alloc.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=riscv64 -mattr=+experimental-v -verify-machineinstrs \
+# RUN:     -start-after finalize-isel -stop-after prologepilog < %s | FileCheck %s
+
+--- |
+  define void @mask_reg_alloc() {
+    ret void
+  }
+...
+---
+name: mask_reg_alloc
+tracksRegLiveness: true
+body:             |
+  bb.0 (%ir-block.0):
+    liveins: $v0, $v1, $v2, $v3
+    ; CHECK-LABEL: name: mask_reg_alloc
+    ; CHECK: liveins: $v0, $v1, $v2, $v3
+    ; CHECK: renamable $v25 = PseudoVMERGE_VIM_M1 killed renamable $v2, 1, killed renamable $v0, $noreg, -1, implicit $vl, implicit $vtype
+    ; CHECK: renamable $v0 = COPY killed renamable $v1
+    ; CHECK: renamable $v26 = PseudoVMERGE_VIM_M1 killed renamable $v3, 1, killed renamable $v0, $noreg, -1, implicit $vl, implicit $vtype
+    ; CHECK: renamable $v0 = PseudoVADD_VV_M1 killed renamable $v25, killed renamable $v26, $noreg, -1, implicit $vl, implicit $vtype
+    ; CHECK: PseudoRET implicit $v0
+    %0:vr = COPY $v0
+    %1:vr = COPY $v1
+    %2:vr = COPY $v2
+    %3:vr = COPY $v3
+    %4:vmv0 = COPY %0
+    %5:vrnov0 = PseudoVMERGE_VIM_M1 killed %2, 1, %4, $noreg, -1, implicit $vl, implicit $vtype
+    %6:vmv0 = COPY %1
+    %7:vrnov0 = PseudoVMERGE_VIM_M1 killed %3, 1, %6, $noreg, -1, implicit $vl, implicit $vtype
+    %8:vr = PseudoVADD_VV_M1 killed %5, killed %7, $noreg, -1, implicit $vl, implicit $vtype
+    $v0 = COPY %8
+    PseudoRET implicit $v0
+...