mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
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:
parent
1a5d641de3
commit
54b66c8c20
|
|
@ -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_;
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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_);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user