[LoongArch] Support lowering FrameIndex

Ensure callee-saved registers are accessed relative to the stackpointer.

Differential Revision: https://reviews.llvm.org/D128430

llvm/lib/Target/LoongArch/LoongArchFrameLowering.cpp

@@ -53,3 +53,32 @@ void LoongArchFrameLowering::emitEpilogue(MachineFunction &MF,
                                           MachineBasicBlock &MBB) const {
   // TODO: Implement this when we have function calls
 }
+
+StackOffset LoongArchFrameLowering::getFrameIndexReference(
+    const MachineFunction &MF, int FI, Register &FrameReg) const {
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+  const TargetRegisterInfo *RI = MF.getSubtarget().getRegisterInfo();
+
+  // Callee-saved registers should be referenced relative to the stack
+  // pointer (positive offset), otherwise use the frame pointer (negative
+  // offset).
+  const auto &CSI = MFI.getCalleeSavedInfo();
+  int MinCSFI = 0;
+  int MaxCSFI = -1;
+  StackOffset Offset =
+      StackOffset::getFixed(MFI.getObjectOffset(FI) - getOffsetOfLocalArea() +
+                            MFI.getOffsetAdjustment());
+
+  if (CSI.size()) {
+    MinCSFI = CSI[0].getFrameIdx();
+    MaxCSFI = CSI[CSI.size() - 1].getFrameIdx();
+  }
+
+  FrameReg = RI->getFrameRegister(MF);
+  if ((FI >= MinCSFI && FI <= MaxCSFI) || !hasFP(MF)) {
+    FrameReg = LoongArch::R3;
+    Offset += StackOffset::getFixed(MFI.getStackSize());
+  }
+
+  return Offset;
+}

llvm/lib/Target/LoongArch/LoongArchFrameLowering.h

@@ -37,6 +37,9 @@ class LoongArchFrameLowering : public TargetFrameLowering {
     return MBB.erase(MI);
   }
 
+  StackOffset getFrameIndexReference(const MachineFunction &MF, int FI,
+                                     Register &FrameReg) const override;
+
   bool hasFP(const MachineFunction &MF) const override;
   bool hasBP(const MachineFunction &MF) const;
 };

llvm/lib/Target/LoongArch/LoongArchISelDAGToDAG.cpp

