Remove GPU dependency from ProfileObserver (#17592)

Summary:
Remove GPU dependency and register ProfileObserver.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17592

Reviewed By: ezyang

Differential Revision: D14265801

Pulled By: mdschatz

fbshipit-source-id: f98c0c32653c64a8b087c58ece4f864dfbe1d4b8

caffe2/observers/CMakeLists.txt

@@ -1,26 +1,16 @@
 if(USE_OBSERVERS)
   message(STATUS "Include Observer library")
-  set(GLOB profile_observer_files profile_observer_*.cc)
   set(Caffe2_CONTRIB_OBSERVERS_CPU_SRC
+    "${CMAKE_CURRENT_SOURCE_DIR}/profile_observer.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/time_observer.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/runcnt_observer.cc"
   )
-  set(Caffe2_CONTRIB_OBSERVERS_GPU_SRC
-    "${CMAKE_CURRENT_SOURCE_DIR}/profile_observer_gpu.cc"
-  )
 
   set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} ${Caffe2_CONTRIB_OBSERVERS_CPU_SRC})
   set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} PARENT_SCOPE)
 
-  set(Caffe2_GPU_SRCS ${Caffe2_GPU_SRCS} ${Caffe2_CONTRIB_OBSERVERS_GPU_SRC})
-  set(Caffe2_GPU_SRCS ${Caffe2_GPU_SRCS} PARENT_SCOPE)
-
   # ---[ CPU test files
   file(GLOB tmp *_test.cc)
   set(Caffe2_CPU_TEST_SRCS ${Caffe2_CPU_TEST_SRCS} ${tmp})
   set(Caffe2_CPU_TEST_SRCS ${Caffe2_CPU_TEST_SRCS} PARENT_SCOPE)
-  exclude(Caffe2_CPU_TEST_SRCS "${Caffe2_CPU_TEST_SRCS}" ${profile_observer_files})
-
-  # ---[ GPU test files
-  set(Caffe2_GPU_TEST_SRCS "${CMAKE_CURRENT_SOURCE_DIR}/profile_observer_test.cc")
 endif()

caffe2/observers/profile_observer.cc

@@ -14,8 +14,8 @@
  * limitations under the License.
  */
 
-#include "caffe2/core/logging.h"
 #include "profile_observer.h"
+#include "caffe2/core/logging.h"
 
 namespace caffe2 {
 
@@ -64,52 +64,11 @@ void ProfileOperatorObserver::Dump() const {
 }
 
 void ProfileOperatorObserver::Start() {
-  auto cudaOp = dynamic_cast_if_rtti<const Operator<CUDAContext>*>(subject_);
-  if (cudaOp) {
-    auto context = cudaOp->getContext();
-    int device;
-    cudaGetDevice(&device);
-
-    cudaSetDevice(context->device_id());
-    cudaEventCreate(&start_);
-    cudaEventRecord(start_, context->cuda_stream());
-
-    cudaSetDevice(device);
-
-    cudaError_t error = cudaGetLastError();
-    if (error != cudaSuccess) {
-      CAFFE_THROW("Encountered CUDA error Start: ", cudaGetErrorString(error));
-    }
-  } else {
   start_time_ = timer_.MilliSeconds();
 }
-}
 
 void ProfileOperatorObserver::Stop() {
-  auto cudaOp = dynamic_cast_if_rtti<const Operator<CUDAContext>*>(subject_);
-  if (cudaOp) {
-    auto context = cudaOp->getContext();
-    int device;
-    cudaGetDevice(&device);
-
-    cudaSetDevice(context->device_id());
-    cudaEventCreate(&stop_);
-    cudaEventRecord(stop_, context->cuda_stream());
-    cudaEventSynchronize(stop_);
-    cudaEventElapsedTime(&run_time_, start_, stop_);
-    cudaEventDestroy(start_);
-    cudaEventDestroy(stop_);
-
-    cudaSetDevice(device);
-
-    cudaError_t error = cudaGetLastError();
-    if (error != cudaSuccess) {
-      CAFFE_THROW("Encountered CUDA error Stop: ", cudaGetErrorString(error));
-    }
-  } else {
   run_time_ = timer_.MilliSeconds() - start_time_;
-  }
-
   Dump();
 }
 

caffe2/observers/profile_observer.h

@@ -19,7 +19,6 @@
 #include <unordered_map>
 
 #include "caffe2/core/common.h"
-#include "caffe2/core/context_gpu.h"
 #include "caffe2/core/event.h"
 #include "caffe2/core/net.h"
 #include "caffe2/core/observer.h"
@@ -45,11 +44,10 @@ class ProfileCounter {
   Timer timer_;
   float start_time_ = 0.0f;
   float run_time_ = 0.0f;
-  cudaEvent_t start_;
-  cudaEvent_t stop_;
 };
 
-class ProfileOperatorObserver : public ProfileCounter,
+class CAFFE2_API ProfileOperatorObserver final
+    : public ProfileCounter,
       public ObserverBase<OperatorBase> {
  public:
   explicit ProfileOperatorObserver(OperatorBase* subject) = delete;
@@ -96,7 +94,7 @@ class ProfileOperatorObserver : public ProfileCounter,
   void Stop() override;
 };
 
-class ProfileObserver final : public OperatorAttachingNetObserver<
+class CAFFE2_API ProfileObserver final : public OperatorAttachingNetObserver<
                                              ProfileOperatorObserver,
                                              ProfileObserver> {
  public:

caffe2/python/pybind_state.cc

@@ -13,6 +13,7 @@
 #include "caffe2/core/operator.h"
 #include "caffe2/core/stats.h"
 #include "caffe2/core/transform.h"
+#include "caffe2/observers/profile_observer.h"
 #include "caffe2/observers/runcnt_observer.h"
 #include "caffe2/observers/time_observer.h"
 #include "caffe2/onnx/backend.h"
@@ -1136,6 +1137,7 @@ void addGlobalMethods(py::module& m) {
     }                                                         \
   }
 
+        REGISTER_PYTHON_EXPOSED_OBSERVER(ProfileObserver);
         REGISTER_PYTHON_EXPOSED_OBSERVER(TimeObserver);
 #undef REGISTER_PYTHON_EXPOSED_OBSERVER