AMDGPU: Error on stack size overflow

[rampitec]

LGTM

[arsenm]

r318189

lib/Target/AMDGPU/AMDGPUAsmPrinter.h

 class AMDGPUAsmPrinter final : public AsmPrinter {
 private:
   // Track resource usage for callee functions.
   struct SIFunctionResourceInfo {
     // Track the number of explicitly used VGPRs. Special registers reserved at
     // the end are tracked separately.
     int32_t NumVGPR = 0;
     int32_t NumExplicitSGPR = 0;
-    uint32_t PrivateSegmentSize = 0;
+    uint64_t PrivateSegmentSize = 0;
     bool UsesVCC = false;
     bool UsesFlatScratch = false;
     bool HasDynamicallySizedStack = false;
     bool HasRecursion = false;
 
     int32_t getTotalNumSGPRs(const SISubtarget &ST) const;
   };
 
   // Track resource usage for kernels / entry functions.
   struct SIProgramInfo {
     // Fields set in PGM_RSRC1 pm4 packet.
     uint32_t VGPRBlocks = 0;
     uint32_t SGPRBlocks = 0;
     uint32_t Priority = 0;
     uint32_t FloatMode = 0;
     uint32_t Priv = 0;
     uint32_t DX10Clamp = 0;
     uint32_t DebugMode = 0;
     uint32_t IEEEMode = 0;
-    uint32_t ScratchSize = 0;
+    uint64_t ScratchSize = 0;
 
     uint64_t ComputePGMRSrc1 = 0;
 
     // Fields set in PGM_RSRC2 pm4 packet.
     uint32_t LDSBlocks = 0;
     uint32_t ScratchBlocks = 0;
 
     uint64_t ComputePGMRSrc2 = 0;
private:
   /// can correctly setup the GPU state.
   void EmitProgramInfoR600(const MachineFunction &MF);
   void EmitProgramInfoSI(const MachineFunction &MF,
                          const SIProgramInfo &KernelInfo);
   void EmitPALMetadata(const MachineFunction &MF,
                        const SIProgramInfo &KernelInfo);
   void emitCommonFunctionComments(uint32_t NumVGPR,
                                   uint32_t NumSGPR,
-                                  uint32_t ScratchSize,
+                                  uint64_t ScratchSize,
                                   uint64_t CodeSize);
 
 public:
   explicit AMDGPUAsmPrinter(TargetMachine &TM,
                             std::unique_ptr<MCStreamer> Streamer);
 
   StringRef getPassName() const override;
 

lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

for (unsigned I = 0, E = Tuple->getNumOperands() & -2; I != E; I += 2) {
     PALMetadataMap[Key->getZExtValue()] = Val->getZExtValue();
   }
 }
 
 // Print comments that apply to both callable functions and entry points.
 void AMDGPUAsmPrinter::emitCommonFunctionComments(
   uint32_t NumVGPR,
   uint32_t NumSGPR,
-  uint32_t ScratchSize,
+  uint64_t ScratchSize,
   uint64_t CodeSize) {
   OutStreamer->emitRawComment(" codeLenInByte = " + Twine(CodeSize), false);
   OutStreamer->emitRawComment(" NumSgprs: " + Twine(NumSGPR), false);
   OutStreamer->emitRawComment(" NumVgprs: " + Twine(NumVGPR), false);
   OutStreamer->emitRawComment(" ScratchSize: " + Twine(ScratchSize), false);
 }
 
 bool AMDGPUAsmPrinter::runOnMachineFunction(MachineFunction &MF) {
if (!FrameInfo.hasCalls() && !FrameInfo.hasTailCall()) {
     Info.NumExplicitSGPR = HighestSGPRReg == AMDGPU::NoRegister ? 0 :
       TRI.getHWRegIndex(HighestSGPRReg) + 1;
 
     return Info;
   }
 
   int32_t MaxVGPR = -1;
   int32_t MaxSGPR = -1;
-  uint32_t CalleeFrameSize = 0;
+  uint64_t CalleeFrameSize = 0;
 
   for (const MachineBasicBlock &MBB : MF) {
     for (const MachineInstr &MI : MBB) {
       // TODO: Check regmasks? Do they occur anywhere except calls?
       for (const MachineOperand &MO : MI.operands()) {
         unsigned Width = 0;
         bool IsSGPR = false;
 
for (const MachineInstr &MI : MBB) {
           // conservative guesses.
 
           // 48 SGPRs - vcc, - flat_scr, -xnack
           int MaxSGPRGuess = 47 - getNumExtraSGPRs(ST, true,
                                                    ST.hasFlatAddressSpace());
           MaxSGPR = std::max(MaxSGPR, MaxSGPRGuess);
           MaxVGPR = std::max(MaxVGPR, 23);
 
-          CalleeFrameSize = std::max(CalleeFrameSize, 16384u);
+          CalleeFrameSize = std::max(CalleeFrameSize, UINT64_C(16384));
           Info.UsesVCC = true;
           Info.UsesFlatScratch = ST.hasFlatAddressSpace();
           Info.HasDynamicallySizedStack = true;
         } else {
           // We force CodeGen to run in SCC order, so the callee's register
           // usage etc. should be the cumulative usage of all callees.
           auto I = CallGraphResourceInfo.find(Callee);
           assert(I != CallGraphResourceInfo.end() &&
void AMDGPUAsmPrinter::getSIProgramInfo(SIProgramInfo &ProgInfo,
 
   ProgInfo.NumVGPR = Info.NumVGPR;
   ProgInfo.NumSGPR = Info.NumExplicitSGPR;
   ProgInfo.ScratchSize = Info.PrivateSegmentSize;
   ProgInfo.VCCUsed = Info.UsesVCC;
   ProgInfo.FlatUsed = Info.UsesFlatScratch;
   ProgInfo.DynamicCallStack = Info.HasDynamicallySizedStack || Info.HasRecursion;
 
+  if (!isUInt<32>(ProgInfo.ScratchSize)) {
+    DiagnosticInfoStackSize DiagStackSize(*MF.getFunction(),
+                                          ProgInfo.ScratchSize, DS_Error);
+    MF.getFunction()->getContext().diagnose(DiagStackSize);
+  }
+
   const SISubtarget &STM = MF.getSubtarget<SISubtarget>();
   const SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
   const SIInstrInfo *TII = STM.getInstrInfo();
   const SIRegisterInfo *RI = &TII->getRegisterInfo();
 
   unsigned ExtraSGPRs = getNumExtraSGPRs(STM,
                                          ProgInfo.VCCUsed,
                                          ProgInfo.FlatUsed);

test/CodeGen/AMDGPU/stack-size-overflow.ll

+; RUN: not llc -march=amdgcn < %s 2>&1 | FileCheck -check-prefix=ERROR %s
+; RUN: not llc -march=amdgcn < %s | FileCheck -check-prefix=GCN %s
+
+declare void @llvm.memset.p0i8.i32(i8* nocapture, i8, i32, i32, i1) #1
+
+; ERROR: error: stack size limit exceeded (4294967296) in stack_size_limit
+; GCN: ; ScratchSize: 4294967296
+define amdgpu_kernel void @stack_size_limit() #0 {
+entry:
+  %alloca = alloca [1073741823 x i32], align 4
+  %bc = bitcast [1073741823 x i32]* %alloca to i8*
+  call void @llvm.memset.p0i8.i32(i8* %bc, i8 9, i32 1073741823, i32 1, i1 true)
+  ret void
+}