[RISCV][WIP] Enable the local stack allocation pass for RISC-V.

[craig.topper]

Summary

While investigating PR58027, I noticed that ARM and AArch64 use
this pass. It creates a virtual base register for the stack accesses
and prevents an emergency spill slot from being needed.

I've done no testing of this other than PR58027. No lit tests were
affected.

Does this patch seem like something we should pursue?

[jrtc27]

llvm/lib/Target/RISCV/RISCVFrameLowering.h

77

Bit too Arm-flavoured :)

llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp

390–395

These should at least be based on XLEN and FLEN calculations rather than a bunch of hard-coded constants, but can they be calculated purely based on the ABI's callee saved register sets?

Also shouldn't the D/F checks be ABI-based rather than ISA-based? They're all caller-saved in the soft-float ABI, and similarly for the high parts of wider-than-ABI_FLEN registers.

404

Maybe also fudged by XLEN?

[craig.topper]

llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp

390–395

Thanks I had never looked into the specifics of the soft float ABI.

I'll definitely look into improving this code if we think it's worthwhile to enable this pass.

[asb]

Hmm, I hadn't realised that this pass might allow us to avoid the need for an emergency spill slot. To answer your question about whether we should try to enable this - I think a definite yes, it should improve generated code slightly too.

D98101 also started to enable it, but the original author unfortunately ran out of time on addressing review feedback - it's probably worth reviewing that implementation and the review comments to see what might be relevant for this patch.

[StephenFan]

In D134851#3823338, @asb wrote:

Hmm, I hadn't realised that this pass might allow us to avoid the need for an emergency spill slot. To answer your question about whether we should try to enable this - I think a definite yes, it should improve generated code slightly too.

D98101 also started to enable it, but the original author unfortunately ran out of time on addressing review feedback - it's probably worth reviewing that implementation and the review comments to see what might be relevant for this patch.

Glad to hear that this pass fixes some potential bugs. I pre-committed a patch D134884 to add a local stack slot allocation test. I'm happy to continue maintaining that patch. Are you interested in taking a look at that patch? Or do you @craig.topper want to continue with your current patch?

[craig.topper]

Abandon in favor of pushing forward on D98101 I completely forgot that existed despite reviewing it over a year ago.

llvm/lib/Target/RISCV/RISCVFrameLowering.h

 //===-- RISCVFrameLowering.h - Define frame lowering for RISCV -*- C++ -*--===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
public:
   bool canUseAsPrologue(const MachineBasicBlock &MBB) const override;
   bool canUseAsEpilogue(const MachineBasicBlock &MBB) const override;
 
   bool enableShrinkWrapping(const MachineFunction &MF) const override;
 
   bool isSupportedStackID(TargetStackID::Value ID) const override;
   TargetStackID::Value getStackIDForScalableVectors() const override;
 
+  bool isStackIdSafeForLocalArea(unsigned StackId) const override {
+    // We don't support putting SVE objects into the pre-allocated local
+    // frame block at the moment.
+    return StackId != TargetStackID::ScalableVector;
+  }
+
 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;
   void adjustStackForRVV(MachineFunction &MF, MachineBasicBlock &MBB,
                          MachineBasicBlock::iterator MBBI, const DebugLoc &DL,
                          int64_t Amount, MachineInstr::MIFlag Flag) const;
   std::pair<int64_t, Align>
   assignRVVStackObjectOffsets(MachineFunction &MF) const;
 };
 } // namespace llvm
 #endif

llvm/lib/Target/RISCV/RISCVRegisterInfo.h

 //===-- RISCVRegisterInfo.h - RISCV Register Information Impl ---*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
struct RISCVRegisterInfo : public RISCVGenRegisterInfo {
   bool requiresRegisterScavenging(const MachineFunction &MF) const override {
     return true;
   }
 
   bool requiresFrameIndexScavenging(const MachineFunction &MF) const override {
     return true;
   }
 
+  bool requiresVirtualBaseRegisters(const MachineFunction &MF) const override {
+    return true;
+  }
+
+  bool needsFrameBaseReg(MachineInstr *MI, int64_t Offset) const override;
+  bool isFrameOffsetLegal(const MachineInstr *MI, Register BaseReg,
+                          int64_t Offset) const override;
+  Register materializeFrameBaseRegister(MachineBasicBlock *MBB, int FrameIdx,
+                                        int64_t Offset) const override;
+  void resolveFrameIndex(MachineInstr &MI, Register BaseReg,
+                         int64_t Offset) const override;
+
   const TargetRegisterClass *
   getPointerRegClass(const MachineFunction &MF,
                      unsigned Kind = 0) const override {
     return &RISCV::GPRRegClass;
   }
 
   const TargetRegisterClass *
   getLargestLegalSuperClass(const TargetRegisterClass *RC,

llvm/lib/Target/RISCV/RISCVRegisterInfo.cpp

 //===-- RISCVRegisterInfo.cpp - RISCV Register Information ------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
if (VLENBSized > 0) {
     Ops.push_back(dwarf::DW_OP_minus);
   }
 }
 
