From f92db6d3fff13bdacdf9b24660eb3f3158c58a17 Mon Sep 17 00:00:00 2001
From: Anshil Gandhi <Anshil.Gandhi@amd.com>
Date: Fri, 15 Oct 2021 15:16:32 -0600
Subject: [PATCH] [HIP] Relax conditions for address space cast in builtin args

Allow (implicit) address space casting between LLVM-equivalent
target address spaces.

Reviewed By: yaxunl, tra

Differential Revision: https://reviews.llvm.org/D111734
---
 clang/lib/Sema/SemaExpr.cpp                   | 10 +++++++---
 .../builtins-unsafe-atomics-gfx90a.cu         | 20 +++++++++++++++++++
 .../builtins-unsafe-atomics-gfx90a.cu         | 12 +++++++++++
 3 files changed, 39 insertions(+), 3 deletions(-)
 create mode 100644 clang/test/CodeGenCUDA/builtins-unsafe-atomics-gfx90a.cu
 create mode 100644 clang/test/SemaCUDA/builtins-unsafe-atomics-gfx90a.cu

diff --git a/clang/lib/Sema/SemaExpr.cpp b/clang/lib/Sema/SemaExpr.cpp
index 94b44714b530..472b15b9ea06 100644
--- a/clang/lib/Sema/SemaExpr.cpp
+++ b/clang/lib/Sema/SemaExpr.cpp
@@ -6545,9 +6545,13 @@ ExprResult Sema::BuildCallExpr(Scope *Scope, Expr *Fn, SourceLocation LParenLoc,
         auto ArgPtTy = ArgTy->getPointeeType();
         auto ArgAS = ArgPtTy.getAddressSpace();
 
-        // Only allow implicit casting from a non-default address space pointee
-        // type to a default address space pointee type
-        if (ArgAS != LangAS::Default || ParamAS == LangAS::Default)
+        // Add address space cast if target address spaces are different
+        bool NeedImplicitASC = 
+          ParamAS != LangAS::Default &&       // Pointer params in generic AS don't need special handling.
+          ( ArgAS == LangAS::Default  ||      // We do allow implicit conversion from generic AS 
+                                              // or from specific AS which has target AS matching that of Param.
+          getASTContext().getTargetAddressSpace(ArgAS) == getASTContext().getTargetAddressSpace(ParamAS));
+        if (!NeedImplicitASC)
           continue;
 
         // First, ensure that the Arg is an RValue.
diff --git a/clang/test/CodeGenCUDA/builtins-unsafe-atomics-gfx90a.cu b/clang/test/CodeGenCUDA/builtins-unsafe-atomics-gfx90a.cu
new file mode 100644
index 000000000000..d15953b3caca
--- /dev/null
+++ b/clang/test/CodeGenCUDA/builtins-unsafe-atomics-gfx90a.cu
@@ -0,0 +1,20 @@
+// RUN: %clang_cc1 -triple amdgcn-amd-amdhsa -target-cpu gfx90a -x hip \
+// RUN:  -aux-triple x86_64-unknown-linux-gnu -fcuda-is-device -emit-llvm %s \
+// RUN:  -o - | FileCheck %s
+
+#define __device__ __attribute__((device))
+typedef __attribute__((address_space(3))) float *LP;
+
+// CHECK-LABEL: test_ds_atomic_add_f32
+// CHECK: %[[ADDR_ADDR:.*]] = alloca float*, align 8, addrspace(5)
+// CHECK: %[[ADDR_ADDR_ASCAST_PTR:.*]] = addrspacecast float* addrspace(5)* %[[ADDR_ADDR]] to float**
+// CHECK: store float* %addr, float** %[[ADDR_ADDR_ASCAST_PTR]], align 8
+// CHECK: %[[ADDR_ADDR_ASCAST:.*]] = load float*, float** %[[ADDR_ADDR_ASCAST_PTR]], align 8
+// CHECK: %[[AS_CAST:.*]] = addrspacecast float* %[[ADDR_ADDR_ASCAST]] to float addrspace(3)*
+// CHECK: %3 = call contract float @llvm.amdgcn.ds.fadd.f32(float addrspace(3)* %[[AS_CAST]]
+// CHECK: %4 = load float*, float** %rtn.ascast, align 8
+// CHECK: store float %3, float* %4, align 4
+__device__ void test_ds_atomic_add_f32(float *addr, float val) {
+  float *rtn;
+  *rtn = __builtin_amdgcn_ds_faddf((LP)addr, val, 0, 0, 0);
+}
diff --git a/clang/test/SemaCUDA/builtins-unsafe-atomics-gfx90a.cu b/clang/test/SemaCUDA/builtins-unsafe-atomics-gfx90a.cu
new file mode 100644
index 000000000000..6f1484c68ec7
--- /dev/null
+++ b/clang/test/SemaCUDA/builtins-unsafe-atomics-gfx90a.cu
@@ -0,0 +1,12 @@
+// RUN: %clang_cc1 -triple amdgcn-amd-amdhsa -target-cpu gfx90a -x hip \
+// RUN:  -aux-triple x86_64-unknown-linux-gnu -fcuda-is-device %s \
+// RUN:  -fsyntax-only -verify
+// expected-no-diagnostics
+
+#define __device__ __attribute__((device))
+typedef __attribute__((address_space(3))) float *LP;
+
+__device__ void test_ds_atomic_add_f32(float *addr, float val) {
+  float *rtn;
+  *rtn = __builtin_amdgcn_ds_faddf((LP)addr, val, 0, 0, 0);
+}