[Polly] [PPCGCodeGeneration] Allow intrinsics within kernels.

[bollu]

Summary

• In D33414, if any function call was found within a kernel, we would bail out.
• This is an over-approximation. This patch changes this by allowing intrinsics within kernels.
• This introduces an additional step when creating a separate llvm::Module for a kernel. We now have to copy the intrinsics from the original llvm::Module to the kernel-specific module.
• Note that this maybe too lax. Perhaps we are allowing intrinsics which should not be allowed through. This can be debated in the review.

[bollu]

the TODO blocks are commentary where I would like some discussion before merging this revision.

test/GPGPU/intrinsic-copied-into-kernel.ll

12

Just realised this: I should also check for the declare of sqrt. Will add.

[bollu]

@singam-sanjay: We had discussed trying this out on the Julia code with debug intrinsics yesterday. Were you able to try it out? Thanks!

[singam-sanjay]

@bollu The patch did prevent passing the address of the debug intrinsic to the kernel, but couldn't guard against buggy metadata,

DICompileUnit not listed in llvm.dbg.cu
!11 = distinct !DICompileUnit(language: DW_LANG_C89, file: !3, producer: "julia", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, enums: !12)

I'll have to look closer at how *GPUModule is being produced to understand this error.

[bollu]

I lack context, I think. Where is wrong metadata coming from?
Also, FWIW, I would not add code that guards against illegal metadata.

[singam-sanjay]

@bollu Please refer this comment on my patch. I observed a pattern with functions that produced the buggy metadata, simple copy kernels (A[i][j][k] = B[i][j][k]) didn't produce buggy metadata but kernels that had computation (A[i][j][k] +=/-=/*= B[i][j][k]) resulted in Modules that didn't list the DICompileUnit in llvm.dbg.cu.

I'm not sure how a module that's built independent from source IR carries metadata. Do you think there'd be a function to explicitly list the DICompileUnit in llvm.dbg.cu ?

[grosser]

I believe the Julia problem consists of two (partially independent) bugs. If this patch fixes a part of your problem (the intrinsics), then we should likely commit it and thing about the metadata afterwards.

[singam-sanjay]

Yes, we can handle the metadata problem in a different patch.

[bollu]

• [NFC] typo fix.
• Try to retreive pre-existing function. If this does not exist, call Function::Create.
• [NFC] add a debug message if we are bailing out from PPCGCodeGeneration due to an unrecognised function in a kernel.

[bollu]

@grosser, ping.

[bollu]

• [NFC] documented replaceKernelFunctions.

[bollu]

• Changed approach as discussed offline with Tobias. It was decided to use the BlockGenerator functionality to setup correct kernel code generation, rather than fixing generated code.

• General patch cleanup.

[grosser]

Some more comments.

lib/CodeGen/PPCGCodeGeneration.cpp

1139

In llvm we start function names with lowercase letters.

Also, please document what the parameter F does.

Also, the comment is not useful. Please describe in one sentence what the function does, and then in a small paragraph what the criteria which a function that is listed here must fullfill.

1141

Please make this an explicit list of intrinsics. Maybe start with sqrt only to extablish the infrastructure.

2721

support
start with lowercase letters (it seems I missed this in the last review)

Also, in general we comment functions as follow:

/// Brief one-line comment
///
/// Longer paragraph explaining what is happening here.

2833

Do not add braces.

[grosser]

lib/CodeGen/PPCGCodeGeneration.cpp

401

Create new function DECLARATIONS ...

408

Should be IslNodeBuilder::ValueMap

1136

Start with a comment describing the function.

1138

Drop TODO.

1141

Please add a comment that we match any llvm.sqrt intrinsics by doing a string match.

[bollu]

• [NFC] style fixes.

[bollu]

• Fix function capitalization, spelling and comments

[bollu]

lib/CodeGen/PPCGCodeGeneration.cpp

1136

This is a member function, so it is already documented at GPUNodeBuilder declaration.

2833

I do need braces because the check contains (potentially) two statements in the body?

[grosser]

