Disable c10 dispatcher on mobile (#17078)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17078

This prevents caffe2 operators from being expsoed to c10 on mobile,
which in turn causes the whole c10 dispatcher to be stripped away
and saves binary size.

We probably want to re-enable the c10 dispatcher for mobile,
but for now this is ok.

Reviewed By: ezyang

Differential Revision: D14077972

fbshipit-source-id: e4dd3e3b60cdfbde91fe0d24102c1d9708d3e5c4

diff --git a/c10/util/Metaprogramming.h b/c10/util/Metaprogramming.h
index aa3cf07..8fa7724 100644 (file)
--- a/c10/util/Metaprogramming.h
+++ b/c10/util/Metaprogramming.h
@@ -51,7 +51,7 @@ struct function_traits<Result (Args...)> {
 
 template <typename Functor>
 struct infer_function_traits {
-  using type = function_traits<detail::strip_class_t<decltype(&Functor::operator())>>;
+  using type = function_traits<c10::guts::detail::strip_class_t<decltype(&Functor::operator())>>;
 };
 
 template <typename Result, typename... Args>

diff --git a/caffe2/core/c10_operator.h b/caffe2/core/c10_operator.h
index d154189..92aef48 100644 (file)
--- a/caffe2/core/c10_operator.h
+++ b/caffe2/core/c10_operator.h
@@ -133,6 +133,7 @@ inline c10::FunctionSchema make_function_schema_for_c10(const char* OperatorName
  * - calling C10_REGISTER_CAFFE2_OPERATOR_CUDA is optional and can be omitted if
  * you don't want to expose the operator for CUDA operations.
  */
+#if !C10_MOBILE
 #define C10_DECLARE_CAFFE2_OPERATOR(OperatorName) \
   namespace caffe2 {                              \
   namespace _c10_ops {                            \
@@ -208,3 +209,11 @@ inline c10::FunctionSchema make_function_schema_for_c10(const char* OperatorName
           at::DeviceType::HIP>>()                                              \
       .dispatchKey(CUDATensorId());                                            \
   }
+
+#else
+// Don't use c10 dispatcher on mobile because of binary size
+#define C10_DECLARE_CAFFE2_OPERATOR(OperatorName)
+#define C10_REGISTER_CAFFE2_OPERATOR_CPU(OperatorName, Inputs, Outputs, OperatorClass)
+#define C10_REGISTER_CAFFE2_OPERATOR_CUDA(OperatorName, OperatorClass)
+#define C10_REGISTER_CAFFE2_OPERATOR_HIP(OperatorName, OperatorClass)
+#endif

diff --git a/caffe2/core/operator_c10wrapper.h b/caffe2/core/operator_c10wrapper.h
index 48d10f9..192a73d 100644 (file)
--- a/caffe2/core/operator_c10wrapper.h
+++ b/caffe2/core/operator_c10wrapper.h
@@ -179,6 +179,8 @@ C10_DECLARE_REGISTRY(
     const OperatorDef&,
     Workspace*);
 
+// TODO Also register c10 operators on mobile
+#if !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, NumOutputParameters)  \
@@ -218,4 +220,10 @@ C10_DECLARE_REGISTRY(
           NumOutputParameters,                                                     \
           std::tuple<__VA_ARGS__>>)
 
+#else
+#define REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(OperatorHandle, Name, NumOutputParameters)
+#define REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_WITH_PARAMETERS(OperatorHandle, Name, NumOutputParameters, ...)
+#define REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_WITH_ARRAY_INPUT(OperatorHandle, Name, NumOutputParameters)
+#define REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_WITH_ARRAY_INPUT_AND_PARAMETERS(OperatorHandle, Name, NumOutputParameters, ...)
+#endif
 } // namespace caffe2