diff --git a/llvm/lib/SYCLLowerIR/LowerWGLocalMemory.cpp b/llvm/lib/SYCLLowerIR/LowerWGLocalMemory.cpp
index e6afe27d58bae..a856d81d6b17f 100644
--- a/llvm/lib/SYCLLowerIR/LowerWGLocalMemory.cpp
+++ b/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) {
diff --git a/sycl/test-e2e/DeviceImageDependencies/free_function_kernels.cpp b/sycl/test-e2e/DeviceImageDependencies/free_function_kernels.cpp
index 674bcb6af977d..2f1bbcd086e86 100644
--- a/sycl/test-e2e/DeviceImageDependencies/free_function_kernels.cpp
+++ b/sycl/test-e2e/DeviceImageDependencies/free_function_kernels.cpp
@@ -1,7 +1,7 @@
 // Ensure -fsycl-allow-device-dependencies can work with free function kernels.
 
 // REQUIRES: aspect-usm_shared_allocations
-// RUN: %{build} --save-temps -o %t.out -fsycl-allow-device-image-dependencies
+// RUN: %{build} -o %t.out -fsycl-allow-device-image-dependencies
 // RUN: %{run} %t.out
 
 #include <iostream>