GPGPU: Add support for shared memory

llvm-svn: 277721

1 files changed, 90 insertions, 5 deletions

diff --git a/polly/lib/CodeGen/PPCGCodeGeneration.cpp b/polly/lib/CodeGen/PPCGCodeGeneration.cpp
index 46deae38a8a..751925a29a9 100644
--- a/polly/lib/CodeGen/PPCGCodeGeneration.cpp
+++ b/polly/lib/CodeGen/PPCGCodeGeneration.cpp
@@ -76,6 +76,10 @@ static cl::opt<bool> FastMath("polly-acc-fastmath",
                               cl::desc("Allow unsafe math optimizations"),
                               cl::Hidden, cl::init(false), cl::ZeroOrMore,
                               cl::cat(PollyCategory));
+static cl::opt<bool> SharedMemory("polly-acc-use-shared",
+                                  cl::desc("Use shared memory"), cl::Hidden,
+                                  cl::init(false), cl::ZeroOrMore,
+                                  cl::cat(PollyCategory));
 
 static cl::opt<std::string>
     CudaVersion("polly-acc-cuda-version",
@@ -155,6 +159,9 @@ private:
   /// The current GPU context.
   Value *GPUContext;
 
+  /// The set of isl_ids allocated in the kernel
+  std::vector<isl_id *> KernelIds;
+
   /// A module containing GPU code.
   ///
   /// This pointer is only set in case we are currently generating GPU code.
@@ -230,6 +237,14 @@ private:
   Value *createLaunchParameters(ppcg_kernel *Kernel, Function *F,
                                 SetVector<Value *> SubtreeValues);
 
+  /// Create declarations for kernel variable.
+  ///
+  /// This includes shared memory declarations.
+  ///
+  /// @param Kernel        The kernel definition to create variables for.
+  /// @param FN            The function into which to generate the variables.
+  void createKernelVariables(ppcg_kernel *Kernel, Function *FN);
+
   /// Create GPU kernel.
   ///
   /// Code generate the kernel described by @p KernelStmt.
@@ -280,6 +295,11 @@ private:
   /// @param The kernel to generate the intrinsic functions for.
   void insertKernelIntrinsics(ppcg_kernel *Kernel);
 
+  /// Create a global-to-shared or shared-to-global copy statement.
+  ///
+  /// @param CopyStmt The copy statement to generate code for
+  void createKernelCopy(ppcg_kernel_stmt *CopyStmt);
+
   /// Create code for a ScopStmt called in @p Expr.
   ///
   /// @param Expr The expression containing the call.
@@ -714,7 +734,7 @@ void GPUNodeBuilder::createUser(__isl_take isl_ast_node *UserStmt) {
     isl_ast_node_free(UserStmt);
     return;
   case ppcg_kernel_copy:
-    // TODO: Create kernel copy stmt
+    createKernelCopy(KernelStmt);
     isl_ast_expr_free(Expr);
     isl_ast_node_free(UserStmt);
     return;
@@ -729,6 +749,22 @@ void GPUNodeBuilder::createUser(__isl_take isl_ast_node *UserStmt) {
   isl_ast_node_free(UserStmt);
   return;
 }
+void GPUNodeBuilder::createKernelCopy(ppcg_kernel_stmt *KernelStmt) {
+  isl_ast_expr *LocalIndex = isl_ast_expr_copy(KernelStmt->u.c.local_index);
+  LocalIndex = isl_ast_expr_address_of(LocalIndex);
+  Value *LocalAddr = ExprBuilder.create(LocalIndex);
+  isl_ast_expr *Index = isl_ast_expr_copy(KernelStmt->u.c.index);
+  Index = isl_ast_expr_address_of(Index);
+  Value *GlobalAddr = ExprBuilder.create(Index);
+
+  if (KernelStmt->u.c.read) {
+    LoadInst *Load = Builder.CreateLoad(GlobalAddr, "shared.read");
+    Builder.CreateStore(Load, LocalAddr);
+  } else {
+    LoadInst *Load = Builder.CreateLoad(LocalAddr, "shared.write");
+    Builder.CreateStore(Load, GlobalAddr);
+  }
+}
 
 void GPUNodeBuilder::createScopStmt(isl_ast_expr *Expr,
                                     ppcg_kernel_stmt *KernelStmt) {
@@ -1042,6 +1078,11 @@ void GPUNodeBuilder::createKernel(__isl_take isl_ast_node *KernelStmt) {
   createCallLaunchKernel(GPUKernel, GridDimX, GridDimY, BlockDimX, BlockDimY,
                          BlockDimZ, Parameters);
   createCallFreeKernel(GPUKernel);
+
+  for (auto Id : KernelIds)
+    isl_id_free(Id);
+
+  KernelIds.clear();
 }
 
 /// Compute the DataLayout string for the NVPTX backend.
@@ -1114,7 +1155,7 @@ GPUNodeBuilder::createKernelFunctionDecl(ppcg_kernel *Kernel,
     LocalArrays.push_back(Val);
 
     isl_ast_build_free(Build);
-    isl_id_free(Id);
+    KernelIds.push_back(Id);
     IDToSAI[Id] = SAIRep;
     Arg++;
   }
@@ -1199,6 +1240,48 @@ void GPUNodeBuilder::prepareKernelArguments(ppcg_kernel *Kernel, Function *FN) {
   }
 }
 
+void GPUNodeBuilder::createKernelVariables(ppcg_kernel *Kernel, Function *FN) {
+  Module *M = Builder.GetInsertBlock()->getParent()->getParent();
+
+  for (int i = 0; i < Kernel->n_var; ++i) {
+    struct ppcg_kernel_var &Var = Kernel->var[i];
+    isl_id *Id = isl_space_get_tuple_id(Var.array->space, isl_dim_set);
+    Type *EleTy = ScopArrayInfo::getFromId(Id)->getElementType();
+
+    SmallVector<const SCEV *, 4> Sizes;
+    isl_val *V0 = isl_vec_get_element_val(Var.size, 0);
+    long Bound = isl_val_get_num_si(V0);
+    isl_val_free(V0);
+    Sizes.push_back(S.getSE()->getConstant(Builder.getInt64Ty(), Bound));
+
+    ArrayType *ArrayTy = ArrayType::get(EleTy, Bound);
+    for (unsigned int j = 1; j < Var.array->n_index; ++j) {
+      isl_val *Val = isl_vec_get_element_val(Var.size, j);
+      Bound = isl_val_get_num_si(Val);
+      isl_val_free(Val);
+      Sizes.push_back(S.getSE()->getConstant(Builder.getInt64Ty(), Bound));
+      ArrayTy = ArrayType::get(ArrayTy, Bound);
+    }
+
+    assert(Var.type == ppcg_access_shared && "Only shared memory supported");
+
+    GlobalVariable *SharedVar = new GlobalVariable(
+        *M, ArrayTy, false, GlobalValue::InternalLinkage, 0, Var.name, nullptr,
+        GlobalValue::ThreadLocalMode::NotThreadLocal, 3);
+    SharedVar->setAlignment(EleTy->getPrimitiveSizeInBits() / 8);
+    ConstantAggregateZero *Zero = ConstantAggregateZero::get(ArrayTy);
+    SharedVar->setInitializer(Zero);
+
+    Id = isl_id_alloc(S.getIslCtx(), Var.name, nullptr);
+    IDToValue[Id] = SharedVar;
+    const ScopArrayInfo *SAI = S.getOrCreateScopArrayInfo(
+        SharedVar, EleTy, Sizes, ScopArrayInfo::MK_Array);
+    LocalArrays.push_back(SharedVar);
+    KernelIds.push_back(Id);
+    IDToSAI[Id] = SAI;
+  }
+}
+
 void GPUNodeBuilder::createKernelFunction(ppcg_kernel *Kernel,
                                           SetVector<Value *> &SubtreeValues) {
 
@@ -1222,6 +1305,7 @@ void GPUNodeBuilder::createKernelFunction(ppcg_kernel *Kernel,
   ScopDetection::markFunctionAsInvalid(FN);
 
   prepareKernelArguments(Kernel, FN);
+  createKernelVariables(Kernel, FN);
   insertKernelIntrinsics(Kernel);
 }
 
@@ -1328,8 +1412,8 @@ public:
     Options->tile_size = 32;
 
     Options->use_private_memory = false;
-    Options->use_shared_memory = false;
-    Options->max_shared_memory = 0;
+    Options->use_shared_memory = SharedMemory;
+    Options->max_shared_memory = 48 * 1024;
 
     Options->target = PPCG_TARGET_CUDA;
     Options->openmp = false;
@@ -1510,9 +1594,10 @@ public:
       isl_map *Universe = isl_map_universe(Space);
       Access->tagged_access =
           isl_map_domain_product(Acc->getAccessRelation(), Universe);
-      Access->exact_write = Acc->isWrite();
+      Access->exact_write = !Acc->isMayWrite();
       Access->ref_id = Acc->getId();
       Access->next = Accesses;
+      Access->n_index = Acc->getScopArrayInfo()->getNumberOfDimensions();
       Accesses = Access;
     }