[SYCL] Add ITT annotations for SPIR functions as well

Previously if was done one for SPIR kernels

Signed-off-by: Dmitry Sidorov <[email protected]>

Author: MrSidims

llvm/lib/Transforms/Instrumentation/SPIRITTAnnotations.cpp

@@ -138,11 +138,17 @@ ModulePass *llvm::createSPIRITTAnnotationsPass() {
 
 namespace {
 
-// Check for calling convention of a function.
-bool isSPIRKernel(Function &F) {
+// Check for calling convention of a function. Return true if it's SPIR kernel.
+inline bool isSPIRKernel(Function &F) {
   return F.getCallingConv() == CallingConv::SPIR_KERNEL;
 }
 
+// Check for calling convention of a function. Return true if it's SPIR
+// function.
+inline bool isSPIRFunction(Function &F) {
+  return F.getCallingConv() == CallingConv::SPIR_FUNC;
+}
+
 Instruction *emitCall(Module &M, Type *RetTy, StringRef FunctionName,
                       ArrayRef<Value *> Args, Instruction *InsertBefore) {
   SmallVector<Type *, 8> ArgTys(Args.size());
@@ -240,20 +246,24 @@ PreservedAnalyses SPIRITTAnnotationsPass::run(Module &M,
       SPIRV_GROUP_FMAX,      SPIRV_GROUP_UMAX, SPIRV_GROUP_SMAX};
 
   for (Function &F : M) {
-    // Annotate only SPIR kernels
-    if (F.isDeclaration() || !isSPIRKernel(F))
+    // Annotate only SPIR kernels and functions
+    bool IsSPIRKernel = isSPIRKernel(F);
+    bool IsSPIRFunction = isSPIRFunction(F);
+    if (F.isDeclaration() || !(isSPIRKernel || isSPIRFunction))
       continue;
 
     // At the beggining of a kernel insert work item start annotation
     // instruction.
-    IRModified |= insertSimpleInstrumentationCall(M, ITT_ANNOTATION_WI_START,
-                                                  &*inst_begin(F));
+    if (IsSPIRKernel)
+      IRModified |= insertSimpleInstrumentationCall(M, ITT_ANNOTATION_WI_START,
+                                                    &*inst_begin(F));
 
     for (BasicBlock &BB : F) {
       // Insert Finish instruction before return instruction
-      if (ReturnInst *RI = dyn_cast<ReturnInst>(BB.getTerminator()))
-        IRModified |=
-            insertSimpleInstrumentationCall(M, ITT_ANNOTATION_WI_FINISH, RI);
+      if (IsSPIRKernel)
+        if (ReturnInst *RI = dyn_cast<ReturnInst>(BB.getTerminator()))
+          IRModified |=
+              insertSimpleInstrumentationCall(M, ITT_ANNOTATION_WI_FINISH, RI);
       for (Instruction &I : BB) {
         CallInst *CI = dyn_cast<CallInst>(&I);
         if (!CI)