[RISCV] Add regalloc hints for Zcb instructions. (#78949)

This hints the register allocator to use the same register for source
and destination to enable more compression.

2 files changed, 120 insertions, 6 deletions

diff --git a/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp b/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
index 730838e..d7bb46a 100644
--- a/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
+++ b/llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp
@@ -753,6 +753,7 @@ bool RISCVRegisterInfo::getRegAllocationHints(
     SmallVectorImpl<MCPhysReg> &Hints, const MachineFunction &MF,
     const VirtRegMap *VRM, const LiveRegMatrix *Matrix) const {
   const MachineRegisterInfo *MRI = &MF.getRegInfo();
+  auto &Subtarget = MF.getSubtarget<RISCVSubtarget>();
 
   bool BaseImplRetVal = TargetRegisterInfo::getRegAllocationHints(
       VirtReg, Order, Hints, MF, VRM, Matrix);
@@ -776,7 +777,7 @@ bool RISCVRegisterInfo::getRegAllocationHints(
 
   // This is all of the compressible binary instructions. If an instruction
   // needs GPRC register class operands \p NeedGPRC will be set to true.
-  auto isCompressible = [](const MachineInstr &MI, bool &NeedGPRC) {
+  auto isCompressible = [&Subtarget](const MachineInstr &MI, bool &NeedGPRC) {
     NeedGPRC = false;
     switch (MI.getOpcode()) {
     default:
@@ -789,9 +790,16 @@ bool RISCVRegisterInfo::getRegAllocationHints(
     case RISCV::SUBW:
       NeedGPRC = true;
       return true;
-    case RISCV::ANDI:
+    case RISCV::ANDI: {
       NeedGPRC = true;
-      return MI.getOperand(2).isImm() && isInt<6>(MI.getOperand(2).getImm());
+      if (!MI.getOperand(2).isImm())
+        return false;
+      int64_t Imm = MI.getOperand(2).getImm();
+      if (isInt<6>(Imm))
+        return true;
+      // c.zext.b
+      return Subtarget.hasStdExtZcb() && Imm == 255;
+    }
     case RISCV::SRAI:
     case RISCV::SRLI:
       NeedGPRC = true;
@@ -802,6 +810,24 @@ bool RISCVRegisterInfo::getRegAllocationHints(
     case RISCV::ADDI:
     case RISCV::ADDIW:
       return MI.getOperand(2).isImm() && isInt<6>(MI.getOperand(2).getImm());
+    case RISCV::MUL:
+    case RISCV::SEXT_B:
+    case RISCV::SEXT_H:
+    case RISCV::ZEXT_H_RV32:
+    case RISCV::ZEXT_H_RV64:
+      // c.mul, c.sext.b, c.sext.h, c.zext.h
+      NeedGPRC = true;
+      return Subtarget.hasStdExtZcb();
+    case RISCV::ADD_UW:
+      // c.zext.w
+      NeedGPRC = true;
+      return Subtarget.hasStdExtZcb() && MI.getOperand(2).isReg() &&
+             MI.getOperand(2).getReg() == RISCV::X0;
+    case RISCV::XORI:
+      // c.not
+      NeedGPRC = true;
+      return Subtarget.hasStdExtZcb() && MI.getOperand(2).isImm() &&
+             MI.getOperand(2).getImm() == -1;
     }
   };
 
@@ -823,13 +849,15 @@ bool RISCVRegisterInfo::getRegAllocationHints(
     bool NeedGPRC;
     if (isCompressible(MI, NeedGPRC)) {
       if (OpIdx == 0 && MI.getOperand(1).isReg()) {
-        if (!NeedGPRC || isCompressibleOpnd(MI.getOperand(2)))
+        if (!NeedGPRC || MI.getNumExplicitOperands() < 3 ||
+            MI.getOpcode() == RISCV::ADD_UW ||
+            isCompressibleOpnd(MI.getOperand(2)))
           tryAddHint(MO, MI.getOperand(1), NeedGPRC);
         if (MI.isCommutable() && MI.getOperand(2).isReg() &&
             (!NeedGPRC || isCompressibleOpnd(MI.getOperand(1))))
           tryAddHint(MO, MI.getOperand(2), NeedGPRC);
-      } else if (OpIdx == 1 &&
-                 (!NeedGPRC || isCompressibleOpnd(MI.getOperand(2)))) {
+      } else if (OpIdx == 1 && (!NeedGPRC || MI.getNumExplicitOperands() < 3 ||
+                                isCompressibleOpnd(MI.getOperand(2)))) {
         tryAddHint(MO, MI.getOperand(0), NeedGPRC);
       } else if (MI.isCommutable() && OpIdx == 2 &&
                  (!NeedGPRC || isCompressibleOpnd(MI.getOperand(1)))) {
diff --git a/llvm/test/CodeGen/RISCV/zcb-regalloc-hints.ll b/llvm/test/CodeGen/RISCV/zcb-regalloc-hints.ll
new file mode 100644
index 0000000..545d6c6
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/zcb-regalloc-hints.ll
@@ -0,0 +1,86 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc < %s -mtriple=riscv64 -mattr=+m,+zba,+zbb,+zcb | FileCheck %s
+
+define i64 @c_not(i64 %x, i64 %y, i64 %z) {
+; CHECK-LABEL: c_not:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    not a1, a1
+; CHECK-NEXT:    li a0, 1234
+; CHECK-NEXT:    mul a0, a0, a1
+; CHECK-NEXT:    ret
+  %a = xor i64 %y, -1
+  %b = mul i64 %a, 1234
+  ret i64 %b
+}
+
+define i64 @c_mul(i64 %x, i64 %y, i64 %z, i64 %w) {
+; CHECK-LABEL: c_mul:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    mul a1, a1, a2
+; CHECK-NEXT:    lui a0, 1
+; CHECK-NEXT:    or a0, a0, a1
+; CHECK-NEXT:    ret
+  %a = mul i64 %y, %z
+  %b = or i64 %a, 4096
+  ret i64 %b
+}
+
+define i64 @c_sext_b(i64 %x, i8 %y, i64 %z) {
+; CHECK-LABEL: c_sext_b:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    sext.b a1, a1
+; CHECK-NEXT:    lui a0, 1
+; CHECK-NEXT:    or a0, a0, a1
+; CHECK-NEXT:    ret
+  %a = sext i8 %y to i64
+  %b = or i64 %a, 4096
+  ret i64 %b
+}
+
+define i64 @c_sext_h(i64 %x, i16 %y, i64 %z) {
+; CHECK-LABEL: c_sext_h:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    sext.h a1, a1
+; CHECK-NEXT:    lui a0, 1
+; CHECK-NEXT:    or a0, a0, a1
+; CHECK-NEXT:    ret
+  %a = sext i16 %y to i64
+  %b = or i64 %a, 4096
+  ret i64 %b
+}
+
+define i64 @c_zext_b(i64 %x, i8 %y, i64 %z) {
+; CHECK-LABEL: c_zext_b:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    andi a1, a1, 255
+; CHECK-NEXT:    lui a0, 1
+; CHECK-NEXT:    or a0, a0, a1
+; CHECK-NEXT:    ret
+  %a = zext i8 %y to i64
+  %b = or i64 %a, 4096
+  ret i64 %b
+}
+
+define i64 @c_zext_h(i64 %x, i16 %y) {
+; CHECK-LABEL: c_zext_h:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    zext.h a1, a1
+; CHECK-NEXT:    lui a0, 4096
+; CHECK-NEXT:    or a0, a0, a1
+; CHECK-NEXT:    ret
+  %a = zext i16 %y to i64
+  %b = or i64 %a, 16777216
+  ret i64 %b
+}
+
+define i64 @c_zext_w(i64 %x, i32 %y) {
+; CHECK-LABEL: c_zext_w:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    zext.w a1, a1
+; CHECK-NEXT:    li a0, 1234
+; CHECK-NEXT:    mul a0, a0, a1
+; CHECK-NEXT:    ret
+  %a = zext i32 %y to i64
+  %b = mul i64 %a, 1234
+  ret i64 %b
+}