mpl/ze: use level zero tracing API to implement zeMemFree callback

Author: zhenggb72

src/mpi/init/init_gpu.c

@@ -22,7 +22,7 @@ int MPII_init_gpu(void)
     if (MPIR_CVAR_ENABLE_GPU) {
         int debug_summary = 0;
         if (MPIR_CVAR_DEBUG_SUMMARY) {
-            debug_summary = (MPIR_Process.rank == 0);
+            debug_summary = (MPIR_Process.rank == 0) ? MPIR_CVAR_DEBUG_SUMMARY : 0;
         }
 
         bool specialized_cache =

src/mpl/src/gpu/mpl_gpu_ze.c

@@ -14,6 +14,16 @@
 
 MPL_SUPPRESS_OSX_HAS_NO_SYMBOLS_WARNING;
 
+/* 1: loader  0: driver */
+#define USELOADERTRACING    1
+
+#if USELOADERTRACING == 1
+#include "level_zero/layers/zel_tracing_api.h"
+#include "level_zero/layers/zel_tracing_register_cb.h"
+#else
+#include <level_zero/zet_api.h>
+#endif
+
 #ifdef MPL_HAVE_ZE
 #include <dirent.h>
 #if defined(MPL_HAVE_DRM_I915_DRM_H)
@@ -211,6 +221,14 @@ static pFnzexMemGetIpcHandles zexMemGetIpcHandles = NULL;
 static pFnzexMemOpenIpcHandles zexMemOpenIpcHandles = NULL;
 /* *INDENT-ON* */
 
+/* tracing */
+#if USELOADERTRACING == 1
+static zel_tracer_handle_t freeTracer;
+#else
+static zet_tracer_exp_handle_t freeTracer;
+#endif
+static int callbacksEnabled = 0;
+
 /* Backend-specific functions */
 /* *INDENT-OFF* */
 typedef ze_result_t (*pFnzexDriverImportExternalPointer)(ze_driver_handle_t, void *, size_t);
@@ -235,6 +253,8 @@ static void MPL_event_pool_destroy(void);
 #ifdef ZE_PCI_PROPERTIES_EXT_NAME
 static int search_physical_devices(ze_pci_address_ext_t pci);
 #endif
+static void onEnterMemFree(ze_mem_free_params_t * tracerParams, ze_result_t result,
+                           void *traceUserData, void **tracerInstanceUserData);
 
 /* For zeMemFree callbacks */
 static gpu_free_hook_s *free_hook_chain = NULL;
@@ -648,7 +668,8 @@ int MPL_gpu_init(int debug_summary)
             ipc_max_entries[i] = max_cache_entries;
         }
 
-        MPL_gpu_free_hook_register(MPL_ze_ipc_remove_cache_handle);
+        if (!callbacksEnabled)
+            MPL_gpu_free_hook_register(MPL_ze_ipc_remove_cache_handle);
     } else {
         MPL_gpu_info.specialized_cache = false;
     }
@@ -661,9 +682,11 @@ int MPL_gpu_init(int debug_summary)
 
     mypid = getpid();
 
-    MPL_initlock_lock(&free_hook_mutex);
-    gpu_mem_hook_init();
-    MPL_initlock_unlock(&free_hook_mutex);
+    if (!callbacksEnabled) {
+        MPL_initlock_lock(&free_hook_mutex);
+        gpu_mem_hook_init();
+        MPL_initlock_unlock(&free_hook_mutex);
+    }
     gpu_initialized = 1;
 
     if (MPL_gpu_info.debug_summary) {
@@ -1217,6 +1240,45 @@ static int gpu_ze_init_driver(void)
                                             (void **) &zexDriverGetHostPointerBaseAddress);
     ZE_ERR_CHECK(ret);
 
+    ret =
+        zeDriverGetExtensionFunctionAddress(ze_driver_handle, "zexMemGetIpcHandles",
+                                            (void **) &zexMemGetIpcHandles);
+    if (ZE_RESULT_SUCCESS != ret)
+        zexMemGetIpcHandles = NULL;
+
+    ret =
+        zeDriverGetExtensionFunctionAddress(ze_driver_handle, "zexMemOpenIpcHandles",
+                                            (void **) &zexMemOpenIpcHandles);
+    if (ZE_RESULT_SUCCESS != ret)
+        zexMemOpenIpcHandles = NULL;
+
+    /* tracing */
+    if (likely(MPL_gpu_info.specialized_cache)) {
+#if USELOADERTRACING == 1
+        zel_tracer_desc_t tracerDesc = { ZEL_STRUCTURE_TYPE_TRACER_DESC, NULL, NULL };
+        ret = zelTracerCreate(&tracerDesc, &freeTracer);
+        if (ret == ZE_RESULT_SUCCESS) {
+            callbacksEnabled = 1;
+            zelTracerMemFreeRegisterCallback(freeTracer, ZEL_REGISTER_PROLOGUE, onEnterMemFree);
+            zelTracerSetEnabled(freeTracer, true);
+            if (MPL_gpu_info.debug_summary)
+                printf("ZE loader tracing is enabled.\n");
+        }
+#else
+        zet_tracer_exp_desc_t tracerDesc = { ZET_STRUCTURE_TYPE_TRACER_EXP_DESC, NULL, NULL };
+        ret = zetTracerExpCreate(ze_context, &tracerDesc, &freeTracer);
+        if (ret == ZE_RESULT_SUCCESS) {
+            callbacksEnabled = 1;
+            ze_callbacks_t prologCbs = { };
+            prologCbs.Mem.pfnFreeCb = onEnterMemFree;
+            zetTracerExpSetPrologues(freeTracer, &prologCbs);
+            zetTracerExpSetEnabled(freeTracer, true);
+            if (MPL_gpu_info.debug_summary)
+                printf("ZE driver tracing is enabled. \n");
+        }
+#endif
+    }
+
   fn_exit:
     MPL_free(all_drivers);
     return ret_error;
@@ -1386,6 +1448,17 @@ int MPL_gpu_finalize(void)
         MPL_free(ipc_lru_mapped_order_head);
         MPL_free(ipc_lru_mapped_order_tail);
         MPL_free(ipc_curr_mapped_entries);
+
+        /* tracing */
+        if (callbacksEnabled) {
+#if USELOADERTRACING == 1
+            zelTracerSetEnabled(freeTracer, false);
+            zelTracerDestroy(freeTracer);
+#else
+            zetTracerExpSetEnabled(freeTracer, false);
+            zetTracerExpDestroy(freeTracer);
+#endif
+        }
     }
 
     MPL_free(ipc_max_entries);
@@ -2754,6 +2827,17 @@ ze_result_t ZE_APICALL zeMemFree(ze_context_handle_t hContext, void *dptr)
     return (result);
 }
 
+/* tracing */
+static void onEnterMemFree(ze_mem_free_params_t * tracerParams, ze_result_t result,
+                           void *traceUserData, void **tracerInstanceUserData)
+{
+    if (MPL_gpu_info.debug_summary == 2) {
+        printf("GPU memory being freed: %lx\n", (unsigned long int) *tracerParams->pptr);
+    }
+    MPL_ze_ipc_remove_cache_handle(*tracerParams->pptr);
+    *tracerInstanceUserData = NULL;
+}
+
 /* ZE-Specific Functions */
 
 #ifdef ZE_PCI_PROPERTIES_EXT_NAME