[RISCV] Allow shrink wrapping for RISC-V

[lewis-revill]

Summary

Enabling shrink wrapping requires ensuring the insertion point of the epilogue is correct for MBBs without a terminator, in which case the instruction to adjust the stack pointer is the last instruction in the block.

[lewis-revill]

I've reworked this patch to safely work around the implementation of save/restore via libcalls. These place constraints on the save/restore points that can be chosen by the shrink wrapping algorithm.

[luismarques]

lib/Target/RISCV/RISCVFrameLowering.cpp

181 ↗

(On Diff #208047)

Document why we advance exactly one instruction.

546 ↗

(On Diff #208047)

This patch is using this function but it seems like D62686 no longer introduces it.

[shiva0217]

lib/Target/RISCV/RISCVFrameLowering.cpp

171 ↗

(On Diff #208047)

We might need to consider that there're two branches at the end of the basic block. Last non-debug instruction is PseudoBr which is a terminator. So the epilogue will be inserted between the branches.

BEQ killed renamable $x10, $x0, %bb.3
PseudoBR %bb.2

Replacing by

MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator();

could handle the case. and

if (MBBI == MBB.end())
  MBBI = MBB.getLastNonDebugInstr();

to handle that there is no terminator in the basic block.

[lewis-revill]

Rebased, and modified the logic to find the insertion point for the epilogue code.

[lewis-revill]

Fix git error whereby logic from D62686 was mistakenly included in this patch

[lewis-revill]

Rebased and updated to incorporate changes to parent revision

[lewis-revill]

Rebased due to updated dependency.

[asb]

I think this patch is orthogonal to save/restore via libcalls, and we can probably land it sooner. Could you please rebase so it's standalone?

[lewis-revill]

Sure, that's sensible. I'll keep a separate patch for ensuring compatibility with save/restore for if/when it's required.

[apazos]

There are some additional perennial test suite failures when applying this patch and enabling -mllvm -enable-shrink-wrap.
When Lewis updates the patch to be standalone, we can verify it again.

[lewis-revill]

Make this patch independent of save/restore compatibility

[apazos]

llvm/lib/Target/RISCV/RISCVFrameLowering.cpp

228

I have been verifying the pending patches at Oz, Os, and also O2. This is helping with uncovering issues.
You can run it with 'llc test.ll -enable-shrink-wrap

With the latest standalone patch, it seems we have a few crashes left. Below is a bugpoint reduced test:

define dso_local void @test() local_unnamed_addr {
entry:

br i1 undef, label %T.exit, label %for.body.i

for.body.i: ; preds = %for.body.i, %entry

store i32 0, i32* undef
%incdec.ptr.i.i = getelementptr inbounds i32, i32* null, i32 1
%cmp.i.i = icmp eq i32* undef, undef
br i1 %cmp.i.i, label %T.exit.loopexit, label %for.body.i

T.exit.loopexit: ; preds = %for.body.i

%0 = ptrtoint i32* %incdec.ptr.i.i to i32
br label %T.exit

T.exit: ; preds = %T.exit.loopexit, %entry

ret void

}

[shiva0217]

llvm/lib/Target/RISCV/RISCVFrameLowering.cpp

228

It seems that shrink wrapping may choose an empty basic block to insert epilogue. We might need to add empty block detection for DL initialization and MBBI advance. Something like:

DebugLoc DL = !MBB.empty() ? MBBI->getDebugLoc() : DebugLoc();
if (!MBB.empty() && !MBBI->isTerminator())
   MBBI = std::next(MBBI);

[lewis-revill]

llvm/lib/Target/RISCV/RISCVFrameLowering.cpp

228

Thanks Ana, it looks like this is a case of the shrink wrapping pass choosing a basic block which is empty as the place to insert the prologue. I didn't realise it was possible, but I'll push an updated patch that takes this into account.

[lewis-revill]

Use MBB.end() and a blank DebugLoc for epilogue insertion when a basic block is empty.

[apazos]

Lewis, this patch is not applying clean on top of https://reviews.llvm.org/D62686. Can you please rebase?

[lewis-revill]

In D62190#1712206, @apazos wrote:

Lewis, this patch is not applying clean on top of https://reviews.llvm.org/D62686. Can you please rebase?

Hi Ana, I've made this patch independent of the save/restore patch as per Alex's request. Unfortunately that means that neither will rebase cleanly on top of the other until one is committed.
I could recreate the old patch with the save/restore compatibility included (currently split into D68644) implemented on top of save/restore, if that is worth doing?

Off topic, this makes me think Phabricator needs some way to add multiple alternative diffs for applying a given patch on top of other patches...

[apazos]

Thanks Lewis, I applied the patch and I am not detecting any new failure with this feature.

[apazos]

Lewis, I am not seeing any code change.

Btw to turn it on by default on a the target you need to add enableShrinkWrapping implementation in RISCVFrameLowering.h

+ /// Returns true if the target will correctly handle shrink wrapping.
+ bool enableShrinkWrapping(const MachineFunction &MF) const override {
+ return true;
+ }
+

[lewis-revill]

So I've been looking into the benchmarks that I've been running to figure out why code doesn't change with this enabled.

I found that GCC's shrink wrapping does make a difference, and that comes from a loop which looks like the following:

int foo(int rpt) {
  int i;
  int r;

  for (i = 0; i < rpt; ++i) {
    srand(0);
    r = ...;
  }

  return r;
}

GCC appears to convert this into the equivalent of:

int foo(int rpt) {
  if (rpt == 0)
    return 0;

  int i;
  int r;

  for (i = 0; i < rpt; ++i) {
    srand(0);
    r = ...;
  }

  return r;
}

And is then able to shrink wrap around the if statement. Sure enough, if I feed the modified function above into LLVM, we are able to shrink wrap it as well.

(I'm not certain if this conversion is sane from GCC, but I haven't properly looked into it.)

So personally I'm not concerned about the fact that shrink wrapping doesn't make a difference in my benchmarks; I can show that it changes C code we expect it to change... I'm not so sure about the lack of differences with SPEC though.

[lewis-revill]

Rebased.

[lewis-revill]

Rebased

[asb]

I note that the TargetFrameLowering hooks canUseAsPrologue and canUseAsEpilogue are both called by the shrink wrapper. They default to true, but targets may need to override this for correctness. Looking at e.g. AArch64, I see it overrides canUseAsPrologue and returns false in the case that the function needs stack realignment and there are no scratch registers available. Are you certain that no such case is needed for RISC-V?

[lewis-revill]

In D62190#1755159, @asb wrote:

I note that the TargetFrameLowering hooks canUseAsPrologue and canUseAsEpilogue are both called by the shrink wrapper. They default to true, but targets may need to override this for correctness. Looking at e.g. AArch64, I see it overrides canUseAsPrologue and returns false in the case that the function needs stack realignment and there are no scratch registers available. Are you certain that no such case is needed for RISC-V?

So I see that the stack realignment implementation for larger values does indeed require the use of a scratch register, so that implies we should attempt to prevent the situation where we don't have one available. However the scratch register is created as a virtual GPR register, rather than attempting to scavenge a real free register. I can imagine that if this was changed, we could directly perform the same check in canUseAsPrologue to guarantee a register being available in a much more transparent and obvious way. EG the way in which AArch64 uses the function findScratchNonCalleeSaveRegister.

[lenary]

In D62190#1755258, @lewis-revill wrote:

In D62190#1755159, @asb wrote:

I note that the TargetFrameLowering hooks canUseAsPrologue and canUseAsEpilogue are both called by the shrink wrapper. They default to true, but targets may need to override this for correctness. Looking at e.g. AArch64, I see it overrides canUseAsPrologue and returns false in the case that the function needs stack realignment and there are no scratch registers available. Are you certain that no such case is needed for RISC-V?

So I see that the stack realignment implementation for larger values does indeed require the use of a scratch register, so that implies we should attempt to prevent the situation where we don't have one available. However the scratch register is created as a virtual GPR register, rather than attempting to scavenge a real free register. I can imagine that if this was changed, we could directly perform the same check in canUseAsPrologue to guarantee a register being available in a much more transparent and obvious way. EG the way in which AArch64 uses the function findScratchNonCalleeSaveRegister.

I'm working on a patch for the realignment implementation based off issues that arose with implementing the ilp32e ABI. I will update this thread when that patch is ready.

[apazos]

Lewis, try rebase again, not applying cleaning, nor https://reviews.llvm.org/D62686

[lewis-revill]

Rebased

[lewis-revill]

Rebased.

[apazos]

Lewis, same question, is the patch final? It would be good to merge it before the 10.0 release branch creation on Jan 15th

[lewis-revill]

This patch is yes, so long as it applies cleanly. So if it's agreed that it should be landed then I can do that as soon as I get a chance.

[luismarques]

Thanks for the patch Lewis. Overall it LGTM, and it would be great to merge this in time for the LLVM 10.0 release.

The patch would benefit from having more extensive tests, although those could be provided by a follow-up patch (but please do follow up). For instance, the current tests don't exercise the if (!MBB.empty()) else condition. It would also be great to include tests that illustrate the impact on debug instructions (even if that has been an area where testing has been limited in the past). See for instance my inline comment for a possible issue regarding correctly updating the debug locations.

llvm/lib/Target/RISCV/RISCVFrameLowering.cpp

233

This declaration shadows the DL above; I assume you meant to update the previously declared DL.

[lewis-revill]

Quick fix of shadow variable assignment.

llvm/lib/Target/RISCV/RISCVFrameLowering.h

bool spillCalleeSavedRegisters(MachineBasicBlock &MBB,
                                  const std::vector<CalleeSavedInfo> &CSI,
                                  const TargetRegisterInfo *TRI) const override;
   bool
   restoreCalleeSavedRegisters(MachineBasicBlock &MBB,
                               MachineBasicBlock::iterator MI,
                               std::vector<CalleeSavedInfo> &CSI,
                               const TargetRegisterInfo *TRI) const override;
 
+  bool canUseAsPrologue(const MachineBasicBlock &MBB) const override;
+  bool canUseAsEpilogue(const MachineBasicBlock &MBB) const override;
+
 protected:
   const RISCVSubtarget &STI;
 
 private:
   void determineFrameLayout(MachineFunction &MF) const;
   void adjustReg(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
                  const DebugLoc &DL, Register DestReg, Register SrcReg,
                  int64_t Val, MachineInstr::MIFlag Flag) const;
 };
 }
 #endif

llvm/lib/Target/RISCV/RISCVFrameLowering.cpp

for (auto &CS : CSI)
     if (CS.getFrameIdx() >= 0)
       NonLibcallCSI.push_back(CS);
 
   return NonLibcallCSI;
 }
 
 void RISCVFrameLowering::emitPrologue(MachineFunction &MF,
                                       MachineBasicBlock &MBB) const {
-  assert(&MF.front() == &MBB && "Shrink-wrapping not yet supported");
-
   MachineFrameInfo &MFI = MF.getFrameInfo();
   auto *RVFI = MF.getInfo<RISCVMachineFunctionInfo>();
   const RISCVRegisterInfo *RI = STI.getRegisterInfo();
   const RISCVInstrInfo *TII = STI.getInstrInfo();
   MachineBasicBlock::iterator MBBI = MBB.begin();
 
   if (RI->needsStackRealignment(MF) && MFI.hasVarSizedObjects()) {
     report_fatal_error(
if (RI->needsStackRealignment(MF)) {
             .addImm(ShiftAmount);
       }
     }
   }
 }
 
 void RISCVFrameLowering::emitEpilogue(MachineFunction &MF,
                                       MachineBasicBlock &MBB) const {
-  MachineBasicBlock::iterator MBBI = MBB.getLastNonDebugInstr();
   const RISCVRegisterInfo *RI = STI.getRegisterInfo();
   MachineFrameInfo &MFI = MF.getFrameInfo();
   auto *RVFI = MF.getInfo<RISCVMachineFunctionInfo>();
-  DebugLoc DL = MBBI->getDebugLoc();
   const RISCVInstrInfo *TII = STI.getInstrInfo();
   Register FPReg = getFPReg(STI);
   Register SPReg = getSPReg(STI);
 
+  // Get the insert location for the epilogue. If there were no terminators in
+  // the block, get the last instruction.
+  MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator();
+  if (MBBI == MBB.end())
+    MBBI = MBB.getLastNonDebugInstr();
+  DebugLoc DL = MBBI->getDebugLoc();
+
+  // If this is not a terminator, the actual insert location should be after the
+  // last instruction.
+  if (!MBBI->isTerminator())
+    MBBI = std::next(MBBI);
+
   // If callee-saved registers are saved via libcall, place stack adjustment
   // before this call.
   while (MBBI != MBB.begin() &&
          std::prev(MBBI)->getFlag(MachineInstr::FrameDestroy))
     --MBBI;
 
   const std::vector<CalleeSavedInfo> &CSI =
       getNonLibcallCSI(MFI.getCalleeSavedInfo());
if (RestoreLibCall) {
     if (MI != MBB.end() && MI->getOpcode() == RISCV::PseudoRET) {
       NewMI->copyImplicitOps(*MF, *MI);
       MI->eraseFromParent();
     }
   }
 
   return true;
 }
