[clang][CodeGen] Fix address space for sret

[yaxunl]

Summary

sret is returned through temporary variables allocated on stack,
therefore it should use alloca address space.

Currently clang use default address space for sret pointers. This
causes inefficient code generated for AMDGPU backend since
alloca address space is 32 bit whereas generic pointer is 64 bit.
It also causes assertions where alloca address space is expected.

This patch uses alloca address space for sret pointers. It is NFC
for targets alloca address space of which is default address space.

[yaxunl]

fix invalid bitcast due to sret pointer passed to ctor

[yaxunl]

cast return value to default address space since it is expected. also fix debug info

[rjmccall]

I think this is intentional; requiring the indirect-result parameter to be in the alloca address space would prevent direct initialization of non-temporary memory, which is an important optimization in C++.

[yaxunl]

In D101389#2724636, @rjmccall wrote:

I think this is intentional; requiring the indirect-result parameter to be in the alloca address space would prevent direct initialization of non-temporary memory, which is an important optimization in C++.

You mean situations like this? https://godbolt.org/z/KnPs6znK8

Address of a global variable is directly passed as the sret arg to the function returning a struct, instead of creating a temporary struct variable and passing its address as the sret arg?

However, this is forbidden in CUDA/HIP: https://godbolt.org/z/1hjsrn9Tn

So can we assume sret arg is always in alloca addr space for CUDA/HIP?

[rjmccall]

In D101389#2724700, @yaxunl wrote:

In D101389#2724636, @rjmccall wrote:

I think this is intentional; requiring the indirect-result parameter to be in the alloca address space would prevent direct initialization of non-temporary memory, which is an important optimization in C++.

You mean situations like this? https://godbolt.org/z/KnPs6znK8

Address of a global variable is directly passed as the sret arg to the function returning a struct, instead of creating a temporary struct variable and passing its address as the sret arg?

That's one specific case of copy-elision, but there are others, e.g.: https://godbolt.org/z/boTeMseaM

I suppose you could disable general copy-elision. But don't you have aggressive optimizations to strength-reduce generic pointer operations?

[yaxunl]

In D101389#2724757, @rjmccall wrote:

In D101389#2724700, @yaxunl wrote:

In D101389#2724636, @rjmccall wrote:

I think this is intentional; requiring the indirect-result parameter to be in the alloca address space would prevent direct initialization of non-temporary memory, which is an important optimization in C++.

You mean situations like this? https://godbolt.org/z/KnPs6znK8

Address of a global variable is directly passed as the sret arg to the function returning a struct, instead of creating a temporary struct variable and passing its address as the sret arg?

That's one specific case of copy-elision, but there are others, e.g.: https://godbolt.org/z/boTeMseaM

I suppose you could disable general copy-elision. But don't you have aggressive optimizations to strength-reduce generic pointer operations?

Thanks for the example. I agree that sret should be in default addr space.

clang/lib/CodeGen/CGCall.cpp

