[NVPTX] Fix integer overflow for alloca arrays

[AlexM]

Summary

Similar to D146767

[tra]

llvm/test/CodeGen/NVPTX/local-variable-big.ll

9

Such a large allocation will certainly cause all sorts of other issues -- we can't allocate that much on a GPU, as it would have to be done for thousands of threads simultaneously.
Also the stack register on NVIDIA GPUs is 32-bit.

PTXAS will accept large values, but appears to overflow internally, so the SASS it produces is broken. With valies more than 2GB (signed overflow?) it does not allocate anything on stack at all. With larger valies NVCC just crashes. https://godbolt.org/z/En716v36P

Ideally we should error out and tell the user that we can't produce sensible code for values that do not fit in 31 bits.

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

void NVPTXAsmPrinter::setAndEmitFunctionVirtualRegisters(
 
   // Map the global virtual register number to a register class specific
   // virtual register number starting from 1 with that class.
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
   //unsigned numRegClasses = TRI->getNumRegClasses();
 
   // Emit the Fake Stack Object
   const MachineFrameInfo &MFI = MF.getFrameInfo();
-  int NumBytes = (int) MFI.getStackSize();
+  int64_t NumBytes = MFI.getStackSize();
   if (NumBytes) {
     O << "\t.local .align " << MFI.getMaxAlign().value() << " .b8 \t"
       << DEPOTNAME << getFunctionNumber() << "[" << NumBytes << "];\n";
     if (static_cast<const NVPTXTargetMachine &>(MF.getTarget()).is64Bit()) {
       O << "\t.reg .b64 \t%SP;\n";
       O << "\t.reg .b64 \t%SPL;\n";
     } else {
       O << "\t.reg .b32 \t%SP;\n";

llvm/test/CodeGen/NVPTX/local-variable-big.ll

+; RUN: llc < %s -march=nvptx | FileCheck %s
+; RUN: llc < %s -march=nvptx64 | FileCheck %s
+; RUN: %if ptxas %{ llc < %s -march=nvptx | %ptxas-verify %}
+; RUN: %if ptxas %{ llc < %s -march=nvptx64 | %ptxas-verify %}
+
+; Make sure there is no integer overflow in the compiler when large
+; local arrays are used.
+;
+; CHECK: .local .align 1 .b8  __local_depot0[4831838208];
+;
+define i32 @large_data() {
+  %arr = alloca [4831838208 x i8]
+  %p = getelementptr inbounds [4831838208 x i8], [4831838208 x i8]* %arr, i64 0, i64 0
+  %call = call i32 @use(ptr %p)
+  ret i32 %call
+}
+
+declare i32 @use(ptr %p)