Remove StaticContext (pytorch#12547)

Summary:
Pull Request resolved: pytorch#12547

Pull Request resolved: pytorch#12305

Remove StaticContext from context_base.h

Reviewed By: dzhulgakov

Differential Revision: D10073519

fbshipit-source-id: 350beec3c54365edef338318ce58229ccb825a98

Author: jerryzh168

aten/src/ATen/core/Allocator.h

@@ -121,6 +121,7 @@ struct Allocator {
   }
 };
 
+// Question: is this still needed?
 struct CAFFE2_API InefficientStdFunctionContext {
   std::unique_ptr<void, std::function<void(void*)>> ptr_;
   InefficientStdFunctionContext(

aten/src/ATen/core/TensorImpl.h

@@ -411,17 +411,6 @@ struct CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
     return storage_.device();
   }
 
-  /**
-   * The static context of a tensor intuitively represents the device
-   * type of a tensor; e.g., a CPU tensor is associated with the
-   * GetCPUStaticContext().  This method replaces the former Context template
-   * parameter which was previously used to identify the device type
-   * of a tensor.
-   */
-  at::BaseStaticContext* GetStaticContext() const {
-    return ::caffe2::get_static_context(device_type());
-  }
-
   /**
    * @brief Copies the data from a source tensor, with a contex provided to
    * carry out the underlying memcpy operation.  This method respects
@@ -761,20 +750,19 @@ struct CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
         return storage_.data();
       }
       const at::Allocator* allocator = storage_.allocator();
-      // TODO: Get rid of StaticContext
       AT_ENFORCE(
           allocator == nullptr,
-          "Allocator is not used within Caffe2 functions, please use StaticContext instead.");
+          "Allocator in storage_ is not used within Caffe2 functions. \
+           we are using global function to get the allocator based on device \
+           type.");
       allocator = caffe2::GetAllocator(storage_.device_type());
       if (meta.ctor()) {
         // For types that need placement new, we will call it, as well as
         // making sure that when the data is freed, it calls the right
         // destruction procedure.
         auto size = numel_;
         auto dtor = data_type_.dtor();
-        auto data_ptr = allocator->allocate(
-            numel_ * storage_.itemsize()); // Removing this can get rid of
-                                           // InefficientStdFunctionContext
+        auto data_ptr = allocator->allocate(numel_ * storage_.itemsize());
         storage_.set_data_ptr(PlacementDeleteContext::makeDataPtr(
             std::move(data_ptr), dtor, size, storage_.device()));
         data_type_.ctor()(storage_.data(), numel_);

aten/src/ATen/core/context_base.cpp

@@ -14,20 +14,5 @@ C10_DEFINE_TYPED_REGISTRY(
 namespace caffe2 {
 
 // TODO: rename context.h -> context_cpu.h & context_base.h -> context.h
-StaticContextMap& GetStaticContexts() {
-  static StaticContextMap static_contexts;
-  return static_contexts;
-}
-
-void set_static_context(at::DeviceType t, BaseStaticContext* ptr) {
-  auto& static_contexts = GetStaticContexts();
-  static_contexts[t] = ptr;
-}
-
-BaseStaticContext* get_static_context(at::DeviceType t) {
-  auto* ptr = GetStaticContexts()[t];
-  AT_ASSERTM(ptr, "StaticContext for ", t, " is not registered yet.");
-  return ptr;
-}
 
 } // namespace caffe2

aten/src/ATen/core/context_base.h

@@ -14,24 +14,12 @@
 
 namespace caffe2 {
 class Event;
-class DeviceOption;
 
 } // namespace caffe2
 namespace at {
 
 class BaseContext;
 
-/* BaseStaticContext defines the interface for static context, which contains
-   functions that are invoked statically before in Tensor class, e.g. New,
-   We will merge this with Allocator later.
- */
-class CAFFE2_API BaseStaticContext {
- public:
-  virtual ~BaseStaticContext() noexcept {}
-
-  virtual DeviceType GetDeviceType() = 0;
-};
-
 /**
  * Virtual interface for the Context class in Caffe2.
  *
@@ -44,8 +32,6 @@ class CAFFE2_API BaseContext {
  public:
   virtual ~BaseContext() noexcept {}
 
-  virtual BaseStaticContext* GetStaticContext() const = 0;
-
   virtual Device device() const = 0;
 
   /* Sorry for the naming, will get rid of this in future diff */
@@ -192,23 +178,5 @@ inline std::unique_ptr<at::BaseContext> CreateContext(
 namespace caffe2 {
 
 using at::BaseContext;
-using at::BaseStaticContext;
-
-using StaticContextMap = std::unordered_map<at::DeviceType, BaseStaticContext*>;
-CAFFE2_API StaticContextMap& GetStaticContexts();
-CAFFE2_API void set_static_context(at::DeviceType t, BaseStaticContext* ptr);
-CAFFE2_API BaseStaticContext* get_static_context(at::DeviceType t);
-
-template <at::DeviceType t>
-struct StaticContextFunctionRegisterer {
-  explicit StaticContextFunctionRegisterer(BaseStaticContext* ptr) {
-    set_static_context(t, ptr);
-  }
-};
-
-#define REGISTER_STATIC_CONTEXT(t, f)                                \
-  namespace {                                                        \
-  static StaticContextFunctionRegisterer<t> g_static_context_##d(f); \
-  }
 
 } // namespace caffe2

caffe2/core/context.cc

@@ -28,11 +28,4 @@ uint32_t RandomNumberSeed() {
       kPrime2 * tv_sec + kPrime3 * tv_usec;
 }
 
-BaseStaticContext* GetCPUStaticContext() {
-  static CPUStaticContext context;
-  return &context;
-}
-
-REGISTER_STATIC_CONTEXT(CPU, GetCPUStaticContext());
-
 } // namespace caffe2

caffe2/core/context.h

@@ -20,8 +20,6 @@ C10_DECLARE_bool(caffe2_report_cpu_memory_usage);
 
 namespace caffe2 {
 
-CAFFE2_API BaseStaticContext* GetCPUStaticContext();
-
 /**
  * A function to generate a random number seed that is unique in a best-effort
  * basis, using an ever-incrementing seed and the current time.
@@ -55,14 +53,6 @@ class CAFFE2_API CPUContext final : public BaseContext {
 
   ~CPUContext() noexcept override {}
 
-  BaseStaticContext* GetStaticContext() const override {
-    return GetCPUStaticContext();
-  }
-
-  static BaseStaticContext* StaticContext() {
-    return GetCPUStaticContext();
-  }
-
   inline void SwitchToDevice(int /*stream_id*/) override {}
 
   using BaseContext::SwitchToDevice;
@@ -187,16 +177,6 @@ inline void CPUContext::CopyBytes<CPUContext, CPUContext>(
   memcpy(dst, src, nbytes);
 }
 
-// TODO(jerryzh): merge CPUStaticContext with Allocator
-class CAFFE2_API CPUStaticContext : public BaseStaticContext {
- public:
-
-  DeviceType GetDeviceType() override {
-    return CPU;
-  }
-
-};
-
 }  // namespace caffe2
 
 #endif  // CAFFE2_CORE_CONTEXT_H_

