Recognize free function kernels that use local work group memory and mark them with the appropriate attribute

Author: lbushi25

llvm/lib/SYCLLowerIR/LowerWGLocalMemory.cpp

@@ -71,10 +71,10 @@ sycl::getKernelNamesUsingImplicitLocalMem(const Module &M) {
       return -1;
     };
     llvm::for_each(M.functions(), [&](const Function &F) {
-      if (F.getCallingConv() == CallingConv::SPIR_KERNEL &&
-          F.hasFnAttribute(WORK_GROUP_STATIC_ATTR)) {
+      if (F.getCallingConv() == CallingConv::SPIR_KERNEL) {
         int ArgPos = GetArgumentPos(F);
-        SPIRKernelNames.emplace_back(F.getName(), ArgPos);
+        if (ArgPos >= 0)
+          SPIRKernelNames.emplace_back(F.getName(), ArgPos);
       }
     });
   }
@@ -184,11 +184,34 @@ lowerDynamicLocalMemCallDirect(CallInst *CI, Triple TT,
 
 static void lowerLocalMemCall(Function *LocalMemAllocFunc,
                               std::function<void(CallInst *CI)> TransformCall) {
+  static SmallPtrSet<Function *, 16> FuncsCache;
   SmallVector<CallInst *, 4> DelCalls;
   for (User *U : LocalMemAllocFunc->users()) {
     auto *CI = cast<CallInst>(U);
     TransformCall(CI);
     DelCalls.push_back(CI);
+    if (!FuncsCache.insert(CI->getFunction()).second)
+      continue; // We have already traversed call graph from this function
+
+    SmallVector<Function *, 8> WorkList;
+    WorkList.push_back(CI->getFunction());
+    while (!WorkList.empty()) {
+      auto *F = WorkList.back();
+      WorkList.pop_back();
+
+      // Mark kernel as using scrach memory if it isn't marked already
+      if (F->getCallingConv() == CallingConv::SPIR_KERNEL &&
+          !F->hasFnAttribute(WORK_GROUP_STATIC_ATTR))
+        F->addFnAttr(WORK_GROUP_STATIC_ATTR);
+
+      for (auto *FU : F->users()) {
+        if (auto *UCI = dyn_cast<CallInst>(FU)) {
+          if (FuncsCache.insert(UCI->getFunction()).second)
+            WorkList.push_back(UCI->getFunction());
+        } // Even though there could be other uses of a Function, we don't
+          // care about them because we are only concerned about call graph
+      }
+    }
   }
 
   for (auto *CI : DelCalls) {