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

[asavonic]

Summary

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

[asavonic]

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

18

Ptxas doesn't like global variables in local address space, so I will remove this case.
ptxas error: Module-scoped variables in .local state space are not allowed with ABI

[asavonic]

• Added ptxas RUN lines.
• Removed test case with .local global variable because ptxas does not support them.

[tra]

Good catch. Thank you for fixing this.

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

18

We could use __const__ AS instead.

For testing local vars we should probably add a function with a local variable.

[asavonic]

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

18

For testing local vars we should probably add a function with a local variable.

Local variables are lowered as __local_depot followed by as a sequence of casts:

	mov.u64 	%SPL, __local_depot0;
	cvta.local.u64 	%SP, %SPL;
	add.u64 	%rd1, %SP, 0;
	{ .reg .b64 %tmp;
	  cvta.to.local.u64 	%tmp, %rd1;
	  cvt.u32.u64 	%r1, %tmp;             <--- r1 is our 32-bit local pointer
        }

So I added a test, but it fails because of a different issue:

declare void @use(i8 addrspace(5)* %local)
define void @test() {
  %v = alloca i8
  %cast = addrspacecast i8* %v to i8 addrspace(5)*
  call void @use(i8 addrspace(5)* %cast)
  ret void
}

PTX:
.extern .func use
(
	.param .b64 use_param_0   <--- pointer is 64-bit instead of 32
)
[...]
.param .b32 param0;
st.param.b32 [param0+0], %r1;    <--- call assumes a 32-bit pointer
call.uni 
use, 
(
param0
);

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

So either the call should pass the pointer as 64-bit, or the use function declaration is wrong. I'm not sure. Should --nvptx-short-ptr change function ABI?

[tra]

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

13

BTW, the global variables themselves can all be in the global AS. we only care about the *type* of these variables which is what we're testing -- how we lower pointers in local/const/shared AS.

https://godbolt.org/z/7rx1jo18v

@s = local_unnamed_addr addrspace(1) global i32 addrspace(3)* null, align 8
@c = local_unnamed_addr addrspace(1) global i32 addrspace(4)* null, align 8
@l = local_unnamed_addr addrspace(1) global i32 addrspace(5)* null, align 8

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

NVPTXAsmPrinter::getPTXFundamentalTypeStr(Type *Ty, bool useB4PTR) const {
   }
   case Type::HalfTyID:
     // fp16 is stored as .b16 for compatibility with pre-sm_53 PTX assembly.
     return "b16";
   case Type::FloatTyID:
     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
         return "u64";
     else if (useB4PTR)
       return "b32";
     else
       return "u32";
+  }
   default:
     break;
   }
   llvm_unreachable("unexpected type");
 }
 
 void NVPTXAsmPrinter::emitPTXGlobalVariable(const GlobalVariable *GVar,
                                             raw_ostream &O,

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,CHECK-SHORT-LOCAL
+
+; 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
+
+; CHECK-DEFAULT: .visible .local .align 8 .u64 l
+; CHECK-DEFAULT-32: .visible .local .align 8 .u32 l
+; CHECK-SHORT-LOCAL: .visible .local .align 8 .u32 l
+@l = local_unnamed_addr addrspace(5) global i32 addrspace(5)* null, align 8