[HIP] Fix pointer type kernel arg for amdgpu

[yaxunl]

Summary

amdgpu target prefers pointer type kernel arg in default address space
to be coerced to device address space for better performance.

This patch fixes that.

[yaxunl]

add a test for non-kernel function.

[tra]

clang/lib/CodeGen/CGCall.cpp

1172

What's the purpose of this change?

clang/test/CodeGenCUDA/amdgpu-kernel-arg-pointer-type.cu

16–17

Interesting. Clang used to crash on explicit address space attribute in the past. I'm glad it works now.

[yaxunl]

BTW Michael Liao will create another patch which will supersede this patch. That patch contains similar changes and also handles pointers in a byval struct or array.

clang/lib/CodeGen/CGCall.cpp

1172

this is supposed to cast the coerced type to the original type. Previously, no coerced type changes the address space, therefore a bitcast is sufficient. Now we coerce a type to different address space, therefore we need a pointer cast.

[yaxunl]

Please review Michael's patch https://reviews.llvm.org/D69826 which supersedes this one. Thanks.

clang/lib/CodeGen/CGCall.cpp

static llvm::Value *CoerceIntOrPtrToIntOrPtr(llvm::Value *Val,
                                              llvm::Type *Ty,
                                              CodeGenFunction &CGF) {
   if (Val->getType() == Ty)
     return Val;
 
   if (isa<llvm::PointerType>(Val->getType())) {
     // If this is Pointer->Pointer avoid conversion to and from int.
     if (isa<llvm::PointerType>(Ty))
-      return CGF.Builder.CreateBitCast(Val, Ty, "coerce.val");
+      return CGF.Builder.CreatePointerCast(Val, Ty, "coerce.val");
 
     // Convert the pointer to an integer so we can play with its width.
     Val = CGF.Builder.CreatePtrToInt(Val, CGF.IntPtrTy, "coerce.val.pi");
   }
 
   llvm::Type *DestIntTy = Ty;
   if (isa<llvm::PointerType>(DestIntTy))
     DestIntTy = CGF.IntPtrTy;

clang/lib/CodeGen/TargetInfo.cpp

ABIArgInfo AMDGPUABIInfo::classifyKernelArgumentType(QualType Ty) const {
   Ty = useFirstFieldIfTransparentUnion(Ty);
 
   // TODO: Can we omit empty structs?
 
   // Coerce single element structs to its element.
   if (const Type *SeltTy = isSingleElementStruct(Ty, getContext()))
     return ABIArgInfo::getDirect(CGT.ConvertType(QualType(SeltTy, 0)));
 
+  // Coerce pointer type kernel arguments in default address space to
+  // device address space for HIP.
+  QualType PointeeTy;
+  if (getContext().getLangOpts().HIP) {
+    if (auto *PT = Ty->getAs<PointerType>()) {
+      if (PT->getPointeeType().getAddressSpace() == LangAS::Default) {
+        PointeeTy = PT->getPointeeType();
+      }
+    } else if (auto *RT = Ty->getAs<ReferenceType>()) {
+      if (RT->getPointeeType().getAddressSpace() == LangAS::Default) {
+        PointeeTy = RT->getPointeeType();
+      }
+    }
+
+    if (PointeeTy != QualType()) {
+      return ABIArgInfo::getDirect(
+        CGT.ConvertType(PointeeTy)
+            ->getPointerTo(
+                getContext().getTargetAddressSpace(LangAS::cuda_device)));
+    }
+  }
   // If we set CanBeFlattened to true, CodeGen will expand the struct to its
   // individual elements, which confuses the Clover OpenCL backend; therefore we
   // have to set it to false here. Other args of getDirect() are just defaults.
   return ABIArgInfo::getDirect(nullptr, 0, nullptr, false);
 }
 
 ABIArgInfo AMDGPUABIInfo::classifyArgumentType(QualType Ty,
                                                unsigned &NumRegsLeft) const {

clang/test/CodeGenCUDA/amdgpu-kernel-arg-pointer-type.cu

+// RUN: %clang_cc1 -triple amdgcn-amd-amdhsa -fcuda-is-device \
+// RUN:     -emit-llvm -x hip %s -o - | FileCheck %s
+#include "Inputs/cuda.h"
+// CHECK: define amdgpu_kernel void  @_Z7kernel1Pi(i32 addrspace(1)* %x.coerce)
+__global__ void kernel1(int *x) {
+  x[0]++;
+}
+
+// CHECK: define amdgpu_kernel void  @_Z7kernel2Ri(i32 addrspace(1)* dereferenceable(4) %x.coerce)
+__global__ void kernel2(int &x) {
+  x++;
+}
+
+// CHECK: define amdgpu_kernel void  @_Z7kernel3PU3AS2iPU3AS1i(i32 addrspace(2)* %x, i32 addrspace(1)* %y)
+__global__ void kernel3(__attribute__((address_space(2))) int *x,
+                        __attribute__((address_space(1))) int *y) {
+  y[0] = x[0];
+}