[LoongArch] Implement LoongArchRegisterInfo::canRealignStack() (#76913)

This patch fixes the crash issue in the test:
CodeGen/LoongArch/can-not-realign-stack.ll

Register allocator may spill virtual registers to the stack, which    
introduces stack alignment requirements (when the size of spilled     
    registers exceeds the default alignment size of the stack). If a  
function does not have stack alignment requirements before register   
allocation, registers used for stack alignment will not be preserved. 

Therefore, we should implement `canRealignStack()` to inform the      
register allocator whether it is allowed to perform stack realignment 
operations.

3 files changed, 75 insertions, 5 deletions

diff --git a/llvm/lib/Target/LoongArch/LoongArchRegisterInfo.cpp b/llvm/lib/Target/LoongArch/LoongArchRegisterInfo.cpp
index 257b947a3ce4..092b5f1fb442 100644
--- a/llvm/lib/Target/LoongArch/LoongArchRegisterInfo.cpp
+++ b/llvm/lib/Target/LoongArch/LoongArchRegisterInfo.cpp
@@ -15,6 +15,7 @@
 #include "LoongArch.h"
 #include "LoongArchInstrInfo.h"
 #include "LoongArchSubtarget.h"
+#include "MCTargetDesc/LoongArchBaseInfo.h"
 #include "MCTargetDesc/LoongArchMCTargetDesc.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
@@ -194,3 +195,25 @@ bool LoongArchRegisterInfo::eliminateFrameIndex(MachineBasicBlock::iterator II,
   MI.getOperand(FIOperandNum + 1).ChangeToImmediate(Offset.getFixed());
   return false;
 }
+
+bool LoongArchRegisterInfo::canRealignStack(const MachineFunction &MF) const {
+  if (!TargetRegisterInfo::canRealignStack(MF))
+    return false;
+
+  const MachineRegisterInfo *MRI = &MF.getRegInfo();
+  const LoongArchFrameLowering *TFI = getFrameLowering(MF);
+
+  // Stack realignment requires a frame pointer.  If we already started
+  // register allocation with frame pointer elimination, it is too late now.
+  if (!MRI->canReserveReg(LoongArch::R22))
+    return false;
+
+  // We may also need a base pointer if there are dynamic allocas or stack
+  // pointer adjustments around calls.
+  if (TFI->hasReservedCallFrame(MF))
+    return true;
+
+  // A base pointer is required and allowed.  Check that it isn't too late to
+  // reserve it.
+  return MRI->canReserveReg(LoongArchABI::getBPReg());
+}
diff --git a/llvm/lib/Target/LoongArch/LoongArchRegisterInfo.h b/llvm/lib/Target/LoongArch/LoongArchRegisterInfo.h
index 7e8f26b14097..d1e40254c297 100644
--- a/llvm/lib/Target/LoongArch/LoongArchRegisterInfo.h
+++ b/llvm/lib/Target/LoongArch/LoongArchRegisterInfo.h
@@ -51,6 +51,7 @@ struct LoongArchRegisterInfo : public LoongArchGenRegisterInfo {
   bool requiresFrameIndexScavenging(const MachineFunction &MF) const override {
     return true;
   }
+  bool canRealignStack(const MachineFunction &MF) const override;
 };
 } // end namespace llvm
 
diff --git a/llvm/test/CodeGen/LoongArch/can-not-realign-stack.ll b/llvm/test/CodeGen/LoongArch/can-not-realign-stack.ll
index 526821076498..af24ae64b7c7 100644
--- a/llvm/test/CodeGen/LoongArch/can-not-realign-stack.ll
+++ b/llvm/test/CodeGen/LoongArch/can-not-realign-stack.ll
@@ -1,14 +1,60 @@
-; REQUIRES: expensive_checks
-; RUN: llc --mtriple=loongarch64 --frame-pointer=none --mattr=+lasx < %s
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc --mtriple=loongarch64 --frame-pointer=none --mattr=+lasx < %s | FileCheck %s
 
