diff options
Diffstat (limited to 'src/cpu/kernels/CpuActivationKernel.h')
-rw-r--r-- | src/cpu/kernels/CpuActivationKernel.h | 10 |
1 files changed, 10 insertions, 0 deletions
diff --git a/src/cpu/kernels/CpuActivationKernel.h b/src/cpu/kernels/CpuActivationKernel.h index d856a9357f..fe2d783059 100644 --- a/src/cpu/kernels/CpuActivationKernel.h +++ b/src/cpu/kernels/CpuActivationKernel.h @@ -73,6 +73,15 @@ public: void run_op(ITensorPack &tensors, const Window &window, const ThreadInfo &info) override; const char *name() const override; + /** Get the preferred dimension in which the scheduler splits the work into multiple jobs. + * + * @return The split dimension hint. + */ + size_t get_split_dimension_hint() const + { + return _split_dimension; + } + struct ActivationKernel { const char *name; @@ -85,6 +94,7 @@ public: private: ActivationLayerInfo _act_info{}; ActivationKernelPtr _run_method{ nullptr }; + size_t _split_dimension{ Window::DimY }; std::string _name{}; }; } // namespace kernels |