Use the same args size (without tail padding) for AMDGPU

Author: kevinsala

clang/lib/CodeGen/CGCUDANV.cpp

@@ -340,7 +340,6 @@ Address CGNVCUDARuntime::prepareKernelArgsLLVMOffload(CodeGenFunction &CGF,
 
   auto *Int64Ty = CGF.Builder.getInt64Ty();
   KernelLaunchParamsTypes.push_back(Int64Ty);
-  KernelLaunchParamsTypes.push_back(Int64Ty);
   KernelLaunchParamsTypes.push_back(PtrTy);
   KernelLaunchParamsTypes.push_back(PtrTy);
 
@@ -352,25 +351,21 @@ Address CGNVCUDARuntime::prepareKernelArgsLLVMOffload(CodeGenFunction &CGF,
       KernelLaunchParamsTy, CharUnits::fromQuantity(16),
       "kernel_launch_params");
 
-  auto KernelArgsSize = CGM.getDataLayout().getTypeAllocSize(KernelArgsTy);
-
-  // Avoid accounting the tail padding for CUDA.
-  auto KernelArgsSizeNoTailPadding = llvm::TypeSize::getZero();
+  // Avoid accounting the tail padding for the kernel arguments.
+  auto KernelArgsSize = llvm::TypeSize::getZero();
   if (auto N = KernelArgsTy->getNumElements()) {
     auto *SL = CGM.getDataLayout().getStructLayout(KernelArgsTy);
-    KernelArgsSizeNoTailPadding = SL->getElementOffset(N - 1);
-    KernelArgsSizeNoTailPadding += CGM.getDataLayout().getTypeAllocSize(
+    KernelArgsSize += SL->getElementOffset(N - 1);
+    KernelArgsSize += CGM.getDataLayout().getTypeAllocSize(
         KernelArgsTy->getElementType(N - 1));
   }
 
   CGF.Builder.CreateStore(llvm::ConstantInt::get(Int64Ty, KernelArgsSize),
                           CGF.Builder.CreateStructGEP(KernelLaunchParams, 0));
-  CGF.Builder.CreateStore(llvm::ConstantInt::get(Int64Ty, KernelArgsSizeNoTailPadding),
-                          CGF.Builder.CreateStructGEP(KernelLaunchParams, 1));
   CGF.Builder.CreateStore(KernelArgs.emitRawPointer(CGF),
-                          CGF.Builder.CreateStructGEP(KernelLaunchParams, 2));
+                          CGF.Builder.CreateStructGEP(KernelLaunchParams, 1));
   CGF.Builder.CreateStore(llvm::Constant::getNullValue(PtrTy),
-                          CGF.Builder.CreateStructGEP(KernelLaunchParams, 3));
+                          CGF.Builder.CreateStructGEP(KernelLaunchParams, 2));
 
   for (unsigned i = 0; i < Args.size(); ++i) {
     auto *ArgVal = CGF.Builder.CreateLoad(CGF.GetAddrOfLocalVar(Args[i]));

offload/include/Shared/APITypes.h

@@ -119,10 +119,8 @@ static_assert(sizeof(KernelArgsTy) ==
 
 /// Flat array of kernel launch parameters and their total size.
 struct KernelLaunchParamsTy {
-  /// Size of the Data array.
+  /// Size of the Data array without the tail padding.
   size_t Size = 0;
-  /// Size of the Data array without tail padding.
-  size_t SizeNoTailPadding = 0;
   /// Flat array of kernel parameters.
   void *Data = nullptr;
   /// Ptrs to the Data entries. Only strictly required for the host plugin.

offload/plugins-nextgen/amdgpu/src/rtl.cpp

@@ -3658,11 +3658,6 @@ Error AMDGPUKernelTy::launchImpl(GenericDeviceTy &GenericDevice,
                                  KernelArgsTy &KernelArgs,
                                  KernelLaunchParamsTy LaunchParams,
                                  AsyncInfoWrapperTy &AsyncInfoWrapper) const {
-  if (ArgsSize != LaunchParams.Size &&
-      ArgsSize > LaunchParams.Size + getImplicitArgsSize())
-    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
-                         "invalid kernel arguments size");
-
   AMDGPUPluginTy &AMDGPUPlugin =
       static_cast<AMDGPUPluginTy &>(GenericDevice.Plugin);
   AMDHostDeviceTy &HostDevice = AMDGPUPlugin.getHostDevice();

offload/plugins-nextgen/common/src/PluginInterface.cpp

@@ -573,7 +573,7 @@ KernelLaunchParamsTy GenericKernelTy::prepareArgs(
   }
 
   size_t ArgsSize = sizeof(void *) * NumArgs;
-  return KernelLaunchParamsTy{ArgsSize, ArgsSize, &Args[0], &Ptrs[0]};
+  return KernelLaunchParamsTy{ArgsSize, &Args[0], &Ptrs[0]};
 }
 
 uint32_t GenericKernelTy::getNumThreads(GenericDeviceTy &GenericDevice,

offload/plugins-nextgen/cuda/src/rtl.cpp

@@ -1414,7 +1414,7 @@ Error CUDAKernelTy::launchImpl(GenericDeviceTy &GenericDevice,
 
   void *Config[] = {CU_LAUNCH_PARAM_BUFFER_POINTER, LaunchParams.Data,
                     CU_LAUNCH_PARAM_BUFFER_SIZE,
-                    reinterpret_cast<void *>(&LaunchParams.SizeNoTailPadding),
+                    reinterpret_cast<void *>(&LaunchParams.Size),
                     CU_LAUNCH_PARAM_END};
 
   // If we are running an RPC server we want to wake up the server thread