caffe2/core/context_gpu.cu

@@ -424,19 +424,7 @@ struct DefaultCUDAAllocator final : public at::Allocator {
   }
 };
 
-at::Allocator* GetCUDAAllocator() {
-  return GetAllocator(CUDA);
-}
-
 static DefaultCUDAAllocator g_cuda_alloc;
-
 REGISTER_ALLOCATOR(CUDA, &g_cuda_alloc);
 
-BaseStaticContext* GetCUDAStaticContext() {
-  static CUDAStaticContext context;
-  return &context;
-}
-
-REGISTER_STATIC_CONTEXT(CUDA, GetCUDAStaticContext());
-
 }  // namespace caffe2

caffe2/core/context_gpu.h

@@ -135,8 +135,6 @@ class CAFFE2_CUDA_API ThreadLocalCUDAObjects {
 #endif // CAFFE2_USE_CUDNN
 };
 
-CAFFE2_CUDA_API BaseStaticContext* GetCUDAStaticContext();
-
 class CAFFE2_CUDA_API CUDAContext final : public BaseContext {
  public:
   // The default cuda context constructor.
@@ -152,14 +150,6 @@ class CAFFE2_CUDA_API CUDAContext final : public BaseContext {
     FinishDeviceComputation();
   }
 
-  BaseStaticContext* GetStaticContext() const override {
-    return GetCUDAStaticContext();
-  }
-
-  static BaseStaticContext* StaticContext() {
-    return GetCUDAStaticContext();
-  }
-
   inline void SwitchToDevice(int stream_id) override {
     set_stream_id(stream_id);
     CaffeCudaSetDevice(gpu_id_);
@@ -389,13 +379,6 @@ struct CAFFE2_CUDA_API PinnedCPUAllocator final : public at::Allocator {
   DefaultCPUAllocator baseAllocator_;
 };
 
-class CAFFE2_CUDA_API CUDAStaticContext final : public BaseStaticContext {
- public:
-  DeviceType GetDeviceType() override {
-    return CUDA;
-  }
-};
-
 using TensorCUDA = Tensor;
 
 }  // namespace caffe2

caffe2/core/hip/context_hip.cc

@@ -430,11 +430,4 @@ struct DefaultHIPAllocator final : public at::Allocator {
 static DefaultHIPAllocator g_hip_alloc;
 REGISTER_ALLOCATOR(HIP, &g_hip_alloc);
 
-BaseStaticContext* GetHIPStaticContext() {
-  static HIPStaticContext context;
-  return &context;
-}
-
-REGISTER_STATIC_CONTEXT(HIP, GetHIPStaticContext());
-
 } // namespace caffe2

caffe2/core/hip/context_hip.h

@@ -120,8 +120,6 @@ class ThreadLocalHIPObjects {
   vector<miopenHandle_t> miopen_handles_[CAFFE2_COMPILE_TIME_MAX_HIP_GPUS];
 };
 
-BaseStaticContext* GetHIPStaticContext();
-
 class HIPContext final : public BaseContext {
  public:
   // The default HIP context constructor.
@@ -137,14 +135,6 @@ class HIPContext final : public BaseContext {
     FinishDeviceComputation();
   }
 
-  BaseStaticContext* GetStaticContext() const override {
-    return GetHIPStaticContext();
-  }
-
-  static BaseStaticContext* StaticContext() {
-    return GetHIPStaticContext();
-  }
-
   inline void SwitchToDevice(int stream_id) override {
     set_stream_id(stream_id);
     CaffeHipSetDevice(gpu_id_);
@@ -378,13 +368,6 @@ struct PinnedCPUAllocator final : public at::Allocator {
   DefaultCPUAllocator baseAllocator_;
 };
 
-class HIPStaticContext final : public BaseStaticContext {
- public:
-  DeviceType GetDeviceType() override {
-    return HIP;
-  }
-};
-
 typedef Tensor TensorHIP;
 
 } // namespace caffe2

caffe2/core/tensor.h

@@ -115,10 +115,6 @@ class CAFFE2_API Tensor final {
     return x;
   }
 
-  BaseStaticContext* GetStaticContext() const {
-    return impl_.get()->GetStaticContext();
-  }
-
   DeviceType GetDeviceType() const {
     return impl_->device_type();
   }

caffe2/ideep/utils/ideep_context.h

@@ -8,8 +8,6 @@
 
 namespace caffe2 {
 
-BaseStaticContext* GetIDEEPStaticContext();
-
 class IDEEPContext final : public BaseContext {
  public:
   typedef std::mt19937 rand_gen_type;
@@ -25,14 +23,6 @@ class IDEEPContext final : public BaseContext {
 
   ~IDEEPContext() noexcept override {}
 
-  BaseStaticContext* GetStaticContext() const override {
-    return GetIDEEPStaticContext();
-  }
-
-  static BaseStaticContext* StaticContext() {
-    return GetIDEEPStaticContext();
-  }
-
   inline void SwitchToDevice(int /*stream_id*/) {}
   using BaseContext::SwitchToDevice;
 
@@ -178,12 +168,4 @@ inline void IDEEPContext::CopyBytes<IDEEPContext, CPUContext>(
   memcpy(dst, src, nbytes);
 }
 
-class IDEEPStaticContext : public BaseStaticContext {
- public:
-  DeviceType GetDeviceType() override {
-    return IDEEP;
-  }
-
-};
-
 } // namespace caffe2

caffe2/ideep/utils/ideep_register.cc

@@ -30,11 +30,4 @@ REGISTER_EVENT_ERROR_MESSAGE_FUNCTION(IDEEP, EventErrorMessageCPU);
 REGISTER_EVENT_SET_FINISHED_FUNCTION(IDEEP, EventSetFinishedCPU);
 REGISTER_EVENT_RESET_FUNCTION(IDEEP, EventResetCPU);
 
-C10_EXPORT BaseStaticContext* GetIDEEPStaticContext() {
-  static IDEEPStaticContext context;
-  return &context;
-}
-
-REGISTER_STATIC_CONTEXT(IDEEP, GetIDEEPStaticContext());
-
 } // namespace caffe2

caffe2/mkl/utils/mkl_context.cc

@@ -23,11 +23,4 @@ REGISTER_EVENT_ERROR_MESSAGE_FUNCTION(MKLDNN, EventErrorMessageCPU);
 REGISTER_EVENT_SET_FINISHED_FUNCTION(MKLDNN, EventSetFinishedCPU);
 REGISTER_EVENT_RESET_FUNCTION(MKLDNN, EventResetCPU);
 
-BaseStaticContext* GetMKLStaticContext() {
-  static MKLStaticContext context;
-  return &context;
-}
-
-REGISTER_STATIC_CONTEXT(MKLDNN, GetMKLStaticContext());
-
 } // namespace caffe2