Add NVTX support via DrHook (Refactored)

CMakeLists.txt

@@ -1,4 +1,5 @@
 # (C) Copyright 2020- ECMWF.
+# (C) Copyright 2024- Meteo-France.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
@@ -50,6 +51,12 @@ ecbuild_add_option( FEATURE WARNINGS
                     DEFAULT ON
                     DESCRIPTION "Add warnings to compiler" )
 
+
+ecbuild_add_option( FEATURE DR_HOOK_NVTX
+                    DEFAULT ${DEFAULT_DR_HOOK_NVTX}
+                    DESCRIPTION "Support for NVTX in DR_HOOK"
+                    REQUIRED_PACKAGES NVTX )
+
 ecbuild_find_package( NAME Realtime QUIET )
 
 ### Sources

cmake/FindNVTX.cmake

@@ -0,0 +1,67 @@
+# (C) Copyright 2024- ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+if(CMAKE_C_COMPILER_ID STREQUAL "PGI" OR CMAKE_C_COMPILER_ID STREQUAL "NVHPC" )
+
+    set (DEFAULT_DR_HOOK_NVTX ON)
+
+    if( ${CMAKE_VERSION} VERSION_LESS "3.25" )
+        find_package(CUDAToolkit COMPONENTS CUDA::nvToolsExt)
+
+        find_path(NVTX_ROOT
+                NAMES include/nvToolsExt.h
+                HINTS ${CUDAToolkit_LIBRARY_DIR}/..
+                )
+
+        find_library(NVTX_LIBRARIES
+                NAMES libnvToolsExt.so nvToolsExt
+                HINTS ${NVTX_ROOT}/lib ${NVTX_ROOT}/lib64
+                )
+
+        find_path(NVTX_INCLUDE_DIRS
+                NAMES nvToolsExt.h
+                HINTS ${NVTX_ROOT}/include
+                )
+
+        include(FindPackageHandleStandardArgs)
+        find_package_handle_standard_args(NVTX DEFAULT_MSG
+                NVTX_LIBRARIES
+                NVTX_INCLUDE_DIRS
+                )
+
+        mark_as_advanced(
+                NVTX_LIBRARIES
+                NVTX_INCLUDE_DIRS
+                )
+    else()
+        find_package(CUDAToolkit COMPONENTS CUDA::nvtx3)
+
+        find_path(NVTX_ROOT
+                NAMES include/nvtx3/nvToolsExt.h
+                HINTS ${CUDAToolkit_LIBRARY_DIR}/..
+                )
+
+        find_path(NVTX_INCLUDE_DIRS
+                NAMES nvToolsExt.h
+                HINTS ${NVTX_ROOT}/include/nvtx3
+                )
+
+        include(FindPackageHandleStandardArgs)
+        find_package_handle_standard_args(NVTX DEFAULT_MSG
+                NVTX_INCLUDE_DIRS
+                )
+
+        mark_as_advanced(
+                NVTX_INCLUDE_DIRS
+                )
+    endif()
+
+
+else ()
+    set (DEFAULT_DR_HOOK_NVTX OFF)
+endif ()

cmake/project_summary.cmake

@@ -23,5 +23,10 @@ ecbuild_info( "MPI (export MPI_HOME to correct MPI implementation)" )
 ecbuild_info( "    MPI_Fortran_INCLUDE_DIRS  : [${MPI_Fortran_INCLUDE_DIRS}]" )
 ecbuild_info( "    MPI_Fortran_LIBRARIES     : [${MPI_Fortran_LIBRARIES}]" )
 ecbuild_info( "    MPIEXEC                   : [${MPIEXEC}]" )
+
+if( HAVE_DR_HOOK_NVTX )
+ecbuild_info( "NVTX_LIBRARIES                : [${NVTX_LIBRARIES}]" )
+endif()
+
 ecbuild_info( "---------------------------------------------------------" )
 

src/fiat/CMakeLists.txt

@@ -1,4 +1,5 @@
 # (C) Copyright 2020- ECMWF.
+# (C) Copyright 2024- Meteo-France.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
@@ -44,6 +45,9 @@ endif()
 configure_file( ${CMAKE_CURRENT_SOURCE_DIR}/library/version.c.in ${CMAKE_CURRENT_BINARY_DIR}/version.c @ONLY )
 
 ecbuild_list_add_pattern( LIST fiat_src GLOB *.c *.F* *.cc )
