Expose c10 cuda ops to caffe2 (pytorch#18036)

Summary:
Pull Request resolved: pytorch#18036

- Add macros to export c10 cuda operators to caffe2 frontend
- Instead of having a separate caffe2 registry for the c10 operator wrappers, use the existing caffe2 registries

Reviewed By: ezyang

Differential Revision: D14467495

fbshipit-source-id: 7715ed2e38d2bbe16f1446ae82c17193a3fabcb9

Author: smessmer

caffe2/core/operator.cc

@@ -119,7 +119,7 @@ GlobalEnginePrefType& g_global_engine_pref() {
   return *g_global_engine_pref_;
 }
 
-unique_ptr<OperatorBase> TryCreateC2Operator(
+unique_ptr<OperatorBase> TryCreateOperator(
     const string& key,
     const OperatorDef& operator_def,
     Workspace* ws) {
@@ -143,24 +143,6 @@ unique_ptr<OperatorBase> TryCreateC2Operator(
   }
 }
 
-unique_ptr<OperatorBase> TryCreateC10Operator(
-    const string& key,
-    const OperatorDef& operator_def,
-    Workspace* ws) {
-  return C10OperatorRegistry()->Create(key, operator_def, ws);
-}
-
-unique_ptr<OperatorBase> TryCreateOperator(
-    const string& key,
-    const OperatorDef& operator_def,
-    Workspace* ws) {
-  if (auto op = TryCreateC10Operator(key, operator_def, ws)) {
-    return op;
-  } else {
-    return TryCreateC2Operator(key, operator_def, ws);
-  }
-}
-
 unique_ptr<OperatorBase> _CreateOperator(
     const OperatorDef& operator_def,
     Workspace* ws) {
@@ -726,11 +708,6 @@ std::set<std::string> GetRegisteredOperators() {
     all_keys.emplace(name);
   }
 
-  // C10 operators
-  for (const auto& name : C10OperatorRegistry()->Keys()) {
-    all_keys.emplace(name);
-  }
-
   return all_keys;
 }
 

caffe2/core/operator_c10wrapper.h

@@ -172,14 +172,14 @@ class C10OperatorWrapper final : public Operator<Context> {
       const std::string& name,
       const c10::optional<IValue>& default_value) {
     if (default_value.has_value()) {
-      return OperatorBase::GetSingleArgument<T>(name, default_value->to<T>());
+      return this->template GetSingleArgument<T>(name, default_value->to<T>());
     } else {
       AT_CHECK(
-          OperatorBase::HasSingleArgumentOfType<T>(name),
+          this->template HasSingleArgumentOfType<T>(name),
           "Error in caffe2->c10 wrapper: Expected argument '",
           name,
           "' missing or wrong type.");
-      return OperatorBase::GetSingleArgument<T>(name, 0);
+      return this->template GetSingleArgument<T>(name, 0);
     }
   }
 
@@ -211,22 +211,22 @@ createC10OperatorWrapper(const c10::OperatorHandle& op_handle) {
 
 } // namespace detail
 
-C10_DECLARE_REGISTRY(
-    C10OperatorRegistry,
-    OperatorBase,
-    const OperatorDef&,
-    Workspace*);
-
 // TODO Also register c10 operators on mobile
 #ifndef C10_MOBILE
 // TODO Currently we only register the CPU variant. This is going to be fixed
 //      once the tensor detemplatization lands.
-#define REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(OperatorHandle, Name) \
-  C10_REGISTER_CREATOR(                                                 \
-      C10OperatorRegistry,                                              \
-      Name,                                                             \
-      detail::createC10OperatorWrapper<CPUContext>(OperatorHandle))
+#define REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_CPU(OperatorHandle, Name) \
+  REGISTER_CPU_OPERATOR_CREATOR(                                            \
+      Name, detail::createC10OperatorWrapper<CPUContext>(OperatorHandle))
+#define REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_CUDA(OperatorHandle, Name) \
+  REGISTER_CUDA_OPERATOR_CREATOR(                                            \
+      Name, detail::createC10OperatorWrapper<CUDAContext>(OperatorHandle))
+#define REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_HIP(OperatorHandle, Name) \
+  REGISTER_HIP_OPERATOR_CREATOR(                                            \
+      Name, detail::createC10OperatorWrapper<HIPContext>(OperatorHandle))
 #else
-#define REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(OperatorHandle, Name)
+#define REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_CPU(OperatorHandle, Name)
+#define REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_CUDA(OperatorHandle, Name)
+#define REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_HIP(OperatorHandle, Name)
 #endif
 } // namespace caffe2

caffe2/operators/CMakeLists.txt

@@ -60,7 +60,6 @@ set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} ${CMAKE_CURRENT_LIST_DIR}/experimental/c1
 set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} ${CMAKE_CURRENT_LIST_DIR}/experimental/c10/cpu/enforce_finite_cpu.cc)
 set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} ${CMAKE_CURRENT_LIST_DIR}/experimental/c10/cpu/add_cpu.cc)
 set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} ${CMAKE_CURRENT_LIST_DIR}/experimental/c10/schemas/sigmoid.cc)