lib/CodeGen/PPCGCodeGeneration.cpp

375

Please document the new parameter as well.

[bollu]

• [NFC] document getFunctionsFromRawSubtreeValues.

[grosser]

LGTM. Please document the new parameter.

lib/CodeGen/PPCGCodeGeneration.cpp

private:
   /// @param Direction The direction in which to transfer data.
   void createDataTransfer(__isl_take isl_ast_node *TransferStmt,
                           enum DataDirection Direction);
 
   /// Find llvm::Values referenced in GPU kernel.
   ///
   /// @param Kernel The kernel to scan for llvm::Values
   ///
-  /// @returns A set of values referenced by the kernel.
-  SetVector<Value *> getReferencesInKernel(ppcg_kernel *Kernel);
+  /// @returns A pair, whose first element contains the set of
+  ///          values referenced by the kernel, and whose
+  ///          second element contains the set of functions
+  ///          referenced by the kernel. All functions in the
+  ///          second set satisfy isValidFunctionInKernel.
+  std::pair<SetVector<Value *>, SetVector<Function *>>
+  getReferencesInKernel(ppcg_kernel *Kernel);
 
   /// Compute the sizes of the execution grid for a given kernel.
   ///
   /// @param Kernel The kernel to compute grid sizes for.
   ///
   /// @returns A tuple with grid sizes for X and Y dimension
   std::tuple<Value *, Value *> getGridSizes(ppcg_kernel *Kernel);
 
private:
   ///
   /// Create a kernel function located in a newly created module that can serve
   /// as target for device code generation. Set the Builder to point to the
   /// start block of this newly created function.
   ///
   /// @param Kernel The kernel to generate code for.
   /// @param SubtreeValues The set of llvm::Values referenced by this kernel.
   void createKernelFunction(ppcg_kernel *Kernel,
-                            SetVector<Value *> &SubtreeValues);
+                            SetVector<Value *> &SubtreeValues,
+                            SetVector<Function *> &SubtreeFunctions);
 
   /// Create the declaration of a kernel function.
   ///
   /// The kernel function takes as arguments:
   ///
   ///   - One i8 pointer for each external array reference used in the kernel.
   ///   - Host iterators
   ///   - Parameters
   ///   - Other LLVM Value references (TODO)
   ///
   /// @param Kernel The kernel to generate the function declaration for.
   /// @param SubtreeValues The set of llvm::Values referenced by this kernel.
   ///
   /// @returns The newly declared function.
   Function *createKernelFunctionDecl(ppcg_kernel *Kernel,
                                      SetVector<Value *> &SubtreeValues);
 
   /// Insert intrinsic functions to obtain thread and block ids.
   ///
   /// @param The kernel to generate the intrinsic functions for.
   void insertKernelIntrinsics(ppcg_kernel *Kernel);
 
+  /// Setup the creation of functions referenced by the GPU kernel.
+  ///
+  /// 1. Create new functions in GPUModule which are the same as
+  /// SubtreeFunctions.
+  ///
+  /// 2. Populate IslNodeBuilder::BlockGen::ValueMap with mappings from
+  /// old functions (that come from the original module) to new functions
+  /// (that are created within GPUModule). That way, we generate references
+  /// to the correct function (in GPUModule) in BlockGenerator.
+  ///
+  /// @see IslNodeBuilder::BlockGen::ValueMap
+  /// @see BlockGenerator::getNewValue
+  /// @see GPUNodeBuilder::getReferencesInKernel.
+  ///
+  /// @param SubtreeFunctions The set of llvm::Functions referenced by
+  ///                         this kernel.
+  void setupKernelSubtreeFunctions(SetVector<Function *> SubtreeFunctions);
+
   /// 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.
isl_bool collectReferencesInGPUStmt(__isl_keep isl_ast_node *Node, void *User) {
   auto Stmt = (ScopStmt *)KernelStmt->u.d.stmt->stmt;
   isl_id_free(Id);
 
   addReferencesFromStmt(Stmt, User, false /* CreateScalarRefs */);
 
   return isl_bool_true;
 }
 
