[LoongArch] Don't merge FrameIndex accesses into [F]{LD,ST}X

[xen0n]

Summary

Otherwise eliminateFrameIndex cannot figure out how to fixup the stack
offset with its stateless logic, because there wouldn't be an immediate
slot for it to trivially write to, and it may not be easy to transform
the surrounding code to make it work.

This fixes a fairly common crash when compiling moderately complex code with
Clang.

[SixWeining]

Thanks for the fix.

I tried to write a simpler test which can reproduce the crash.

define i8 @test(i64 %i) {
  %1 = alloca ptr 
  %2 = getelementptr inbounds i8, ptr %1, i64 %i
  %3 = load i8, ptr %2
  ret i8 %3
}

BTW. Seems @wangleiat has another approach to fix this issue. @wangleiat Do you have any comments?

[xen0n]

In D135251#3844467, @SixWeining wrote:

Thanks for the fix.

I tried to write a simpler test which can reproduce the crash.

define i8 @test(i64 %i) {
  %1 = alloca ptr 
  %2 = getelementptr inbounds i8, ptr %1, i64 %i
  %3 = load i8, ptr %2
  ret i8 %3
}

BTW. Seems @wangleiat has another approach to fix this issue. @wangleiat Do you have any comments?

Nice. I simply minimized my cases with bugpoint --run-llc, although I suspected alloca may be the culprit but I didn't try to remove the array structure.

As for which approach to choose for fixing this bug, I'm open to any better proposal. Maybe we can wait for @wangleiat to post his patch to compare.

[wangleiat]

In D135251#3844567, @xen0n wrote:

In D135251#3844467, @SixWeining wrote:

Thanks for the fix.

I tried to write a simpler test which can reproduce the crash.

define i8 @test(i64 %i) {
  %1 = alloca ptr 
  %2 = getelementptr inbounds i8, ptr %1, i64 %i
  %3 = load i8, ptr %2
  ret i8 %3
}

BTW. Seems @wangleiat has another approach to fix this issue. @wangleiat Do you have any comments?

Nice. I simply minimized my cases with bugpoint --run-llc, although I suspected alloca may be the culprit but I didn't try to remove the array structure.

As for which approach to choose for fixing this bug, I'm open to any better proposal. Maybe we can wait for @wangleiat to post his patch to compare.

This would be better.

[xen0n]

include @SixWeining's cleaner reproducer as an additional test case

[SixWeining]

LGTM except a small nit. Thanks.

llvm/test/CodeGen/LoongArch/ldx-stx-sp-1.ll

3

Nit: --mattr=+f

[xen0n]

address @SixWeining's comment

llvm/lib/Target/LoongArch/LoongArchISelDAGToDAG.h

public:
   bool runOnMachineFunction(MachineFunction &MF) override {
     Subtarget = &MF.getSubtarget<LoongArchSubtarget>();
     return SelectionDAGISel::runOnMachineFunction(MF);
   }
 
   void Select(SDNode *Node) override;
 
   bool SelectBaseAddr(SDValue Addr, SDValue &Base);
+  bool selectNonFIBaseAddr(SDValue Addr, SDValue &Base);
 
   bool selectShiftMask(SDValue N, unsigned ShiftWidth, SDValue &ShAmt);
   bool selectShiftMaskGRLen(SDValue N, SDValue &ShAmt) {
     return selectShiftMask(N, Subtarget->getGRLen(), ShAmt);
   }
   bool selectShiftMask32(SDValue N, SDValue &ShAmt) {
     return selectShiftMask(N, 32, ShAmt);
   }

llvm/lib/Target/LoongArch/LoongArchISelDAGToDAG.cpp

bool LoongArchDAGToDAGISel::SelectBaseAddr(SDValue Addr, SDValue &Base) {
   if (auto *FIN = dyn_cast<FrameIndexSDNode>(Addr))
     Base =
         CurDAG->getTargetFrameIndex(FIN->getIndex(), Subtarget->getGRLenVT());
   else
     Base = Addr;
   return true;
 }
 