 unsigned
 RISCVRegisterInfo::getRegisterCostTableIndex(const MachineFunction &MF) const {
   return MF.getSubtarget<RISCVSubtarget>().hasStdExtC() ? 1 : 0;
 }
+
+/// needsFrameBaseReg - Returns true if the instruction's frame index
+/// reference would be better served by a base register other than FP
+/// or SP. Used by LocalStackFrameAllocation to determine which frame index
+/// references it should create new base registers for.
+bool RISCVRegisterInfo::needsFrameBaseReg(MachineInstr *MI,
+                                          int64_t Offset) const {
+  for (unsigned i = 0; !MI->getOperand(i).isFI(); ++i)
+    assert(i < MI->getNumOperands() &&
+           "Instr doesn't have FrameIndex operand!");
+
+  // It's the load/store FI references that cause issues, as it can be difficult
+  // to materialize the offset if it won't fit in the literal field. Estimate
+  // based on the size of the local frame and some conservative assumptions
+  // about the rest of the stack frame (note, this is pre-regalloc, so
+  // we don't know everything for certain yet) whether this offset is likely
+  // to be out of range of the immediate. Return true if so.
+
+  // We only generate virtual base registers for loads and stores, so
+  // return false for everything else.
+  if (!MI->mayLoad() && !MI->mayStore())
+    return false;
+
+  // Without a virtual base register, if the function has variable sized
+  // objects, all fixed-size local references will be via the frame pointer,
+  // Approximate the offset and see if it's legal for the instruction.
+  // Note that the incoming offset is based on the SP value at function entry,
+  // so it'll be negative.
+  MachineFunction &MF = *MI->getParent()->getParent();
+  const RISCVFrameLowering *TFI = getFrameLowering(MF);
+  MachineFrameInfo &MFI = MF.getFrameInfo();
+
+  // Estimate an offset from the frame pointer.
+  // Conservatively assume all GPR callee-saved registers get pushed.
+  // RA, S0-S11.
+  auto &Subtarget = MF.getSubtarget<RISCVSubtarget>();
+  int64_t FPOffset = Offset - (Subtarget.is64Bit() ? 104 : 52);
+  // Consider FS0-FS11 if present.
+  if (Subtarget.hasStdExtD())
+    FPOffset -= 96;
+  else if (Subtarget.hasStdExtF())
+    FPOffset -= 48;
+
+  // Estimate an offset from the stack pointer.
+  // The incoming offset is relating to the SP at the start of the function,
+  // but when we access the local it'll be relative to the SP after local
+  // allocation, so adjust our SP-relative offset by that allocation size.
+  Offset += MFI.getLocalFrameSize();
+  // Assume that we'll have at least some spill slots allocated.
+  // FIXME: This is a total SWAG number. We should run some statistics
+  //        and pick a real one.
+  Offset += 128; // 128 bytes of spill slots
+
+  // If there is a frame pointer, try using it.
+  // The FP is only available if there is no dynamic realignment. We
+  // don't know for sure yet whether we'll need that, so we guess based
+  // on whether there are any local variables that would trigger it.
+  if (TFI->hasFP(MF) && isFrameOffsetLegal(MI, RISCV::X8, FPOffset))
+    return false;
+
+  // If we can reference via the stack pointer or base pointer, try that.
+  // FIXME: This (and the code that resolves the references) can be improved
+  //        to only disallow SP relative references in the live range of
+  //        the VLA(s). In practice, it's unclear how much difference that
+  //        would make, but it may be worth doing.
+  if (isFrameOffsetLegal(MI, RISCV::X2, Offset))
+    return false;
+
+  // If even offset 0 is illegal, we don't want a virtual base register.
+  if (!isFrameOffsetLegal(MI, RISCV::X2, 0))
+    return false;
+
+  // The offset likely isn't legal; we want to allocate a virtual base register.
+  return true;
+}
+
+bool RISCVRegisterInfo::isFrameOffsetLegal(const MachineInstr *MI,
+                                           Register BaseReg,
+                                           int64_t Offset) const {
+  unsigned i = 0;
+  for (; !MI->getOperand(i).isFI(); ++i)
+    assert(i + 1 < MI->getNumOperands() &&
+           "Instr doesn't have FrameIndex operand!");
+
+  Offset += MI->getOperand(i + 1).getImm();
+
+  return isInt<12>(Offset);
+}
+
+/// materializeFrameBaseRegister - Insert defining instruction(s) for BaseReg to
+/// be a pointer to FrameIdx at the beginning of the basic block.
+Register RISCVRegisterInfo::materializeFrameBaseRegister(MachineBasicBlock *MBB,
+                                                         int FrameIdx,
+                                                         int64_t Offset) const {
+  MachineBasicBlock::iterator Ins = MBB->begin();
+  DebugLoc DL; // Defaults to "unknown"
+  if (Ins != MBB->end())
+    DL = Ins->getDebugLoc();
+
+  const MachineFunction &MF = *MBB->getParent();
+  MachineRegisterInfo &MRI = MBB->getParent()->getRegInfo();
+  const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
+  const MCInstrDesc &MCID = TII.get(RISCV::ADDI);
+  Register BaseReg = MRI.createVirtualRegister(&RISCV::GPRRegClass);
+
+  BuildMI(*MBB, Ins, DL, MCID, BaseReg).addFrameIndex(FrameIdx).addImm(Offset);
+
+  return BaseReg;
+}
+
+void RISCVRegisterInfo::resolveFrameIndex(MachineInstr &MI, Register BaseReg,
+                                          int64_t Offset) const {
+  unsigned i = 0;
+  while (!MI.getOperand(i).isFI()) {
+    ++i;
+    assert(i < MI.getNumOperands() && "Instr doesn't have FrameIndex operand!");
+  }
+
+  Offset += MI.getOperand(i + 1).getImm();
+  assert(isInt<12>(Offset));
+  MI.getOperand(i).ChangeToRegister(BaseReg, false);
+  MI.getOperand(i + 1).ChangeToImmediate(Offset);
+}