Merge pull request #172 from sony/feature/20190624-nvtx-support

Add NVTX utilities in Python

.gitignore

@@ -23,6 +23,7 @@
 /python/setup.cfg.in.configured
 /python/src/nnabla_ext/cuda/_version.py
 /python/src/nnabla_ext/cuda/init.cpp
+/python/src/nnabla_ext/cuda/nvtx.cpp
 /python/src/nnabla_ext/cuda/libnnabla_cuda.so
 /python/src/nnabla_ext/cuda/nnabla_cuda.dll
 /python/src/nnabla_ext/cuda/nnabla_cuda.lib

include/nbla/cuda/nvtx.hpp

@@ -0,0 +1,34 @@
+// Copyright (c) 2017 Sony Corporation. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef __NBLA_NVTX_HPP__
+#define __NBLA_NVTX_HPP__
+
+#include "nbla/cuda/defs.hpp"
+
+#include <string>
+
+namespace nbla {
+using std::string;
+
+NBLA_CUDA_API void nvtx_mark_A(string msg);
+
+NBLA_CUDA_API void nvtx_range_push_A(string msg);
+
+NBLA_CUDA_API void nvtx_range_push_with_C(string msg);
+
+NBLA_CUDA_API void nvtx_range_pop();
+}
+
+#endif //__NBLA_NVTX_HPP__

python/setup.py

@@ -143,6 +143,9 @@ def cuda_config(root_dir, cuda_lib, ext_opts):
         Extension(cuda_pkg + '.init',
                   [join(path_cuda_pkg, 'init.pyx')],
                   **cuda_ext_opts),
+        Extension(cuda_pkg + '.nvtx',
+                  [join(path_cuda_pkg, 'nvtx.pyx')],
+                  **cuda_ext_opts),
     ]
     return ExtConfig(package_dir, packages, package_data,
                      ext_modules, cuda_ext_opts)

python/src/nnabla_ext/cuda/nvtx.pyx

@@ -0,0 +1,34 @@
+# Copyright (c) 2017 Sony Corporation. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from libcpp.string cimport string
+
+cdef extern from "nbla/cuda/nvtx.hpp" namespace "nbla":
+    void nvtx_mark_A(string) except +
+    void nvtx_range_push_A(string) except +
+    void nvtx_range_push_with_C(string) except +
+    void nvtx_range_pop() except +
+
+
+def mark_A(str msg):
+    nvtx_mark_A(msg)
+
+def range_push(str msg, coloring=True):
+    if coloring:
+        nvtx_range_push_with_C(msg)
+    else:
+        nvtx_range_push_A(msg)
+
+def range_pop():
+    nvtx_range_pop()

src/nbla/cuda/CMakeLists.txt

@@ -45,6 +45,26 @@ if(WITH_NCCL)
 	  )
 endif()
 
+# nvprof support
+option(WITH_NVTX "Build with nvtx API for nvprof support" OFF)
+if(WITH_NVTX)
+    find_library(CUDA_NVTX_LIBRARY nvToolsExt
+            HINTS
+            "${CUDA_TOOLKIT_ROOT_DIR}/lib64"
+            "${CUDA_TOOLKIT_ROOT_DIR}/lib"
+            "${CUDA_TOOLKIT_ROOT_DIR}"
+            ENV CUDA_PATH
+            ENV CUDA_LIB_PATH)
+    if (NOT CUDA_NVTX_LIBRARY)
+        message(WARNING "nvToolsExt-libs is not found. Build without nvToolsExt.")
+    else()
+        list(APPEND NBLA_CUDA_LINKER_LIBS ${CUDA_NVTX_LIBRARY})
+        add_definitions(-DWITH_NVTX)
+        message("NVToolsExt: " ${CUDA_NVTX_LIBRARY})
+    endif()
+endif()
+
+
 ####################################################################################################
 # cuDNN
 

src/nbla/cuda/nvtx.cpp

@@ -0,0 +1,95 @@
+// Copyright (c) 2017 Sony Corporation. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "nbla/cuda/nvtx.hpp"
+
+#include <string>
+
+using std::string;
+
+////////////////////////////////////////
+// Define internal use only functions
+////////////////////////////////////////
+
+#ifdef WITH_NVTX
+#include <nvToolsExt.h>
+#include <vector>
+
+using std::vector;
+
+const vector<uint32_t> colors = {
+    0x00ff0000, // red
+    0x0000ff00, // green
+    0x000000ff, // blue
+    0x00dc143c, // crimson
+    0x00006400, // darkgreen
+    0x0000ffff, // cyan
+    0x00ffa500, // orange
+    0x008b008b, // darkmagenta
+    0x00000080  // navy
+};
+static int color_id = 0;
+
+inline void nvtxRangePushWithColorChange(const char *msg) {
+  nvtxEventAttributes_t eventAttrib = {0};
+  eventAttrib.version = NVTX_VERSION;
+  eventAttrib.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
+  eventAttrib.colorType = NVTX_COLOR_ARGB;
+  eventAttrib.color = colors[color_id++];
+  color_id %= colors.size();
+  eventAttrib.messageType = NVTX_MESSAGE_TYPE_ASCII;
+  eventAttrib.message.ascii = msg;
+
+  nvtxRangePushEx(&eventAttrib);
+}
+
+#else
+
+inline void nvtxMarkA(const char *) {}
+
+//
+// inline void nvtxMarkEx(...) {}
+//
+// inline int nvtxRangeStartA(const char *) { return 0; }
+//
+// inline int nvtxRangeStartEx(...) { return 0; }
+//
+// inline void nvtxRangeEnd(nvtxRangeId_t) {}
+//
+inline void nvtxRangePushA(const char *) {}
+//
+// inline void nvtxRangePushEx(...) {}
+//
+
+inline void nvtxRangePushWithColorChange(const char *) {}
+
+inline void nvtxRangePop() {}
+
+#endif // WITH_NVTX
+
+////////////////////////////////////////
+// Define functions used by cython
+////////////////////////////////////////
+
+namespace nbla {
+void nvtx_mark_A(string msg) { nvtxMarkA(msg.c_str()); };
+
+void nvtx_range_push_A(string msg) { nvtxRangePushA(msg.c_str()); };
+
+void nvtx_range_push_with_C(string msg) {
+  nvtxRangePushWithColorChange(msg.c_str());
+}
+
+void nvtx_range_pop() { nvtxRangePop(); };
+}