+
+bool RISCVFrameLowering::canUseAsPrologue(const MachineBasicBlock &MBB) const {
+  MachineBasicBlock *TmpMBB = const_cast<MachineBasicBlock *>(&MBB);
+  const auto *RVFI = MBB.getParent()->getInfo<RISCVMachineFunctionInfo>();
+
+  if (!RVFI->useSaveRestoreLibCalls())
+    return true;
+
+  // Inserting a call to a __riscv_save libcall requires the use of the register
+  // t0 (X5) to hold the return address. Therefore if this register is already
+  // used we can't insert the call.
+
+  RegScavenger RS;
+  RS.enterBasicBlock(*TmpMBB);
+  return !RS.isRegUsed(RISCV::X5);
+}
+
+bool RISCVFrameLowering::canUseAsEpilogue(const MachineBasicBlock &MBB) const {
+  MachineBasicBlock *TmpMBB = const_cast<MachineBasicBlock *>(&MBB);
+  const auto *RVFI = MBB.getParent()->getInfo<RISCVMachineFunctionInfo>();
+
+  if (!RVFI->useSaveRestoreLibCalls())
+    return true;
+
+  // Using the __riscv_restore libcalls to restore CSRs requires a tail call.
+  // This means if we still need to continue executing code within this function
+  // the restore cannot take place in this basic block.
+
+  if (MBB.succ_size() > 1)
+    return false;
+
+  MachineBasicBlock *SuccMBB =
+      MBB.succ_empty() ? TmpMBB->getFallThrough() : *MBB.succ_begin();
+
+  // Doing a tail call should be safe if there are no successors, because either
+  // we have a returning block or the end of the block is unreachable, so the
+  // restore will be eliminated regardless.
+  if (!SuccMBB)
+    return true;
+
+  // The successor can only contain a return, since we would effectively be
+  // replacing the successor with our own tail return at the end of our block.
+  return SuccMBB->isReturnBlock() && SuccMBB->size() == 1;
+}

