[RISCV] Fold stack reload into sext.w by using lw instead of ld.

[craig.topper]

Summary

We can use lw to load 4 bytes from the stack and sign extend them
instead of loading all 8 bytes.

[asb]

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

648

It would probably be a minor improvement in readability to define a isSEXTW helper and use it here and in RISCVSExtWRemoval.

650

I think you could ditch this and replace the .addMemOperand below with .setMemRefs(MI.memoperands())

[craig.topper]

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

650

MI is the sext.w, it doesn’t have a memoperand and the load hasn’t been created yet.

[asb]

LGTM.

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

650

Ah yes, of course it is. Sorry for the confusion.

[craig.topper]

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

650

Should I check for little endian target here? For big endian I think we need to use 4 for the offset.

[craig.topper]

Add isSEXT_W helper.
Add check for big endian.

[frasercrmck]

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

652

Maybe it's just me but I read "sext.w+load" the other way round than it's intended here. Maybe I'm interpreting "+" as "followed-by"?

[asb]

Good thought re big-endian. This looks good to me, modulo tweaking the comment as Fraser suggested.

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

652

You're right, load+sext.w would seem clearer. Or maybe just use S-expressions to remove any ambiguity.

[craig.topper]

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

652

I think i wrong it that way because of "load from stack". I rewrote to "load from stack followed by sext.w" and remove the +

[reames]

llvm/test/CodeGen/RISCV/stack-folding.ll

24

Hm, looking at these spills sparks a question for discussion.

Do we have infrastructure to narrow the spill itself? If we know only four bytes are used, can we use a narrower store? And maybe even a smaller stack slot? I think the answer is here is no, but maybe it's worth some discussion?

There's two ways of narrowing this I see.

The first is the use of the narrow type - though that's been lowered so I really mean the presence of the sext.w. This could be phrased as some time of demanded bits analysis. Though, given how late we spill, maybe we don't have anything like that.

The second is exploiting the fact that only the i1 comparison value is actually needed. We could hoist the compare up to the site of the spill, and then store only the i1 - probably i8 in practice - result to the stack. This has all the general speculation safety and profitability concerns.

Anyways, just a thought that occurred.

llvm/lib/Target/RISCV/RISCVInstrInfo.h

void storeRegToStackSlot(MachineBasicBlock &MBB,
                            const TargetRegisterClass *RC,
                            const TargetRegisterInfo *TRI) const override;
 
   void loadRegFromStackSlot(MachineBasicBlock &MBB,
                             MachineBasicBlock::iterator MBBI, Register DstReg,
                             int FrameIndex, const TargetRegisterClass *RC,
                             const TargetRegisterInfo *TRI) const override;
 
+  using TargetInstrInfo::foldMemoryOperandImpl;
+  MachineInstr *foldMemoryOperandImpl(MachineFunction &MF, MachineInstr &MI,
+                                      ArrayRef<unsigned> Ops,
+                                      MachineBasicBlock::iterator InsertPt,
+                                      int FrameIndex,
+                                      LiveIntervals *LIS = nullptr,
+                                      VirtRegMap *VRM = nullptr) const override;
+
   // Materializes the given integer Val into DstReg.
   void movImm(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
               const DebugLoc &DL, Register DstReg, uint64_t Val,
               MachineInstr::MIFlag Flag = MachineInstr::NoFlags) const;
 
   unsigned getInstSizeInBytes(const MachineInstr &MI) const override;
 
   bool analyzeBranch(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
Register getVLENFactoredAmount(
       MachineInstr::MIFlag Flag = MachineInstr::NoFlags) const;
 
 protected:
   const RISCVSubtarget &STI;
 };
 
 namespace RISCV {
 
+// Returns true if this is the sext.w pattern, addiw rd, rs1, 0.
+bool isSEXT_W(const MachineInstr &MI);
+
 // Returns true if the given MI is an RVV instruction opcode for which we may
 // expect to see a FrameIndex operand.
 bool isRVVSpill(const MachineInstr &MI);
 
 Optional<std::pair<unsigned, unsigned>> isRVVSpillForZvlsseg(unsigned Opcode);
 
 bool isFaultFirstLoad(const MachineInstr &MI);
 

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

if (IsScalableVector) {
 
     BuildMI(MBB, I, DL, get(Opcode), DstReg)
         .addFrameIndex(FI)
         .addImm(0)
         .addMemOperand(MMO);
   }
 }
 