@@ -33,13 +33,14 @@ void LoongArchDAGToDAGISel::Select(SDNode *Node) {
   unsigned Opcode = Node->getOpcode();
   MVT GRLenVT = Subtarget->getGRLenVT();
   SDLoc DL(Node);
+  MVT VT = Node->getSimpleValueType(0);
 
   switch (Opcode) {
   default:
     break;
   case ISD::Constant: {
     int64_t Imm = cast<ConstantSDNode>(Node)->getSExtValue();
-    if (Imm == 0 && Node->getSimpleValueType(0) == GRLenVT) {
+    if (Imm == 0 && VT == GRLenVT) {
       SDValue New = CurDAG->getCopyFromReg(CurDAG->getEntryNode(), DL,
                                            LoongArch::R0, GRLenVT);
       ReplaceNode(Node, New.getNode());
@@ -59,6 +60,15 @@ void LoongArchDAGToDAGISel::Select(SDNode *Node) {
 
     ReplaceNode(Node, Result);
     return;
+  }
+  case ISD::FrameIndex: {
+    SDValue Imm = CurDAG->getTargetConstant(0, DL, GRLenVT);
+    int FI = cast<FrameIndexSDNode>(Node)->getIndex();
+    SDValue TFI = CurDAG->getTargetFrameIndex(FI, VT);
+    unsigned ADDIOp =
+        Subtarget->is64Bit() ? LoongArch::ADDI_D : LoongArch::ADDI_W;
+    ReplaceNode(Node, CurDAG->getMachineNode(ADDIOp, DL, VT, TFI, Imm));
+    return;
   }
     // TODO: Add selection nodes needed later.
   }
@@ -67,6 +77,17 @@ void LoongArchDAGToDAGISel::Select(SDNode *Node) {
   SelectCode(Node);
 }
 
+bool LoongArchDAGToDAGISel::SelectBaseAddr(SDValue Addr, SDValue &Base) {
+  // If this is FrameIndex, select it directly. Otherwise just let it get
+  // selected to a register independently.
+  if (auto *FIN = dyn_cast<FrameIndexSDNode>(Addr))
+    Base =
+        CurDAG->getTargetFrameIndex(FIN->getIndex(), Subtarget->getGRLenVT());
+  else
+    Base = Addr;
+  return true;
+}
+
 bool LoongArchDAGToDAGISel::selectShiftMask(SDValue N, unsigned ShiftWidth,
                                             SDValue &ShAmt) {
   // Shift instructions on LoongArch only read the lower 5 or 6 bits of the

llvm/lib/Target/LoongArch/LoongArchISelDAGToDAG.h

@@ -38,6 +38,8 @@ class LoongArchDAGToDAGISel : public SelectionDAGISel {
 
   void Select(SDNode *Node) override;
 
+  bool SelectBaseAddr(SDValue Addr, SDValue &Base);
+
   bool selectShiftMask(SDValue N, unsigned ShiftWidth, SDValue &ShAmt);
   bool selectShiftMaskGRLen(SDValue N, SDValue &ShAmt) {
     return selectShiftMask(N, Subtarget->getGRLen(), ShAmt);

llvm/lib/Target/LoongArch/LoongArchInstrInfo.td

@@ -175,6 +175,8 @@ def call_symbol : Operand<iPTR> {
   let ParserMatchClass = CallSymbol;
 }
 
+def BaseAddr : ComplexPattern<iPTR, 1, "SelectBaseAddr">;
+
 //===----------------------------------------------------------------------===//
 // Instruction Formats
 //===----------------------------------------------------------------------===//
@@ -751,9 +753,9 @@ def : Pat<(loongarch_bstrpick GPR:$rj, uimm6:$msbd, uimm6:$lsbd),
 /// Loads
 
 multiclass LdPat<PatFrag LoadOp, LAInst Inst, ValueType vt = GRLenVT> {
-  def : Pat<(vt (LoadOp GPR:$rj)), (Inst GPR:$rj, 0)>;
-  def : Pat<(vt (LoadOp (add GPR:$rj, simm12:$imm12))),
-            (Inst GPR:$rj, simm12:$imm12)>;
+  def : Pat<(vt (LoadOp BaseAddr:$rj)), (Inst BaseAddr:$rj, 0)>;
+  def : Pat<(vt (LoadOp (add BaseAddr:$rj, simm12:$imm12))),
+            (Inst BaseAddr:$rj, simm12:$imm12)>;
 }
 
 defm : LdPat<sextloadi8, LD_B>;
@@ -774,10 +776,10 @@ defm : LdPat<load, LD_D, i64>;
 
 multiclass StPat<PatFrag StoreOp, LAInst Inst, RegisterClass StTy,
                  ValueType vt> {
-  def : Pat<(StoreOp (vt StTy:$rd), GPR:$rj),
-            (Inst StTy:$rd, GPR:$rj, 0)>;
-  def : Pat<(StoreOp (vt StTy:$rd), (add GPR:$rj, simm12:$imm12)),
-            (Inst StTy:$rd, GPR:$rj, simm12:$imm12)>;
+  def : Pat<(StoreOp (vt StTy:$rd), BaseAddr:$rj),
+            (Inst StTy:$rd, BaseAddr:$rj, 0)>;
+  def : Pat<(StoreOp (vt StTy:$rd), (add BaseAddr:$rj, simm12:$imm12)),
+            (Inst StTy:$rd, BaseAddr:$rj, simm12:$imm12)>;
 }
 
 defm : StPat<truncstorei8, ST_B, GPR, GRLenVT>;

llvm/test/CodeGen/LoongArch/frame.ll

@@ -0,0 +1,30 @@
+; RUN: llc --mtriple=loongarch64 < %s | FileCheck %s
+
+%struct.key_t = type { i32, [16 x i8] }
+
+;; FIXME: prologue and epilogue insertion must be implemented to complete this
+;; test
+
+define i32 @test() nounwind {
+; CHECK-LABEL: test:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    st.d $ra, $sp, 24 # 8-byte Folded Spill
+; CHECK-NEXT:    st.w $zero, $sp, 16
+; CHECK-NEXT:    st.d $zero, $sp, 8
+; CHECK-NEXT:    st.d $zero, $sp, 0
+; CHECK-NEXT:    addi.d $a0, $sp, 0
+; CHECK-NEXT:    ori $a0, $a0, 4
+; CHECK-NEXT:    bl test1
+; CHECK-NEXT:    move $a0, $zero
+; CHECK-NEXT:    ld.d $ra, $sp, 24 # 8-byte Folded Reload
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %key = alloca %struct.key_t, align 4
+  call void @llvm.memset.p0i8.i64(ptr %key, i8 0, i64 20, i1 false)
+  %1 = getelementptr inbounds %struct.key_t, ptr %key, i64 0, i32 1, i64 0
+  call void @test1(ptr %1)
+  ret i32 0
+}
+
+declare void @llvm.memset.p0i8.i64(ptr, i8, i64, i1)
+
+declare void @test1(ptr)