[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();

clang/lib/CodeGen/CGDecl.cpp

if (EmitDebugInfo && HaveInsertPoint()) {
     Address DebugAddr = address;
     bool UsePointerValue = NRVO && ReturnValuePointer.isValid();
     DI->setLocation(D.getLocation());
 
     // If NRVO, use a pointer to the return address.
     if (UsePointerValue)
       DebugAddr = ReturnValuePointer;
 
-    (void)DI->EmitDeclareOfAutoVariable(&D, DebugAddr.getPointer(), Builder,
-                                        UsePointerValue);
+    // Local variables are casted to default address space if the alloca address
+    // space is different. Need to strip casts to get the real variables.
+    (void)DI->EmitDeclareOfAutoVariable(
+        &D, DebugAddr.getPointer()->stripPointerCasts(), Builder,
   }
 
   if (D.hasAttr<AnnotateAttr>() && HaveInsertPoint())
     EmitVarAnnotations(&D, address.getPointer());
 
   // Make sure we call @llvm.lifetime.end.
   if (emission.useLifetimeMarkers())
     EHStack.pushCleanup<CallLifetimeEnd>(NormalEHLifetimeMarker,

clang/lib/CodeGen/CodeGenFunction.cpp

if (RetTy->isVoidType()) {
     // Tell the epilog emitter to autorelease the result.  We do this
     // now so that various specialized functions can suppress it
     // during their IR-generation.
     if (getLangOpts().ObjCAutoRefCount &&
         !CurFnInfo->isReturnsRetained() &&
         RetTy->isObjCRetainableType())
       AutoreleaseResult = true;
   }
+  // Alloca address space may be different than default address space. Return
+  // value is expected to be in default address space. This is no-op if they
+  // are the same.
+  if (ReturnValue.isValid()) {
+    ReturnValue =
+        Address(getTargetHooks().performAddrSpaceCast(
+                    *this, ReturnValue.getPointer(), LangAS::Default,
+                    getASTAllocaAddressSpace(),
+                    ReturnValue.getPointer()
+                        ->getType()
+                        ->getPointerElementType()
+                        ->getPointerTo(getContext().getTargetAddressSpace(
+                            LangAS::Default)),
+                    /*non-null*/ true),
+                ReturnValue.getAlignment());
+  }
 
   EmitStartEHSpec(CurCodeDecl);
 
   PrologueCleanupDepth = EHStack.stable_begin();
 
   // Emit OpenMP specific initialization of the device functions.
   if (getLangOpts().OpenMP && CurCodeDecl)
     CGM.getOpenMPRuntime().emitFunctionProlog(*this, CurCodeDecl);

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 -debug-info-kind=limited \
+// 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;
+
+__device__ void callee(A *a);
+
+// CHECK-LABEL: @_Z5func1v(
+// CHECK-SAME: %struct.A addrspace(5)* noalias sret(%struct.A) align 4 %[[RET:.*]])
+// CHECK: %x = alloca [100 x i32], align 16, addrspace(5)
+// CHECK: %x.ascast = addrspacecast [100 x i32] addrspace(5)* %x to [100 x i32]*
+// CHECK: %p = alloca %struct.A*, align 8, addrspace(5)
+// CHECK: %p.ascast = addrspacecast %struct.A* addrspace(5)* %p to %struct.A**
+// CHECK: %[[RET_CAST:.*]] = addrspacecast %struct.A addrspace(5)* %[[RET]] to %struct.A*
+// CHECK: call void @llvm.dbg.declare(metadata %struct.A addrspace(5)* %[[RET]]
+// CHECK: call void @_ZN1AC1Ev(%struct.A* nonnull dereferenceable(400) %[[RET_CAST]])
+// CHECK: call void @llvm.dbg.declare(metadata [100 x i32] addrspace(5)* %x
+// CHECK: call void @_Z6calleeP1A(%struct.A* %[[RET_CAST]])
+// CHECK: %[[RET_CAST2:.*]] = bitcast %struct.A* %[[RET_CAST]] to i8*
+// CHECK: call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 4 %[[RET_CAST2]], i8* align 16 %{{.*}}, i64 400, i1 false)
+// CHECK: call void @llvm.dbg.declare(metadata %struct.A* addrspace(5)* %p
+// CHECK: store %struct.A* %[[RET_CAST]], %struct.A** %p.ascast
+__device__ A func1() {
+  A a;
+  int x[100];
+  callee(&a);
+  __builtin_memcpy(&a, x, 400);
+  A *p = &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 %[[RET:.*]])
+// CHECK: %[[CAST1:.*]] = addrspacecast %struct.A addrspace(5)* %[[RET]] to %struct.A*
+// CHECK: %[[CAST2:.*]] = addrspacecast %struct.A* %[[CAST1]] to %struct.A addrspace(5)*
+// CHECK: call void @_Z5func1v(%struct.A addrspace(5)* sret(%struct.A) align 4 %[[CAST2]])
+__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();
+}