Fix shared_ptr binary size in op registration (#26869)

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

Having a lot of shared_ptr<Functor> cost us ~1.1MB of binary size in libtorch.so.
This PR fixes that.
ghstack-source-id: 90842812

Test Plan: measure libtorch.so size

Differential Revision: D17595674

fbshipit-source-id: 05151047ee8e85c05205b7510a33915ba98bab58
This commit is contained in:
Sebastian Messmer 2019-09-26 16:53:13 -07:00 committed by Facebook Github Bot
parent 1a5d641de3
commit 54b66c8c20
4 changed files with 42 additions and 36 deletions

View File

@ -159,10 +159,10 @@ public:
* > public:
* > Tensor operator()(Tensor a, Tensor b) {...}
* > };
* > KernelFunction func = KernelFunction::makeFromUnboxedFunctor(std::make_shared<MyFunctor>());
* > KernelFunction func = KernelFunction::makeFromUnboxedFunctor(guts::make_unique<MyFunctor>());
*/
template<bool AllowLegacyTypes = false, class KernelFunctor>
static KernelFunction makeFromUnboxedFunctor(std::shared_ptr<KernelFunctor> kernelFunctor) {
static KernelFunction makeFromUnboxedFunctor(std::unique_ptr<OperatorKernel> kernelFunctor) {
static_assert(guts::is_functor<KernelFunctor>::value, "Tried to call KernelFunction::makeFromUnboxedFunctor<KernelFunctor> but the argument is not a functor.");
static_assert(std::is_base_of<OperatorKernel, KernelFunctor>::value, "Tried to call KernelFunction::makeFromUnboxedFunctor<KernelFunctor>, but the functor doesn't inherit from c10::OperatorKernel. Please have the functor inherit from it.");
@ -191,11 +191,11 @@ public:
* > Tensor operator()(Tensor a, Tensor b) {...}
* > };
* > KernelFunction func = KernelFunction::makeFromUnboxedFunctor([] {
* > return std::make_shared<MyFunctor>();
* > return guts::make_unique<MyFunctor>();
* > });
*/
template<class KernelFunctor, bool AllowLegacyTypes = false>
static KernelFunction makeFromUnboxedFunctorFactory(std::function<std::shared_ptr<KernelFunctor>()> kernelFunctorFactory) {
static KernelFunction makeFromUnboxedFunctorFactory(std::function<std::unique_ptr<OperatorKernel>()> kernelFunctorFactory) {
static_assert(guts::is_functor<KernelFunctor>::value, "Tried to call KernelFunction::makeFromUnboxedFunctor<KernelFunctor> but the argument is not a functor.");
static_assert(std::is_base_of<OperatorKernel, KernelFunctor>::value, "Tried to call KernelFunction::makeFromUnboxedFunctor<KernelFunctor>, but the functor doesn't inherit from c10::OperatorKernel. Please have the functor inherit from it.");
@ -223,10 +223,10 @@ public:
* > public:
* > Tensor operator()(Tensor a, Tensor b) {...}
* > };
* > KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor(std::make_shared<MyFunctor>());
* > KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor(guts::make_unique<MyFunctor>());
*/
template<class KernelFunctor>
static KernelFunction makeFromUnboxedOnlyFunctor(std::shared_ptr<KernelFunctor> kernelFunctor) {
static KernelFunction makeFromUnboxedOnlyFunctor(std::unique_ptr<OperatorKernel> kernelFunctor) {
// TODO We want to get rid of kernels that have only an unboxed function pointer.
// All kernels should have a boxed pointer.
@ -259,8 +259,8 @@ public:
static_assert(!std::is_same<FuncType, BoxedKernelFunction>::value, "Tried to call KernelFunction::makeFromUnboxedFunction with a boxed function pointer. Please use KernelFunction::makeFromBoxedFunction instead.");
static_assert(func != nullptr, "Kernel function cannot be nullptr");
return makeFromUnboxedFunctor<AllowLegacyTypes>(
std::make_shared<typename detail::WrapKernelFunction<FuncType, func>::type>()
return makeFromUnboxedFunctor<AllowLegacyTypes, typename detail::WrapKernelFunction<FuncType, func>::type>(
guts::make_unique_base<OperatorKernel, typename detail::WrapKernelFunction<FuncType, func>::type>()
);
}
@ -288,8 +288,8 @@ public:
static_assert(!std::is_same<FuncType, BoxedKernelFunction>::value, "Tried to call KernelFunction::makeFromUnboxedOnlyFunction with a boxed function pointer. Please use KernelFunction::makeFromBoxedFunction instead.");
static_assert(func != nullptr, "Kernel function cannot be nullptr");
return makeFromUnboxedOnlyFunctor(
std::make_shared<typename detail::WrapKernelFunction<FuncType, func>::type>()
return makeFromUnboxedOnlyFunctor<typename detail::WrapKernelFunction<FuncType, func>::type> (
guts::make_unique_base<OperatorKernel, typename detail::WrapKernelFunction<FuncType, func>::type>()
);
}
@ -310,8 +310,8 @@ public:
static_assert(!std::is_same<FuncType, BoxedKernelFunction>::value, "Tried to call KernelFunction::makeFromUnboxedRuntimeFunction with a boxed function pointer. Please use KernelFunction::makeFromBoxedFunction instead.");
TORCH_INTERNAL_ASSERT(func != nullptr, "Kernel function cannot be nullptr");
return makeFromUnboxedFunctor<AllowLegacyTypes>(
std::make_shared<detail::WrapRuntimeKernelFunctor<guts::decay_t<FuncType>>>(func)
return makeFromUnboxedFunctor<AllowLegacyTypes, detail::WrapRuntimeKernelFunctor<guts::decay_t<FuncType>>>(
guts::make_unique_base<OperatorKernel, detail::WrapRuntimeKernelFunctor<guts::decay_t<FuncType>>>(func)
);
}
@ -321,8 +321,8 @@ public:
static_assert(!std::is_same<FuncType, BoxedKernelFunction>::value, "Tried to call KernelFunction::makeFromUnboxedRuntimeFunction with a boxed function pointer. Please use KernelFunction::makeFromBoxedFunction instead.");
TORCH_INTERNAL_ASSERT(func != nullptr, "Kernel function cannot be nullptr");
return makeFromUnboxedOnlyFunctor(
std::make_shared<detail::WrapRuntimeKernelFunctor<guts::decay_t<FuncType>>>(func)
return makeFromUnboxedOnlyFunctor<detail::WrapRuntimeKernelFunctor<guts::decay_t<FuncType>>>(
guts::make_unique_base<OperatorKernel, detail::WrapRuntimeKernelFunctor<guts::decay_t<FuncType>>>(func)
);
}
@ -338,14 +338,14 @@ public:
static KernelFunction makeFromUnboxedLambda(Lambda&& lambda) {
static_assert(guts::is_functor<guts::decay_t<Lambda>>::value, "Tried to call KernelFunction::makeFromUnboxedLambda with a non-lambda type.");
return makeFromUnboxedFunctor<AllowLegacyTypes>(
std::make_shared<detail::WrapRuntimeKernelFunctor<guts::decay_t<Lambda>>>(std::forward<Lambda>(lambda))
return makeFromUnboxedFunctor<AllowLegacyTypes, detail::WrapRuntimeKernelFunctor<guts::decay_t<Lambda>>>(
guts::make_unique_base<OperatorKernel, detail::WrapRuntimeKernelFunctor<guts::decay_t<Lambda>>>(std::forward<Lambda>(lambda))
);
}
private:
explicit KernelFunction(std::function<std::shared_ptr<OperatorKernel>()> functorFactory, std::shared_ptr<OperatorKernel> functor, BoxedKernelFunction* boxed_kernel_func, void* unboxed_kernel_func)
explicit KernelFunction(std::function<std::unique_ptr<OperatorKernel>()> functorFactory, std::unique_ptr<OperatorKernel> functor, BoxedKernelFunction* boxed_kernel_func, void* unboxed_kernel_func)
: functorFactory_(std::move(functorFactory))
, functor_(std::move(functor))
, boxed_kernel_func_(boxed_kernel_func)
@ -372,7 +372,7 @@ private:
// initialization time yet (SIOF). So these register with a
// functorFactory_ instead of a functor_ and will be initialized
// on the first call to the KernelFunction.
std::function<std::shared_ptr<OperatorKernel>()> functorFactory_;
std::function<std::unique_ptr<OperatorKernel>()> functorFactory_;
mutable std::shared_ptr<OperatorKernel> functor_;
BoxedKernelFunction* boxed_kernel_func_;

View File

@ -55,14 +55,14 @@ struct unboxed_functor_without_return final : OperatorKernel {
};
struct unboxed_functor_with_return_factory final {
std::shared_ptr<unboxed_functor_with_return> operator()() {
return std::make_shared<unboxed_functor_with_return>();
std::unique_ptr<OperatorKernel> operator()() {
return c10::guts::make_unique<unboxed_functor_with_return>();
}
};
struct unboxed_functor_without_return_factory final {
std::shared_ptr<unboxed_functor_without_return> operator()() {
return std::make_shared<unboxed_functor_without_return>();
std::unique_ptr<OperatorKernel> operator()() {
return c10::guts::make_unique<unboxed_functor_without_return>();
}
};
@ -199,32 +199,32 @@ TEST(KernelFunctionTest, givenBoxedFunction_withoutReturn_whenCallingUnboxedOnly
}
TEST(KernelFunctionTest, givenUnboxedFunctor_withReturn_whenCallingBoxed_thenWorks) {
KernelFunction func = KernelFunction::makeFromUnboxedFunctor(std::make_shared<kernels::unboxed_functor_with_return>());
KernelFunction func = KernelFunction::makeFromUnboxedFunctor<false, kernels::unboxed_functor_with_return>(std::unique_ptr<OperatorKernel>(c10::guts::make_unique<kernels::unboxed_functor_with_return>()));
kernels::expectBoxedCallingWithReturnWorks(func);
}
TEST(KernelFunctionTest, givenUnboxedFunctor_withoutReturn_whenCallingBoxed_thenWorks) {
KernelFunction func = KernelFunction::makeFromUnboxedFunctor(std::make_shared<kernels::unboxed_functor_without_return>());
KernelFunction func = KernelFunction::makeFromUnboxedFunctor<false, kernels::unboxed_functor_without_return>(std::unique_ptr<OperatorKernel>(c10::guts::make_unique<kernels::unboxed_functor_without_return>()));
kernels::expectBoxedCallingWithoutReturnWorks(func);
}
TEST(KernelFunctionTest, givenUnboxedFunctor_withReturn_whenCallingUnboxed_thenWorks) {
KernelFunction func = KernelFunction::makeFromUnboxedFunctor(std::make_shared<kernels::unboxed_functor_with_return>());
KernelFunction func = KernelFunction::makeFromUnboxedFunctor<false, kernels::unboxed_functor_with_return>(std::unique_ptr<OperatorKernel>(c10::guts::make_unique<kernels::unboxed_functor_with_return>()));
kernels::expectUnboxedCallingWithReturnWorks(func);
}
TEST(KernelFunctionTest, givenUnboxedFunctor_withoutReturn_whenCallingUnboxed_thenWorks) {
KernelFunction func = KernelFunction::makeFromUnboxedFunctor(std::make_shared<kernels::unboxed_functor_without_return>());
KernelFunction func = KernelFunction::makeFromUnboxedFunctor<false, kernels::unboxed_functor_without_return>(std::unique_ptr<OperatorKernel>(c10::guts::make_unique<kernels::unboxed_functor_without_return>()));
kernels::expectUnboxedCallingWithoutReturnWorks(func);
}
TEST(KernelFunctionTest, givenUnboxedFunctor_withReturn_whenCallingUnboxedOnly_thenWorks) {
KernelFunction func = KernelFunction::makeFromUnboxedFunctor(std::make_shared<kernels::unboxed_functor_with_return>());
KernelFunction func = KernelFunction::makeFromUnboxedFunctor<false, kernels::unboxed_functor_with_return>(std::unique_ptr<OperatorKernel>(c10::guts::make_unique<kernels::unboxed_functor_with_return>()));
kernels::expectUnboxedOnlyCallingWithReturnWorks(func);
}
TEST(KernelFunctionTest, givenUnboxedFunctor_withoutReturn_whenCallingUnboxedOnly_thenWorks) {
KernelFunction func = KernelFunction::makeFromUnboxedFunctor(std::make_shared<kernels::unboxed_functor_without_return>());
KernelFunction func = KernelFunction::makeFromUnboxedFunctor<false, kernels::unboxed_functor_without_return>(std::unique_ptr<OperatorKernel>(c10::guts::make_unique<kernels::unboxed_functor_without_return>()));
kernels::expectUnboxedOnlyCallingWithoutReturnWorks(func);
}
@ -259,32 +259,32 @@ TEST(KernelFunctionTest, givenUnboxedFunctorFactory_withoutReturn_whenCallingUnb
}
TEST(KernelFunctionTest, givenUnboxedOnlyFunctor_withReturn_whenCallingBoxed_thenFails) {
KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor(std::make_shared<kernels::unboxed_functor_with_return>());
KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor<kernels::unboxed_functor_with_return>(std::unique_ptr<OperatorKernel>(c10::guts::make_unique<kernels::unboxed_functor_with_return>()));
kernels::expectBoxedCallingFailsWith(func, "Tried to call KernelFunction::callBoxed() on a KernelFunction that can only be called with KernelFunction::callUnboxed()");
}
TEST(KernelFunctionTest, givenUnboxedOnlyFunctor_withoutReturn_whenCallingBoxed_thenFails) {
KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor(std::make_shared<kernels::unboxed_functor_without_return>());
KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor<kernels::unboxed_functor_without_return>(std::unique_ptr<OperatorKernel>(c10::guts::make_unique<kernels::unboxed_functor_without_return>()));
kernels::expectBoxedCallingFailsWith(func, "Tried to call KernelFunction::callBoxed() on a KernelFunction that can only be called with KernelFunction::callUnboxed()");
}
TEST(KernelFunctionTest, givenUnboxedOnlyFunctor_withReturn_whenCallingUnboxed_thenWorks) {
KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor(std::make_shared<kernels::unboxed_functor_with_return>());
KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor<kernels::unboxed_functor_with_return>(std::unique_ptr<OperatorKernel>(c10::guts::make_unique<kernels::unboxed_functor_with_return>()));
kernels::expectUnboxedCallingWithReturnWorks(func);
}
TEST(KernelFunctionTest, givenUnboxedOnlyFunctor_withoutReturn_whenCallingUnboxed_thenWorks) {
KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor(std::make_shared<kernels::unboxed_functor_without_return>());
KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor<kernels::unboxed_functor_without_return>(std::unique_ptr<OperatorKernel>(c10::guts::make_unique<kernels::unboxed_functor_without_return>()));
kernels::expectUnboxedCallingWithoutReturnWorks(func);
}
TEST(KernelFunctionTest, givenUnboxedOnlyFunctor_withReturn_whenCallingUnboxedOnly_thenWorks) {
KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor(std::make_shared<kernels::unboxed_functor_with_return>());
KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor<kernels::unboxed_functor_with_return>(std::unique_ptr<OperatorKernel>(c10::guts::make_unique<kernels::unboxed_functor_with_return>()));
kernels::expectUnboxedOnlyCallingWithReturnWorks(func);
}
TEST(KernelFunctionTest, givenUnboxedOnlyFunctor_withoutReturn_whenCallingUnboxedOnly_thenWorks) {
KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor(std::make_shared<kernels::unboxed_functor_without_return>());
KernelFunction func = KernelFunction::makeFromUnboxedOnlyFunctor<kernels::unboxed_functor_without_return>(std::unique_ptr<OperatorKernel>(c10::guts::make_unique<kernels::unboxed_functor_without_return>()));
kernels::expectUnboxedOnlyCallingWithoutReturnWorks(func);
}

View File

@ -270,9 +270,9 @@ namespace detail {
explicit constexpr KernelFactory(Args... args)
: constructor_parameters_(std::move(args)...) {}
std::shared_ptr<KernelFunctor> operator()() const {
std::unique_ptr<OperatorKernel> operator()() const {
return guts::apply(
[] (const Args&... params) {return std::make_shared<KernelFunctor>(params...); },
[] (const Args&... params) -> std::unique_ptr<OperatorKernel> {return guts::make_unique_base<OperatorKernel, KernelFunctor>(params...); },
constructor_parameters_);
}

View File

@ -92,6 +92,12 @@ make_unique(Args&&...) = delete;
#endif
template <typename Base, typename Child, typename... Args>
typename std::enable_if<!std::is_array<Base>::value && !std::is_array<Base>::value && std::is_base_of<Base, Child>::value, std::unique_ptr<Base>>::type
make_unique_base(Args&&... args) {
return std::unique_ptr<Base>(new Child(c10::guts::forward<Args>(args)...));
}
#ifdef __cpp_lib_integer_sequence