CodeGenTypes::GetFunctionType(const CGFunctionInfo &FI) {
 
   ClangToLLVMArgMapping IRFunctionArgs(getContext(), FI, true);
   SmallVector<llvm::Type*, 8> ArgTypes(IRFunctionArgs.totalIRArgs());
 
   // Add type for sret argument.
   if (IRFunctionArgs.hasSRetArg()) {
     QualType Ret = FI.getReturnType();
     llvm::Type *Ty = ConvertType(Ret);
-    unsigned AddressSpace = Context.getTargetAddressSpace(Ret);
+    unsigned AddressSpace =
+        Context.getTargetAddressSpace(CGM.getASTAllocaAddressSpace());
     ArgTypes[IRFunctionArgs.getSRetArgNo()] =
         llvm::PointerType::get(Ty, AddressSpace);
   }
 
   // Add type for inalloca argument.
   if (IRFunctionArgs.hasInallocaArg()) {
     auto ArgStruct = FI.getArgStruct();
     assert(ArgStruct);
if (!ReturnValue.isNull()) {
       SRetPtr = CreateMemTemp(RetTy, "tmp", &SRetAlloca);
       if (HaveInsertPoint() && ReturnValue.isUnused()) {
         uint64_t size =
             CGM.getDataLayout().getTypeAllocSize(ConvertTypeForMem(RetTy));
         UnusedReturnSizePtr = EmitLifetimeStart(size, SRetAlloca.getPointer());
       }
     }
     if (IRFunctionArgs.hasSRetArg()) {
-      IRCallArgs[IRFunctionArgs.getSRetArgNo()] = SRetPtr.getPointer();
+      IRCallArgs[IRFunctionArgs.getSRetArgNo()] =
+          getTargetHooks().performAddrSpaceCast(
+              *this, SRetPtr.getPointer(), LangAS::Default,
+              getASTAllocaAddressSpace(),
+              SRetPtr.getPointer()
+                  ->getType()
+                  ->getPointerElementType()
+                  ->getPointerTo(getContext().getTargetAddressSpace(
+                      getASTAllocaAddressSpace())),
+              /*non-null*/ true);
+
     } else if (RetAI.isInAlloca()) {
       Address Addr =
           Builder.CreateStructGEP(ArgMemory, RetAI.getInAllocaFieldIndex());
       Builder.CreateStore(SRetPtr.getPointer(), Addr);
     }
   }
 
   Address swiftErrorTemp = Address::invalid();
case ABIArgInfo::Direct: {
           Builder.CreateStore(errorValue, swiftErrorTemp);
         }
 
         // We might have to widen integers, but we should never truncate.
         if (ArgInfo.getCoerceToType() != V->getType() &&
             V->getType()->isIntegerTy())
           V = Builder.CreateZExt(V, ArgInfo.getCoerceToType());
 
-        // If the argument doesn't match, perform a bitcast to coerce it.  This
-        // can happen due to trivial type mismatches.
+        // If the argument doesn't match, perform a bitcast or addrspacecast
+        // to coerce it.  This can happen due to trivial type mismatches.
         if (FirstIRArg < IRFuncTy->getNumParams() &&
-            V->getType() != IRFuncTy->getParamType(FirstIRArg))
-          V = Builder.CreateBitCast(V, IRFuncTy->getParamType(FirstIRArg));
+            V->getType() != IRFuncTy->getParamType(FirstIRArg)) {
+          V = Builder.CreatePointerBitCastOrAddrSpaceCast(
+              V, IRFuncTy->getParamType(FirstIRArg));
+        }
 
         IRCallArgs[FirstIRArg] = V;
         break;
       }
 
       // FIXME: Avoid the conversion through memory if possible.
       Address Src = Address::invalid();
       if (!I->isAggregate()) {

clang/test/CodeGenCUDA/amdgpu-sret.cu

+// REQUIRES: amdgpu-registered-target
+
+// RUN: %clang_cc1 -triple amdgcn-amd-amdhsa -target-cpu gfx906 \
+// RUN:   -aux-triple x86_64-unknown-linux-gnu -fcuda-is-device \
+// RUN:   -emit-llvm -o - -x hip %s \
+// RUN:   | FileCheck %s
+
+// Check no assertion with debug info.
+
+// RUN: %clang_cc1 -triple amdgcn-amd-amdhsa -target-cpu gfx906 \
+// RUN:   -aux-triple x86_64-unknown-linux-gnu -fcuda-is-device \
+// RUN:   -S -o %t.s -x hip %s \
+// RUN:   -debug-info-kind=limited
+
+#include "Inputs/cuda.h"
+ 
+struct A {
+  int x[100];
+  __device__ A();
+};
+
+struct B {
+  int x[100];
+};
+
+__device__ B b;
+
+// CHECK-LABEL: @_Z5func1v(
+// CHECK-SAME: %struct.A addrspace(5)* noalias sret(%struct.A) align 4 %[[ARG:.*]])
+// CHECK: %[[CAST:.*]] = addrspacecast %struct.A addrspace(5)* %[[ARG]] to %struct.A*
+// CHECK: call void @_ZN1AC1Ev(%struct.A* nonnull dereferenceable(400) %[[CAST]])
+__device__ A func1() {
+  A a;
+  return a;
+}
+
+// CHECK-LABEL: @_Z6func1av(%struct.B addrspace(5)* noalias sret(%struct.B) align 4 
+__device__ B func1a() {
+  B b;
+  return b;
+}
+
+// Check returning the return value again.
+
+// CHECK-LABEL: @_Z5func2v(
+// CHECK-SAME: %struct.A addrspace(5)* noalias sret(%struct.A) align 4 %[[ARG:.*]])
+// CHECK: call void @_Z5func1v(%struct.A addrspace(5)* sret(%struct.A) align 4 %[[ARG]])
+__device__ A func2() {
+  A a = func1();
+  return a;
+}
+
+// Check assigning the return value to a global variable.
+
+// CHECK-LABEL: @_Z5func3v(
+// CHECK: %[[RET:.*]] = alloca %struct.B, align 4, addrspace(5)
+// CHECK: %[[CAST1:.*]] = addrspacecast %struct.B addrspace(5)* %[[RET]] to %struct.B*
+// CHECK: %[[CAST2:.*]] = addrspacecast %struct.B* %[[CAST1]] to %struct.B addrspace(5)*
+// CHECK: call void @_Z6func1av(%struct.B addrspace(5)* sret(%struct.B) align 4 %[[CAST2]]
+// CHECK: %[[CAST3:.*]] = bitcast %struct.B* %[[CAST1]] to i8*
+// CHECK: call void @llvm.memcpy.p0i8.p0i8.i64{{.*}}@b{{.*}}%[[CAST3]]
+__device__ void func3() {
+  b = func1a();
+}
+
+// Check assigning the return value to a temporary variable.
+
+// CHECK-LABEL: @_Z5func4v(
+// CHECK: %[[TMP:.*]] = alloca %struct.A, align 4, addrspace(5)
+// CHECK: %[[TMP_CAST1:.*]] = addrspacecast %struct.A addrspace(5)* %[[TMP]] to %struct.A*
+// CHECK: %[[RET:.*]] = alloca %struct.A, align 4, addrspace(5)
+// CHECK: %[[RET_CAST1:.*]] = addrspacecast %struct.A addrspace(5)* %[[RET]] to %struct.A*
+// CHECK: call void @_ZN1AC1Ev(%struct.A* nonnull dereferenceable(400) %[[TMP_CAST1]])
+// CHECK: %[[RET_CAST2:.*]] = addrspacecast %struct.A* %[[RET_CAST1]] to %struct.A addrspace(5)*
+// CHECK: call void @_Z5func1v(%struct.A addrspace(5)* sret(%struct.A) align 4 %[[RET_CAST2]]
+// CHECK: %[[TMP_CAST2:.*]] = bitcast %struct.A* %[[TMP_CAST1]] to i8*
+// CHECK: %[[RET_CAST3:.*]] = bitcast %struct.A* %[[RET_CAST1]] to i8*
+// CHECK: call void @llvm.memcpy.p0i8.p0i8.i64{{.*}}%[[TMP_CAST2]]{{.*}}%[[RET_CAST3]]
+__device__ void func4() {
+  A a;
+  a = func1();
+}