From 696b0253e597a38edb948daf3278adc52a69b004 Mon Sep 17 00:00:00 2001
From: dangqingqing <dangqingqing@baidu.com>
Date: Tue, 28 Nov 2017 18:28:35 +0800
Subject: [PATCH] Refine paddle/v2/fluid/profiler.py.

---
 paddle/platform/cuda_profiler.h               |  8 +-
 python/paddle/v2/fluid/profiler.py            | 78 ++++++-------------
 python/paddle/v2/fluid/tests/test_profiler.py |  2 +-
 3 files changed, 30 insertions(+), 58 deletions(-)

diff --git a/paddle/platform/cuda_profiler.h b/paddle/platform/cuda_profiler.h
index c096ce37c5..b6311cb23d 100644
--- a/paddle/platform/cuda_profiler.h
+++ b/paddle/platform/cuda_profiler.h
@@ -29,10 +29,10 @@ void CudaProfilerInit(std::string output_file, std::string output_mode,
   memcpy(buf.data(), tmpl.data(), tmpl.size());
   auto result = mktemp(buf.data());
   PADDLE_ENFORCE(strlen(result) != 0);
-  std::string config = result;
+  std::string config_file = result;
 
   {
-    std::ofstream ofs(config, std::ios::out | std::ios::trunc);
+    std::ofstream ofs(config_file, std::ios::out | std::ios::trunc);
     PADDLE_ENFORCE(ofs.is_open(), "ofstream: ", ofs.rdstate());
     for (const auto& line : config_flags) {
       ofs << line << std::endl;
@@ -42,12 +42,12 @@ void CudaProfilerInit(std::string output_file, std::string output_mode,
   PADDLE_ENFORCE(output_mode == "kvp" || output_mode == "csv");
   cudaOutputMode_t mode = output_mode == "csv" ? cudaCSV : cudaKeyValuePair;
   PADDLE_ENFORCE(
-      cudaProfilerInitialize(config.c_str(), output_file.c_str(), mode));
+      cudaProfilerInitialize(config_file.c_str(), output_file.c_str(), mode));
 }
 
 void CudaProfilerStart() { PADDLE_ENFORCE(cudaProfilerStart()); }
 
-void CudaProfilerStop() { PADDLE_ENFORCE((cudaProfilerStop())); }
+void CudaProfilerStop() { PADDLE_ENFORCE(cudaProfilerStop()); }
 
 }  // namespace platform
 }  // namespace paddle
diff --git a/python/paddle/v2/fluid/profiler.py b/python/paddle/v2/fluid/profiler.py
index f31d6f0a61..2dbba72c64 100644
--- a/python/paddle/v2/fluid/profiler.py
+++ b/python/paddle/v2/fluid/profiler.py
@@ -1,9 +1,9 @@
 import paddle.v2.fluid.core as core
-import subprocess
+from contextlib import contextmanager
 
 __all__ = ['CudaProfiler']
 
-NV_FLAGS = [
+NVPROF_CONFIG = [
     "gpustarttimestamp",
     "gpuendtimestamp",
     "gridsize3d",
@@ -14,61 +14,33 @@ NV_FLAGS = [
 ]
 
 
-def nvporf_init(output_file, output_mode=None, flags=None):
-    """
-    Initialize the CUDA profiler.
-    This methods must be called before nvprof_start.
-
-    :param output_file: The output file name.
-    :type output_file: string
-    :param output_mode: The output mode has Key-Value pair format and
-                        Comma separated values format.
-                        It should be 'kv' or 'csv'.
-    :type output_mode: string
+@contextmanager
+def cuda_profiler(output_file, output_mode=None, config=None):
+    """The CUDA profiler.
+    This fuctions is used to profile CUDA program by CUDA runtime application
+    programming interface. The profiling result will be written into
+    `output_file` with Key-Value pair format or Comma separated values format.
+    The user can set the output mode by `output_mode` argument and set the
+    counters/options for profiling by `config` argument. The default config
+    caontains 'gpustarttimestamp', 'gpustarttimestamp', 'gridsize3d',
+    'threadblocksize', 'streamid', 'enableonstart 0', 'conckerneltrace'.
+
+    Args:
+        output_file (string) : The output file name, the result will be
+            written into this file.
+        output_mode (string) : The output mode has Key-Value pair format and
+            Comma separated values format. It should be 'kv' or 'csv'.
+        config (string) : The profiler options and counters can refer to
+            "Compute Command Line Profiler User Guide".
     """
     if output_mode is None:
         output_mode = 'csv'
     if output_mode not in ['kv', 'csv']:
         raise ValueError("The output mode must be 'key-value' or 'csv'.")
-    flags = NV_FLAGS if flags is None else flags
-    core.nvprof_init(output_file, output_mode, flags)
-
-
-def nvporf_start():
-    """
-    Enables profiler collection by the active CUDA profiling tool.
-    """
+    config = NVPROF_CONFIG if config is None else config
+    core.nvprof_init(output_file, output_mode, config)
+    # Enables profiler collection by the active CUDA profiling tool.
     core.nvprof_start()
-
-
-def nvporf_stop():
-    """
-    Disables profiler collection.
-    """
+    yield
+    # Disables profiler collection.
     core.nvprof_stop()
-
-
-class CudaProfiler(object):
-    def __init__(self, output_file, output_mode=None, flags=None, enabled=True):
-        self.enabled = enabled
-        if not self.enabled:
-            return
-        self.entered = False
-        self.out_file = output_file
-        nvporf_init(output_file, output_mode, flags)
-
-    def __enter__(self):
-        if not self.enabled:
-            return
-        if self.entered:
-            raise RuntimeError("The profiler traces are not reentrant")
-        self.entered = True
-        nvporf_start()
-        return self
-
-    def __exit__(self, exc_type, exc_value, tb):
-        if exc_value is not None:
-            raise exc_value
-        if not self.enabled:
-            return
-        nvporf_stop()
diff --git a/python/paddle/v2/fluid/tests/test_profiler.py b/python/paddle/v2/fluid/tests/test_profiler.py
index 1fec5c99bf..e8f24251b9 100644
--- a/python/paddle/v2/fluid/tests/test_profiler.py
+++ b/python/paddle/v2/fluid/tests/test_profiler.py
@@ -18,7 +18,7 @@ class TestProfiler(unittest.TestCase):
         exe = fluid.Executor(place)
         exe.run(fluid.default_startup_program())
 
-        with profiler.CudaProfiler("cuda_profiler.txt", 'csv') as nvprof:
+        with profiler.cuda_profiler("cuda_profiler.txt", 'csv') as nvprof:
             for i in range(epoc):
                 input = np.random.random(dshape).astype("float32")
                 exe.run(fluid.default_main_program(), feed={'data': input})
-- 
GitLab