[RISCV] Extend sext.w removal pass to remove unused sign-extensions

[mohammed-nurulhoque]

Summary

Current pass recursively searches definitions backwards to determine if a sext.w is redundant because the value is already in sign-extended form. An instruction propagates a sign-extended definition if giving it a sign-extended input yields a sign-extended output.

This patch adds recursively searching uses forwards to determine if bits 63:32 of sext.w result are not used, so it can be removed. An isntruction propagates sign-extended use if the lower word of its output is only dependent on the lower word of its input.
It also adds a few W instructions to the backwards search.

I have a pending patch dependent on this one that "fixes" other instructions (e.g. ADD -> ADDW) to make a sext.w redundant.

[mohammed-nurulhoque]

Rebase on latest main

[craig.topper]

llvm/test/CodeGen/RISCV/sextw-removal.ll

506

The loop calculate the same value every time. The loop optimizers in the middle can remove it.

Nothing about this test looks a problem that is unique to RISCV. The shifts are redundant on any target and are visible to the middle end. Why shouldn't we clean it up in the middle end?

520

This function is unused

[mohammed-nurulhoque]

Thanks for your comments. As I mentioned, this patch is actually first half of a larger patch which I now uploaded as a signle patch here https://reviews.llvm.org/D119928 (didn't request reviews there yet). The main thing is transforming instructions to their W variant to eliminate sext.w occurences. With that change, checking for unused sext.w was just cheap to add because we already have the isAllUsesReadW function.

I originally broke it into 2 parts because I thought it would be too big for a single review, but I'm finding it hard to isolate a test case for the first part (this review) that wouldn't be optimized out by the middle-end. The full patch linked above has more convincing tests cases. Please do check the full patch and let me know if you think I should proceed with the full patch as a single review and cancel this one.

[craig.topper]

llvm/lib/Target/RISCV/RISCVSExtWRemoval.cpp

116

The register form is handled without knowing anything about the shift amount operand. So why does the immediate value matter?

120

Similar question.

298

Can't changing bit 31 break the sign extension?

[mohammed-nurulhoque]

llvm/lib/Target/RISCV/RISCVSExtWRemoval.cpp

116

In the register form, we check recursively if the uses of SLL only read the lower word.
We can do the same for SLLI, but if the immediate is >= 32, there's a stronger conclusion: we don't have to check the uses because the upper bit are already lost.
If the condition is false, we do the same as the register form.

120

Same here, AND with an immadiate that has leading zeros, we know we don't have to check the uses further. The upper bits are already lost.
Otherwise, we do the same as the register form.

298

true. Updating

[mohammed-nurulhoque]

a polite ping for this review.

[craig.topper]

llvm/lib/Target/RISCV/RISCVSExtWRemoval.cpp

63

const auto -> const auto *. LLVM coding guidelines want to see the *

149

Only one of the operands of ADD_UW uses all bits. Can we know which operand is being used?

162

SLLI_UW could be treated the same as W instructions. It only uses the lower 32 bits of the first operand.

180

Can you put the changes to this function in a separate patch with tests?

396

Combine this into a single if using &&

llvm/test/CodeGen/RISCV/sextw-removal.ll

506

I still don't think this is a robust test. The middle end can make it not a loop so it would never make it to the backend in this form. And it doesn't demonstrate a RISCV specific problem This code is a problem for any target so why is a RISC-V specific solution that correct solution?

[mohammed-nurulhoque]

Thank you for your feedback. I decided to abandon this revision and split its content to one that has just the changes in isSignExtendingOpW and one that has the forward searching https://reviews.llvm.org/D119928. I will request reviewing for these 2 soon.

llvm/lib/Target/RISCV/RISCVSExtWRemoval.cpp

 char RISCVSExtWRemoval::ID = 0;
 INITIALIZE_PASS(RISCVSExtWRemoval, DEBUG_TYPE, "RISCV sext.w Removal", false,
                 false)
 
 FunctionPass *llvm::createRISCVSExtWRemovalPass() {
   return new RISCVSExtWRemoval();
 }
 