+
+ecbuild_list_exclude_pattern( LIST fiat_src REGEX drhook/extensions/*)
+
 set( fiat_src ${fiat_src} PARENT_SCOPE )
 
 ecbuild_add_library( TARGET fiat
@@ -65,6 +69,19 @@ ecbuild_add_library( TARGET fiat
         $<INSTALL_INTERFACE:include/fiat>
 )
 
+if (HAVE_DR_HOOK_NVTX)
+    # Files from within DrHook
+    ecbuild_list_add_pattern( LIST fiat_nvtx_src GLOB *.c SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/drhook/extensions/nvtx)
+    target_sources(fiat PRIVATE ${fiat_nvtx_src})
+    target_include_directories(fiat PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/drhook/extensions/nvtx)
+
+    # Files defined externally
+    target_include_directories( fiat PRIVATE ${NVTX_INCLUDE_DIRS} )
+    target_compile_definitions( fiat PRIVATE DR_HOOK_HAVE_NVTX=1 )
+    if( CMAKE_VERSION VERSION_LESS 3.25 )
+        target_link_libraries     ( fiat PRIVATE ${NVTX_LIBRARIES} )
+    endif()
+endif()
 
 if( ${CMAKE_SYSTEM_NAME} MATCHES "Darwin" )
   # Following should not be necessary; 

src/fiat/drhook/drhook.c

@@ -82,6 +82,9 @@ static int backtrace(void **buffer, int size) { return 0; }
 #include <sys/stat.h>
 #include <fcntl.h>
 #include <dlfcn.h>
+#ifdef DR_HOOK_HAVE_NVTX
+#include "dr_hook_nvtx.h"
+#endif
 
 #include "ec_get_cycles.h"
 static long long int *thread_cycles = NULL;
@@ -319,6 +322,13 @@ static int callpath_indent = callpath_indent_default;
 #define callpath_depth_default 50
 static int callpath_depth = callpath_depth_default;
 static int callpath_packed = 0;
+static int opt_nvtx = 0;
+#define nvtx_SCC_default 10
+static int opt_nvtx_SCC = nvtx_SCC_default;
+#define nvtx_SWT_default 0.0001
+static double opt_nvtx_SWT = nvtx_SWT_default;
+static int opt_strict_regions = 0;
+static int opt_silent = 0;
 
 static int opt_calltrace = 0;
 static int opt_funcenter = 0;
@@ -471,6 +481,9 @@ typedef struct drhook_key_t {
   long long int mem_maxhwm, mem_maxrss, mem_maxstk, mem_maxpagdelta;
   long long int paging_in;
   unsigned long long int alloc_count, free_count;
+#if defined(DR_HOOK_HAVE_NVTX)
+  unsigned long long int skipped_nvtx_calls;
+#endif
   struct drhook_key_t *next;
 } drhook_key_t;
 
@@ -2195,16 +2208,18 @@ process_options()
 
   tid = drhook_oml_get_thread_num();
 
-  int silent = 0;
   env = getenv("DR_HOOK_SILENT");
-  silent = env ? atoi(env) : silent;
+  if (env) {
+    opt_silent = atoi(env);
+  }
 
   env = getenv("DR_HOOK_SHOW_PROCESS_OPTIONS");
-  ienv = env ? atoi(env) : silent ? 0 : 1;
+  ienv = env ? atoi(env) : opt_silent ? 0 : 1;
   if (ienv == -1 || ienv == myproc) fp = stderr;
   if (fp) pfx = PREFIX(tid);
 
   if(fp) fprintf(fp,"[EC_DRHOOK:hostname:myproc:omltid:pid:unixtid] [YYYYMMDD:HHMMSS:walltime] [function@file:lineno] -- Max OpenMP threads = %d\n",drhook_oml_get_max_threads());
+  OPTPRINT(fp,"%s %s [%s@%s:%d] DR_HOOK_SILENT=%d\n",pfx,TIMESTR(tid),FFL,opt_silent);
 
   OPTPRINT(fp,"%s %s [%s@%s:%d] fp = %p\n",pfx,TIMESTR(tid),FFL,(void*)fp);
 
@@ -2478,6 +2493,48 @@ process_options()
     OPTPRINT(fp,"%s %s [%s@%s:%d] DR_HOOK_GENCORE_SIGNAL=%d\n",pfx,TIMESTR(tid),FFL,opt_gencore_signal);
   }
 
+  env = getenv("DR_HOOK_STRICT_REGIONS");
+  int strict_regions_opt_touched = 0;
+  if (env) {
+    opt_strict_regions = atoi(env);
+    strict_regions_opt_touched = 1;
+  }
+
+  env = getenv("DR_HOOK_NVTX");
+  if (env) {
+    opt_nvtx = atoi(env);
+    opt_strict_regions = opt_strict_regions || opt_nvtx;
+    strict_regions_opt_touched = 1;
+    opt_walltime = 1;
+    opt_calls = 1;
+    OPTPRINT(fp,"%s %s [%s@%s:%d] DR_HOOK_NVTX=%d\n",pfx,TIMESTR(tid),FFL,opt_nvtx);
+  }
+
+  if (strict_regions_opt_touched)
+    OPTPRINT(fp,"%s %s [%s@%s:%d] DR_HOOK_STRICT_REGIONS=%d\n",pfx,TIMESTR(tid),FFL,opt_strict_regions);
+
+  if (opt_nvtx) {
+    env = getenv("DR_HOOK_NVTX_SPAM_CALL_COUNT");
+    if (env) {
+      opt_nvtx_SCC = atoi(env);
+
+      if (opt_nvtx_SCC < 0)
+        opt_nvtx_SCC = nvtx_SCC_default;
+
+      OPTPRINT(fp,"%s %s [%s@%s:%d] DR_HOOK_NVTX_SPAM_CALL_COUNT=%d\n",pfx,TIMESTR(tid),FFL,opt_nvtx_SCC);
+    }
+
+    env = getenv("DR_HOOK_NVTX_SPAM_WT");
+    if (env) {
+      opt_nvtx_SWT = atof(env);
+
+      if (opt_nvtx_SWT < 0)
+        opt_nvtx_SWT = nvtx_SWT_default;
+
+      OPTPRINT(fp, "%s %s [%s@%s:%g] DR_HOOK_NVTX_SPAM_WT=%g\n", pfx, TIMESTR(tid), FFL, nvtx_SWT_default);
+    }
+  }
+
   newline = 0;
   env = getenv("DR_HOOK_OPT");
   if (env) {
@@ -2786,6 +2843,18 @@ getkey(int tid, const char *name, int name_len,
           keyptr->calls++;
           keyptr->status++;
         }
+#if defined(DR_HOOK_HAVE_NVTX)
+        // Helps filter out wrapper calls that may be noise
+        if (opt_nvtx && drhook_oml_get_thread_num() == 1){
+          if (keyptr->calls > opt_nvtx_SCC && keyptr->delta_wall_all < opt_nvtx_SWT) {
+            if (!opt_silent)
+              fprintf(stderr,"DRHOOK:NVTX: Skipping opening of region %s\n", keyptr->name);
+            keyptr->skipped_nvtx_calls++;
+          }
+          else
+            dr_hook_nvtx_start(keyptr->name);
+        }
+#endif
         insert_calltree(tid, keyptr);
         break; /* for (;;) */
       }