+bool LoongArchDAGToDAGISel::selectNonFIBaseAddr(SDValue Addr, SDValue &Base) {
+  // If this is FrameIndex, don't select it.
+  if (isa<FrameIndexSDNode>(Addr))
+    return false;
+  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
   // shift amount. If there is an AND on the shift amount, we can bypass it if
   // it doesn't affect any of those bits.
   if (N.getOpcode() == ISD::AND && isa<ConstantSDNode>(N.getOperand(1))) {
     const APInt &AndMask = N->getConstantOperandAPInt(1);
 

llvm/lib/Target/LoongArch/LoongArchInstrInfo.td

 
 // Return an immediate subtracted from 32.
 def ImmSubFrom32 : SDNodeXForm<imm, [{
   return CurDAG->getTargetConstant(32 - N->getZExtValue(), SDLoc(N),
                                    N->getValueType(0));
 }]>;
 
 def BaseAddr : ComplexPattern<iPTR, 1, "SelectBaseAddr">;
+def NonFIBaseAddr : ComplexPattern<iPTR, 1, "selectNonFIBaseAddr">;
 
 def fma_nsz : PatFrag<(ops node:$fj, node:$fk, node:$fa),
                       (fma node:$fj, node:$fk, node:$fa), [{
   return N->getFlags().hasNoSignedZeros();
 }]>;
 
 //===----------------------------------------------------------------------===//
 // Instruction Formats

 defm : LdPat<extloadi32, LD_W, i64>;
 defm : LdPat<zextloadi32, LD_WU, i64>;
 defm : LdPat<load, LD_D, i64>;
 } // Predicates = [IsLA64]
 
 // LA64 register-register-addressed loads
 let Predicates = [IsLA64] in {
 class RegRegLdPat<PatFrag LoadOp, LAInst Inst, ValueType vt>
-  : Pat<(vt (LoadOp (add BaseAddr:$rj, GPR:$rk))),
-        (Inst BaseAddr:$rj, GPR:$rk)>;
+  : Pat<(vt (LoadOp (add NonFIBaseAddr:$rj, GPR:$rk))),
+        (Inst NonFIBaseAddr:$rj, GPR:$rk)>;
 
 def : RegRegLdPat<extloadi8, LDX_B, i64>;
 def : RegRegLdPat<sextloadi8, LDX_B, i64>;
 def : RegRegLdPat<zextloadi8, LDX_BU, i64>;
 def : RegRegLdPat<extloadi16, LDX_H, i64>;
 def : RegRegLdPat<sextloadi16, LDX_H, i64>;
 def : RegRegLdPat<zextloadi16, LDX_HU, i64>;
 def : RegRegLdPat<extloadi32, LDX_W, i64>;

 def : Pat<(store (i64 GPR:$rd), (AddLike BaseAddr:$rj, simm14_lsl2:$imm14)),
           (STPTR_D GPR:$rd, BaseAddr:$rj, simm14_lsl2:$imm14)>;
 } // Predicates = [IsLA64]
 
 // LA64 register-register-addressed stores
 let Predicates = [IsLA64] in {
 class RegRegStPat<PatFrag StoreOp, LAInst Inst, RegisterClass StTy,
                   ValueType vt>
-  : Pat<(StoreOp (vt StTy:$rd), (add BaseAddr:$rj, GPR:$rk)),
-        (Inst StTy:$rd, BaseAddr:$rj, GPR:$rk)>;
+  : Pat<(StoreOp (vt StTy:$rd), (add NonFIBaseAddr:$rj, GPR:$rk)),
+        (Inst StTy:$rd, NonFIBaseAddr:$rj, GPR:$rk)>;
 
 def : RegRegStPat<truncstorei8, STX_B, GPR, i64>;
 def : RegRegStPat<truncstorei16, STX_H, GPR, i64>;
 def : RegRegStPat<truncstorei32, STX_W, GPR, i64>;
 def : RegRegStPat<store, STX_D, GPR, i64>;
 } // Predicates = [IsLA64]
 
 /// Atomic loads and stores

llvm/lib/Target/LoongArch/LoongArchRegisterInfo.cpp

void LoongArchRegisterInfo::eliminateFrameIndex(MachineBasicBlock::iterator II,
                                                 unsigned FIOperandNum,
                                                 RegScavenger *RS) const {
   // TODO: this implementation is a temporary placeholder which does just
   // enough to allow other aspects of code generation to be tested.
 
   assert(SPAdj == 0 && "Unexpected non-zero SPAdj value");
 
   MachineInstr &MI = *II;
+  assert(MI.getOperand(FIOperandNum + 1).isImm() &&
+         "Unexpected FI-consuming insn");
+
   MachineBasicBlock &MBB = *MI.getParent();
   MachineFunction &MF = *MI.getParent()->getParent();
   MachineRegisterInfo &MRI = MF.getRegInfo();
   const LoongArchSubtarget &STI = MF.getSubtarget<LoongArchSubtarget>();
   const LoongArchInstrInfo *TII = STI.getInstrInfo();
   const TargetFrameLowering *TFI = MF.getSubtarget().getFrameLowering();
   DebugLoc DL = MI.getDebugLoc();
 

llvm/test/CodeGen/LoongArch/ldx-stx-sp-1.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc --mtriple=loongarch64 --mattr=+f < %s | FileCheck %s
+
+;; This should not crash the code generator, but the indexed loads/stores
+;; should still be present (the important part is that [f]{ld,st}x shouldn't
+;; take an $sp argument).
+
+define i8 @test_load_i(i64 %i) {
+; CHECK-LABEL: test_load_i:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    addi.d $sp, $sp, -16
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    addi.d $a1, $sp, 8
+; CHECK-NEXT:    ldx.b $a0, $a0, $a1
+; CHECK-NEXT:    addi.d $sp, $sp, 16
+; CHECK-NEXT:    ret
+  %1 = alloca ptr
+  %2 = getelementptr inbounds i8, ptr %1, i64 %i
+  %3 = load i8, ptr %2
+  ret i8 %3
+}
+
+define float @test_load_f(i64 %i) {
+; CHECK-LABEL: test_load_f:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    addi.d $sp, $sp, -16
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    slli.d $a0, $a0, 2
+; CHECK-NEXT:    addi.d $a1, $sp, 8
+; CHECK-NEXT:    fldx.s $fa0, $a0, $a1
+; CHECK-NEXT:    addi.d $sp, $sp, 16
+; CHECK-NEXT:    ret
+  %1 = alloca ptr
+  %2 = getelementptr inbounds float, ptr %1, i64 %i
+  %3 = load float, ptr %2
+  ret float %3
+}
+
+define void @test_store_i(i64 %i, i8 %v) {
+; CHECK-LABEL: test_store_i:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    addi.d $sp, $sp, -16
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    addi.d $a2, $sp, 8
+; CHECK-NEXT:    stx.b $a1, $a0, $a2
+; CHECK-NEXT:    addi.d $sp, $sp, 16
+; CHECK-NEXT:    ret
+  %1 = alloca ptr
+  %2 = getelementptr inbounds i8, ptr %1, i64 %i
+  store i8 %v, ptr %2, align 1
+  ret void
+}
+
+define void @test_store_f(i64 %i, float %v) {
+; CHECK-LABEL: test_store_f:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    addi.d $sp, $sp, -16
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    slli.d $a0, $a0, 2
+; CHECK-NEXT:    addi.d $a1, $sp, 8
+; CHECK-NEXT:    fstx.s $fa0, $a0, $a1
+; CHECK-NEXT:    addi.d $sp, $sp, 16
+; CHECK-NEXT:    ret
+  %1 = alloca ptr
+  %2 = getelementptr inbounds float, ptr %1, i64 %i
+  store float %v, ptr %2, align 4
+  ret void
+}

llvm/test/CodeGen/LoongArch/ldx-stx-sp-2.ll

+; RUN: llc --mtriple=loongarch32 < %s
+; RUN: llc --mtriple=loongarch64 < %s
+
+;; This should not crash the code generator.
+
+@.str.2 = external dso_local unnamed_addr constant [69 x i8], align 1
+
+define dso_local void @main() {
+entry:
+  %n0 = alloca [2 x [3 x i32]], align 4
+  %0 = load i32, ptr poison, align 4
+  %idxprom15 = sext i32 %0 to i64
+  %arrayidx16 = getelementptr inbounds [2 x [3 x i32]], ptr %n0, i64 0, i64 %idxprom15
+  %arrayidx17 = getelementptr inbounds [3 x i32], ptr %arrayidx16, i64 0, i64 0
+  %1 = load i32, ptr %arrayidx17, align 4
+  call void (ptr, ...) @printf(ptr noundef @.str.2, i32 noundef signext %1)
+  ret void
+}
+
+declare void @printf(ptr, ...)

llvm/test/CodeGen/LoongArch/ldx-stx-sp-3.ll

+; RUN: llc --mtriple=loongarch32 < %s
+; RUN: llc --mtriple=loongarch64 < %s
+
+;; This should not crash the code generator.
+
+define void @_ZN12_GLOBAL__N_111DumpVisitorclIN4llvm16itanium_demangle8FoldExprEEEvPKT_() {
+entry:
+  %ref.tmp6.i.i = alloca [4 x i8], align 1
+  br label %for.cond.i.i
+
+for.cond.i.i:                                     ; preds = %for.body.i.i, %entry
+  %__begin0.0.add.i.i = add nuw nsw i64 poison, 1
+  br label %for.body.i.i
+
+for.body.i.i:                                     ; preds = %for.cond.i.i
+  %__begin0.0.ptr.i.i = getelementptr inbounds i8, ptr %ref.tmp6.i.i, i64 %__begin0.0.add.i.i
+  %0 = load i8, ptr %__begin0.0.ptr.i.i, align 1
+  %tobool18.not.i.i = icmp eq i8 %0, 0
+  br i1 %tobool18.not.i.i, label %for.cond.i.i, label %exit
+
+exit: ; preds = %for.body.i.i
+  unreachable
+}