-; XFAIL: *
+;; This test is checking that when a function allows stack realignment and
+;; realignment needs were not detected before register allocation (at this
+;; point, fp is not preserved), but realignment is required during register
+;; allocation, the stack should not undergo realignment.
 
-;; FIXME: This test will crash with expensive check. The subsequent patch will
-;; address and fix this issue.
+;; Ensure that the `bstrins.d $sp, $zero, n, 0` instruction is not generated.
+;; n = log2(realign_size) - 1
 
 %struct.S = type { [64 x i16] }
 
 define dso_local noundef signext i32 @main() nounwind {
+; CHECK-LABEL: main:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    addi.d $sp, $sp, -272
+; CHECK-NEXT:    st.d $ra, $sp, 264 # 8-byte Folded Spill
+; CHECK-NEXT:    st.d $fp, $sp, 256 # 8-byte Folded Spill
+; CHECK-NEXT:    pcalau12i $a0, %pc_hi20(.LCPI0_0)
+; CHECK-NEXT:    addi.d $a0, $a0, %pc_lo12(.LCPI0_0)
+; CHECK-NEXT:    xvld $xr0, $a0, 0
+; CHECK-NEXT:    xvst $xr0, $sp, 96 # 32-byte Folded Spill
+; CHECK-NEXT:    pcalau12i $a0, %pc_hi20(.LCPI0_1)
+; CHECK-NEXT:    addi.d $a0, $a0, %pc_lo12(.LCPI0_1)
+; CHECK-NEXT:    xvld $xr1, $a0, 0
+; CHECK-NEXT:    xvst $xr1, $sp, 64 # 32-byte Folded Spill
+; CHECK-NEXT:    xvst $xr1, $sp, 224
+; CHECK-NEXT:    xvst $xr0, $sp, 192
+; CHECK-NEXT:    pcalau12i $a0, %pc_hi20(.LCPI0_2)
+; CHECK-NEXT:    addi.d $a0, $a0, %pc_lo12(.LCPI0_2)
+; CHECK-NEXT:    xvld $xr0, $a0, 0
+; CHECK-NEXT:    xvst $xr0, $sp, 32 # 32-byte Folded Spill
+; CHECK-NEXT:    xvst $xr0, $sp, 160
+; CHECK-NEXT:    pcalau12i $a0, %pc_hi20(.LCPI0_3)
+; CHECK-NEXT:    addi.d $a0, $a0, %pc_lo12(.LCPI0_3)
+; CHECK-NEXT:    xvld $xr0, $a0, 0
+; CHECK-NEXT:    xvst $xr0, $sp, 0 # 32-byte Folded Spill
+; CHECK-NEXT:    xvst $xr0, $sp, 128
+; CHECK-NEXT:    addi.d $fp, $sp, 128
+; CHECK-NEXT:    move $a0, $fp
+; CHECK-NEXT:    bl %plt(foo)
+; CHECK-NEXT:    xvld $xr0, $sp, 64 # 32-byte Folded Reload
+; CHECK-NEXT:    xvst $xr0, $sp, 224
+; CHECK-NEXT:    xvld $xr0, $sp, 96 # 32-byte Folded Reload
+; CHECK-NEXT:    xvst $xr0, $sp, 192
+; CHECK-NEXT:    xvld $xr0, $sp, 32 # 32-byte Folded Reload
+; CHECK-NEXT:    xvst $xr0, $sp, 160
+; CHECK-NEXT:    xvld $xr0, $sp, 0 # 32-byte Folded Reload
+; CHECK-NEXT:    xvst $xr0, $sp, 128
+; CHECK-NEXT:    move $a0, $fp
+; CHECK-NEXT:    bl %plt(bar)
+; CHECK-NEXT:    move $a0, $zero
+; CHECK-NEXT:    ld.d $fp, $sp, 256 # 8-byte Folded Reload
+; CHECK-NEXT:    ld.d $ra, $sp, 264 # 8-byte Folded Reload
+; CHECK-NEXT:    addi.d $sp, $sp, 272
+; CHECK-NEXT:    ret
 entry:
   %s = alloca %struct.S, align 2
   call void @llvm.lifetime.start.p0(i64 128, ptr nonnull %s)