GlobalISel: Don't expand stacksave/stackrestore in IRTranslator

[arsenm]

Summary

In some (likely invalid edge cases anyway), it's not correct to
directly copy the stack pointer register.

[nitinjohnraj]

LGTM

[arsenm]

1ca0808db2c6393f4a79dcb57b89627eeefffffc

llvm/include/llvm/CodeGen/GlobalISel/LegalizerHelper.h

public:
   LegalizeResult lowerFMad(MachineInstr &MI);
   LegalizeResult lowerIntrinsicRound(MachineInstr &MI);
   LegalizeResult lowerFFloor(MachineInstr &MI);
   LegalizeResult lowerMergeValues(MachineInstr &MI);
   LegalizeResult lowerUnmergeValues(MachineInstr &MI);
   LegalizeResult lowerExtractInsertVectorElt(MachineInstr &MI);
   LegalizeResult lowerShuffleVector(MachineInstr &MI);
   LegalizeResult lowerDynStackAlloc(MachineInstr &MI);
+  LegalizeResult lowerStackSave(MachineInstr &MI);
+  LegalizeResult lowerStackRestore(MachineInstr &MI);
   LegalizeResult lowerExtract(MachineInstr &MI);
   LegalizeResult lowerInsert(MachineInstr &MI);
   LegalizeResult lowerSADDO_SSUBO(MachineInstr &MI);
   LegalizeResult lowerAddSubSatToMinMax(MachineInstr &MI);
   LegalizeResult lowerAddSubSatToAddoSubo(MachineInstr &MI);
   LegalizeResult lowerShlSat(MachineInstr &MI);
   LegalizeResult lowerBswap(MachineInstr &MI);
   LegalizeResult lowerBitreverse(MachineInstr &MI);

llvm/include/llvm/Support/TargetOpcodes.def

 HANDLE_TARGET_OPCODE(G_BLOCK_ADDR)
 
 /// Generic jump table address
 HANDLE_TARGET_OPCODE(G_JUMP_TABLE)
 
 /// Generic dynamic stack allocation.
 HANDLE_TARGET_OPCODE(G_DYN_STACKALLOC)
 
+/// Generic stack pointer save.
+HANDLE_TARGET_OPCODE(G_STACKSAVE)
+
+/// Generic stack pointer restore.
+HANDLE_TARGET_OPCODE(G_STACKRESTORE)
+
 /// Strict floating point instructions.
 HANDLE_TARGET_OPCODE(G_STRICT_FADD)
 HANDLE_TARGET_OPCODE(G_STRICT_FSUB)
 HANDLE_TARGET_OPCODE(G_STRICT_FMUL)
 HANDLE_TARGET_OPCODE(G_STRICT_FDIV)
 HANDLE_TARGET_OPCODE(G_STRICT_FREM)
 HANDLE_TARGET_OPCODE(G_STRICT_FMA)
 HANDLE_TARGET_OPCODE(G_STRICT_FSQRT)

llvm/include/llvm/Target/GenericOpcodes.td

 }
 
 def G_DYN_STACKALLOC : GenericInstruction {
   let OutOperandList = (outs ptype0:$dst);
   let InOperandList = (ins type1:$size, i32imm:$align);
   let hasSideEffects = true;
 }
 
