caffe2::DeviceType -> at::DeviceType (pytorch#11254)

Summary:
Pull Request resolved: pytorch#11254
Previously we use DeviceType in caffe2.proto directly, but it's an `enum` and have implicit conversion to int, which does not have type safety, e.g. we have to explicitly check for a device type is valid in event.h:
```
template <int d>
struct EventCreateFunctionRegisterer {
  explicit EventCreateFunctionRegisterer(EventCreateFunction f) {
    static_assert(d < MaxDeviceTypes, "");
    Event::event_creator_[d] = f;
  }
};
```
at::DeviceType is an `enum class`, and it does not have implicit conversion to int, and provides better type safety guarantees. In this diff we have done the following refactor(taking CPU as an example):

    1. caffe2::DeviceType → caffe2::DeviceTypeProto
    2. caffe2::CPU → caffe2::PROTO_CPU
    3. caffe2::DeviceType = at::DeviceType
    4. caffe2::CPU = at::DeviceType::CPU

codemod -d caffe2/caffe2 --extensions h,cc,cpp 'device_type\(\), ' 'device_type(), PROTO_'
+ some manual changes

In short, after this diff, in c++, caffe2::CPU refers to the at::DeviceType::CPU and the old proto caffe2::CPU will be caffe2::PROTO_CPU.
In python side, we have a temporary workaround that alias `caffe2_pb2.CPU = caffe2_pb2.PROOT_CPU` to make the change easier to review and this will be removed later.

Reviewed By: ezyang

Differential Revision: D9545704

fbshipit-source-id: 461a28a4ca74e616d3ee183a607078a717fd38a7

Author: jerryzh168

binaries/benchmark_helper.cc

@@ -68,7 +68,7 @@ bool backendCudaSet(const string& backend) {
 void setDeviceType(caffe2::NetDef* net_def, caffe2::DeviceType& run_dev) {
   for (int j = 0; j < net_def->op_size(); j++) {
     caffe2::OperatorDef* op = net_def->mutable_op(j);
-    op->mutable_device_option()->set_device_type(run_dev);
+    op->mutable_device_option()->set_device_type(caffe2::TypeToProto(run_dev));
   }
 }
 

binaries/core_overhead_benchmark_gpu.cc

@@ -167,7 +167,7 @@ static void BM_OperatorCreationCPU(benchmark::State& state) {
   OperatorDef def;
   Workspace ws;
   def.set_type("DummyEmpty");
-  def.mutable_device_option()->set_device_type(CPU);
+  def.mutable_device_option()->set_device_type(PROTO_CPU);
   while (state.KeepRunning()) {
     op = CreateOperator(def, &ws);
   }
@@ -180,7 +180,7 @@ static void BM_OperatorCreationCUDA(benchmark::State& state) {
   OperatorDef def;
   Workspace ws;
   def.set_type("DummyEmpty");
-  def.mutable_device_option()->set_device_type(CUDA);
+  def.mutable_device_option()->set_device_type(PROTO_CUDA);
   while (state.KeepRunning()) {
     op = CreateOperator(def, &ws);
   }

binaries/print_registered_core_operators.cc

@@ -52,8 +52,8 @@ int main(int argc, char** argv) {
   for (const auto& pair : *caffe2::gDeviceTypeRegistry()) {
     std::cout << "Device type " << pair.first
 #ifndef CAFFE2_USE_LITE_PROTO
-              << " (" << caffe2::DeviceType_Name(
-                             static_cast<caffe2::DeviceType>(pair.first))
+              << " ("
+              << at::DeviceTypeName(static_cast<caffe2::DeviceType>(pair.first))
               << ")"
 #endif
               << std::endl;

caffe2/contrib/nervana/nervana_fc_op_gpu_test.cc

@@ -19,7 +19,7 @@ namespace {
 static void AddConstInput(const std::vector<int>& shape, const float value,
                           const string& name, Workspace* ws) {
   DeviceOption option;
-  option.set_device_type(CUDA);
+  option.set_device_type(PROTO_CUDA);
   CUDAContext context(option);
   Blob* blob = ws->CreateBlob(name);
   auto* tensor = blob->GetMutableTensor(CUDA);
@@ -43,7 +43,7 @@ TEST(NervanaFullyConnectedTest, Test) {
   def.add_input("W");
   def.add_input("B");
   def.add_output("Y");
-  def.mutable_device_option()->set_device_type(CUDA);
+  def.mutable_device_option()->set_device_type(PROTO_CUDA);
   def.set_engine("NERVANA");
   AddConstInput(std::vector<int>{5, 10}, 1., "X", &ws);
   AddConstInput(std::vector<int>{6, 10}, 1., "W", &ws);

caffe2/contrib/opencl/context.h

@@ -36,7 +36,7 @@ class OpenCLContext final {
  public:
   explicit OpenCLContext();
   explicit OpenCLContext(const DeviceOption& option) {
-    DCHECK_EQ(option.device_type(), OPENCL);
+    DCHECK_EQ(option.device_type(), PROTO_OPENCL);
     OpenCLContext();
   }
   ~OpenCLContext() {}

caffe2/core/blob_gpu_test.cc

@@ -193,7 +193,7 @@ TEST(TensorTest, TensorSerializationMultiDevices) {
       EXPECT_EQ(tensor_proto.float_data(i), i);
     }
     EXPECT_TRUE(tensor_proto.has_device_detail());
-    EXPECT_EQ(tensor_proto.device_detail().device_type(), CUDA);
+    EXPECT_EQ(tensor_proto.device_detail().device_type(), PROTO_CUDA);
     EXPECT_EQ(tensor_proto.device_detail().cuda_gpu_id(), gpu_id);
     // Test if the restored blob is still of the same device.
     blob.Reset();

caffe2/core/blob_test.cc

@@ -883,7 +883,7 @@ TYPED_TEST(TypedTensorTest, BigTensorSerialization) {
 
   {
     DeviceOption option;
-    option.set_device_type(CPU);
+    option.set_device_type(PROTO_CPU);
     Argument db_type_arg = MakeArgument<string>("db_type", "vector_db");
     Argument absolute_path_arg = MakeArgument<bool>("absolute_path", true);
     Argument db_source_arg = MakeArgument<string>("db", db_source);
@@ -996,7 +996,7 @@ TEST(ContentChunks, Serialization) {
 
   {
     DeviceOption option;
-    option.set_device_type(CPU);
+    option.set_device_type(PROTO_CPU);
     Argument db_type_arg = MakeArgument<string>("db_type", "vector_db");
     Argument absolute_path_arg = MakeArgument<bool>("absolute_path", true);
     Argument db_source_arg = MakeArgument<string>("db", db_source);

caffe2/core/context.h

@@ -48,7 +48,7 @@ class CAFFE2_API CPUContext final : public BaseContext {
       : random_seed_(
             option.has_random_seed() ? option.random_seed()
                                      : RandomNumberSeed()) {
-    CAFFE_ENFORCE_EQ(option.device_type(), CPU);
+    CAFFE_ENFORCE_EQ(option.device_type(), PROTO_CPU);
   }
 
   ~CPUContext() noexcept override {}

caffe2/core/context_base.cc

@@ -3,20 +3,20 @@
 namespace caffe2 {
 
 // TODO: rename context.h -> context_cpu.h & context_base.h -> context.h
-std::array<BaseStaticContext*, COMPILE_TIME_MAX_DEVICE_TYPES>&
-GetStaticContexts() {
-  static std::array<BaseStaticContext*, COMPILE_TIME_MAX_DEVICE_TYPES>
-      static_contexts;
+StaticContextMap& GetStaticContexts() {
+  static StaticContextMap static_contexts;
   return static_contexts;
 }
 
-void set_static_context(int d, BaseStaticContext* ptr) {
+void set_static_context(DeviceType t, BaseStaticContext* ptr) {
   auto& static_contexts = GetStaticContexts();
-  static_contexts[d] = ptr;
+  static_contexts[t] = ptr;
 }
 
-BaseStaticContext* get_static_context(int d) {
-  return GetStaticContexts()[d];
+BaseStaticContext* get_static_context(DeviceType t) {
+  auto* ptr = GetStaticContexts()[t];
+  CAFFE_ENFORCE(ptr, "StaticContext is not registered yet.");
+  return ptr;
 }
 
 } // namespace caffe2

caffe2/core/context_base.h

@@ -37,7 +37,7 @@ class CAFFE2_API BaseStaticContext {
    * current context and the a data pointer
    */
   virtual void ExtractDeviceOption(DeviceOption* device, const void* /*data*/) {
-    device->set_device_type(GetDeviceType());
+    device->set_device_type(TypeToProto(GetDeviceType()));
   }
 };
 
@@ -169,22 +169,21 @@ class CAFFE2_API BaseContext {
   }
 };
 
-CAFFE2_API std::array<BaseStaticContext*, COMPILE_TIME_MAX_DEVICE_TYPES>&
-GetStaticContexts();
-CAFFE2_API void set_static_context(int d, BaseStaticContext* ptr);
-CAFFE2_API BaseStaticContext* get_static_context(int d);
+using StaticContextMap = CaffeMap<DeviceType, BaseStaticContext*>;
+CAFFE2_API StaticContextMap& GetStaticContexts();
+CAFFE2_API void set_static_context(DeviceType t, BaseStaticContext* ptr);
+CAFFE2_API BaseStaticContext* get_static_context(DeviceType t);
 
-template <int d>
+template <DeviceType t>
 struct StaticContextFunctionRegisterer {
   explicit StaticContextFunctionRegisterer(BaseStaticContext* ptr) {
-    static_assert(d < COMPILE_TIME_MAX_DEVICE_TYPES, "");
-    set_static_context(d, ptr);
+    set_static_context(t, ptr);
   }
 };
 
-#define REGISTER_STATIC_CONTEXT(d, f)                                \
+#define REGISTER_STATIC_CONTEXT(t, f)                                \
   namespace {                                                        \
-  static StaticContextFunctionRegisterer<d> g_static_context_##d(f); \
+  static StaticContextFunctionRegisterer<t> g_static_context_##d(f); \
   }
 
 } // namespace caffe2

caffe2/core/context_gpu.cu

@@ -257,7 +257,7 @@ CUDAContext::CUDAContext(const DeviceOption& option)
           option.has_random_seed() ? option.random_seed()
                                    : RandomNumberSeed()) {
   static Caffe2CudaInitializerHelper g_cuda_initializer_;
-  DCHECK_EQ(option.device_type(), CUDA);
+  DCHECK_EQ(option.device_type(), PROTO_CUDA);
 }
 
 // shared mutex to lock out alloc / free during NCCL launches

caffe2/core/context_gpu.h

@@ -403,7 +403,7 @@ class CAFFE2_API CUDAStaticContext final : public BaseStaticContext {
   }
 
   void ExtractDeviceOption(DeviceOption* device, const void* data) override {
-    device->set_device_type(GetDeviceType());
+    device->set_device_type(TypeToProto(GetDeviceType()));
     device->set_cuda_gpu_id(GetGPUIDForPointer(data));
   }