Extract CudaProfiler from test_harness and graduate to cuda.h

Author: nicolasvasilache

include/tc/core/cuda/cuda.h

@@ -27,6 +27,7 @@
 #include <stdexcept>
 
 #include <cuda.h>
+#include <cuda_profiler_api.h>
 #include <cuda_runtime.h>
 
 #include <glog/logging.h>
@@ -70,6 +71,8 @@
 
 namespace tc {
 
+DECLARE_bool(use_nvprof);
+
 struct WithDevice {
   WithDevice(size_t g) : newGpu(g) {
     int dev;
@@ -111,4 +114,17 @@ class CudaGPUInfo {
   std::vector<size_t> sharedMemSizes_;
 };
 
+struct CudaProfiler {
+  CudaProfiler() {
+    if (FLAGS_use_nvprof) {
+      cudaProfilerStart();
+    }
+  }
+  ~CudaProfiler() {
+    if (FLAGS_use_nvprof) {
+      cudaProfilerStop();
+    }
+  }
+};
+
 } // namespace tc

src/core/cuda/cuda.cc

@@ -26,6 +26,8 @@
 #include "tc/core/flags.h"
 
 namespace tc {
+DEFINE_bool(use_nvprof, false, "Start / stop nvprof");
+
 namespace {
 
 std::tuple<std::vector<std::string>, std::vector<size_t>> init() {

test/test_harness.h

@@ -15,7 +15,6 @@
  */
 #pragma once
 
-#include <cuda_profiler_api.h>
 #include <gtest/gtest.h>
 #include <mutex>
 #include <string>
@@ -32,23 +31,8 @@
 #include "tc/c2/tc_op.h"
 #include "tc/core/cuda/cuda.h"
 
-DEFINE_bool(use_nvprof, false, "Start / stop nvprof");
-
 namespace caffe2 {
 
-struct CudaProfiler {
-  CudaProfiler() {
-    if (FLAGS_use_nvprof) {
-      cudaProfilerStart();
-    }
-  }
-  ~CudaProfiler() {
-    if (FLAGS_use_nvprof) {
-      cudaProfilerStop();
-    }
-  }
-};
-
 caffe2::TensorCPU context2tensor(caffe2::CPUContext& ctx) {
   return caffe2::TensorCPU();
 }
@@ -315,7 +299,7 @@ struct TestHarness {
 
     void RunReference() {
       ASSERT_TRUE(net_ref.get());
-      CudaProfiler p;
+      tc::CudaProfiler p;
       ASSERT_TRUE(net_ref->Run());
     }
 
@@ -326,7 +310,7 @@ struct TestHarness {
 
     void Run() {
       ASSERT_TRUE(op_test.get());
-      CudaProfiler p;
+      tc::CudaProfiler p;
       ASSERT_TRUE(op_test->Run());
     }
 
@@ -406,7 +390,7 @@ struct TestHarness {
       unique_ptr<OperatorBase> op_g(CreateOperator(g_op, &w));
       ASSERT_TRUE(op_g.get());
       {
-        CudaProfiler p;
+        tc::CudaProfiler p;
         ASSERT_TRUE(op_g->Run());
       }
     }
@@ -424,7 +408,7 @@ struct TestHarness {
     unique_ptr<NetBase> ref_net(CreateNet(ref_net_def, &w1));
     ASSERT_TRUE(ref_net.get());
     {
-      CudaProfiler p;
+      tc::CudaProfiler p;
       ASSERT_TRUE(ref_net->Run());
     }
 
@@ -433,7 +417,7 @@ struct TestHarness {
     unique_ptr<NetBase> net(CreateNet(net_def, &w2));
     ASSERT_TRUE(net.get());
     {
-      CudaProfiler p;
+      tc::CudaProfiler p;
       ASSERT_TRUE(net->Run());
     }
 
@@ -467,7 +451,7 @@ struct TestHarness {
     unique_ptr<NetBase> net(CreateNet(net_def, &w1));
     ASSERT_TRUE(net.get());
     {
-      CudaProfiler p;
+      tc::CudaProfiler p;
       ASSERT_TRUE(net->Run());
     }
     RunGradient(w1, *net_def.mutable_op()->Mutable(0));
@@ -477,7 +461,7 @@ struct TestHarness {
     unique_ptr<OperatorBase> op(CreateOperator(op_def, &w2));
     ASSERT_TRUE(op.get());
     {
-      CudaProfiler p;
+      tc::CudaProfiler p;
       ASSERT_TRUE(op->Run());
     }
     OperatorDef def = op_def;