From 9e586217f4e331bca241df00648ca0b71edd3ef6 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper@sifive.com>
Date: Fri, 16 Dec 2022 23:25:20 -0800
Subject: [PATCH] [RISCV] Support the short-forward-branch predicatd ops in
 RISCVSExtWRemoval.

---
 llvm/lib/Target/RISCV/RISCVInstrInfo.td     |  2 ++
 llvm/lib/Target/RISCV/RISCVSExtWRemoval.cpp | 41 +++++++++++++++++++++++++++--
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVInstrInfo.td b/llvm/lib/Target/RISCV/RISCVInstrInfo.td
index 51b320a..fb890e1 100644
--- a/llvm/lib/Target/RISCV/RISCVInstrInfo.td
+++ b/llvm/lib/Target/RISCV/RISCVInstrInfo.td
@@ -1378,6 +1378,7 @@ def PseudoCCXOR : Pseudo<(outs GPR:$dst),
                   Sched<[WriteSFB, ReadSFB, ReadSFB, ReadSFB, ReadSFB, ReadSFB]>;
 
 // RV64I instructions
+let IsSignExtendingOpW = 1 in {
 def PseudoCCADDW : Pseudo<(outs GPR:$dst),
                           (ins GPR:$lhs, GPR:$rhs, ixlenimm:$cc,
                            GPR:$falsev, GPR:$rs1, GPR:$rs2), []>,
@@ -1387,6 +1388,7 @@ def PseudoCCSUBW : Pseudo<(outs GPR:$dst),
                            GPR:$falsev, GPR:$rs1, GPR:$rs2), []>,
                    Sched<[WriteSFB, ReadSFB, ReadSFB, ReadSFB, ReadSFB, ReadSFB]>;
 }
+}
 
 multiclass SelectCC_GPR_rrirr<RegisterClass valty> {
   let usesCustomInserter = 1 in
diff --git a/llvm/lib/Target/RISCV/RISCVSExtWRemoval.cpp b/llvm/lib/Target/RISCV/RISCVSExtWRemoval.cpp
index 73b9d67..076c787 100644
--- a/llvm/lib/Target/RISCV/RISCVSExtWRemoval.cpp
+++ b/llvm/lib/Target/RISCV/RISCVSExtWRemoval.cpp
@@ -222,6 +222,30 @@ static bool hasAllWUsers(const MachineInstr &OrigMI, MachineRegisterInfo &MRI) {
           return false;
         Worklist.push_back(UserMI);
         break;
+      case RISCV::PseudoCCADD:
+      case RISCV::PseudoCCSUB:
+      case RISCV::PseudoCCAND:
+      case RISCV::PseudoCCOR:
+      case RISCV::PseudoCCXOR:
+        // Instruction either returns operand 4 or performs ADD/SUB/AND/OR/XOR
+        // of operand 5 and 6 and returns that result. If only lower word of the
+        // result is used, then only the lower result of these operands is used.
+        if (OpIdx != 4 && OpIdx != 5 && OpIdx != 6)
+          return false;
+        Worklist.push_back(UserMI);
+        break;
+      case RISCV::PseudoCCADDW:
+      case RISCV::PseudoCCSUBW:
+        // Instruction either returns operand 4 or performs an ADDW/SUBW of
+        // operand 5 and 6 and returns that result.
+        // Only the lower word of operand 5 and 6 is ever needed.
+        if (OpIdx == 5 || OpIdx == 6)
+          return true;
+        // For operand 4 we need to check the users of this instruction.
+        if (OpIdx != 4)
+          return false;
+        Worklist.push_back(UserMI);
+        break;
       }
     }
   }
@@ -383,20 +407,33 @@ static bool isSignExtendedW(Register SrcReg, MachineRegisterInfo &MRI,
     case RISCV::MIN:
     case RISCV::MINU:
     case RISCV::PseudoCCMOVGPR:
+    case RISCV::PseudoCCAND:
+    case RISCV::PseudoCCOR:
+    case RISCV::PseudoCCXOR:
     case RISCV::PHI: {
       // If all incoming values are sign-extended, the output of AND, OR, XOR,
       // MIN, MAX, or PHI is also sign-extended.
 
       // The input registers for PHI are operand 1, 3, ...
       // The input registers for PseudoCCMOVGPR are 4 and 5.
+      // The input registers for PseudoCCAND/OR/XOR are 4, 5, and 6.
       // The input registers for others are operand 1 and 2.
       unsigned B = 1, E = 3, D = 1;
-      if (MI->getOpcode() == RISCV::PHI) {
+      switch (MI->getOpcode()) {
+      case RISCV::PHI:
         E = MI->getNumOperands();
         D = 2;
-      } else if (MI->getOpcode() == RISCV::PseudoCCMOVGPR) {
+        break;
+      case RISCV::PseudoCCMOVGPR:
         B = 4;
         E = 6;
+        break;
+      case RISCV::PseudoCCAND:
+      case RISCV::PseudoCCOR:
+      case RISCV::PseudoCCXOR:
+        B = 4;
+        E = 7;
+        break;
       }
 
       for (unsigned I = B; I != E; I += D) {
-- 
2.7.4