FIXES

jdoerfert · jdoerfert · commit 038c7d57bfed · 2024-06-07T11:45:24.000-07:00
diff --git a/clang/lib/CodeGen/CGCUDANV.cpp b/clang/lib/CodeGen/CGCUDANV.cpp
@@ -15,6 +15,7 @@
 #include "CGCXXABI.h"
 #include "CodeGenFunction.h"
 #include "CodeGenModule.h"
+#include "clang/AST/CharUnits.h"
 #include "clang/AST/Decl.h"
 #include "clang/Basic/Cuda.h"
 #include "clang/CodeGen/CodeGenABITypes.h"
@@ -138,6 +139,9 @@ class CGNVCUDARuntime : public CGCUDARuntime {
     return DummyFunc;
   }
 
+  Address prepareKernelArgs(CodeGenFunction &CGF, FunctionArgList &Args);
+  Address prepareKernelArgsLLVMOffload(CodeGenFunction &CGF,
+                                       FunctionArgList &Args);
   void emitDeviceStubBodyLegacy(CodeGenFunction &CGF, FunctionArgList &Args);
   void emitDeviceStubBodyNew(CodeGenFunction &CGF, FunctionArgList &Args);
   std::string getDeviceSideName(const NamedDecl *ND) override;
@@ -322,12 +326,30 @@ void CGNVCUDARuntime::emitDeviceStub(CodeGenFunction &CGF,
     emitDeviceStubBodyLegacy(CGF, Args);
 }
 
