[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

11 ↗

(On Diff #102326)

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

1120 ↗

(On Diff #102326)

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.

1122 ↗

(On Diff #102326)

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

2704 ↗

(On Diff #102326)

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.

2816 ↗

(On Diff #102326)

Do not add braces.

[grosser]

lib/CodeGen/PPCGCodeGeneration.cpp

401 ↗

(On Diff #103923)

Create new function DECLARATIONS ...

408 ↗

(On Diff #103923)

Should be IslNodeBuilder::ValueMap

1136 ↗

(On Diff #103923)

Start with a comment describing the function.

1138 ↗

(On Diff #103923)

Drop TODO.

1122 ↗

(On Diff #102326)

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 ↗

(On Diff #103923)

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

2816 ↗

(On Diff #102326)

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

[grosser]

lib/CodeGen/PPCGCodeGeneration.cpp

374 ↗

(On Diff #103928)

Please document the new parameter as well.

[bollu]

• [NFC] document getFunctionsFromRawSubtreeValues.

[grosser]

LGTM. Please document the new parameter.

polly/trunk/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 kernel function.
   ///
   /// 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.
+  /// @param SubtreeFunctions The set of llvm::Functions 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 function declarations in GPUModule which are the same as
+  /// SubtreeFunctions.
+  ///
+  /// 2. Populate IslNodeBuilder::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::ValueMap
+  /// @see BlockGenerator::GlobalMap
+  /// @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) {
+/// Check if F is a function that we can code-generate in a GPU kernel.
+static bool isValidFunctionInKernel(llvm::Function *F) {
+  assert(F && "F is an invalid pointer");
+  // We string compare against the name of the function to allow
+  // all variants of the intrinsic "llvm.sqrt.*"
+  return F->isIntrinsic() && 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); }
+
+/// Return `Function`s from `RawSubtreeValues`.
+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;
+  // Note: { ValidSubtreeValues, ValidSubtreeFunctions } partitions
+  // SubtreeValues. This is important, because we should not lose any
+  // SubtreeValues in the process of constructing the
+  // "ValidSubtree{Values, Functions} sets. Nor should the set
+  // ValidSubtree{Values, Functions} have any common element.
+  auto ValidSubtreeValuesIt =
+      make_filter_range(SubtreeValues, isValidSubtreeValue);
+  SetVector<Value *> ValidSubtreeValues(ValidSubtreeValuesIt.begin(),
+                                        ValidSubtreeValuesIt.end());
+  SetVector<Function *> ValidSubtreeFunctions(
+      getFunctionsFromRawSubtreeValues(SubtreeValues));
+
+  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.");
+    assert(ValueMap.find(Fn) == ValueMap.end() &&
+           "Fn already present in ValueMap");
+    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)
+  /// Check if the basic block contains a function we cannot codegen for GPU
+  /// kernels.
+  ///
+  /// If this basic block does something with a `Function` other than calling
+  /// a function that we support in a kernel, return true.
+  bool containsInvalidKernelFunctionInBllock(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 (containsInvalidKernelFunctionInBllock(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 (containsInvalidKernelFunctionInBllock(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);
 

polly/trunk/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.
+
+; REQUIRES: pollyacc
+
+; 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 }
+