fixed matching of kp_sampler with kokkosp_end check

kokkos · Sep 7, 2023 · 4b3bda1 · 4b3bda1
1 parent 761fcc9
commit 4b3bda1
Showing 1 changed file with 12 additions and 8 deletions.
diff --git a/common/kokkos-sampler/kp_sampler_skip.cpp b/common/kokkos-sampler/kp_sampler_skip.cpp
@@ -172,15 +172,19 @@ void kokkosp_begin_parallel_for(const char* name, const uint32_t devID,
   }
 }
 
-void kokkosp_end_parallel_for(const uint64_t kID) {
-    if (tool_verbosity > 0) {
-      printf("KokkosP: sample %llu calling child-end function...\n",
-             (unsigned long long)(kID));
-    }
+void kokkosp_end_parallel_for(const uint64_t kID) { 
+
     if (NULL != endForCallee) {
-      uint64_t retrievedNestedkID = infokIDSample.at(kID);
-      (*endForCallee)(retrievedNestedkID);
-    }
+     if ( !(infokIDSample.find(kID) == infokIDSample.end()) )
+      {
+       uint64_t retrievedNestedkID = infokIDSample.at(kID);
+       if (tool_verbosity > 0) { 
+       printf("KokkosP: sample %llu calling child-end function...\n",
+             (unsigned long long)(kID));
+        }
+       (*endForCallee)(retrievedNestedkID);
+       }
+     }
 }
 
 void kokkosp_begin_parallel_scan(const char* name, const uint32_t devID,