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
--- a/c10/util/Metaprogramming.h
+++ b/c10/util/Metaprogramming.h
@@ -51,7 +51,7 @@
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
--- a/caffe2/core/c10_operator.h
+++ b/caffe2/core/c10_operator.h
@@ -133,6 +133,7 @@
* - 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 @@
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
--- a/caffe2/core/operator_c10wrapper.h
+++ b/caffe2/core/operator_c10wrapper.h
@@ -179,6 +179,8 @@
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 @@
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
