GPGPU: Pass subtree values correctly to the kernel

llvm-svn: 277697

2 files changed, 37 insertions, 6 deletions

diff --git a/polly/lib/CodeGen/PPCGCodeGeneration.cpp b/polly/lib/CodeGen/PPCGCodeGeneration.cpp
index 0d6e6ca078d..388ee8fdf05 100644
--- a/polly/lib/CodeGen/PPCGCodeGeneration.cpp
+++ b/polly/lib/CodeGen/PPCGCodeGeneration.cpp
@@ -221,12 +221,14 @@ private:
 
   /// Create kernel launch parameters.
   ///
-  /// @param Kernel The kernel to create parameters for.
-  /// @param F      The kernel function that has been created.
+  /// @param Kernel        The kernel to create parameters for.
+  /// @param F             The kernel function that has been created.
+  /// @param SubtreeValues The set of llvm::Values referenced by this kernel.
   ///
   /// @returns A stack allocated array with pointers to the parameter
   ///          values that are passed to the kernel.
-  Value *createLaunchParameters(ppcg_kernel *Kernel, Function *F);
+  Value *createLaunchParameters(ppcg_kernel *Kernel, Function *F,
+                                SetVector<Value *> SubtreeValues);
 
   /// Create GPU kernel.
   ///
@@ -879,8 +881,9 @@ GPUNodeBuilder::getBlockSizes(ppcg_kernel *Kernel) {
   return std::make_tuple(Sizes[0], Sizes[1], Sizes[2]);
 }
 
-Value *GPUNodeBuilder::createLaunchParameters(ppcg_kernel *Kernel,
-                                              Function *F) {
+Value *
+GPUNodeBuilder::createLaunchParameters(ppcg_kernel *Kernel, Function *F,
+                                       SetVector<Value *> SubtreeValues) {
   Type *ArrayTy = ArrayType::get(Builder.getInt8PtrTy(),
                                  std::distance(F->arg_begin(), F->arg_end()));
 
@@ -948,6 +951,19 @@ Value *GPUNodeBuilder::createLaunchParameters(ppcg_kernel *Kernel,
     Index++;
   }
 
+  for (auto Val : SubtreeValues) {
+    Instruction *Param = new AllocaInst(
+        Val->getType(), Launch + "_param_" + std::to_string(Index),
+        EntryBlock->getTerminator());
+    Builder.CreateStore(Val, Param);
+    Value *Slot = Builder.CreateGEP(
+        Parameters, {Builder.getInt64(0), Builder.getInt64(Index)});
+    Value *ParamTyped =
+        Builder.CreatePointerCast(Param, Builder.getInt8PtrTy());
+    Builder.CreateStore(ParamTyped, Slot);
+    Index++;
+  }
+
   auto Location = EntryBlock->getTerminator();
   return new BitCastInst(Parameters, Builder.getInt8PtrTy(),
                          Launch + "_params_i8ptr", Location);
@@ -1003,7 +1019,7 @@ void GPUNodeBuilder::createKernel(__isl_take isl_ast_node *KernelStmt) {
     S.invalidateScopArrayInfo(BasePtr, ScopArrayInfo::MK_Array);
   LocalArrays.clear();
 
-  Value *Parameters = createLaunchParameters(Kernel, F);
+  Value *Parameters = createLaunchParameters(Kernel, F, SubtreeValues);
 
   std::string ASMString = finalizeKernelFunction();
   std::string Name = "kernel_" + std::to_string(Kernel->id);
diff --git a/polly/test/GPGPU/scalar-parameter.ll b/polly/test/GPGPU/scalar-parameter.ll
index 38b38f4054e..7edc70b552c 100644
--- a/polly/test/GPGPU/scalar-parameter.ll
+++ b/polly/test/GPGPU/scalar-parameter.ll
@@ -2,6 +2,10 @@
 ; RUN: -disable-output < %s | \
 ; RUN: FileCheck -check-prefix=CODE %s
 
+; RUN: opt %loadPolly -polly-codegen-ppcg \
+; RUN: -S < %s | \
+; RUN: FileCheck -check-prefix=IR %s
+
 ; REQUIRES: pollyacc
 
 ; CODE: Code
@@ -301,6 +305,17 @@ bb7:                                              ; preds = %bb1
   ret void
 }
 
+; IR-LABEL: @i8
+
+; IR: %1 = call i8* @polly_getDevicePtr(i8* %p_dev_array_MemRef_A)
+; IR-NEXT: store i8* %1, i8** %polly_launch_0_param_0
+; IR-NEXT: %2 = getelementptr [2 x i8*], [2 x i8*]* %polly_launch_0_params, i64 0, i64 0
+; IR-NEXT: %3 = bitcast i8** %polly_launch_0_param_0 to i8*
+; IR-NEXT: store i8* %3, i8** %2
+; IR-NEXT: store i8 %b, i8* %polly_launch_0_param_1
+; IR-NEXT: %4 = getelementptr [2 x i8*], [2 x i8*]* %polly_launch_0_params, i64 0, i64 1
+; IR-NEXT: store i8* %polly_launch_0_param_1, i8** %4
+
 ; CODE: Code
 ; CODE-NEXT: ====
 ; CODE-NEXT: # host