From ce43e2f07477a166712b7ee4a89fbc65d2ed56ad Mon Sep 17 00:00:00 2001
From: Hugh Delaney <hugh.delaney@codeplay.com>
Date: Wed, 4 Jan 2023 11:52:08 -0800
Subject: [PATCH] [llvm][CUDA] Allow NVVMREflect to process OpenCL-specific
 __nvvm_reflect_ocl()

OpenCL requires constant string arguments to be in a particular address space,
so OpenCL sources can't use the regular `__nvvm_reflect()`.

Allow NVVMReflect pass to accept an Open_CL specific variant with a constant
string in a non-default address space.

Differential Revision: https://reviews.llvm.org/D139213
---
 llvm/lib/Target/NVPTX/NVVMReflect.cpp       |  5 ++++-
 llvm/test/CodeGen/NVPTX/nvvm-reflect-ocl.ll | 20 ++++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/NVPTX/nvvm-reflect-ocl.ll

diff --git a/llvm/lib/Target/NVPTX/NVVMReflect.cpp b/llvm/lib/Target/NVPTX/NVVMReflect.cpp
index 3f3c496..7ff5a5e 100644
--- a/llvm/lib/Target/NVPTX/NVVMReflect.cpp
+++ b/llvm/lib/Target/NVPTX/NVVMReflect.cpp
@@ -40,6 +40,7 @@
 #include <sstream>
 #include <string>
 #define NVVM_REFLECT_FUNCTION "__nvvm_reflect"
+#define NVVM_REFLECT_OCL_FUNCTION "__nvvm_reflect_ocl"
 
 using namespace llvm;
 
@@ -78,7 +79,8 @@ static bool runNVVMReflect(Function &F, unsigned SmVersion) {
   if (!NVVMReflectEnabled)
     return false;
 
-  if (F.getName() == NVVM_REFLECT_FUNCTION) {
+  if (F.getName() == NVVM_REFLECT_FUNCTION ||
+      F.getName() == NVVM_REFLECT_OCL_FUNCTION) {
     assert(F.isDeclaration() && "_reflect function should not have a body");
     assert(F.getReturnType()->isIntegerTy() &&
            "_reflect's return type should be integer");
@@ -119,6 +121,7 @@ static bool runNVVMReflect(Function &F, unsigned SmVersion) {
       continue;
     Function *Callee = Call->getCalledFunction();
     if (!Callee || (Callee->getName() != NVVM_REFLECT_FUNCTION &&
+                    Callee->getName() != NVVM_REFLECT_OCL_FUNCTION &&
                     Callee->getIntrinsicID() != Intrinsic::nvvm_reflect))
       continue;
 
diff --git a/llvm/test/CodeGen/NVPTX/nvvm-reflect-ocl.ll b/llvm/test/CodeGen/NVPTX/nvvm-reflect-ocl.ll
new file mode 100644
index 0000000..7a5a5a7
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/nvvm-reflect-ocl.ll
@@ -0,0 +1,20 @@
+; Verify that __nvvm_reflect_ocl() is replaced with an appropriate value
+;
+; RUN: opt %s -S -passes='default<O2>' -mtriple=nvptx64 \
+; RUN:   | FileCheck %s --check-prefixes=COMMON,SM20
+; RUN: opt %s -S -passes='default<O2>' -mtriple=nvptx64 -mcpu=sm_35 \
+; RUN:   | FileCheck %s --check-prefixes=COMMON,SM35
+
+@"$str" = private addrspace(4) constant [12 x i8] c"__CUDA_ARCH\00"
+
+declare i32 @__nvvm_reflect_ocl(ptr addrspace(4) noundef)
+
+; COMMON-LABEL: @foo
+define i32 @foo(float %a, float %b) {
+; COMMON-NOT: call i32 @__nvvm_reflect_ocl
+  %reflect = tail call i32 @__nvvm_reflect_ocl(ptr addrspace(4) noundef getelementptr inbounds ([12 x i8], [12 x i8] addrspace(4)* @"$str", i64 0, i64 0))
+; SM20: ret i32 200
+; SM35: ret i32 350
+  ret i32 %reflect
+}
+
-- 
2.7.4