-SetVector<Value *> GPUNodeBuilder::getReferencesInKernel(ppcg_kernel *Kernel) {
+// By this point, we have only allowed "safe" function calls which we know
+// will be lowered. For now, we only allow intrinsics within kernel functions
+// TODO: is this too lax? We maybe allowing intrinsics that cannot be lowered.
+static bool IsValidFunctionInKernel(llvm::Function *F) {
+  assert(F && "F is a invalid pointer");
+  return F->getName().startswith("llvm.sqrt");
+}
+
+// Do not take `Function` as a subtree value.
+// We try to take the reference of all subtree values and pass them along
+// to the kernel from the host. Taking an address of any function and
+// trying to pass along is nonsensical. Only allow `Value`s that are not
+// `Function`s.
+static bool isValidSubtreeValue(llvm::Value *V) { return !isa<Function>(V); }
+
+// TODO: Consider converting to filter + map. Not sure how to use the
+// `map_iterator` in llvm's STLExtras.
+static SetVector<Function *>
+getFunctionsFromRawSubtreeValues(SetVector<Value *> RawSubtreeValues) {
+  SetVector<Function *> SubtreeFunctions;
+  for (Value *It : RawSubtreeValues) {
+    Function *F = dyn_cast<Function>(It);
+    if (F) {
+      assert(IsValidFunctionInKernel(F) && "Code should have bailed out by "
+                                           "this point if an invalid function "
+                                           "were present in a kernel");
+      SubtreeFunctions.insert(F);
+    }
+  }
+  return SubtreeFunctions;
+}
+
+std::pair<SetVector<Value *>, SetVector<Function *>>
+GPUNodeBuilder::getReferencesInKernel(ppcg_kernel *Kernel) {
   SetVector<Value *> SubtreeValues;
   SetVector<const SCEV *> SCEVs;
   SetVector<const Loop *> Loops;
   SubtreeReferences References = {
       LI, SE, S, ValueMap, SubtreeValues, SCEVs, getBlockGenerator()};
 
   for (const auto &I : IDToValue)
     SubtreeValues.insert(I.second);
GPUNodeBuilder::getReferencesInKernel(ppcg_kernel *Kernel) {
   for (long i = 0; i < isl_space_dim(Kernel->space, isl_dim_set); i++) {
     isl_id *Id = isl_space_get_dim_id(Kernel->space, isl_dim_set, i);
     assert(IDToValue.count(Id));
     Value *Val = IDToValue[Id];
     SubtreeValues.remove(Val);
     isl_id_free(Id);
   }
 
-  return SubtreeValues;
+  // TODO: Consider creating both SubtreeValue and SubtreeFunction in one
+  // loop since they are mutually exclusive? This makes it harder to read.
+  auto ValidSubtreeValuesIt =
+      make_filter_range(SubtreeValues, isValidSubtreeValue);
+  SetVector<Value *> ValidSubtreeValues(ValidSubtreeValuesIt.begin(),
+                                        ValidSubtreeValuesIt.end());
+
+  SetVector<Function *> ValidSubtreeFunctions(
+      getFunctionsFromRawSubtreeValues(SubtreeValues));
+
+#ifndef NDEBUG
+  // Note that { ValidSubtreeValues, ValidSubtreeFunctions } partions
+  // SubtreeValues. We are checking for that property here.
+
+  // Check that ValidSubtreeValues U ValidSubtreeFunctions == SubtreeValues
+  {
+    auto ValidSubtreeValuesUnion = SetVector<Value *>(
+        ValidSubtreeValues.begin(), ValidSubtreeValues.end());
+    ValidSubtreeValuesUnion.set_union(ValidSubtreeFunctions);
+    assert(ValidSubtreeValuesUnion == SubtreeValues);
+  }
+
+// TODO: Check ValidSubtreeValues intersection ValidSubtreeFunctions == empty.
+// llvm::SetVector does not support set_intersect.
+// llvm::set_intersect from llvm/ADT/SetOperations.h does not work over
+// llvm::SetVector, unfortunately.
+#endif
+  return std::make_pair(ValidSubtreeValues, ValidSubtreeFunctions);
 }
 
 void GPUNodeBuilder::clearDominators(Function *F) {
   DomTreeNode *N = DT.getNode(&F->getEntryBlock());
   std::vector<BasicBlock *> Nodes;
   for (po_iterator<DomTreeNode *> I = po_begin(N), E = po_end(N); I != E; ++I)
     Nodes.push_back(I->getBlock());
 
for (int i = 0; i < NumArgs; i++) {
     Index++;
   }
 
   auto Location = EntryBlock->getTerminator();
   return new BitCastInst(Parameters, Builder.getInt8PtrTy(),
                          Launch + "_params_i8ptr", Location);
 }
 
+void GPUNodeBuilder::setupKernelSubtreeFunctions(
+    SetVector<Function *> SubtreeFunctions) {
+  for (auto Fn : SubtreeFunctions) {
+    const std::string ClonedFnName = Fn->getName();
+    Function *Clone = GPUModule->getFunction(ClonedFnName);
+    if (!Clone)
+      Clone =
+          Function::Create(Fn->getFunctionType(), GlobalValue::ExternalLinkage,
+                           ClonedFnName, GPUModule.get());
+    assert(Clone && "Expected cloned function to be initialized.");
+#ifndef NDEBUG
+    assert(ValueMap.find(Fn) == ValueMap.end() && "Fn already present in"
+                                                  " ValueMap");
+#endif
+    ValueMap[Fn] = Clone;
+  }
+}
 void GPUNodeBuilder::createKernel(__isl_take isl_ast_node *KernelStmt) {
   isl_id *Id = isl_ast_node_get_annotation(KernelStmt);
   ppcg_kernel *Kernel = (ppcg_kernel *)isl_id_get_user(Id);
   isl_id_free(Id);
   isl_ast_node_free(KernelStmt);
 
   if (Kernel->n_grid > 1)
     DeepestParallel =
         std::max(DeepestParallel, isl_space_dim(Kernel->space, isl_dim_set));
   else
     DeepestSequential =
         std::max(DeepestSequential, isl_space_dim(Kernel->space, isl_dim_set));
 
   Value *BlockDimX, *BlockDimY, *BlockDimZ;
   std::tie(BlockDimX, BlockDimY, BlockDimZ) = getBlockSizes(Kernel);
 
-  SetVector<Value *> SubtreeValues = getReferencesInKernel(Kernel);
+  SetVector<Value *> SubtreeValues;
+  SetVector<Function *> SubtreeFunctions;
+  std::tie(SubtreeValues, SubtreeFunctions) = getReferencesInKernel(Kernel);
 
   assert(Kernel->tree && "Device AST of kernel node is empty");
 
   Instruction &HostInsertPoint = *Builder.GetInsertPoint();
   IslExprBuilder::IDToValueTy HostIDs = IDToValue;
   ValueMapT HostValueMap = ValueMap;
   BlockGenerator::AllocaMapTy HostScalarMap = ScalarMap;
   ScalarMap.clear();
 
   SetVector<const Loop *> Loops;
 
   // Create for all loops we depend on values that contain the current loop
   // iteration. These values are necessary to generate code for SCEVs that
   // depend on such loops. As a result we need to pass them to the subfunction.
   for (const Loop *L : Loops) {
     const SCEV *OuterLIV = SE.getAddRecExpr(SE.getUnknown(Builder.getInt64(0)),
                                             SE.getUnknown(Builder.getInt64(1)),
                                             L, SCEV::FlagAnyWrap);
     Value *V = generateSCEV(OuterLIV);
     OutsideLoopIterations[L] = SE.getUnknown(V);
     SubtreeValues.insert(V);
   }
 
-  createKernelFunction(Kernel, SubtreeValues);
+  createKernelFunction(Kernel, SubtreeValues, SubtreeFunctions);
+  setupKernelSubtreeFunctions(SubtreeFunctions);
 
   create(isl_ast_node_copy(Kernel->tree));
 
   finalizeKernelArguments(Kernel);
   Function *F = Builder.GetInsertBlock()->getParent();
   addCUDAAnnotations(F->getParent(), BlockDimX, BlockDimY, BlockDimZ);
   clearDominators(F);
   clearScalarEvolution(F);
for (int i = 0; i < Kernel->n_var; ++i) {
     Id = isl_id_alloc(S.getIslCtx(), Var.name, nullptr);
     IDToValue[Id] = Allocation;
     LocalArrays.push_back(Allocation);
     KernelIds.push_back(Id);
     IDToSAI[Id] = SAI;
   }
 }
 
-void GPUNodeBuilder::createKernelFunction(ppcg_kernel *Kernel,
-                                          SetVector<Value *> &SubtreeValues) {
+void GPUNodeBuilder::createKernelFunction(
+    ppcg_kernel *Kernel, SetVector<Value *> &SubtreeValues,
+    SetVector<Function *> &SubtreeFunctions) {
   std::string Identifier = "kernel_" + std::to_string(Kernel->id);
   GPUModule.reset(new Module(Identifier, Builder.getContext()));
 
   switch (Arch) {
   case GPUArch::NVPTX64:
     if (Runtime == GPURuntime::CUDA)
       GPUModule->setTargetTriple(Triple::normalize("nvptx64-nvidia-cuda"));
     else if (Runtime == GPURuntime::OpenCL)
public:
   __isl_give isl_ast_expr *
   createSufficientComputeCheck(Scop &S, __isl_keep isl_ast_build *Build) {
     auto Iterations = getNumberOfIterations(S, Build);
     auto *MinComputeVal = isl_val_int_from_si(S.getIslCtx(), MinCompute);
     auto *MinComputeExpr = isl_ast_expr_from_val(MinComputeVal);
     return isl_ast_expr_ge(Iterations, MinComputeExpr);
   }
 
-  /// Check whether the Block contains any Function value.
-  bool ContainsFnPtrValInBlock(const BasicBlock *BB) {
-    for (const Instruction &Inst : *BB)
+  // If this basic block does something with a `Function` other than calling
+  // a function that we support in a kernel, return true.
+  bool ContainsInvalidKernelFunctionInBlock(const BasicBlock *BB) {
+    for (const Instruction &Inst : *BB) {
+      const CallInst *Call = dyn_cast<CallInst>(&Inst);
+      if (Call && IsValidFunctionInKernel(Call->getCalledFunction())) {
+        continue;
+      }
+
       for (Value *SrcVal : Inst.operands()) {
         PointerType *p = dyn_cast<PointerType>(SrcVal->getType());
         if (!p)
           continue;
         if (isa<FunctionType>(p->getElementType()))
           return true;
       }
+    }
     return false;
   }
 
-  /// Return whether the Scop S has functions.
-  bool ContainsFnPtr(const Scop &S) {
+  /// Return whether the Scop S uses functions in a way that we do not support.
+  bool ContainsInvalidKernelFunction(const Scop &S) {
     for (auto &Stmt : S) {
       if (Stmt.isBlockStmt()) {
-        if (ContainsFnPtrValInBlock(Stmt.getBasicBlock()))
+        if (ContainsInvalidKernelFunctionInBlock(Stmt.getBasicBlock()))
           return true;
       } else {
         assert(Stmt.isRegionStmt() &&
                "Stmt was neither block nor region statement");
         for (const BasicBlock *BB : Stmt.getRegion()->blocks())
-          if (ContainsFnPtrValInBlock(BB))
+          if (ContainsInvalidKernelFunctionInBlock(BB))
             return true;
       }
     }
     return false;
   }
 
   /// Generate code for a given GPU AST described by @p Root.
   ///
public:
   bool runOnScop(Scop &CurrentScop) override {
     S = &CurrentScop;
     LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
     DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
     DL = &S->getRegion().getEntry()->getModule()->getDataLayout();
     RI = &getAnalysis<RegionInfoPass>().getRegionInfo();
 
-    // We currently do not support functions inside kernels, as code
-    // generation will need to offload function calls to the kernel.
-    // This may lead to a kernel trying to call a function on the host.
+    // We currently do not support functions other than intrinsics inside
+    // kernels, as code generation will need to offload function calls to the
+    // kernel. This may lead to a kernel trying to call a function on the host.
     // This also allows us to prevent codegen from trying to take the
     // address of an intrinsic function to send to the kernel.
-    if (ContainsFnPtr(CurrentScop))
+    if (ContainsInvalidKernelFunction(CurrentScop)) {
+      DEBUG(
+          dbgs()
+              << "Scop contains function which cannot be materialised in a GPU "
+                 "kernel. bailing out.\n";);
       return false;
+    }
 
     auto PPCGScop = createPPCGScop();
     auto PPCGProg = createPPCGProg(PPCGScop);
     auto PPCGGen = generateGPU(PPCGScop, PPCGProg);
 
     if (PPCGGen->tree)
       generateCode(isl_ast_node_copy(PPCGGen->tree), PPCGProg);
 

test/GPGPU/intrinsic-copied-into-kernel.ll

+; RUN: opt %loadPolly -analyze -polly-scops < %s | FileCheck %s --check-prefix=SCOP
+; RUN: opt %loadPolly -analyze -polly-codegen-ppcg -polly-acc-dump-kernel-ir < %s | FileCheck %s --check-prefix=KERNEL-IR
+; RUN: opt %loadPolly -S -polly-codegen-ppcg  < %s | FileCheck %s --check-prefix=HOST-IR
+
+; Test that we do recognise and codegen a kernel that has intrinsics.
+
+; Check that we model the kernel as a scop.
+; SCOP:      Function: f
+; SCOP-NEXT:       Region: %entry.split---%for.end
+
+; Check that the intrinsic call is present in the kernel IR.
+; KERNEL-IR:   %p_sqrt = tail call float @llvm.sqrt.f32(float %A.arr.i.val_p_scalar_)
+; KERNEL-IR:   declare float @llvm.sqrt.f32(float) #2
+
+; Check that kernel launch is generated in host IR.
+; the declare would not be generated unless a call to a kernel exists.
+; HOST-IR: declare void @polly_launchKernel(i8*, i32, i32, i32, i32, i32, i8*)
+
+
+; void f(float *A, float *B, int N) {
+;   for(int i = 0; i < N; i++) {
+;       B[i] = sqrt(A[i]);
+;   }
+; }
+
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+
+define void @f(float* %A, float* %B, i32 %N) {
+entry:
+  br label %entry.split
+
+entry.split:                                      ; preds = %entry
+  %cmp1 = icmp sgt i32 %N, 0
+  br i1 %cmp1, label %for.body.lr.ph, label %for.end
+
+for.body.lr.ph:                                   ; preds = %entry.split
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.lr.ph, %for.body
+  %indvars.iv = phi i64 [ 0, %for.body.lr.ph ], [ %indvars.iv.next, %for.body ]
+  %A.arr.i = getelementptr inbounds float, float* %A, i64 %indvars.iv
+  %A.arr.i.val = load float, float* %A.arr.i, align 4
+  ; Call to intrinsic that should be part of the kernel.
+  %sqrt = tail call float @llvm.sqrt.f32(float %A.arr.i.val)
+  %B.arr.i = getelementptr inbounds float, float* %B, i64 %indvars.iv
+  store float %sqrt, float* %B.arr.i, align 4
+
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %wide.trip.count = zext i32 %N to i64
+  %exitcond = icmp ne i64 %indvars.iv.next, %wide.trip.count
+  br i1 %exitcond, label %for.body, label %for.cond.for.end_crit_edge
+
+for.cond.for.end_crit_edge:                       ; preds = %for.body
+  br label %for.end
+
+for.end:                                          ; preds = %for.cond.for.end_crit_edge, %entry.split
+  ret void
+}
+
+; Function Attrs: nounwind readnone
+declare float @llvm.sqrt.f32(float) #0
+
+attributes #0 = { nounwind readnone }
+