+MachineInstr *RISCVInstrInfo::foldMemoryOperandImpl(
+    MachineFunction &MF, MachineInstr &MI, ArrayRef<unsigned> Ops,
+    MachineBasicBlock::iterator InsertPt, int FrameIndex, LiveIntervals *LIS,
+    VirtRegMap *VRM) const {
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+
+  // The below optimizations narrow the load so they are only valid for little
+  // endian.
+  // TODO: Support big endian by adding an offset into the frame object?
+  if (MF.getDataLayout().isBigEndian())
+    return nullptr;
+
+  // Fold load from stack followed by sext.w into lw.
+  // TODO: Fold with sext.b, sext.h, zext.b, zext.h, zext.w?
+  if (Ops.size() == 1 && Ops[0] == 1 && RISCV::isSEXT_W(MI)) {
+    MachineMemOperand *MMO = MF.getMachineMemOperand(
+        MachinePointerInfo::getFixedStack(MF, FrameIndex),
+        MachineMemOperand::MOLoad, MFI.getObjectSize(FrameIndex),
+        MFI.getObjectAlign(FrameIndex));
+
+    Register DstReg = MI.getOperand(0).getReg();
+    return BuildMI(*MI.getParent(), InsertPt, MI.getDebugLoc(), get(RISCV::LW),
+                   DstReg)
+        .addFrameIndex(FrameIndex)
+        .addImm(0)
+        .addMemOperand(MMO);
+  }
+
+  return nullptr;
+}
+
 void RISCVInstrInfo::movImm(MachineBasicBlock &MBB,
                             MachineBasicBlock::iterator MBBI,
                             const DebugLoc &DL, Register DstReg, uint64_t Val,
                             MachineInstr::MIFlag Flag) const {
   Register SrcReg = RISCV::X0;
 
   if (!STI.is64Bit() && !isInt<32>(Val))
     report_fatal_error("Should only materialize 32-bit constants for RV32");
BuildMI(MBB, II, DL, get(RISCV::MUL), VL)
         .addReg(VL, RegState::Kill)
         .addReg(N, RegState::Kill)
         .setMIFlag(Flag);
   }
 
   return VL;
 }
 
+// Returns true if this is the sext.w pattern, addiw rd, rs1, 0.
+bool RISCV::isSEXT_W(const MachineInstr &MI) {
+  return MI.getOpcode() == RISCV::ADDIW && MI.getOperand(1).isReg() &&
+         MI.getOperand(2).isImm() && MI.getOperand(2).getImm() == 0;
+}
+
 static bool isRVVWholeLoadStore(unsigned Opcode) {
   switch (Opcode) {
   default:
     return false;
   case RISCV::VS1R_V:
   case RISCV::VS2R_V:
   case RISCV::VS4R_V:
   case RISCV::VS8R_V:

llvm/lib/Target/RISCV/RISCVSExtWRemoval.cpp

bool RISCVSExtWRemoval::runOnMachineFunction(MachineFunction &MF) {
 
   // Replacing instructions invalidates the MI iterator
   // we collect the candidates, then iterate over them separately.
   for (MachineBasicBlock &MBB : MF) {
     for (auto I = MBB.begin(), IE = MBB.end(); I != IE;) {
       MachineInstr *MI = &*I++;
 
       // We're looking for the sext.w pattern ADDIW rd, rs1, 0.
-      if (MI->getOpcode() != RISCV::ADDIW || !MI->getOperand(2).isImm() ||
-          MI->getOperand(2).getImm() != 0 || !MI->getOperand(1).isReg())
+      if (!RISCV::isSEXT_W(*MI))
         continue;
 
       // Input should be a virtual register.
       Register SrcReg = MI->getOperand(1).getReg();
       if (!SrcReg.isVirtual())
         continue;
 
       SExtWRemovalCands.insert(MI);

llvm/test/CodeGen/RISCV/stack-folding.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=riscv64 | FileCheck %s
+
+; Make sure we emit an lw for the stack reload in 'truebb'.
+define i1 @foo(i64 %x, i32 %y) nounwind {
+; CHECK-LABEL: foo:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    addi sp, sp, -144
+; CHECK-NEXT:    sd ra, 136(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd gp, 128(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd tp, 120(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s0, 112(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s1, 104(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s2, 96(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s3, 88(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s4, 80(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s5, 72(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s6, 64(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s7, 56(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s8, 48(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s9, 40(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s10, 32(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd s11, 24(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd a1, 8(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    sd a0, 16(sp) # 8-byte Folded Spill
+; CHECK-NEXT:    #APP
+; CHECK-NEXT:    #NO_APP
+; CHECK-NEXT:    ld a0, 16(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    beqz a0, .LBB0_2
+; CHECK-NEXT:  # %bb.1: # %falsebb
+; CHECK-NEXT:    li a0, 0
+; CHECK-NEXT:    j .LBB0_3
+; CHECK-NEXT:  .LBB0_2: # %truebb
+; CHECK-NEXT:    lw a0, 8(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    slti a0, a0, 0
+; CHECK-NEXT:  .LBB0_3: # %falsebb
+; CHECK-NEXT:    ld ra, 136(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld gp, 128(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld tp, 120(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s0, 112(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s1, 104(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s2, 96(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s3, 88(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s4, 80(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s5, 72(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s6, 64(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s7, 56(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s8, 48(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s9, 40(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s10, 32(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    ld s11, 24(sp) # 8-byte Folded Reload
+; CHECK-NEXT:    addi sp, sp, 144
+; CHECK-NEXT:    ret
+  tail call void asm sideeffect "", "~{x1},~{x3},~{x4},~{x5},~{x6},~{x7},~{x8},~{x9},~{x10},~{x11},~{x12},~{x13},~{x14},~{x15},~{x16},~{x17},~{x18},~{x19},~{x20},~{x21},~{x22},~{x23},~{x24},~{x25},~{x26},~{x27},~{x28},~{x29},~{x30},~{x31}"()
+  %a = icmp eq i64 %x, 0
+  br i1 %a, label %truebb, label %falsebb
+truebb:
+  %b = icmp slt i32 %y, 0
+  ret i1 %b
+falsebb:
+  ret i1 0
+}