llvm/test/CodeGen/RISCV/shrinkwrap.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple riscv32 < %s | FileCheck %s -check-prefix=RV32I-NOSW
+; RUN: llc -mtriple riscv32 -enable-shrink-wrap < %s | FileCheck %s -check-prefix=RV32I-SW
+; RUN: llc -mtriple riscv32 -enable-shrink-wrap -mattr=+save-restore < %s \
+; RUN:     | FileCheck %s -check-prefix=RV32I-SW-SR
+
+
+declare void @abort()
+
+define void @eliminate_restore(i32 %n) nounwind {
+; RV32I-NOSW-LABEL: eliminate_restore:
+; RV32I-NOSW:       # %bb.0:
+; RV32I-NOSW-NEXT:    addi sp, sp, -16
+; RV32I-NOSW-NEXT:    sw ra, 12(sp)
+; RV32I-NOSW-NEXT:    addi a1, zero, 32
+; RV32I-NOSW-NEXT:    bgeu a1, a0, .LBB0_2
+; RV32I-NOSW-NEXT:  # %bb.1: # %if.end
+; RV32I-NOSW-NEXT:    lw ra, 12(sp)
+; RV32I-NOSW-NEXT:    addi sp, sp, 16
+; RV32I-NOSW-NEXT:    ret
+; RV32I-NOSW-NEXT:  .LBB0_2: # %if.then
+; RV32I-NOSW-NEXT:    call abort
+;
+; RV32I-SW-LABEL: eliminate_restore:
+; RV32I-SW:       # %bb.0:
+; RV32I-SW-NEXT:    addi a1, zero, 32
+; RV32I-SW-NEXT:    bgeu a1, a0, .LBB0_2
+; RV32I-SW-NEXT:  # %bb.1: # %if.end
+; RV32I-SW-NEXT:    ret
+; RV32I-SW-NEXT:  .LBB0_2: # %if.then
+; RV32I-SW-NEXT:    addi sp, sp, -16
+; RV32I-SW-NEXT:    sw ra, 12(sp)
+; RV32I-SW-NEXT:    call abort
+;
+; RV32I-SW-SR-LABEL: eliminate_restore:
+; RV32I-SW-SR:       # %bb.0:
+; RV32I-SW-SR-NEXT:    addi a1, zero, 32
+; RV32I-SW-SR-NEXT:    bgeu a1, a0, .LBB0_2
+; RV32I-SW-SR-NEXT:  # %bb.1: # %if.end
+; RV32I-SW-SR-NEXT:    ret
+; RV32I-SW-SR-NEXT:  .LBB0_2: # %if.then
+; RV32I-SW-SR-NEXT:    call t0, __riscv_save_0
+; RV32I-SW-SR-NEXT:    call abort
+  %cmp = icmp ule i32 %n, 32
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  call void @abort()
+  unreachable
+
+if.end:
+  ret void
+}
+
+declare void @notdead(i8*)
+
+define void @conditional_alloca(i32 %n) nounwind {
+; RV32I-NOSW-LABEL: conditional_alloca:
+; RV32I-NOSW:       # %bb.0:
+; RV32I-NOSW-NEXT:    addi sp, sp, -16
+; RV32I-NOSW-NEXT:    sw ra, 12(sp)
+; RV32I-NOSW-NEXT:    sw s0, 8(sp)
+; RV32I-NOSW-NEXT:    addi s0, sp, 16
+; RV32I-NOSW-NEXT:    addi a1, zero, 32
+; RV32I-NOSW-NEXT:    bltu a1, a0, .LBB1_2
+; RV32I-NOSW-NEXT:  # %bb.1: # %if.then
+; RV32I-NOSW-NEXT:    addi a0, a0, 15
+; RV32I-NOSW-NEXT:    andi a0, a0, -16
+; RV32I-NOSW-NEXT:    sub a0, sp, a0
+; RV32I-NOSW-NEXT:    mv sp, a0
+; RV32I-NOSW-NEXT:    call notdead
+; RV32I-NOSW-NEXT:  .LBB1_2: # %if.end
+; RV32I-NOSW-NEXT:    addi sp, s0, -16
+; RV32I-NOSW-NEXT:    lw s0, 8(sp)
+; RV32I-NOSW-NEXT:    lw ra, 12(sp)
+; RV32I-NOSW-NEXT:    addi sp, sp, 16
+; RV32I-NOSW-NEXT:    ret
+;
+; RV32I-SW-LABEL: conditional_alloca:
+; RV32I-SW:       # %bb.0:
+; RV32I-SW-NEXT:    addi a1, zero, 32
+; RV32I-SW-NEXT:    bltu a1, a0, .LBB1_2
+; RV32I-SW-NEXT:  # %bb.1: # %if.then
+; RV32I-SW-NEXT:    addi sp, sp, -16
+; RV32I-SW-NEXT:    sw ra, 12(sp)
+; RV32I-SW-NEXT:    sw s0, 8(sp)
+; RV32I-SW-NEXT:    addi s0, sp, 16
+; RV32I-SW-NEXT:    addi a0, a0, 15
+; RV32I-SW-NEXT:    andi a0, a0, -16
+; RV32I-SW-NEXT:    sub a0, sp, a0
+; RV32I-SW-NEXT:    mv sp, a0
+; RV32I-SW-NEXT:    call notdead
+; RV32I-SW-NEXT:    addi sp, s0, -16
+; RV32I-SW-NEXT:    lw s0, 8(sp)
+; RV32I-SW-NEXT:    lw ra, 12(sp)
+; RV32I-SW-NEXT:    addi sp, sp, 16
+; RV32I-SW-NEXT:  .LBB1_2: # %if.end
+; RV32I-SW-NEXT:    ret
+;
+; RV32I-SW-SR-LABEL: conditional_alloca:
+; RV32I-SW-SR:       # %bb.0:
+; RV32I-SW-SR-NEXT:    addi a1, zero, 32
+; RV32I-SW-SR-NEXT:    bltu a1, a0, .LBB1_2
+; RV32I-SW-SR-NEXT:  # %bb.1: # %if.then
+; RV32I-SW-SR-NEXT:    call t0, __riscv_save_1
+; RV32I-SW-SR-NEXT:    mv s0, sp
+; RV32I-SW-SR-NEXT:    addi a0, a0, 15
+; RV32I-SW-SR-NEXT:    andi a0, a0, -16
+; RV32I-SW-SR-NEXT:    sub a0, sp, a0
+; RV32I-SW-SR-NEXT:    mv sp, a0
+; RV32I-SW-SR-NEXT:    call notdead
+; RV32I-SW-SR-NEXT:    mv sp, s0
+; RV32I-SW-SR-NEXT:    tail __riscv_restore_1
+; RV32I-SW-SR-NEXT:  .LBB1_2: # %if.end
+; RV32I-SW-SR-NEXT:    ret
+  %cmp = icmp ule i32 %n, 32
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  %addr = alloca i8, i32 %n
+  call void @notdead(i8* %addr)
+  br label %if.end
+
+if.end:
+  ret void
+}