chore(gpu): add nvtx tool for profiling

backends/tfhe-cuda-backend/build.rs

@@ -39,6 +39,7 @@ fn main() {
         println!("cargo:rustc-link-lib=cudart");
         println!("cargo:rustc-link-search=native=/usr/lib/x86_64-linux-gnu/");
         println!("cargo:rustc-link-lib=stdc++");
+        println!("cargo:rustc-link-lib=nvToolsExt");
     } else {
         panic!(
             "Error: platform not supported, tfhe-cuda-backend not built (only Linux is supported)"

backends/tfhe-cuda-backend/cuda/CMakeLists.txt

@@ -67,7 +67,8 @@ endif()
 
 add_compile_definitions(CUDA_ARCH=${CUDA_ARCH})
 
-# in production, should use -arch=sm_70 --ptxas-options=-v to see register spills -lineinfo for better debugging
+# in production, should use -arch=sm_70 --ptxas-options=-v to see register spills -lineinfo for better debugging to use
+# nvtx when profiling -lnvToolsExt
 set(CMAKE_CUDA_FLAGS
     "${CMAKE_CUDA_FLAGS} -ccbin ${CMAKE_CXX_COMPILER} -O3 \
   -std=c++17 --no-exceptions  --expt-relaxed-constexpr -rdc=true \

backends/tfhe-cuda-backend/cuda/src/utils/helper_profile.cu

@@ -0,0 +1,35 @@
+#include "helper_profile.cuh"
+
+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 cuda_nvtx_label_with_color(const char *name) {
+  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);
+}

backends/tfhe-cuda-backend/cuda/src/utils/helper_profile.cuh

@@ -0,0 +1,11 @@
+#ifndef HELPER_PROFILE
+#define HELPER_PROFILE
+#include <nvToolsExt.h>
+
+void cuda_nvtx_label_with_color(const char *name);
+
+#define PUSH_RANGE(name)                                                       \
+  { cuda_nvtx_label_with_color(name); }
+#define POP_RANGE() nvtxRangePop();
+
+#endif