diff --git a/llvm/lib/SYCLLowerIR/SYCLVirtualFunctionsAnalysis.cpp b/llvm/lib/SYCLLowerIR/SYCLVirtualFunctionsAnalysis.cpp
index 3d9ca035be927..13c2db4918990 100644
--- a/llvm/lib/SYCLLowerIR/SYCLVirtualFunctionsAnalysis.cpp
+++ b/llvm/lib/SYCLLowerIR/SYCLVirtualFunctionsAnalysis.cpp
@@ -71,37 +71,23 @@ void checkKernel(const Function *F, const CallGraphTy &CG) {
 void computeFunctionToKernelsMappingImpl(Function *Kernel, const Function *F,
                                          const CallGraphTy &CG,
                                          FuncToFuncMapTy &Mapping) {
+  Mapping[F].insert(Kernel);
   CallGraphTy::const_iterator It = CG.find(F);
   // It could be that the function itself is a leaf and doesn't call anything
   if (It == CG.end())
     return;
 
-  Mapping[F].insert(Kernel);
-
   const SmallPtrSet<Value *, 8> &Callees = It->getSecond();
   for (const Value *V : Callees)
     if (auto *Callee = dyn_cast<Function>(V))
       computeFunctionToKernelsMappingImpl(Kernel, Callee, CG, Mapping);
 }
 
+// Compute a map from functions used by a kernel to that kernel.
+// For simplicity we also consider a kernel to be using itself.
 void computeFunctionToKernelsMapping(Function *Kernel, const CallGraphTy &CG,
                                      FuncToFuncMapTy &Mapping) {
-  // For simplicity we also consider a kernel to be using itself
-  Mapping[Kernel].insert(Kernel);
-
-  CallGraphTy::const_iterator It = CG.find(Kernel);
-  // It could be that the kernel doesn't call anything
-  if (It == CG.end())
-    return;
-
-  const SmallPtrSet<Value *, 8> &Callees = It->getSecond();
-  for (const Value *V : Callees) {
-    auto *Callee = dyn_cast<Function>(V);
-    if (!Callee)
-      continue;
-    Mapping[Callee].insert(Kernel);
-    computeFunctionToKernelsMappingImpl(Kernel, Callee, CG, Mapping);
-  }
+  computeFunctionToKernelsMappingImpl(Kernel, Kernel, CG, Mapping);
 }
 
 void collectVTablesThatUseFunction(
diff --git a/llvm/test/SYCLLowerIR/SYCLVirtualFunctionsAnalysis/calls-indirectly-propagation-4.ll b/llvm/test/SYCLLowerIR/SYCLVirtualFunctionsAnalysis/calls-indirectly-propagation-4.ll
index 21da86cb120de..71222e2594b9e 100644
--- a/llvm/test/SYCLLowerIR/SYCLVirtualFunctionsAnalysis/calls-indirectly-propagation-4.ll
+++ b/llvm/test/SYCLLowerIR/SYCLVirtualFunctionsAnalysis/calls-indirectly-propagation-4.ll
@@ -35,7 +35,20 @@ entry:
   ret void
 }
 
+define internal spir_func void @helper2(ptr addrspace(1) noundef align 8 %arg) {
+entry:
+  call void @helper(ptr addrspace(1) %arg)
+  ret void
+}
+
+define weak_odr dso_local spir_kernel void @kernel2(ptr addrspace(1) noundef align 8 %_arg_StorageAcc) #2 {
+entry:
+  call void @helper2(ptr addrspace(1) %_arg_StorageAcc)
+  ret void
+}
+
 ; CHECK: @kernel{{.*}} #[[#KERNEL_ATTRS:]]
+; CHECK: @kernel2{{.*}} #[[#KERNEL_ATTRS]]
 ;
 ; CHECK: attributes #[[#KERNEL_ATTRS]] = {{.*}}"calls-indirectly"="set-foo,set-bar"