@@ -2811,7 +2880,12 @@ putkey(int tid, drhook_key_t *keyptr, const char *name, int name_len,
   const int sig = SIGABRT;
   const char sl_name[] = "SIGABRT";
   drhook_calltree_t *treeptr = (tid >= 1 && tid <= numthreads) ? thiscall[tid-1] : NULL;
-  if (!treeptr || !treeptr->active || treeptr->keyptr != keyptr) {
+
+  int regions_mismatch = 0;
+  if (opt_strict_regions)
+    regions_mismatch = strncasecmp(keyptr->name, name, name_len);
+
+  if (!treeptr || !treeptr->active || treeptr->keyptr != keyptr || regions_mismatch) {
     char *pfx = PREFIX(tid);
     char *s;
     unsigned int hash;
@@ -2911,6 +2985,17 @@ putkey(int tid, drhook_key_t *keyptr, const char *name, int name_len,
     if (opt_walltime) keyptr->delta_wall_all   += delta_wall;
     if (opt_cputime)  keyptr->delta_cpu_all    += delta_cpu;
     if (opt_cycles)   keyptr->delta_cycles_all += delta_cycles;
+#if defined(DR_HOOK_HAVE_NVTX)
+    if (opt_nvtx && drhook_oml_get_thread_num() == 1) {
+      if (keyptr->skipped_nvtx_calls > 0) {
+        if (!opt_silent)
+          fprintf(stderr, "DRHOOK:NVTX: Skipping closing of region %s\n", keyptr->name);
+        keyptr->skipped_nvtx_calls--;
+      } else {
+        dr_hook_nvtx_end();
+      }
+    }
+#endif
     remove_calltree(tid, keyptr, &delta_wall, &delta_cpu, &delta_cycles);
   }
 }

src/fiat/drhook/extensions/nvtx/dr_hook_nvtx.c

@@ -0,0 +1,67 @@
+/*
+ * (C) Copyright 2024- ECMWF.
+ * (C) Copyright 2024- Meteo-France.
+ *
+ * This software is licensed under the terms of the Apache Licence Version 2.0
+ * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+ * In applying this licence, ECMWF does not waive the privileges and immunities
+ * granted to it by virtue of its status as an intergovernmental organisation
+ * nor does it submit to any jurisdiction.
+ */
+
+#include <nvToolsExt.h>
+#include <string.h>
+#include <stdlib.h>
+
+#include "dr_hook_nvtx.h"
+
+static uint32_t adler32 (const unsigned char *data)
+{
+  const uint32_t MOD_ADLER = 65521;
+  uint32_t a = 1, b = 0;
+  size_t index;
+
+  for (index = 0; data[index] != 0; ++index)
+    {
+      a = (a + data[index]*2) % MOD_ADLER;
+      b = (b + a) % MOD_ADLER;
+    }
+
+  return (b << 16) | a;
+}
+
+
+void dr_hook_nvtx_start (const char * name)
+{
+  int hash = 0;
+  int color_id = adler32 ((const unsigned char*)name);
+  int r,g,b;
+
+  r=color_id & 0x000000ff;
+  g=(color_id & 0x000ff000) >> 12;
+  b=(color_id & 0x0ff00000) >> 20;
+
+  if (r<64 & g<64 & b<64) 
+    {
+      r=r*3;
+      g=g*3+64;
+      b=b*4;
+    }
+
+  color_id = 0xff000000 | (r << 16) | (g << 8) | (b);
+
+  nvtxEventAttributes_t eventAttrib = {0};
+  eventAttrib.version       = NVTX_VERSION;
+  eventAttrib.size          = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
+  eventAttrib.colorType     = NVTX_COLOR_ARGB;
+  eventAttrib.color         = color_id;
+  eventAttrib.messageType   = NVTX_MESSAGE_TYPE_ASCII;
+  eventAttrib.message.ascii = name;
+
+  nvtxRangePushEx (&eventAttrib);
+
+}
+
+void dr_hook_nvtx_end () {
+  nvtxRangePop ();
+}

src/fiat/drhook/extensions/nvtx/dr_hook_nvtx.h

@@ -0,0 +1,17 @@
+/*
+ * (C) Copyright 2024- ECMWF.
+ *
+ * This software is licensed under the terms of the Apache Licence Version 2.0
+ * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+ * In applying this licence, ECMWF does not waive the privileges and immunities
+ * granted to it by virtue of its status as an intergovernmental organisation
+ * nor does it submit to any jurisdiction.
+ */
+
+#ifndef FIAT_DRHOOK_NVTX_DR_HOOK_NVTX_H
+#define FIAT_DRHOOK_NVTX_DR_HOOK_NVTX_H
+
+void dr_hook_nvtx_start (const char* name);
+void dr_hook_nvtx_end ();
+
+#endif //FIAT_DRHOOK_NVTX_DR_HOOK_NVTX_H

tests/CMakeLists.txt

@@ -39,12 +39,8 @@ add_test(NAME fiat_test_abort_exception_handler
   -P ${CMAKE_CURRENT_SOURCE_DIR}/test_program_output.cmake )
 endif()
 
-
 add_subdirectory( drhook )
 
-
-
-
 # ----------------------------------------------------------------------------------------
 # Test: fiat_test_drhook_no_output
 #