[NVPTX] Fix pointer type for short 32-bit pointers

Global variables used to be printed as u64/b64 even when
-nvptx-short-ptr is set.

Differential Revision: https://reviews.llvm.org/D127668

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

@@ -1365,8 +1365,11 @@ NVPTXAsmPrinter::getPTXFundamentalTypeStr(Type *Ty, bool useB4PTR) const {
     return "f32";
   case Type::DoubleTyID:
     return "f64";
-  case Type::PointerTyID:
-    if (static_cast<const NVPTXTargetMachine &>(TM).is64Bit())
+  case Type::PointerTyID: {
+    unsigned PtrSize = TM.getPointerSizeInBits(Ty->getPointerAddressSpace());
+    assert((PtrSize == 64 || PtrSize == 32) && "Unexpected pointer size");
+
+    if (PtrSize == 64)
       if (useB4PTR)
         return "b64";
       else
@@ -1375,6 +1378,7 @@ NVPTXAsmPrinter::getPTXFundamentalTypeStr(Type *Ty, bool useB4PTR) const {
       return "b32";
     else
       return "u32";
+  }
   default:
     break;
   }

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

@@ -0,0 +1,17 @@
+; 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: %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