+// add uses of MI to the Worklist
+static void addUses(const MachineInstr &MI,
+                    SmallVectorImpl<const MachineInstr *> &Worklist,
+                    MachineRegisterInfo &MRI) {
+  for (auto &UserOp : MRI.reg_operands(MI.getOperand(0).getReg())) {
+    const auto User = UserOp.getParent();
+    if (User == &MI) // ignore the def, current MI
+      continue;
+    Worklist.push_back(User);
+  }
+}
+
+// returns true if all uses of OrigMI only depend on bits 31:0 of its output
+static bool isAllUsesReadW(const MachineInstr &OrigMI,
+                           MachineRegisterInfo &MRI) {
+
+  SmallPtrSet<const MachineInstr *, 4> Visited;
+  SmallVector<const MachineInstr *, 4> Worklist;
+
+  Visited.insert(&OrigMI);
+  addUses(OrigMI, Worklist, MRI);
+
+  while (!Worklist.empty()) {
+    const MachineInstr *MI = Worklist.pop_back_val();
+
+    if (!Visited.insert(MI).second)
+      continue;
+
+    switch (MI->getOpcode()) {
+    case RISCV::ADDIW:
+    case RISCV::ADDW:
+    case RISCV::DIVUW:
+    case RISCV::DIVW:
+    case RISCV::MULW:
+    case RISCV::REMUW:
+    case RISCV::REMW:
+    case RISCV::SLLIW:
+    case RISCV::SLLW:
+    case RISCV::SRAIW:
+    case RISCV::SRAW:
+    case RISCV::SRLIW:
+    case RISCV::SRLW:
+    case RISCV::SUBW:
+    case RISCV::ROLW:
+    case RISCV::RORW:
+    case RISCV::RORIW:
+    case RISCV::CLZW:
+    case RISCV::CTZW:
+    case RISCV::CPOPW:
+    case RISCV::FCVT_S_W:
+    case RISCV::FCVT_S_WU:
+    case RISCV::FCVT_D_W:
+    case RISCV::FCVT_D_WU:
+      continue;
+
+    // these overwrite higher input bits, otherwise the lower word of output
+    // depends only on the lower word of input. So check their uses read W.
+    case RISCV::SLLI:
+      if (MI->getOperand(2).getImm() >= 32)
+        continue;
+      addUses(*MI, Worklist, MRI);
+      continue;
+    case RISCV::ANDI:
+      if (isUInt<11>(MI->getOperand(2).getImm()))
+        continue;
+      addUses(*MI, Worklist, MRI);
+      continue;
+    case RISCV::ORI:
+      if (!isUInt<11>(MI->getOperand(2).getImm()))
+        continue;
+      addUses(*MI, Worklist, MRI);
+      continue;
+
+    case RISCV::BEXTI:
+      if (MI->getOperand(2).getImm() >= 32)
+        return false;
+      continue;
+
+    // For these, lower word of output in these operations, depends only on
+    // the lower word of input. So, we check all uses only read lower word.
+    case RISCV::COPY:
+    case RISCV::PHI:
+
+    case RISCV::ADD:
+    case RISCV::ADDI:
+    case RISCV::AND:
+    case RISCV::MUL:
+    case RISCV::OR:
+    case RISCV::SLL:
+    case RISCV::SUB:
+
+    case RISCV::ADD_UW:
+    case RISCV::ANDN:
+    case RISCV::CLMUL:
+    case RISCV::ORC_B:
+    case RISCV::ORN:
+    case RISCV::SEXT_B:
+    case RISCV::SEXT_H:
+    case RISCV::SH1ADD:
+    case RISCV::SH1ADD_UW:
+    case RISCV::SH2ADD:
+    case RISCV::SH2ADD_UW:
+    case RISCV::SH3ADD:
+    case RISCV::SH3ADD_UW:
+    case RISCV::SLLI_UW:
+    case RISCV::XNOR:
+    case RISCV::XOR:
+    case RISCV::XORI:
+    case RISCV::ZEXT_H_RV64:
+      addUses(*MI, Worklist, MRI);
+      continue;
+    default:
+      return false;
+    }
+  }
+  return true;
+}
+
 // This function returns true if the machine instruction always outputs a value
 // where bits 63:32 match bit 31.
 // TODO: Allocate a bit in TSFlags for the W instructions?
 // TODO: Add other W instructions.
 static bool isSignExtendingOpW(const MachineInstr &MI) {
   switch (MI.getOpcode()) {
   case RISCV::LUI:
   case RISCV::LW:
   case RISCV::ADDW:
   case RISCV::ADDIW:
   case RISCV::SUBW:
   case RISCV::MULW:
   case RISCV::SLLW:
static bool isSignExtendingOpW(const MachineInstr &MI) {
   case RISCV::FCVT_W_H:
   case RISCV::FCVT_WU_H:
   case RISCV::FCVT_W_S:
   case RISCV::FCVT_WU_S:
   case RISCV::FCVT_W_D:
   case RISCV::FCVT_WU_D:
   case RISCV::FMV_X_W:
   // The following aren't W instructions, but are either sign extended from a
-  // smaller size or put zeros in bits 63:31.
+  // smaller size, always outputs a small integer, or put zeros in bits 63:31.
   case RISCV::LBU:
   case RISCV::LHU:
   case RISCV::LB:
   case RISCV::LH:
   case RISCV::SLT:
   case RISCV::SLTI:
   case RISCV::SLTU:
   case RISCV::SLTIU:
   case RISCV::SEXT_B:
   case RISCV::SEXT_H:
   case RISCV::ZEXT_H_RV64:
   case RISCV::FMV_X_H:
+  case RISCV::BEXT:
+  case RISCV::BEXTI:
+  case RISCV::CLZ:
+  case RISCV::CPOP:
+  case RISCV::CTZ:
     return true;
   // shifting right sufficiently makes the value 32-bit sign-extended
   case RISCV::SRAI:
     return MI.getOperand(2).getImm() >= 32;
   case RISCV::SRLI:
     return MI.getOperand(2).getImm() > 32;
   // The LI pattern ADDI rd, X0, imm is sign extended.
   case RISCV::ADDI:
case RISCV::COPY: {
       const MachineInstr *SrcMI = MRI.getVRegDef(SrcReg);
       if (!SrcMI)
         return false;
 
       // Add SrcMI to the worklist.
       Worklist.push_back(SrcMI);
       break;
     }
+
+    // For these We just need to check if the 1st operand is sign extended.
+    case RISCV::BCLRI:
+    case RISCV::BINVI:
+    case RISCV::BSETI:
+      if (MI->getOperand(2).getImm() >= 31)
+        return false;
+      LLVM_FALLTHROUGH;
     case RISCV::REM:
     case RISCV::ANDI:
     case RISCV::ORI:
     case RISCV::XORI: {
       // |Remainder| is always <= |Dividend|. If D is 32-bit, then so is R.
       // DIV doesn't work because of the edge case 0xf..f 8000 0000 / (long)-1
-      // Logical operations use a sign extended 12-bit immediate. We just need
-      // to check if the other operand is sign extended.
+      // Logical operations use a sign extended 12-bit immediate.
       Register SrcReg = MI->getOperand(1).getReg();
       if (!SrcReg.isVirtual())
         return false;
       const MachineInstr *SrcMI = MRI.getVRegDef(SrcReg);
       if (!SrcMI)
         return false;
 
       // Add SrcMI to the worklist.
for (auto I = MBB.begin(), IE = MBB.end(); I != IE;) {
         continue;
 
       // Input should be a virtual register.
       Register SrcReg = MI->getOperand(1).getReg();
       if (!SrcReg.isVirtual())
         continue;
 
       const MachineInstr &SrcMI = *MRI.getVRegDef(SrcReg);
+
+      // If all uses only depend on the lower 32-bits of the result of MI,
+      // then sext.w is unnecessary
+      if (!isAllUsesReadW(*MI, MRI)) {
+        // If all definitions reaching MI sign-extend their output,
+        // then sext.w is redundant
         if (!isSignExtendedW(SrcMI, MRI))
           continue;
+      }
+      // either condition is satisfied, so proceed with deleting the sext.w
 
       Register DstReg = MI->getOperand(0).getReg();
       if (!MRI.constrainRegClass(SrcReg, MRI.getRegClass(DstReg)))
         continue;
 
       LLVM_DEBUG(dbgs() << "Removing redundant sign-extension\n");
       MRI.replaceRegWith(DstReg, SrcReg);
       MRI.clearKillFlags(SrcReg);
       MI->eraseFromParent();
       ++NumRemovedSExtW;
       MadeChange = true;
     }
   }
 
   return MadeChange;
 }

llvm/test/CodeGen/RISCV/sextw-removal.ll

bb2:                                              ; preds = %bb2, %bb
   %i4 = tail call float @baz(i32 signext %i3)
   %i5 = bitcast float %i4 to i32
   %i6 = fcmp oeq float %i4, zeroinitializer
   br i1 %i6, label %bb7, label %bb2
 
 bb7:                                              ; preds = %bb2
   ret void
 }
+
+define i64 @test11(i64 %arg1, i64 %arg2)  {
+; CHECK-LABEL: test11:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    li a2, 254
+; CHECK-NEXT:    mv a3, a0
+; CHECK-NEXT:  .LBB9_1: # %bb2
+; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    mul a3, a3, a1
+; CHECK-NEXT:    add a3, a3, a1
+; CHECK-NEXT:    blt a2, a0, .LBB9_1
+; CHECK-NEXT:  # %bb.2: # %bb7
+; CHECK-NEXT:    sext.w a0, a3
+; CHECK-NEXT:    ret
+;
+; NOREMOVAL-LABEL: test11:
+; NOREMOVAL:       # %bb.0: # %entry
+; NOREMOVAL-NEXT:    li a2, 254
+; NOREMOVAL-NEXT:    mv a3, a0
+; NOREMOVAL-NEXT:  .LBB9_1: # %bb2
+; NOREMOVAL-NEXT:    # =>This Inner Loop Header: Depth=1
+; NOREMOVAL-NEXT:    sext.w a3, a3
+; NOREMOVAL-NEXT:    mul a3, a3, a1
+; NOREMOVAL-NEXT:    add a3, a3, a1
+; NOREMOVAL-NEXT:    blt a2, a0, .LBB9_1
+; NOREMOVAL-NEXT:  # %bb.2: # %bb7
+; NOREMOVAL-NEXT:    sext.w a0, a3
+; NOREMOVAL-NEXT:    ret
+entry:
+  br label %bb2
+
+bb2:                                              ; preds = %bb2, %bb
+  %i1 = phi i64 [ %arg1, %entry ], [ %i5, %bb2 ]
+  %i2 = shl i64 %i1, 32
+  %i3 = ashr i64 %i2, 32
+  %i4 = mul i64 %i3, %arg2
+  %i5 = add i64 %i4, %arg2
+  %i6 = icmp slt i64 %arg1, 255
+  br i1 %i6, label %bb7, label %bb2
+
+bb7:                                              ; preds = %bb2
+  %i7 = shl i64 %i5, 32
+  %i8 = ashr i64 %i7, 32
+  ret i64 %i8
+}