-// CUDA 9.0+ uses new way to launch kernels. Parameters are packed in a local
-// array and kernels are launched using cudaLaunchKernel().
-void CGNVCUDARuntime::emitDeviceStubBodyNew(CodeGenFunction &CGF,
-                                            FunctionArgList &Args) {
-  // Build the shadow stack entry at the very start of the function.
+/// CUDA passes the arguments with a level of indirection. For example, a
+/// (void*, short, void*) is passed as {void **, short *, void **} to the launch
+/// function. For the LLVM/offload launch we flatten the arguments into the
+/// struct directly, thus pass {void *, short, void *}
+Address CGNVCUDARuntime::prepareKernelArgsLLVMOffload(CodeGenFunction &CGF,
+                                                      FunctionArgList &Args) {
+  SmallVector<llvm::Type *> ArgTypes;
+  for (auto &Arg : Args)
+    ArgTypes.push_back(CGF.ConvertTypeForMem(Arg->getType()));
+
+  llvm::StructType *ST = llvm::StructType::create(ArgTypes);
+  Address KernelArgs = CGF.CreateTempAllocaWithoutCast(
+      ST, CharUnits::fromQuantity(16), "kernel_args");
+
+  for (unsigned i = 0; i < Args.size(); ++i) {
+    auto *ArgVal = CGF.Builder.CreateLoad(CGF.GetAddrOfLocalVar(Args[i]));
+    CGF.Builder.CreateStore(ArgVal, CGF.Builder.CreateStructGEP(KernelArgs, i));
+  }
 
+  return KernelArgs;
+}
+
+Address CGNVCUDARuntime::prepareKernelArgs(CodeGenFunction &CGF,
+                                           FunctionArgList &Args) {
   // Calculate amount of space we will need for all arguments.  If we have no
   // args, allocate a single pointer so we still have a valid pointer to the
   // argument array that we can pass to runtime, even if it will be unused.
@@ -342,6 +364,17 @@ void CGNVCUDARuntime::emitDeviceStubBodyNew(CodeGenFunction &CGF,
         VoidVarPtr, CGF.Builder.CreateConstGEP1_32(
                         PtrTy, KernelArgs.emitRawPointer(CGF), i));
   }
+  return KernelArgs;
+}
+
+// CUDA 9.0+ uses new way to launch kernels. Parameters are packed in a local
+// array and kernels are launched using cudaLaunchKernel().
+void CGNVCUDARuntime::emitDeviceStubBodyNew(CodeGenFunction &CGF,
+                                            FunctionArgList &Args) {
+  // Build the shadow stack entry at the very start of the function.
+  Address KernelArgs = CGF.getLangOpts().OffloadViaLLVM
+                           ? prepareKernelArgsLLVMOffload(CGF, Args)
+                           : prepareKernelArgs(CGF, Args);
 
   llvm::BasicBlock *EndBlock = CGF.createBasicBlock("setup.end");
 
diff --git a/clang/lib/Driver/ToolChains/Clang.cpp b/clang/lib/Driver/ToolChains/Clang.cpp
@@ -1126,15 +1126,19 @@ void Clang::AddPreprocessingOptions(Compilation &C, const JobAction &JA,
   }
 
   if (Args.hasArg(options::OPT_foffload_via_llvm)) {
-    CmdArgs.push_back("-include");
+    // Add llvm_wrappers/* to our system include path.  This lets us wrap
+    // standard library headers and other headers.
     SmallString<128> P(D.ResourceDir);
     llvm::sys::path::append(P, "include");
-    llvm::sys::path::append(P, "openmp_wrappers");
+    llvm::sys::path::append(P, "llvm_offload_wrappers");
+    CmdArgs.push_back("-internal-isystem");
+    CmdArgs.push_back(Args.MakeArgString(P));
+
+    CmdArgs.push_back("-include");
     if (JA.isDeviceOffloading(Action::OFK_OpenMP))
-      llvm::sys::path::append(P, "__llvm_offload_device.h");
+      CmdArgs.push_back("__llvm_offload_device.h");
     else
-      llvm::sys::path::append(P, "__llvm_offload_host.h");
-    CmdArgs.push_back(Args.MakeArgString(P));
+      CmdArgs.push_back("__llvm_offload_host.h");
   }
 
   // Add -i* options, and automatically translate to
diff --git a/clang/lib/Headers/CMakeLists.txt b/clang/lib/Headers/CMakeLists.txt
@@ -323,9 +323,12 @@ set(openmp_wrapper_files
   openmp_wrappers/__clang_openmp_device_functions.h
   openmp_wrappers/complex_cmath.h
   openmp_wrappers/new
-  openmp_wrappers/__llvm_offload.h
-  openmp_wrappers/__llvm_offload_host.h
-  openmp_wrappers/__llvm_offload_device.h
+)
+
+set(llvm_offload_wrapper_files
+  llvm_offload_wrappers/__llvm_offload.h
+  llvm_offload_wrappers/__llvm_offload_host.h
+  llvm_offload_wrappers/__llvm_offload_device.h
 )
 
 set(llvm_libc_wrapper_files
@@ -378,7 +381,7 @@ endfunction(clang_generate_header)
 # Copy header files from the source directory to the build directory
 foreach( f ${files} ${cuda_wrapper_files} ${cuda_wrapper_bits_files}
            ${ppc_wrapper_files} ${openmp_wrapper_files} ${zos_wrapper_files} ${hlsl_files}
-           ${llvm_libc_wrapper_files})
+	   ${llvm_libc_wrapper_files} ${llvm_offload_wrapper_files})
   copy_header_to_output_dir(${CMAKE_CURRENT_SOURCE_DIR} ${f})
 endforeach( f )
 
@@ -504,6 +507,7 @@ add_header_target("hlsl-resource-headers" ${hlsl_files})
 add_header_target("opencl-resource-headers" ${opencl_files})
 add_header_target("llvm-libc-resource-headers" ${llvm_libc_wrapper_files})
 add_header_target("openmp-resource-headers" ${openmp_wrapper_files})
+add_header_target("llvm-offload-resource-headers" ${llvm_libc_wrapper_files})
 add_header_target("windows-resource-headers" ${windows_only_files})
 add_header_target("utility-resource-headers" ${utility_files})
 
@@ -545,6 +549,11 @@ install(
   DESTINATION ${header_install_dir}/openmp_wrappers
   COMPONENT clang-resource-headers)
 
+install(
+  FILES ${llvm_offload_wrapper_files}
+  DESTINATION ${header_install_dir}/llvm_offload_wrappers
+  COMPONENT clang-resource-headers)
+
 install(
   FILES ${zos_wrapper_files}
   DESTINATION ${header_install_dir}/zos_wrappers
@@ -707,8 +716,8 @@ install(
   COMPONENT openmp-resource-headers)
 
 install(
-  FILES ${openmp_wrapper_files}
-  DESTINATION ${header_install_dir}/openmp_wrappers
+  FILES ${llvm_offload_wrapper_files}
+  DESTINATION ${header_install_dir}/llvm_offload_wrappers
   EXCLUDE_FROM_ALL
   COMPONENT openmp-resource-headers)
 
diff --git a/clang/lib/Headers/llvm_offload_wrappers/__llvm_offload.h b/clang/lib/Headers/llvm_offload_wrappers/__llvm_offload.h
diff --git a/clang/lib/Headers/llvm_offload_wrappers/__llvm_offload_device.h b/clang/lib/Headers/llvm_offload_wrappers/__llvm_offload_device.h
diff --git a/clang/lib/Headers/llvm_offload_wrappers/__llvm_offload_host.h b/clang/lib/Headers/llvm_offload_wrappers/__llvm_offload_host.h
diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -3268,13 +3268,13 @@ Error AMDGPUKernelTy::launchImpl(GenericDeviceTy &GenericDevice,
                                  uint32_t NumThreads, uint64_t NumBlocks,
                                  KernelArgsTy &KernelArgs, void *Args,
                                  AsyncInfoWrapperTy &AsyncInfoWrapper) const {
+  uint32_t KernelArgsSize = KernelArgs.NumArgs * sizeof(void *);
+
   if (KernelArgs.Flags.IsCUDA) {
-    // For CUDA kernels we compute the number of arguments here.
-    KernelArgs.NumArgs = (ArgsSize - ImplicitArgsSize) / sizeof(void *);
+    // For CUDA kernels we compute the kernel argument size explicitly.
+    KernelArgsSize = ArgsSize - ImplicitArgsSize;
   }
 
-  const uint32_t KernelArgsSize = KernelArgs.NumArgs * sizeof(void *);
-
   if (ArgsSize < KernelArgsSize)
     return Plugin::error("Mismatch of kernel arguments size");
 
@@ -3315,14 +3315,8 @@ Error AMDGPUKernelTy::launchImpl(GenericDeviceTy &GenericDevice,
   // Copy the explicit arguments.
   // TODO: We should expose the args memory manager alloc to the common part as
   // 	   alternative to copying them twice.
-  if (KernelArgs.NumArgs && !KernelArgs.Flags.IsCUDA) {
-    std::memcpy(AllArgs, *static_cast<void **>(Args),
-                sizeof(void *) * KernelArgs.NumArgs);
-  } else {
-    for (uint32_t I = 0; I < KernelArgs.NumArgs; ++I)
-      std::memcpy(advanceVoidPtr(AllArgs, sizeof(void *) * I),
-                  static_cast<void **>(Args)[I], sizeof(void *));
-  }
+  if (KernelArgsSize)
+    std::memcpy(AllArgs, *static_cast<void **>(Args), KernelArgsSize);
 
   AMDGPUDeviceTy &AMDGPUDevice = static_cast<AMDGPUDeviceTy &>(GenericDevice);
 
diff --git a/offload/src/KernelLanguage/API.cpp b/offload/src/KernelLanguage/API.cpp
@@ -58,7 +58,7 @@ int __tgt_target_kernel(void *Loc, int64_t DeviceId, int32_t NumTeams,
                         KernelArgsTy *Args);
 
 unsigned llvmLaunchKernel(const void *func, dim3 gridDim, dim3 blockDim,
-                          void **args, size_t sharedMem, void *stream) {
+                          void *args, size_t sharedMem, void *stream) {
   KernelArgsTy Args = {};
   Args.DynCGroupMem = sharedMem;
   Args.NumTeams[0] = gridDim.x;
@@ -67,7 +67,7 @@ unsigned llvmLaunchKernel(const void *func, dim3 gridDim, dim3 blockDim,
   Args.ThreadLimit[0] = blockDim.x;
   Args.ThreadLimit[1] = blockDim.y;
   Args.ThreadLimit[2] = blockDim.z;
-  Args.ArgPtrs = args;
+  Args.ArgPtrs = &args;
   Args.Flags.IsCUDA = true;
   int rv = __tgt_target_kernel(nullptr, 0, gridDim.x,
                                blockDim.x, func, &Args);
diff --git a/offload/test/lit.cfg b/offload/test/lit.cfg
@@ -66,7 +66,7 @@ def evaluate_bool_env(env):
 config.name = 'libomptarget :: ' + config.libomptarget_current_target
 
 # suffixes: A list of file extensions to treat as test files.
-config.suffixes = ['.c', '.cpp', '.cc', '.f90']
+config.suffixes = ['.c', '.cpp', '.cc', '.f90', '.cu']
 
 # excludes: A list of directories to exclude from the testuites.
 config.excludes = ['Inputs']
diff --git a/offload/test/offloading/CUDA/basic_launch_blocks_and_threads.cu b/offload/test/offloading/CUDA/basic_launch_blocks_and_threads.cu
@@ -22,7 +22,7 @@ int main(int argc, char **argv) {
   int *Ptr = reinterpret_cast<int *>(llvm_omp_target_alloc_shared(4, DevNo));
   *Ptr = 0;
   printf("Ptr %p, *Ptr: %i\n", Ptr, *Ptr);
-  // CHECK: Ptr [[Ptr:0x.*]], *Ptr: 7
+  // CHECK: Ptr [[Ptr:0x.*]], *Ptr: 0
   square<<<7, 6>>>(Ptr);
   printf("Ptr %p, *Ptr: %i\n", Ptr, *Ptr);
   // CHECK: Ptr [[Ptr]], *Ptr: 42
diff --git a/offload/test/offloading/CUDA/basic_launch_multi_arg.cu b/offload/test/offloading/CUDA/basic_launch_multi_arg.cu
@@ -13,8 +13,10 @@ void *llvm_omp_target_alloc_shared(size_t Size, int DeviceNum);
 void llvm_omp_target_free_shared(void *DevicePtr, int DeviceNum);
 }
 
-__global__ void square(int *Dst, int *Src, short Q, short P) {
+__global__ void square(int *Dst, short Q, int *Src, short P) {
   *Dst = (Src[0] + Src[1]) * (Q + P);
+  Src[0] = Q;
+  Src[1] = P;
 }
 
 int main(int argc, char **argv) {
@@ -25,9 +27,11 @@ int main(int argc, char **argv) {
   Src[0] = -2;
   Src[1] = 8;
   printf("Ptr %p, *Ptr: %i\n", Ptr, *Ptr);
+  printf("%i : %i\n", Src[0], Src[1]);
   // CHECK: Ptr [[Ptr:0x.*]], *Ptr: 7
-  square<<<1, 1>>>(Ptr, Src, 3, 4);
+  square<<<1, 1>>>(Ptr, 3, Src, 4);
   printf("Ptr %p, *Ptr: %i\n", Ptr, *Ptr);
+  printf("%i : %i\n", Src[0], Src[1]);
   // CHECK: Ptr [[Ptr]], *Ptr: 42
   llvm_omp_target_free_shared(Ptr, DevNo);
 }