-set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} ${CMAKE_CURRENT_LIST_DIR}/experimental/c10/schemas/layer_norm.cc)
 set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} ${CMAKE_CURRENT_LIST_DIR}/experimental/c10/schemas/filler.cc)
 set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} ${CMAKE_CURRENT_LIST_DIR}/experimental/c10/schemas/expand_dims.cc)
 set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} ${CMAKE_CURRENT_LIST_DIR}/experimental/c10/schemas/mul.cc)

caffe2/operators/experimental/c10/schemas/add.cc

@@ -21,5 +21,7 @@ C10_DEFINE_OP_SCHEMA(Add, FunctionSchema(
 }
 
 namespace caffe2 {
-REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(ops::Add(), C10Add_DontUseThisOpYet)
+REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_CPU(
+    ops::Add(),
+    C10Add_DontUseThisOpYet)
 }

caffe2/operators/experimental/c10/schemas/averaged_loss.cc

@@ -19,7 +19,7 @@ C10_DEFINE_OP_SCHEMA(AveragedLoss, FunctionSchema(
 }
 
 namespace caffe2 {
-REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(
+REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_CPU(
     ops::AveragedLoss(),
     C10AveragedLoss_DontUseThisOpYet)
 }

caffe2/operators/experimental/c10/schemas/batch_gather.cc

@@ -20,7 +20,7 @@ C10_DEFINE_OP_SCHEMA(BatchGather, FunctionSchema(
 }
 
 namespace caffe2 {
-REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(
+REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_CPU(
     ops::BatchGather(),
     C10BatchGather_DontUseThisOpYet)
 }

caffe2/operators/experimental/c10/schemas/batch_matmul.cc

@@ -24,7 +24,7 @@ C10_DEFINE_OP_SCHEMA(BatchMatmul, FunctionSchema(
 
 namespace caffe2 {
 
-REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(
+REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_CPU(
     ops::BatchMatmul(),
     C10BatchMatMul_DontUseThisOpYet)
 }

caffe2/operators/experimental/c10/schemas/cast.cc

@@ -21,5 +21,7 @@ C10_DEFINE_OP_SCHEMA(Cast, FunctionSchema(
 }
 
 namespace caffe2 {
-REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(ops::Cast(), C10Cast_DontUseThisOpYet)
+REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_CPU(
+    ops::Cast(),
+    C10Cast_DontUseThisOpYet)
 }

caffe2/operators/experimental/c10/schemas/concat.cc

@@ -22,7 +22,7 @@ C10_DEFINE_OP_SCHEMA(Concat, FunctionSchema(
 }
 
 namespace caffe2 {
-REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(
+REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_CPU(
     ops::Concat(),
     C10Concat_DontUseThisOpYet)
 }

caffe2/operators/experimental/c10/schemas/enforce_finite.cc

@@ -18,7 +18,7 @@ C10_DEFINE_OP_SCHEMA(EnforceFinite, FunctionSchema(
 }
 
 namespace caffe2 {
-REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(
+REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_CPU(
     ops::EnforceFinite(),
     C10EnforceFinite_DontUseThisOpYet)
 }