diff --git a/llvm/test/tools/sycl-post-link/assert-property-2.ll b/llvm/test/tools/sycl-post-link/assert-property-2.ll
index 8b707c1a4ef4c..cb1e53d129c75 100644
--- a/llvm/test/tools/sycl-post-link/assert-property-2.ll
+++ b/llvm/test/tools/sycl-post-link/assert-property-2.ll
@@ -38,9 +38,26 @@
 ; void G() { common3(); }
 ; void H() { common3(); }
 ;
+; void no_assert_func() {
+;   return;
+; }
+; void common4() {
+;   assert_func();
+;   no_assert_func();
+; }
+; void J() {
+;   common4();
+; }
+;
 ; int main() {
 ;   queue Q;
 ;   Q.submit([&] (handler& CGH) {
+;     CGH.parallel_for<class Kernel9>(range<1>{1}, [=](id<1> i) {
+;       J();
+;     });
+;     CGH.parallel_for<class Kernel10>(range<1>{1}, [=](id<1> i) {
+;       common4();
+;     });
 ;     CGH.parallel_for<class Kernel>(range<1>{1}, [=](id<1> i) {
 ;       A_excl();
 ;       B_incl();
@@ -90,6 +107,43 @@ target triple = "spir64_x86_64-unknown-unknown-sycldevice"
 
 ; CHECK: [SYCL/assert used]
 
+; Function Attrs: convergent noinline norecurse optnone mustprogress
+define dso_local spir_func void @_Z1Jv() #3 {
+entry:
+  call spir_func void @_Z7common4v()
+  ret void
+}
+
+; Function Attrs: convergent noinline norecurse optnone mustprogress
+define dso_local spir_func void @_Z7common4v() #3 {
+entry:
+  call spir_func void @_Z11assert_funcv()
+  call spir_func void @_Z14no_assert_funcv()
+  ret void
+}
+
+; CHECK: _ZTSZZ4mainENKUlRN2cl4sycl7handlerEE_clES2_E7Kernel9
+; Function Attrs: convergent noinline norecurse mustprogress
+define weak_odr dso_local spir_kernel void @_ZTSZZ4mainENKUlRN2cl4sycl7handlerEE_clES2_E7Kernel9() #0 {
+entry:
+  call spir_func void @_Z1Jv()
+  ret void
+}
+
+; CHECK: _ZTSZZ4mainENKUlRN2cl4sycl7handlerEE_clES2_E8Kernel10
+; Function Attrs: convergent noinline norecurse optnone mustprogress
+define weak_odr dso_local spir_kernel void @_ZTSZZ4mainENKUlRN2cl4sycl7handlerEE_clES2_E8Kernel10() #0 {
+entry:
+  call spir_func void @_Z7common4v()
+  ret void
+}
+
+; Function Attrs: convergent noinline norecurse nounwind optnone mustprogress
+define dso_local spir_func void @_Z14no_assert_funcv() #2 {
+entry:
+  ret void
+}
+
 ; Function Attrs: convergent norecurse nounwind mustprogress
 define dso_local spir_func void @_Z6B_inclv() local_unnamed_addr {
 entry:
diff --git a/llvm/tools/sycl-post-link/sycl-post-link.cpp b/llvm/tools/sycl-post-link/sycl-post-link.cpp
index 9a5ea2da3bf0b..e7a4651ba7eab 100644
--- a/llvm/tools/sycl-post-link/sycl-post-link.cpp
+++ b/llvm/tools/sycl-post-link/sycl-post-link.cpp
@@ -316,10 +316,11 @@ static bool hasAssertInFunctionCallGraph(llvm::Function *Func) {
 
       // Return if we've already discovered if there are asserts in the
       // function call graph.
-      if (hasAssertionInCallGraphMap.count(CF)) {
+      auto HasAssert = hasAssertionInCallGraphMap.find(CF);
+      if (HasAssert != hasAssertionInCallGraphMap.end()) {
         // If we know, that this function does not contain assert, we still
         // should investigate another instructions in the function.
-        if (!hasAssertionInCallGraphMap[CF])
+        if (!HasAssert->second)
           continue;
 
         return true;
@@ -343,10 +344,9 @@ static bool hasAssertInFunctionCallGraph(llvm::Function *Func) {
       }
     }
 
-    if (IsLeaf) {
-      // Mark the above functions as ones that definetely do not call assert.
-      for (auto *It : FuncCallStack)
-        hasAssertionInCallGraphMap[It] = false;
+    if (IsLeaf && !FuncCallStack.empty()) {
+      // Mark the leaf function as one that definetely does not call assert.
+      hasAssertionInCallGraphMap[FuncCallStack.back()] = false;
       FuncCallStack.clear();
     }
   }