+def G_STACKSAVE : GenericInstruction {
+  let OutOperandList = (outs ptype0:$dst);
+  let InOperandList = (ins);
+  let hasSideEffects = true;
+}
+
+def G_STACKRESTORE : GenericInstruction {
+  let OutOperandList = (outs);
+  let InOperandList = (ins ptype0:$src);
+  let hasSideEffects = true;
+}
+
 def G_FREEZE : GenericInstruction {
   let OutOperandList = (outs type0:$dst);
   let InOperandList = (ins type0:$src);
   let hasSideEffects = false;
 }
 
 def G_LROUND: GenericInstruction {
   let OutOperandList = (outs type0:$dst);

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

MIRBuilder.buildStore(
         GuardVal, getOrCreateVReg(*Slot),
         *MF->getMachineMemOperand(MachinePointerInfo::getFixedStack(*MF, FI),
                                   MachineMemOperand::MOStore |
                                       MachineMemOperand::MOVolatile,
                                   PtrTy, Align(8)));
     return true;
   }
   case Intrinsic::stacksave: {
-    // Save the stack pointer to the location provided by the intrinsic.
-    Register Reg = getOrCreateVReg(CI);
-    Register StackPtr = MF->getSubtarget()
-                            .getTargetLowering()
-                            ->getStackPointerRegisterToSaveRestore();
-
-    // If the target doesn't specify a stack pointer, then fall back.
-    if (!StackPtr)
-      return false;
-
-    MIRBuilder.buildCopy(Reg, StackPtr);
+    MIRBuilder.buildInstr(TargetOpcode::G_STACKSAVE, {getOrCreateVReg(CI)}, {});
     return true;
   }
   case Intrinsic::stackrestore: {
-    // Restore the stack pointer from the location provided by the intrinsic.
-    Register Reg = getOrCreateVReg(*CI.getArgOperand(0));
-    Register StackPtr = MF->getSubtarget()
-                            .getTargetLowering()
-                            ->getStackPointerRegisterToSaveRestore();
-
-    // If the target doesn't specify a stack pointer, then fall back.
-    if (!StackPtr)
-      return false;
-
-    MIRBuilder.buildCopy(StackPtr, Reg);
+    MIRBuilder.buildInstr(TargetOpcode::G_STACKRESTORE, {},
+                          {getOrCreateVReg(*CI.getArgOperand(0))});
     return true;
   }
   case Intrinsic::cttz:
   case Intrinsic::ctlz: {
     ConstantInt *Cst = cast<ConstantInt>(CI.getArgOperand(1));
     bool isTrailing = ID == Intrinsic::cttz;
     unsigned Opcode = isTrailing
                           ? Cst->isZero() ? TargetOpcode::G_CTTZ

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

LegalizerHelper::lower(MachineInstr &MI, unsigned TypeIdx, LLT LowerHintTy) {
   }
   case G_EXTRACT_VECTOR_ELT:
   case G_INSERT_VECTOR_ELT:
     return lowerExtractInsertVectorElt(MI);
   case G_SHUFFLE_VECTOR:
     return lowerShuffleVector(MI);
   case G_DYN_STACKALLOC:
     return lowerDynStackAlloc(MI);
+  case G_STACKSAVE:
+    return lowerStackSave(MI);
+  case G_STACKRESTORE:
+    return lowerStackRestore(MI);
   case G_EXTRACT:
     return lowerExtract(MI);
   case G_INSERT:
     return lowerInsert(MI);
   case G_BSWAP:
     return lowerBswap(MI);
   case G_BITREVERSE:
     return lowerBitreverse(MI);
LegalizerHelper::lowerDynStackAlloc(MachineInstr &MI) {
   MIRBuilder.buildCopy(SPReg, SPTmp);
   MIRBuilder.buildCopy(Dst, SPTmp);
 
   MI.eraseFromParent();
   return Legalized;
 }
 
 LegalizerHelper::LegalizeResult
+LegalizerHelper::lowerStackSave(MachineInstr &MI) {
+  Register StackPtr = TLI.getStackPointerRegisterToSaveRestore();
+  if (!StackPtr)
+    return UnableToLegalize;
+
+  MIRBuilder.buildCopy(MI.getOperand(0), StackPtr);
+  MI.eraseFromParent();
+  return Legalized;
+}
+
+LegalizerHelper::LegalizeResult
+LegalizerHelper::lowerStackRestore(MachineInstr &MI) {
+  Register StackPtr = TLI.getStackPointerRegisterToSaveRestore();
+  if (!StackPtr)
+    return UnableToLegalize;
+
+  MIRBuilder.buildCopy(StackPtr, MI.getOperand(0));
+  MI.eraseFromParent();
+  return Legalized;
+}
+
+LegalizerHelper::LegalizeResult
 LegalizerHelper::lowerExtract(MachineInstr &MI) {
   auto [DstReg, DstTy, SrcReg, SrcTy] = MI.getFirst2RegLLTs();
   unsigned Offset = MI.getOperand(2).getImm();
 
   // Extract sub-vector or one element
   if (SrcTy.isVector()) {
     unsigned SrcEltSize = SrcTy.getElementType().getSizeInBits();
     unsigned DstSize = DstTy.getSizeInBits();

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp

getActionDefinitionsBuilder(G_CONCAT_VECTORS)
       .legalFor({{v4s32, v2s32}, {v8s16, v4s16}, {v16s8, v8s8}});
 
   getActionDefinitionsBuilder(G_JUMP_TABLE).legalFor({{p0}, {s64}});
 
   getActionDefinitionsBuilder(G_BRJT).legalIf([=](const LegalityQuery &Query) {
     return Query.Types[0] == p0 && Query.Types[1] == s64;
   });
 
-  getActionDefinitionsBuilder(G_DYN_STACKALLOC).lower();
+  getActionDefinitionsBuilder({G_DYN_STACKALLOC,
+                               G_STACKSAVE,
+                               G_STACKRESTORE}).lower();
 
   if (ST.hasMOPS()) {
     // G_BZERO is not supported. Currently it is only emitted by
     // PreLegalizerCombiner for G_MEMSET with zero constant.
     getActionDefinitionsBuilder(G_BZERO).unsupported();
 
     getActionDefinitionsBuilder(G_MEMSET)
         .legalForCartesianProduct({p0}, {s64}, {s64})

llvm/lib/Target/X86/X86LegalizerInfo.cpp

getActionDefinitionsBuilder(G_SELECT)
                   {p0, s32}})
     .widenScalarToNextPow2(0, /*Min=*/8)
     .clampScalar(0, s8, sMaxScalar)
     .clampScalar(1, s32, s32);
 
   // memory intrinsics
   getActionDefinitionsBuilder({G_MEMCPY, G_MEMMOVE, G_MEMSET}).libcall();
 
+  getActionDefinitionsBuilder({G_DYN_STACKALLOC,
+                               G_STACKSAVE,
+                               G_STACKRESTORE}).lower();
+
   // fp intrinsics
   getActionDefinitionsBuilder(G_INTRINSIC_ROUNDEVEN)
       .scalarize(0)
       .minScalar(0, LLT::scalar(32))
       .libcall();
 
   getActionDefinitionsBuilder({G_FREEZE, G_CONSTANT_FOLD_BARRIER})
     .legalFor({s8, s16, s32, s64, p0})

llvm/test/CodeGen/AArch64/GlobalISel/arm64-irtranslator.ll

; CHECK: $w0 = COPY [[ZEXT1]](s32)
   call void %f(i1 true)
   ret void
 }
 
 declare ptr @llvm.stacksave()
 declare void @llvm.stackrestore(ptr)
 define void @test_stacksaverestore() {
   ; CHECK-LABEL: name: test_stacksaverestore
-  ; CHECK: [[SAVE:%[0-9]+]]:_(p0) = COPY $sp
-  ; CHECK-NEXT: $sp = COPY [[SAVE]](p0)
+  ; CHECK: [[SAVE:%[0-9]+]]:_(p0) = G_STACKSAVE
+  ; CHECK-NEXT: G_STACKRESTORE [[SAVE]]
   ; CHECK-NEXT: RET_ReallyLR
   %sp = call ptr @llvm.stacksave()
   call void @llvm.stackrestore(ptr %sp)
   ret void
 }
 
 declare float @llvm.rint.f32(float)
 define float @test_rint_f32(float %x) {

llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir

 # DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: G_BLOCK_ADDR (opcode {{[0-9]+}}): 1 type index, 0 imm indices
 # DEBUG-NEXT: .. the first uncovered type index: 1, OK
 # DEBUG-NEXT: .. the first uncovered imm index: 0, OK
 # DEBUG-NEXT: G_JUMP_TABLE (opcode {{[0-9]+}}): 1 type index, 0 imm indices
 # DEBUG-NEXT: .. the first uncovered type index: 1, OK
 # DEBUG-NEXT: .. the first uncovered imm index: 0, OK
-# DEBUG-NEXT: G_DYN_STACKALLOC (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
+# DEBUG-NEXT: G_DYN_STACKALLOC (opcode [[DYN_STACKALLOC:[0-9]+]]): 2 type indices, 0 imm indices
+# DEBUG-NEXT: .. type index coverage check SKIPPED: user-defined predicate detected
+# DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected
+# DEBUG-NEXT: G_STACKSAVE (opcode {{[0-9]+}}): 1 type index, 0 imm indices
+# DEBUG-NEXT: .. opcode {{[0-9]+}} is aliased to [[DYN_STACKALLOC]]
+# DEBUG-NEXT: .. type index coverage check SKIPPED: user-defined predicate detected
+# DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected
+# DEBUG-NEXT: G_STACKRESTORE (opcode {{[0-9]+}}): 1 type index, 0 imm indices
+# DEBUG-NEXT: .. opcode {{[0-9]+}} is aliased to [[DYN_STACKALLOC]]
 # DEBUG-NEXT: .. type index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: G_STRICT_FADD (opcode {{[0-9]+}}): 1 type index, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: G_STRICT_FSUB (opcode {{[0-9]+}}): 1 type index, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined

llvm/test/CodeGen/AArch64/GlobalISel/stacksave-stackrestore.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; RUN: llc -global-isel=1 -mtriple=aarch64-linux-gnu -o - %s | FileCheck %s
+
+declare void @use_addr(ptr)
+declare ptr @llvm.stacksave.p0()
+declare void @llvm.stackrestore.p0(ptr)
+
+define void @test_scoped_alloca(i64 %n) {
+; CHECK-LABEL: test_scoped_alloca:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    stp x29, x30, [sp, #-32]! // 16-byte Folded Spill
+; CHECK-NEXT:    str x19, [sp, #16] // 8-byte Folded Spill
+; CHECK-NEXT:    mov x29, sp
+; CHECK-NEXT:    .cfi_def_cfa w29, 32
+; CHECK-NEXT:    .cfi_offset w19, -16
+; CHECK-NEXT:    .cfi_offset w30, -24
+; CHECK-NEXT:    .cfi_offset w29, -32
+; CHECK-NEXT:    add x9, x0, #15
+; CHECK-NEXT:    mov x8, sp
+; CHECK-NEXT:    and x9, x9, #0xfffffffffffffff0
+; CHECK-NEXT:    mov x19, sp
+; CHECK-NEXT:    sub x0, x8, x9
+; CHECK-NEXT:    mov sp, x0
+; CHECK-NEXT:    bl use_addr
+; CHECK-NEXT:    mov sp, x19
+; CHECK-NEXT:    mov sp, x29
+; CHECK-NEXT:    ldr x19, [sp, #16] // 8-byte Folded Reload
+; CHECK-NEXT:    ldp x29, x30, [sp], #32 // 16-byte Folded Reload
+; CHECK-NEXT:    ret
+  %sp = call ptr @llvm.stacksave.p0()
+  %addr = alloca i8, i64 %n
+  call void @use_addr(ptr %addr)
+  call void @llvm.stackrestore.p0(ptr %sp)
+  ret void
+}

llvm/test/CodeGen/X86/GlobalISel/stacksave-stackrestore.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; RUN: llc -global-isel=1 -mtriple=x86_64-linux-gnu -o - %s | FileCheck %s
+
+declare void @use_addr(ptr)
+declare ptr @llvm.stacksave.p0()
+declare void @llvm.stackrestore.p0(ptr)
+
+define void @test_scoped_alloca(i64 %n) {
+; CHECK-LABEL: test_scoped_alloca:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    pushq %rbp
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    .cfi_offset %rbp, -16
+; CHECK-NEXT:    movq %rsp, %rbp
+; CHECK-NEXT:    .cfi_def_cfa_register %rbp
+; CHECK-NEXT:    pushq %rbx
+; CHECK-NEXT:    pushq %rax
+; CHECK-NEXT:    .cfi_offset %rbx, -24
+; CHECK-NEXT:    movq %rsp, %rbx
+; CHECK-NEXT:    movq %rsp, %rax
+; CHECK-NEXT:    imulq $1, %rdi, %rcx
+; CHECK-NEXT:    addq $15, %rcx
+; CHECK-NEXT:    andq $-16, %rcx
+; CHECK-NEXT:    subq %rcx, %rax
+; CHECK-NEXT:    movq %rax, %rsp
+; CHECK-NEXT:    movq %rax, %rdi
+; CHECK-NEXT:    callq use_addr
+; CHECK-NEXT:    movq %rbx, %rsp
+; CHECK-NEXT:    leaq -8(%rbp), %rsp
+; CHECK-NEXT:    popq %rbx
+; CHECK-NEXT:    popq %rbp
+; CHECK-NEXT:    .cfi_def_cfa %rsp, 8
+; CHECK-NEXT:    retq
+  %sp = call ptr @llvm.stacksave.p0()
+  %addr = alloca i8, i64 %n
+  call void @use_addr(ptr %addr)
+  call void @llvm.stackrestore.p0(ptr %sp)
+  ret void
+}