AMDGPU/GlobalISel: Handle stacksave/stackrestore

[Pierre-vh]

llvm/test/CodeGen/AMDGPU/stacksave_stackrestore.invalid.ll

5

nit: why does this test use amdpal but the MIR ones don't?
Also is filetype=null needed?

[arsenm]

llvm/test/CodeGen/AMDGPU/stacksave_stackrestore.invalid.ll

5

We emit an error on amdgpu_gfx functions for amdhsa which is really annoying.

-filetype=null saves time by not actually emitting the broken result

[arsenm]

1030483561d6abf0a160bca324509d0c7ef4e398

llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.h

private:
   bool selectG_PTRMASK(MachineInstr &I) const;
   bool selectG_EXTRACT_VECTOR_ELT(MachineInstr &I) const;
   bool selectG_INSERT_VECTOR_ELT(MachineInstr &I) const;
   bool selectBufferLoadLds(MachineInstr &MI) const;
   bool selectGlobalLoadLds(MachineInstr &MI) const;
   bool selectBVHIntrinsic(MachineInstr &I) const;
   bool selectSMFMACIntrin(MachineInstr &I) const;
   bool selectWaveAddress(MachineInstr &I) const;
+  bool selectStackRestore(MachineInstr &MI) const;
 
   std::pair<Register, unsigned> selectVOP3ModsImpl(MachineOperand &Root,
                                                    bool IsCanonicalizing = true,
                                                    bool AllowAbs = true,
                                                    bool OpSel = false) const;
 
   Register copyToVGPRIfSrcFolded(Register Src, unsigned Mods,
                                  MachineOperand Root, MachineInstr *InsertPt,

llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp

void AMDGPUInstructionSelector::setupMF(MachineFunction &MF, GISelKnownBits *KB,
                                         CodeGenCoverage *CoverageInfo,
                                         ProfileSummaryInfo *PSI,
                                         BlockFrequencyInfo *BFI) {
   MRI = &MF.getRegInfo();
   Subtarget = &MF.getSubtarget<GCNSubtarget>();
   InstructionSelector::setupMF(MF, KB, CoverageInfo, PSI, BFI);
 }
 
+// Return the wave level SGPR base address if this is a wave address.
+static Register getWaveAddress(const MachineInstr *Def) {
+  return Def->getOpcode() == AMDGPU::G_AMDGPU_WAVE_ADDRESS
+             ? Def->getOperand(1).getReg()
+             : Register();
+}
+
 bool AMDGPUInstructionSelector::isVCC(Register Reg,
                                       const MachineRegisterInfo &MRI) const {
   // The verifier is oblivious to s1 being a valid value for wavesize registers.
   if (Reg.isPhysical())
     return false;
 
   auto &RegClassOrBank = MRI.getRegClassOrRegBank(Reg);
   const TargetRegisterClass *RC =
const TargetRegisterClass &RC =
       IsVALU ? AMDGPU::VGPR_32RegClass : AMDGPU::SReg_32RegClass;
   if (!RBI.constrainGenericRegister(DstReg, RC, *MRI))
     return false;
 
   MI.eraseFromParent();
   return true;
 }
 
+bool AMDGPUInstructionSelector::selectStackRestore(MachineInstr &MI) const {
+  Register SrcReg = MI.getOperand(0).getReg();
+  if (!RBI.constrainGenericRegister(SrcReg, AMDGPU::SReg_32RegClass, *MRI))
+    return false;
+
+  MachineInstr *DefMI = MRI->getVRegDef(SrcReg);
+  Register SP =
+      Subtarget->getTargetLowering()->getStackPointerRegisterToSaveRestore();
+  Register WaveAddr = getWaveAddress(DefMI);
+  MachineBasicBlock *MBB = MI.getParent();
+  const DebugLoc &DL = MI.getDebugLoc();
+
+  if (!WaveAddr) {
+    WaveAddr = MRI->createVirtualRegister(&AMDGPU::SReg_32RegClass);
+    BuildMI(*MBB, MI, DL, TII.get(AMDGPU::S_LSHR_B32), WaveAddr)
+      .addReg(SrcReg)
+      .addImm(Subtarget->getWavefrontSizeLog2())
+      .setOperandDead(3); // Dead scc
+  }
+
+  BuildMI(*MBB, &MI, DL, TII.get(AMDGPU::COPY), SP)
+    .addReg(WaveAddr);
+
+  MI.eraseFromParent();
+  return true;
+}
+
 bool AMDGPUInstructionSelector::select(MachineInstr &I) {
   if (I.isPHI())
     return selectPHI(I);
 
   if (!I.isPreISelOpcode()) {
     if (I.isCopy())
       return selectCOPY(I);
     return true;
bool AMDGPUInstructionSelector::select(MachineInstr &I) {
   case AMDGPU::G_SBFX:
   case AMDGPU::G_UBFX:
     return selectG_SBFX_UBFX(I);
   case AMDGPU::G_SI_CALL:
     I.setDesc(TII.get(AMDGPU::SI_CALL));
     return true;
   case AMDGPU::G_AMDGPU_WAVE_ADDRESS:
     return selectWaveAddress(I);
+  case AMDGPU::G_STACKRESTORE:
+    return selectStackRestore(I);
   default:
     return selectImpl(I, *CoverageInfo);
   }
   return false;
 }
 
 InstructionSelector::ComplexRendererFns
 AMDGPUInstructionSelector::selectVCSRC(MachineOperand &Root) const {
bool AMDGPUInstructionSelector::isUnneededShiftMask(const MachineInstr &MI,
 
   if (RHS->countr_one() >= ShAmtBits)
     return true;
 
   const APInt &LHSKnownZeros = KB->getKnownZeroes(MI.getOperand(1).getReg());
   return (LHSKnownZeros | *RHS).countr_one() >= ShAmtBits;
 }
 
-// Return the wave level SGPR base address if this is a wave address.
-static Register getWaveAddress(const MachineInstr *Def) {
-  return Def->getOpcode() == AMDGPU::G_AMDGPU_WAVE_ADDRESS
-             ? Def->getOperand(1).getReg()
-             : Register();
-}
-
 InstructionSelector::ComplexRendererFns
 AMDGPUInstructionSelector::selectMUBUFScratchOffset(
     MachineOperand &Root) const {
   Register Reg = Root.getReg();
   const SIMachineFunctionInfo *Info = MF->getInfo<SIMachineFunctionInfo>();
 
   std::optional<DefinitionAndSourceRegister> Def =
     getDefSrcRegIgnoringCopies(Reg, *MRI);

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.h

bool legalizeBufferLoad(MachineInstr &MI, MachineRegisterInfo &MRI,
                           MachineIRBuilder &B, bool IsFormat,
                           bool IsTyped) const;
   bool legalizeBufferAtomic(MachineInstr &MI, MachineIRBuilder &B,
                             Intrinsic::ID IID) const;
 
   bool legalizeBVHIntrinsic(MachineInstr &MI, MachineIRBuilder &B) const;
 
   bool legalizeFPTruncRound(MachineInstr &MI, MachineIRBuilder &B) const;
+  bool legalizeStackSave(MachineInstr &MI, MachineIRBuilder &B) const;
 
   bool legalizeImageIntrinsic(
       MachineInstr &MI, MachineIRBuilder &B,
       GISelChangeObserver &Observer,
       const AMDGPU::ImageDimIntrinsicInfo *ImageDimIntr) const;
 
   bool legalizeSBufferLoad(LegalizerHelper &Helper, MachineInstr &MI) const;
 

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

AMDGPULegalizerInfo::AMDGPULegalizerInfo(const GCNSubtarget &ST_,
 
   getActionDefinitionsBuilder(G_FRAME_INDEX).legalFor({PrivatePtr});
 
   // If the amount is divergent, we have to do a wave reduction to get the
   // maximum value, so this is expanded during RegBankSelect.
   getActionDefinitionsBuilder(G_DYN_STACKALLOC)
     .legalFor({{PrivatePtr, S32}});
 
+  getActionDefinitionsBuilder(G_STACKSAVE)
+    .customFor({PrivatePtr});
+  getActionDefinitionsBuilder(G_STACKRESTORE)
+    .legalFor({PrivatePtr});
+
   getActionDefinitionsBuilder(G_GLOBAL_VALUE)
     .customIf(typeIsNot(0, PrivatePtr));
 
   getActionDefinitionsBuilder(G_BLOCK_ADDR).legalFor({CodePtr});
 
   auto &FPOpActions = getActionDefinitionsBuilder(
     { G_FADD, G_FMUL, G_FMA, G_FCANONICALIZE,
       G_STRICT_FADD, G_STRICT_FMUL, G_STRICT_FMA})
case TargetOpcode::G_BUILD_VECTOR_TRUNC:
     return legalizeBuildVector(MI, MRI, B);
   case TargetOpcode::G_MUL:
     return legalizeMul(Helper, MI);
   case TargetOpcode::G_CTLZ:
   case TargetOpcode::G_CTTZ:
     return legalizeCTLZ_CTTZ(MI, MRI, B);
   case TargetOpcode::G_INTRINSIC_FPTRUNC_ROUND:
     return legalizeFPTruncRound(MI, B);
+  case TargetOpcode::G_STACKSAVE:
+    return legalizeStackSave(MI, B);
   default:
     return false;
   }
 
   llvm_unreachable("expected switch to return");
 }
 
 Register AMDGPULegalizerInfo::getSegmentAperture(
B.buildInstr(Opc)
       .addDef(MI.getOperand(0).getReg())
       .addUse(MI.getOperand(1).getReg());
 
   MI.eraseFromParent();
 
   return true;
 }
 
+bool AMDGPULegalizerInfo::legalizeStackSave(MachineInstr &MI,
+                                            MachineIRBuilder &B) const {
+  const SITargetLowering *TLI = ST.getTargetLowering();
+  Register StackPtr = TLI->getStackPointerRegisterToSaveRestore();
+  Register DstReg = MI.getOperand(0).getReg();
+  B.buildInstr(AMDGPU::G_AMDGPU_WAVE_ADDRESS, {DstReg}, {StackPtr});
+  MI.eraseFromParent();
+  return true;
+}
+
 bool AMDGPULegalizerInfo::legalizeIntrinsic(LegalizerHelper &Helper,
                                             MachineInstr &MI) const {
   MachineIRBuilder &B = Helper.MIRBuilder;
   MachineRegisterInfo &MRI = *B.getMRI();
 
   // Replace the use G_BRCOND with the exec manipulate and branch pseudos.
   auto IntrID = cast<GIntrinsic>(MI).getIntrinsicID();
   switch (IntrID) {

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

   case AMDGPU::G_SEXTLOAD: {
     if (applyMappingLoad(MI, OpdMapper, MRI))
       return;
     break;
   }
   case AMDGPU::G_DYN_STACKALLOC:
     applyMappingDynStackAlloc(MI, OpdMapper, MRI);
     return;
+  case AMDGPU::G_STACKRESTORE: {
+    applyDefaultMapping(OpdMapper);
+    constrainOpWithReadfirstlane(MI, MRI, 0);
+    return;
+  }
   case AMDGPU::G_SBFX:
     applyMappingBFE(OpdMapper, /*Signed*/ true);
     return;
   case AMDGPU::G_UBFX:
     applyMappingBFE(OpdMapper, /*Signed*/ false);
     return;
   case AMDGPU::G_AMDGPU_MAD_U64_U32:
   case AMDGPU::G_AMDGPU_MAD_I64_I32:

llvm/test/CodeGen/AMDGPU/GlobalISel/inst-select-stacksave-stackrestore.invalid.mir

+# RUN: not --crash llc -march=amdgcn -mcpu=gfx1030 -run-pass=instruction-select -o /dev/null %s 2>&1 | FileCheck -check-prefix=ERR %s
+
+# ERR: LLVM ERROR: cannot select: G_STACKRESTORE %{{[0-9]+}}:vgpr(p5) (in function: stackrestore_waveaddress_vgpr)
+
+---
+name:            stackrestore_waveaddress_vgpr
+legalized:       true
+regBankSelected: true
+body: |
+  bb.0:
+    %0:vgpr(p5) = G_AMDGPU_WAVE_ADDRESS $sgpr32
+    G_STACKRESTORE %0
+
+...

llvm/test/CodeGen/AMDGPU/GlobalISel/inst-select-stacksave-stackrestore.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 2
+# RUN: llc -march=amdgcn -mcpu=gfx1030 -run-pass=instruction-select -o - %s | FileCheck -check-prefix=GFX10-WAVE32 %s
+# RUN: llc -march=amdgcn -mcpu=gfx1030 -mattr=+wavefrontsize64 -run-pass=instruction-select -o - %s | FileCheck -check-prefix=GFX10-WAVE64 %s
+
+---
+name:            stackrestore_waveaddress_sgpr
+legalized:       true
+regBankSelected: true
+body: |
+  bb.0:
+    ; GFX10-WAVE32-LABEL: name: stackrestore_waveaddress_sgpr
+    ; GFX10-WAVE32: $sgpr32 = COPY $sgpr32
+    ;
+    ; GFX10-WAVE64-LABEL: name: stackrestore_waveaddress_sgpr
+    ; GFX10-WAVE64: $sgpr32 = COPY $sgpr32
+    %0:sgpr(p5) = G_AMDGPU_WAVE_ADDRESS $sgpr32
+    G_STACKRESTORE %0
+
+...
+
+# Test we aren't special casing the direct register value.
+---
+name:            stackrestore_direct_sp_sgpr
+legalized:       true
+regBankSelected: true
+body: |
+  bb.0:
+    liveins: $sgpr10
+    ; GFX10-WAVE32-LABEL: name: stackrestore_direct_sp_sgpr
+    ; GFX10-WAVE32: liveins: $sgpr10
+    ; GFX10-WAVE32-NEXT: {{  $}}
+    ; GFX10-WAVE32-NEXT: [[COPY:%[0-9]+]]:sreg_32 = COPY $sgpr32
+    ; GFX10-WAVE32-NEXT: [[S_LSHR_B32_:%[0-9]+]]:sreg_32 = S_LSHR_B32 [[COPY]], 5, implicit-def dead $scc
+    ; GFX10-WAVE32-NEXT: $sgpr32 = COPY [[S_LSHR_B32_]]
+    ;
+    ; GFX10-WAVE64-LABEL: name: stackrestore_direct_sp_sgpr
+    ; GFX10-WAVE64: liveins: $sgpr10
+    ; GFX10-WAVE64-NEXT: {{  $}}
+    ; GFX10-WAVE64-NEXT: [[COPY:%[0-9]+]]:sreg_32 = COPY $sgpr32
+    ; GFX10-WAVE64-NEXT: [[S_LSHR_B32_:%[0-9]+]]:sreg_32 = S_LSHR_B32 [[COPY]], 6, implicit-def dead $scc
+    ; GFX10-WAVE64-NEXT: $sgpr32 = COPY [[S_LSHR_B32_]]
+    %0:sgpr(p5) = COPY $sgpr32
+    G_STACKRESTORE %0
+
+...
+
+---
+name:            stackrestore_any_sgpr
+legalized:       true
+regBankSelected: true
+body: |
+  bb.0:
+    liveins: $sgpr10
+    ; GFX10-WAVE32-LABEL: name: stackrestore_any_sgpr
+    ; GFX10-WAVE32: liveins: $sgpr10
+    ; GFX10-WAVE32-NEXT: {{  $}}
+    ; GFX10-WAVE32-NEXT: [[COPY:%[0-9]+]]:sreg_32 = COPY $sgpr10
+    ; GFX10-WAVE32-NEXT: [[S_LSHR_B32_:%[0-9]+]]:sreg_32 = S_LSHR_B32 [[COPY]], 5, implicit-def dead $scc
+    ; GFX10-WAVE32-NEXT: $sgpr32 = COPY [[S_LSHR_B32_]]
+    ;
+    ; GFX10-WAVE64-LABEL: name: stackrestore_any_sgpr
+    ; GFX10-WAVE64: liveins: $sgpr10
+    ; GFX10-WAVE64-NEXT: {{  $}}
+    ; GFX10-WAVE64-NEXT: [[COPY:%[0-9]+]]:sreg_32 = COPY $sgpr10
+    ; GFX10-WAVE64-NEXT: [[S_LSHR_B32_:%[0-9]+]]:sreg_32 = S_LSHR_B32 [[COPY]], 6, implicit-def dead $scc
+    ; GFX10-WAVE64-NEXT: $sgpr32 = COPY [[S_LSHR_B32_]]
+    %0:sgpr(p5) = COPY $sgpr10
+    G_STACKRESTORE %0
+
+...
+

llvm/test/CodeGen/AMDGPU/stacksave_stackrestore.invalid.ll

 ; RUN: split-file %s %t
-; RUN: not --crash llc -mtriple=amdgcn-amd-amdpal -mcpu=gfx1030 -filetype=null %t/stacksave-error.ll 2>&1 | FileCheck -check-prefix=ERR-SAVE %s
-; RUN: not --crash llc -mtriple=amdgcn-amd-amdpal -mcpu=gfx1030 -filetype=null %t/stackrestore-error.ll 2>&1 | FileCheck -check-prefix=ERR-RESTORE %s
+; RUN: not --crash llc -global-isel=0 -mtriple=amdgcn-amd-amdpal -mcpu=gfx1030 -filetype=null %t/stacksave-error.ll 2>&1 | FileCheck -check-prefix=ERR-SAVE-SDAG %s
+; RUN: not --crash llc -global-isel=0 -mtriple=amdgcn-amd-amdpal -mcpu=gfx1030 -filetype=null %t/stackrestore-error.ll 2>&1 | FileCheck -check-prefix=ERR-RESTORE-SDAG %s
+
+; RUN: not --crash llc -global-isel=1 -mtriple=amdgcn-amd-amdpal -mcpu=gfx1030 -filetype=null %t/stacksave-error.ll 2>&1 | FileCheck -check-prefix=ERR-SAVE-GISEL %s
+; RUN: not --crash llc -global-isel=1 -mtriple=amdgcn-amd-amdpal -mcpu=gfx1030 -filetype=null %t/stackrestore-error.ll 2>&1 | FileCheck -check-prefix=ERR-RESTORE-GISEL %s
 
 ; Test that an error is produced if stacksave/stackrestore are used
 ; with the wrong (default) address space.
 
 ;--- stacksave-error.ll
 
 declare ptr @llvm.stacksave.p0()
 
-; ERR-SAVE: LLVM ERROR: Cannot select: {{.+}}: i64,ch = stacksave
+; ERR-SAVE-SDAG: LLVM ERROR: Cannot select: {{.+}}: i64,ch = stacksave
+; ERR-SAVE-GISEL: LLVM ERROR: unable to legalize instruction: %{{[0-9]+}}:_(p0) = G_STACKSAVE (in function: func_store_stacksave)
 define void @func_store_stacksave() {
   %stacksave = call ptr @llvm.stacksave.p0()
   call void asm sideeffect "; use $0", "s"(ptr %stacksave)
   ret void
 }
 
 ;--- stackrestore-error.ll
 
 declare void @llvm.stackrestore.p0(ptr)
 
-; ERR-RESTORE: LLVM ERROR: Cannot select: {{.+}}: ch = stackrestore {{.+}}, {{.+}}
+; ERR-RESTORE-SDAG: LLVM ERROR: Cannot select: {{.+}}: ch = stackrestore {{.+}}, {{.+}}
+; ERR-RESTORE-GISEL: LLVM ERROR: unable to legalize instruction: G_STACKRESTORE %{{[0-9]+}}:_(p0) (in function: func_stacksave_sgpr)
 define amdgpu_gfx void @func_stacksave_sgpr(ptr inreg %stack) {
   call void @llvm.stackrestore.p0(ptr %stack)
   ret void
 }