AMDGPU: Annotate necessity for flat-scratch-init

[arsenm]

Summary

As an approximation of the existing handling to avoid
regressions. Fixes using too many registers with calls
on subtargets with the SGPR allocation bug.

[rampitec]

LGTM

[arsenm]

r308326

lib/Target/AMDGPU/AMDGPUAnnotateKernelFeatures.cpp

static void copyFeaturesToFunction(Function &Parent, const Function &Callee,
   if (handleAttr(Parent, Callee, "amdgpu-queue-ptr"))
     NeedQueuePtr = true;
 
   for (StringRef AttrName : AttrNames)
     handleAttr(Parent, Callee, AttrName);
 }
 
 bool AMDGPUAnnotateKernelFeatures::addFeatureAttributes(Function &F) {
-  bool HasApertureRegs = TM->getSubtarget<AMDGPUSubtarget>(F).hasApertureRegs();
+  const AMDGPUSubtarget &ST = TM->getSubtarget<AMDGPUSubtarget>(F);
+  bool HasFlat = ST.hasFlatAddressSpace();
+  bool HasApertureRegs = ST.hasApertureRegs();
   SmallPtrSet<const Constant *, 8> ConstantExprVisited;
 
   bool Changed = false;
   bool NeedQueuePtr = false;
+  bool HaveCall = false;
   bool IsFunc = !AMDGPU::isEntryFunctionCC(F.getCallingConv());
 
   for (BasicBlock &BB : F) {
     for (Instruction &I : BB) {
       CallSite CS(&I);
       if (CS) {
         Function *Callee = CS.getCalledFunction();
 
         // TODO: Do something with indirect calls.
-        if (!Callee)
+        if (!Callee) {
+          if (!CS.isInlineAsm())
+            HaveCall = true;
           continue;
+        }
 
         Intrinsic::ID IID = Callee->getIntrinsicID();
         if (IID == Intrinsic::not_intrinsic) {
+          HaveCall = true;
           copyFeaturesToFunction(F, *Callee, NeedQueuePtr);
           Changed = true;
         } else {
           bool NonKernelOnly = false;
           StringRef AttrName = intrinsicToAttrName(IID,
                                                    NonKernelOnly, NeedQueuePtr);
           if (!AttrName.empty() && (IsFunc || !NonKernelOnly)) {
             F.addFnAttr(AttrName);
for (BasicBlock &BB : F) {
     }
   }
 
   if (NeedQueuePtr) {
     F.addFnAttr("amdgpu-queue-ptr");
     Changed = true;
   }
 
+  // TODO: We could refine this to captured pointers that could possibly be
+  // accessed by flat instructions. For now this is mostly a poor way of
+  // estimating whether there are calls before argument lowering.
+  if (HasFlat && !IsFunc && HaveCall) {
+    F.addFnAttr("amdgpu-flat-scratch");
+    Changed = true;
+  }
+
   return Changed;
 }
 
 bool AMDGPUAnnotateKernelFeatures::runOnSCC(CallGraphSCC &SCC) {
   Module &M = SCC.getCallGraph().getModule();
   Triple TT(M.getTargetTriple());
 
   bool Changed = false;

lib/Target/AMDGPU/SIFrameLowering.cpp

void SIFrameLowering::emitEntryFunctionPrologue(MachineFunction &MF,
   // register even if there are no stack objects. There could be stores to undef
   // or a constant without an associated object.
 
   // FIXME: We still have implicit uses on SGPR spill instructions in case they
   // need to spill to vector memory. It's likely that will not happen, but at
   // this point it appears we need the setup. This part of the prolog should be
   // emitted after frame indices are eliminated.
 
-  if (MF.getFrameInfo().hasStackObjects() && MFI->hasFlatScratchInit())
+  if (MFI->hasFlatScratchInit())
     emitFlatScratchInit(ST, MF, MBB);
 
   unsigned SPReg = MFI->getStackPtrOffsetReg();
   if (SPReg != AMDGPU::SP_REG) {
     assert(MRI.isReserved(SPReg) && "SPReg used but not reserved");
 
     DebugLoc DL;
-    int64_t StackSize = MF.getFrameInfo().getStackSize();
+    const MachineFrameInfo &FrameInfo = MF.getFrameInfo();
+    int64_t StackSize = FrameInfo.getStackSize();
 
     if (StackSize == 0) {
       BuildMI(MBB, MBB.begin(), DL, TII->get(AMDGPU::COPY), SPReg)
         .addReg(MFI->getScratchWaveOffsetReg());
     } else {
       BuildMI(MBB, MBB.begin(), DL, TII->get(AMDGPU::S_ADD_U32), SPReg)
         .addReg(MFI->getScratchWaveOffsetReg())
         .addImm(StackSize * ST.getWavefrontSize());

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

if (HasStackObjects || MaySpill) {
 
       // HS and GS always have the scratch wave offset in SGPR5 on GFX9.
       if (ST.getGeneration() >= AMDGPUSubtarget::GFX9 &&
           (CC == CallingConv::AMDGPU_HS || CC == CallingConv::AMDGPU_GS))
         PrivateSegmentWaveByteOffsetSystemSGPR = AMDGPU::SGPR5;
     }
   }
 
-  if (ST.isAmdCodeObjectV2(MF)) {
+  bool IsCOV2 = ST.isAmdCodeObjectV2(MF);
+  if (IsCOV2) {
     if (HasStackObjects || MaySpill)
       PrivateSegmentBuffer = true;
 
     if (F->hasFnAttribute("amdgpu-dispatch-ptr"))
       DispatchPtr = true;
 
     if (F->hasFnAttribute("amdgpu-queue-ptr"))
       QueuePtr = true;
 
     if (F->hasFnAttribute("amdgpu-dispatch-id"))
       DispatchID = true;
   } else if (ST.isMesaGfxShader(MF)) {
     if (HasStackObjects || MaySpill)
       ImplicitBufferPtr = true;
   }
 
   if (F->hasFnAttribute("amdgpu-kernarg-segment-ptr"))
     KernargSegmentPtr = true;
 
-  // We don't need to worry about accessing spills with flat instructions.
-  // TODO: On VI where we must use flat for global, we should be able to omit
-  // this if it is never used for generic access.
-  if (HasStackObjects && ST.hasFlatAddressSpace() && ST.isAmdHsaOS() &&
-      isEntryFunction())
+  if (ST.hasFlatAddressSpace() && isEntryFunction() && IsCOV2) {
+    // TODO: This could be refined a lot. The attribute is a poor way of
+    // detecting calls that may require it before argument lowering.
+    if (HasStackObjects || F->hasFnAttribute("amdgpu-flat-scratch"))
       FlatScratchInit = true;
   }
+}
 
 unsigned SIMachineFunctionInfo::addPrivateSegmentBuffer(
   const SIRegisterInfo &TRI) {
   PrivateSegmentBufferUserSGPR = TRI.getMatchingSuperReg(
     getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_128RegClass);
   NumUserSGPRs += 4;
   return PrivateSegmentBufferUserSGPR;
 }

test/CodeGen/AMDGPU/annotate-kernel-features-hsa-call.ll

 }
 
 ; HSA: define void @func_indirect_use_implicitarg_ptr() #14 {
 define void @func_indirect_use_implicitarg_ptr() #1 {
   call void @use_implicitarg_ptr()
   ret void
 }
 
+; HSA: declare void @external.func() #15
+declare void @external.func() #3
+
+; HSA: define internal void @defined.func() #15 {
+define internal void @defined.func() #3 {
+  ret void
+}
+
+; HSA: define void @func_call_external() #15 {
+define void @func_call_external() #3 {
+  call void @external.func()
+  ret void
+}
+
+; HSA: define void @func_call_defined() #15 {
+define void @func_call_defined() #3 {
+  call void @defined.func()
+  ret void
+}
+
+; HSA: define void @func_call_asm() #15 {
+define void @func_call_asm() #3 {
+  call void asm sideeffect "", ""() #3
+  ret void
+}
+
+; HSA: define amdgpu_kernel void @kern_call_external() #16 {
+define amdgpu_kernel void @kern_call_external() #3 {
+  call void @external.func()
+  ret void
+}
+
+; HSA: define amdgpu_kernel void @func_kern_defined() #16 {
+define amdgpu_kernel void @func_kern_defined() #3 {
+  call void @defined.func()
+  ret void
+}
+
 attributes #0 = { nounwind readnone speculatable }
 attributes #1 = { nounwind "target-cpu"="fiji" }
 attributes #2 = { nounwind "target-cpu"="gfx900" }
