Revert "Move CreateContext to global registry (pytorch#11688)" (pytorch#12049)

ezyang · facebook-github-bot · commit d7e11e3aaec8 · 2018-09-25T10:13:43.000-07:00
Summary: This reverts commit 3ae6ee4. Pull Request resolved: pytorch#12049 Differential Revision: D10030954 Pulled By: ezyang fbshipit-source-id: 6ca9de65b707c5b4c68280fc6f1b8e5ad7251efc
diff --git a/aten/src/ATen/core/context_base.cpp b/aten/src/ATen/core/context_base.cpp
@@ -1,16 +1,5 @@
 #include <ATen/core/context_base.h>
 
-namespace at {
-
-AT_DEFINE_TYPED_REGISTRY(
-    ContextRegistry,
-    DeviceType,
-    BaseContext,
-    std::unique_ptr,
-    at::Device);
-
-} // namespace at
-
 namespace caffe2 {
 
 // TODO: rename context.h -> context_cpu.h & context_base.h -> context.h
diff --git a/aten/src/ATen/core/context_base.h b/aten/src/ATen/core/context_base.h
@@ -6,12 +6,11 @@
 #include <memory>
 #include <unordered_map>
 
-#include <ATen/core/ATenGeneral.h>
-#include <ATen/core/Device.h>
+#include <ATen/core/DeviceType.h>
 #include <ATen/core/Error.h>
-#include <ATen/core/Registry.h>
 #include <ATen/core/UniqueVoidPtr.h>
 #include <ATen/core/typeid.h>
+#include <ATen/core/ATenGeneral.h>
 
 namespace caffe2 {
 class Event;
@@ -32,6 +31,11 @@ class AT_CORE_API BaseStaticContext {
 
   virtual std::pair<void*, DeleterFnPtr> New(size_t nbytes) const = 0;
 
+  virtual std::unique_ptr<BaseContext> CreateContext() = 0;
+
+  virtual std::unique_ptr<BaseContext> CreateContext(
+      const caffe2::DeviceOption&) = 0;
+
   virtual DeviceType GetDeviceType() = 0;
 
   /*
@@ -180,22 +184,6 @@ class AT_CORE_API BaseContext {
   }
 };
 
-// Context constructor registry
-AT_DECLARE_TYPED_REGISTRY(
-    ContextRegistry,
-    at::DeviceType,
-    BaseContext,
-    std::unique_ptr,
-    at::Device);
-
-#define REGISTER_CONTEXT(type, ...) \
-  AT_REGISTER_TYPED_CLASS(ContextRegistry, type, __VA_ARGS__)
-
-inline std::unique_ptr<at::BaseContext> CreateContext(
-    const at::Device& device) {
-  return ContextRegistry()->Create(device.type(), device);
-}
-
 } // namespace at
 
 namespace caffe2 {
diff --git a/caffe2/core/blob_serialization.cc b/caffe2/core/blob_serialization.cc
@@ -196,7 +196,7 @@ void TensorSerializer::Serialize(
   const TensorProto::DataType data_type = TypeMetaToDataType(input.meta());
   proto.set_data_type(data_type);
   StoreDeviceDetail(input, &proto);
-  auto uniq_ptr = CreateContext(input.GetDevice());
+  auto uniq_ptr = input.GetStaticContext()->CreateContext();
   // A lot of copypaste is error prone. Should we create a macro for this?
   switch (data_type) {
     case TensorProto_DataType_FLOAT:
@@ -370,7 +370,8 @@ void TensorDeserializer::Deserialize(const BlobProto& blob_proto, Blob* blob) {
 void TensorDeserializer::Deserialize(const TensorProto& proto, Tensor* tensor) {
   // We create a local context for deserializing. Since Caffe2 contexts are
   // usually lightweight, this should not involve too much overhead.
-  auto uniq_ptr = CreateContext(OptionToDevice(proto.device_detail()));
+  auto uniq_ptr =
+      tensor->GetStaticContext()->CreateContext(proto.device_detail());
   auto context = uniq_ptr.get();
   context->SwitchToDevice(0);
   vector<int64_t> dims;
diff --git a/caffe2/core/context.cc b/caffe2/core/context.cc
@@ -5,10 +5,6 @@
 #include <process.h>
 #endif
 
-namespace at {
-
-REGISTER_CONTEXT(DeviceType::CPU, caffe2::CPUContext);
-} // namespace at
 namespace caffe2 {
 
 uint32_t RandomNumberSeed() {
diff --git a/caffe2/core/context.h b/caffe2/core/context.h
@@ -50,8 +50,6 @@ class CAFFE2_API CPUContext final : public BaseContext {
                                      : RandomNumberSeed()) {
     CAFFE_ENFORCE_EQ(option.device_type(), PROTO_CPU);
   }
-  explicit CPUContext(const at::Device& device)
-      : CPUContext(DeviceToOption(device)) {}
 
   ~CPUContext() noexcept override {}
 
@@ -194,6 +192,15 @@ class CAFFE2_API CPUStaticContext : public BaseStaticContext {
     return data_and_deleter;
   }
 
+  std::unique_ptr<BaseContext> CreateContext() override {
+    return caffe2::make_unique<CPUContext>();
+  }
+
+  std::unique_ptr<BaseContext> CreateContext(
+      const DeviceOption& option) override {
+    return caffe2::make_unique<CPUContext>(option);
+  }
+
   DeviceType GetDeviceType() override {
     return CPU;
   }
diff --git a/caffe2/core/context_base.cc b/caffe2/core/context_base.cc
@@ -1,5 +1,4 @@
 #include "context_base.h"
 
 namespace caffe2 {
-
 } // namespace caffe2
diff --git a/caffe2/core/context_base.h b/caffe2/core/context_base.h
@@ -5,5 +5,3 @@
 #include "caffe2/core/common.h"
 #include "caffe2/core/logging.h"
 #include "caffe2/proto/caffe2_pb.h"
-
-namespace caffe2 {} // namespace caffe2
diff --git a/caffe2/core/context_gpu.cu b/caffe2/core/context_gpu.cu
@@ -57,11 +57,6 @@ CAFFE2_DEFINE_int(
     128,
     "The threshold in MB on how frequently to report memory changes");
 
-namespace at {
-
-REGISTER_CONTEXT(DeviceType::CUDA, caffe2::CUDAContext);
-} // namespace at
-
 namespace caffe2 {
 
 ThreadLocalCUDAObjects& CUDAContext::getCudaObjects() {
diff --git a/caffe2/core/context_gpu.h b/caffe2/core/context_gpu.h
@@ -142,8 +142,6 @@ class CAFFE2_CUDA_API CUDAContext final : public BaseContext {
   // The default cuda context constructor.
   explicit CUDAContext(const int gpu_id = -1);
   explicit CUDAContext(const DeviceOption& option);
-  explicit CUDAContext(const at::Device& device)
-      : CUDAContext(DeviceToOption(device)) {}
 
   ~CUDAContext() override {
     if (curand_generator_) {
@@ -387,6 +385,19 @@ class CAFFE2_CUDA_API CUDAStaticContext final : public BaseStaticContext {
  public:
   std::pair<void*, MemoryDeleter> New(size_t nbytes) const override;
 
+  std::unique_ptr<BaseContext> CreateContext() override {
+    return caffe2::make_unique<CUDAContext>();
+  }
+
+  std::unique_ptr<BaseContext> CreateContext(
+      const DeviceOption& option) override {
+    return caffe2::make_unique<CUDAContext>(option);
+  }
+
+  std::unique_ptr<BaseContext> CreateContext(int gpu_id = -1) {
+    return caffe2::make_unique<CUDAContext>(gpu_id);
+  }
+
   DeviceType GetDeviceType() override {
     return CUDA;
   }
diff --git a/caffe2/core/hip/context_hip.cc b/caffe2/core/hip/context_hip.cc
@@ -50,11 +50,6 @@ CAFFE2_DEFINE_int(caffe2_gpu_memory_report_interval_mb,
                   128,
                   "The threshold in MB on how frequently to report memory changes");
 
-namespace at {
-
-REGISTER_CONTEXT(DeviceType::HIP, caffe2::HIPContext);
-} // namespace at
-
 namespace caffe2 {
 
 thread_local ThreadLocalHIPObjects HIPContext::hip_objects_;
@@ -413,12 +408,13 @@ void HIPStaticContext::Delete(void* ptr) {
         g_hip_device_affiliation.erase(it);
         break;
     }
-    case HipMemoryPoolType::THC: {
-      HIP_ENFORCE(g_thc_allocator->Free(ptr));
-      if (FLAGS_caffe2_gpu_memory_tracking) {
-        g_hip_device_affiliation.erase(g_hip_device_affiliation.find(ptr));
-      }
-      break;
+    case HipMemoryPoolType::THC: 
+    {
+        HIP_ENFORCE(g_thc_allocator->Free(ptr));
+        if (FLAGS_caffe2_gpu_memory_tracking) {
+          g_hip_device_affiliation.erase(g_hip_device_affiliation.find(ptr));
+        }
+        break;
     }
     }
 }
diff --git a/caffe2/core/hip/context_hip.h b/caffe2/core/hip/context_hip.h
@@ -127,8 +127,6 @@ class HIPContext final : public BaseContext {
   // The default HIP context constructor.
   explicit HIPContext(const int gpu_id = -1);
   explicit HIPContext(const DeviceOption& option);
-  explicit HIPContext(const at::Device& device)
-      : HIPContext(DeviceToOption(device)) {}
 
   ~HIPContext() override {
     if (hiprand_generator_) {
@@ -376,6 +374,19 @@ class HIPStaticContext final : public BaseStaticContext {
  public:
   std::pair<void*, MemoryDeleter> New(size_t nbytes) const override;
 
+  std::unique_ptr<BaseContext> CreateContext() override {
+    return caffe2::make_unique<HIPContext>();
+  }
+
+  std::unique_ptr<BaseContext> CreateContext(
+      const DeviceOption& option) override {
+    return caffe2::make_unique<HIPContext>(option);
+  }
+
+  std::unique_ptr<BaseContext> CreateContext(int gpu_id = -1) {
+    return caffe2::make_unique<HIPContext>(gpu_id);
+  }
+
   DeviceType GetDeviceType() override {
     return HIP;
   }
diff --git a/caffe2/core/registry.h b/caffe2/core/registry.h
@@ -172,7 +172,7 @@ class Registerer {
       key,                                                                    \
       RegistryName(),                                                         \
       Registerer##RegistryName::DefaultCreator<__VA_ARGS__>,                  \
-      at::demangle_type<__VA_ARGS__>());                                      \
+      at::demangle_type<__VA_ARGS__>());                                           \
   }
 
 // CAFFE_DECLARE_REGISTRY and CAFFE_DEFINE_REGISTRY are hard-wired to use string
diff --git a/caffe2/core/tensor.h b/caffe2/core/tensor.h
@@ -130,12 +130,12 @@ class CAFFE2_API Tensor final {
     return impl_.get()->GetStaticContext();
   }
 
-  DeviceType GetDeviceType() const {
-    return impl_.get()->GetDeviceType();
+  std::unique_ptr<BaseContext> CreateContext() const {
+    return impl_.get()->CreateContext();
   }
 
-  at::Device GetDevice() const {
-    return impl_.get()->GetDevice();
+  DeviceType GetDeviceType() const {
+    return impl_.get()->GetDeviceType();
   }
 
   void CopyFrom(const Tensor& src, BaseContext* context = nullptr) const {
diff --git a/caffe2/core/tensor_impl.cc b/caffe2/core/tensor_impl.cc
@@ -1,5 +1,5 @@
 #include "caffe2/core/tensor_impl.h"
-#include "caffe2/core/context_base.h"
+
 #include "caffe2/core/flags.h"
 
 CAFFE2_DEFINE_bool(
diff --git a/caffe2/core/tensor_impl.h b/caffe2/core/tensor_impl.h
@@ -3,6 +3,7 @@
 #include <ATen/core/DimVector.h>
 #include <ATen/core/TensorImpl.h>
 #include <ATen/core/context_base.h>
+#include <ATen/core/context_base.h>
 
 #include "caffe2/core/allocator.h"
 #include "caffe2/core/common.h"
@@ -111,12 +112,19 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
     return get_static_context(device_type);
   }
 
-  at::DeviceType GetDeviceType() const {
-    return storage_.device_type();
+  /* @brief
+   * Create a context that has the same device_type
+   * as the tensor.
+   * Note that this doesn't support passing in argument
+   * TODO(jerryzh): move this to a global registry
+   * that can create context for us
+   */
+  std::unique_ptr<at::BaseContext> CreateContext() const {
+    return GetStaticContext()->CreateContext();
   }
 
-  at::Device GetDevice() const {
-    return storage_.device();
+  at::DeviceType GetDeviceType() const {
+    return storage_.device_type();
   }
 
   /**
@@ -159,12 +167,8 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
         // knows how to copy between CPU and that context
         if (src.GetDeviceType() != ::at::DeviceType::CPU || GetDeviceType() == ::at::DeviceType::CPU) {
           if (!context) {
-            CreateContext(src.GetDevice())
-                ->CopyBytesToDevice(
-                    nbytes(),
-                    src.raw_data(),
-                    raw_mutable_data(),
-                    GetDeviceType());
+            src.CreateContext()->CopyBytesToDevice(
+                nbytes(), src.raw_data(), raw_mutable_data(), GetDeviceType());
           } else {
             CAFFE_ENFORCE(
                 context->device_type() == src.GetDeviceType(),
@@ -176,8 +180,8 @@ class CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
           // In case source context is CPU, and target context is non-CPU
           // We'll have to create a Context from target and perform the
           // copy using that context
-          CreateContext(GetDevice())
-              ->CopyBytesFromCPU(nbytes(), src.raw_data(), raw_mutable_data());
+          CreateContext()->CopyBytesFromCPU(
+              nbytes(), src.raw_data(), raw_mutable_data());
         }
       }
     }
diff --git a/caffe2/ideep/utils/ideep_context.h b/caffe2/ideep/utils/ideep_context.h
@@ -20,8 +20,6 @@ class IDEEPContext final : public BaseContext {
                                      : RandomNumberSeed()) {
     CAFFE_ENFORCE_EQ(option.device_type(), PROTO_IDEEP);
   }
-  explicit IDEEPContext(const at::Device& device)
-      : IDEEPContext(DeviceToOption(device)) {}
 
   ~IDEEPContext() noexcept override {}
 
@@ -180,6 +178,15 @@ class IDEEPStaticContext : public BaseStaticContext {
     return GetCPUAllocator()->New(nbytes);
   }
 
+  std::unique_ptr<BaseContext> CreateContext() override {
+    return caffe2::make_unique<IDEEPContext>();
+  }
+
+  std::unique_ptr<BaseContext> CreateContext(
+      const DeviceOption& option) override {
+    return caffe2::make_unique<IDEEPContext>(option);
+  }
+
   DeviceType GetDeviceType() override {
     return IDEEP;
   }
diff --git a/caffe2/ideep/utils/ideep_register.cc b/caffe2/ideep/utils/ideep_register.cc
@@ -4,9 +4,6 @@
 #include <ideep_pin_singletons.hpp>
 #include "ideep_context.h"
 
-namespace at {
-REGISTER_CONTEXT(DeviceType::IDEEP, caffe2::IDEEPContext);
-} // namespace at
 namespace caffe2 {
 
 CAFFE_KNOWN_TYPE(ideep::tensor);
diff --git a/caffe2/mkl/utils/mkl_context.cc b/caffe2/mkl/utils/mkl_context.cc
@@ -3,10 +3,6 @@
 #include "mkl_context.h"
 #include "caffe2/core/event_cpu.h"
 
-namespace at {
-
-REGISTER_CONTEXT(DeviceType::MKLDNN, caffe2::MKLContext);
-} // namespace at
 namespace caffe2 {
 
 // MKL events are the same as CPU events
diff --git a/caffe2/mkl/utils/mkl_context.h b/caffe2/mkl/utils/mkl_context.h
@@ -29,8 +29,6 @@ class MKLContext : public BaseContext {
                                      : RandomNumberSeed()) {
     CAFFE_ENFORCE_EQ(option.device_type(), PROTO_MKLDNN);
   }
-  explicit MKLContext(const at::Device& device)
-      : MKLContext(DeviceToOption(device)) {}
 
   ~MKLContext() override {}
 
@@ -157,6 +155,15 @@ class MKLStaticContext : public BaseStaticContext {
     return GetCPUAllocator()->New(nbytes);
   }
 
+  std::unique_ptr<BaseContext> CreateContext() override {
+    return caffe2::make_unique<MKLContext>();
+  }
+
+  std::unique_ptr<BaseContext> CreateContext(
+      const DeviceOption& option) override {
+    return caffe2::make_unique<MKLContext>(option);
+  }
+
   DeviceType GetDeviceType() override {
     return MKLDNN;
   }
diff --git a/caffe2/proto/caffe2_pb.h b/caffe2/proto/caffe2_pb.h
diff --git a/caffe2/python/pybind_state.h b/caffe2/python/pybind_state.h
