[NVPTX] Fix pointer argument declaration for --nvptx-short-ptr

[asavonic]

Summary

When --nvptx-short-ptr is set, local pointers are 32-bit on nvptx64 target.

Before this patch, arguments for a function declaration were always
emitted as b64 regardless of their address space, but they were correctly
set as b32 for the corresponding call instruction:

.extern .func test
(
 .param .b64 test_param_0
)
[...]
 .param .b32 param0;
 st.param.b32 [param0+0], %r1;
 call.uni test, (param0);

This is not supported:

ptxas: Type of argument does not match formal parameter
'test_param_0'

Now short pointers in a function declaration are emitted as b32 if
--nvptx-short-ptr is set.

This patch fixes the bug found in https://reviews.llvm.org/D127668#3587241.

[tra]

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1544

Should there be an else PTySizeInBits = TLI->getPointerTy(DL).getSizeInBits() ?

1592–1593

assert(PTySizeInBits) ?

[asavonic]

• Applied code review comments

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1544

We only use PTySizeInBits when PTy is set. I've now changed the condition below to make it clear.

[tra]

LGTM with a nit.

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1594

The assertion should probably be moved to where we set PTySizeInBits, so it covers if(isKernelFunc) case, too.

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

void NVPTXAsmPrinter::emitFunctionParamList(const Function *F, raw_ostream &O) {
   const auto *TLI = cast<NVPTXTargetLowering>(STI.getTargetLowering());
 
   Function::const_arg_iterator I, E;
   unsigned paramIndex = 0;
   bool first = true;
   bool isKernelFunc = isKernelFunction(*F);
   bool isABI = (STI.getSmVersion() >= 20);
   bool hasImageHandles = STI.hasImageHandles();
-  MVT thePointerTy = TLI->getPointerTy(DL);
 
   if (F->arg_empty()) {
     O << "()\n";
     return;
   }
 
   O << "(\n";
 
if (!PAL.hasParamAttr(paramIndex, Attribute::ByVal)) {
         O << "\t.param .align " << OptimalAlign.value() << " .b8 ";
         printParamName(I, paramIndex, O);
         O << "[" << DL.getTypeAllocSize(Ty) << "]";
 
         continue;
       }
       // Just a scalar
       auto *PTy = dyn_cast<PointerType>(Ty);
+      unsigned PTySizeInBits = 0;
+      if (PTy) {
+        PTySizeInBits =
+            TLI->getPointerTy(DL, PTy->getAddressSpace()).getSizeInBits();
+        assert(PTySizeInBits && "Invalid pointer size");
+      }
+
       if (isKernelFunc) {
         if (PTy) {
           // Special handling for pointer arguments to kernel
-          O << "\t.param .u" << thePointerTy.getSizeInBits() << " ";
+          O << "\t.param .u" << PTySizeInBits << " ";
 
           if (static_cast<NVPTXTargetMachine &>(TM).getDrvInterface() !=
               NVPTX::CUDA) {
             int addrSpace = PTy->getAddressSpace();
             switch (addrSpace) {
             default:
               O << ".ptr ";
               break;
if (!PAL.hasParamAttr(paramIndex, Attribute::ByVal)) {
         printParamName(I, paramIndex, O);
         continue;
       }
       // Non-kernel function, just print .param .b<size> for ABI
       // and .reg .b<size> for non-ABI
       unsigned sz = 0;
       if (isa<IntegerType>(Ty)) {
         sz = cast<IntegerType>(Ty)->getBitWidth();
         sz = promoteScalarArgumentSize(sz);
-      } else if (isa<PointerType>(Ty))
-        sz = thePointerTy.getSizeInBits();
-      else if (Ty->isHalfTy())
+      } else if (PTy) {
+        assert(PTySizeInBits && "Invalid pointer size");
+        sz = PTySizeInBits;
+      } else if (Ty->isHalfTy())
         // PTX ABI requires all scalar parameters to be at least 32
         // bits in size.  fp16 normally uses .b16 as its storage type
         // in PTX, so its size must be adjusted here, too.
         sz = 32;
       else
         sz = Ty->getPrimitiveSizeInBits();
       if (isABI)
         O << "\t.param .b" << sz << " ";

llvm/test/CodeGen/NVPTX/short-ptr.ll

 ; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 | FileCheck %s --check-prefix CHECK-DEFAULT
 ; RUN: llc < %s -march=nvptx -mcpu=sm_20 | FileCheck %s --check-prefix CHECK-DEFAULT-32
-; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -nvptx-short-ptr | FileCheck %s --check-prefixes CHECK-SHORT-SHARED,CHECK-SHORT-CONST
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -nvptx-short-ptr | FileCheck %s --check-prefixes CHECK-SHORT-SHARED,CHECK-SHORT-CONST,CHECK-SHORT-LOCAL
 
 ; RUN: %if ptxas %{ llc < %s -march=nvptx -mcpu=sm_20 | %ptxas-verify %}
 ; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_20 | %ptxas-verify %}
 ; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_20 -nvptx-short-ptr | %ptxas-verify %}
 
 ; CHECK-DEFAULT: .visible .shared .align 8 .u64 s
 ; CHECK-DEFAULT-32: .visible .shared .align 8 .u32 s
 ; CHECK-SHORT-SHARED: .visible .shared .align 8 .u32 s
 @s = local_unnamed_addr addrspace(3) global i32 addrspace(3)* null, align 8
 
 ; CHECK-DEFAULT: .visible .const .align 8 .u64 c
 ; CHECK-DEFAULT-32: .visible .const .align 8 .u32 c
 ; CHECK-SHORT-CONST: .visible .const .align 8 .u32 c
 @c = local_unnamed_addr addrspace(4) global i32 addrspace(4)* null, align 8
+
+declare void @use(i8 %arg);
+
+; CHECK-DEFAULT: .param .b64 test1_param_0
+; CHECK-DEFAULT-32: .param .b32 test1_param_0
+; CHECK-SHORT-LOCAL: .param .b32 test1_param_0
+define void @test1(i8 addrspace(5)* %local) {
+  ; CHECK-DEFAULT: ld.param.u64 %rd{{.*}}, [test1_param_0];
+  ; CHECK-DEFAULT-32:  ld.param.u32 %r{{.*}}, [test1_param_0];
+  ; CHECK-SHORT-LOCAL: ld.param.u32 %r{{.*}}, [test1_param_0];
+  %v = load i8, i8 addrspace(5)* %local
+  call void @use(i8 %v)
+  ret void
+}
+
+define void @test2() {
+  %v = alloca i8
+  %cast = addrspacecast i8* %v to i8 addrspace(5)*
+  ; CHECK-DEFAULT: .param .b64 param0;
+  ; CHECK-DEFAULT: st.param.b64
+  ; CHECK-DEFAULT-32: .param .b32 param0;
+  ; CHECK-DEFAULT-32: st.param.b32
+  ; CHECK-SHORT-LOCAL: .param .b32 param0;
+  ; CHECK-SHORT-LOCAL: st.param.b32
+  call void @test1(i8 addrspace(5)* %cast)
+  ret void
+}