+attributes #3 = { nounwind }
 
 ; HSA: attributes #0 = { nounwind readnone speculatable }
 ; HSA: attributes #1 = { nounwind "amdgpu-work-item-id-x" "target-cpu"="fiji" }
 ; HSA: attributes #2 = { nounwind "amdgpu-work-item-id-y" "target-cpu"="fiji" }
 ; HSA: attributes #3 = { nounwind "amdgpu-work-item-id-z" "target-cpu"="fiji" }
 ; HSA: attributes #4 = { nounwind "amdgpu-work-group-id-x" "target-cpu"="fiji" }
 ; HSA: attributes #5 = { nounwind "amdgpu-work-group-id-y" "target-cpu"="fiji" }
 ; HSA: attributes #6 = { nounwind "amdgpu-work-group-id-z" "target-cpu"="fiji" }
 ; HSA: attributes #7 = { nounwind "amdgpu-dispatch-ptr" "target-cpu"="fiji" }
 ; HSA: attributes #8 = { nounwind "amdgpu-queue-ptr" "target-cpu"="fiji" }
 ; HSA: attributes #9 = { nounwind "amdgpu-dispatch-id" "target-cpu"="fiji" }
 ; HSA: attributes #10 = { nounwind "amdgpu-work-group-id-y" "amdgpu-work-group-id-z" "target-cpu"="fiji" }
 ; HSA: attributes #11 = { nounwind "target-cpu"="fiji" }
 ; HSA: attributes #12 = { nounwind "target-cpu"="gfx900" }
 ; HSA: attributes #13 = { nounwind "amdgpu-queue-ptr" "target-cpu"="gfx900" }
 ; HSA: attributes #14 = { nounwind "amdgpu-kernarg-segment-ptr" "target-cpu"="fiji" }
+; HSA: attributes #15 = { nounwind }
+; HSA: attributes #16 = { nounwind "amdgpu-flat-scratch" }