kp_sampler_skip.cpp: add deep copy support

kokkos · Nov 2, 2023 · 5d05d38 · 5d05d38
1 parent 2ddedef
commit 5d05d38
Showing 1 changed file with 52 additions and 1 deletion.
diff --git a/common/kokkos-sampler/kp_sampler_skip.cpp b/common/kokkos-sampler/kp_sampler_skip.cpp
@@ -27,9 +27,11 @@ static finalizeFunction finalizeProfileLibrary = NULL;
 static beginFunction beginForCallee            = NULL;
 static beginFunction beginScanCallee           = NULL;
 static beginFunction beginReduceCallee         = NULL;
+static beginFunction beginDeepCopyCallee       = NULL;
 static endFunction endForCallee                = NULL;
 static endFunction endScanCallee               = NULL;
 static endFunction endReduceCallee             = NULL;
+static endFunction endDeepCopyCallee           = NULL;
 
 void kokkosp_request_tool_settings(const uint32_t,
                                    Kokkos_Tools_ToolSettings* settings) {
@@ -134,14 +136,19 @@ void kokkosp_init_library(const int loadSeq, const uint64_t interfaceVer,
           (beginFunction)dlsym(childLibrary, "kokkosp_begin_parallel_scan");
       beginReduceCallee =
           (beginFunction)dlsym(childLibrary, "kokkosp_begin_parallel_reduce");
-
+      beginDeepCopyCallee =
+          (beginFunction)dlsym(childLibrary, "kokkosp_begin_deep_copy");
+
       endScanCallee =
           (endFunction)dlsym(childLibrary, "kokkosp_end_parallel_scan");
       endForCallee =
           (endFunction)dlsym(childLibrary, "kokkosp_end_parallel_for");
       endReduceCallee =
           (endFunction)dlsym(childLibrary, "kokkosp_end_parallel_reduce");
+      endDeepCopyCallee =
+          (endFunction)dlsym(childLibrary, "kokkosp_end_deep_copy");
 
+
       initProfileLibrary =
           (initFunction)dlsym(childLibrary, "kokkosp_init_library");
       finalizeProfileLibrary =
@@ -160,12 +167,16 @@ void kokkosp_init_library(const int loadSeq, const uint64_t interfaceVer,
                (beginScanCallee == NULL) ? "no" : "yes");
         printf("KokkosP: begin-parallel-reduce:   %s\n",
                (beginReduceCallee == NULL) ? "no" : "yes");
+        printf("KokkosP: begin-deep-copy:   %s\n",
+               (beginDeepCopyCallee == NULL) ? "no" : "yes");
         printf("KokkosP: end-parallel-for:        %s\n",
                (endForCallee == NULL) ? "no" : "yes");
         printf("KokkosP: end-parallel-scan:       %s\n",
                (endScanCallee == NULL) ? "no" : "yes");
         printf("KokkosP: end-parallel-reduce:     %s\n",
                (endReduceCallee == NULL) ? "no" : "yes");
+        printf("KokkosP: end-deep-copy:     %s\n",
+               (endDeepCopyCallee == NULL) ? "no" : "yes");
       }
     }
   }
@@ -302,6 +313,44 @@ void kokkosp_end_parallel_reduce(const uint64_t kID) {
   }
 }
 
+void kokkosp_begin_deep_copy(const char* name, const uint32_t devID,
+                                   uint64_t* kID) {
+  *kID                          = uniqID++;
+  static uint64_t invocationNum = 0;
+  ++invocationNum;
+  if ((invocationNum % kernelSampleSkip) == 0) {
+    if (tool_verbosity > 0) {
+      printf("KokkosP: sample %llu calling child-begin function...\n",
+             (unsigned long long)(*kID));
+    }
+    if (NULL != beginDeepCopyCallee) {
+      uint64_t nestedkID = 0;
+      if (tool_globFence) {
+        invoke_ktools_fence(0);
+      }
+      (*beginDeepCopyCallee)(name, devID, &nestedkID);
+      infokIDSample.insert({*kID, nestedkID});
+    }
+  }
+}
+
+void kokkosp_end_deep_copy(const uint64_t kID) {
+  if (NULL != endDeepCopyCallee) {
+    if (!(infokIDSample.find(kID) == infokIDSample.end())) {
+      uint64_t retrievedNestedkID = infokIDSample[kID];
+      if (tool_verbosity > 0) {
+        printf("KokkosP: sample %llu calling child-end function...\n",
+               (unsigned long long)(kID));
+      }
+      if (tool_globFence) {
+        invoke_ktools_fence(0);
+      }
+      (*endDeepCopyCallee)(retrievedNestedkID);
+      infokIDSample.erase(kID);
+    }
+  }
+} 
+
 }  // namespace Sampler
 }  // end namespace KokkosTools
 
@@ -319,5 +368,7 @@ EXPOSE_BEGIN_PARALLEL_SCAN(impl::kokkosp_begin_parallel_scan)
 EXPOSE_END_PARALLEL_SCAN(impl::kokkosp_end_parallel_scan)
 EXPOSE_BEGIN_PARALLEL_REDUCE(impl::kokkosp_begin_parallel_reduce)
 EXPOSE_END_PARALLEL_REDUCE(impl::kokkosp_end_parallel_reduce)
+EXPOSE_BEGIN_DEEP_COPY(impl::kokkosp_begin_deep_copy)
+EXPOSE_END_DEEP_COPY(impl::kokkosp_end_deep_copy)
 
 }  // end extern "C"