From 31fccc801fc6219d73ede793f4e52e36ae9249da Mon Sep 17 00:00:00 2001 From: Edward Shogulin Date: Tue, 28 May 2024 11:23:34 +0100 Subject: [PATCH] [CPU] [ARM] JIT GELU Tanh (#23969) ### Details: - *[CPU] [ARM] jit gelu tanh* ### Tickets: - *CVS-138293* --- .../plugin/aarch64/jit_eltwise_emitters.cpp | 91 +++++++++++++++++++ .../plugin/aarch64/jit_eltwise_emitters.hpp | 33 +++++++ .../nodes/executors/aarch64/jit_eltwise.cpp | 1 + .../aarch64/jit_uni_eltwise_generic.cpp | 2 + .../aarch64/jit_uni_eltwise_generic.hpp | 19 ++-- .../single_layer_tests/classes/activation.cpp | 9 +- 6 files changed, 145 insertions(+), 10 deletions(-) diff --git a/src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_eltwise_emitters.cpp b/src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_eltwise_emitters.cpp index aa522a5e69a..dbe951a4dd3 100644 --- a/src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_eltwise_emitters.cpp +++ b/src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_eltwise_emitters.cpp @@ -613,6 +613,97 @@ std::set> jit_gelu_erf_emitter::get_supported_precisi return {{element::f32}}; } +/// GELU_TANH /// +jit_gelu_tanh_emitter::jit_gelu_tanh_emitter(dnnl::impl::cpu::aarch64::jit_generator* host, + dnnl::impl::cpu::aarch64::cpu_isa_t host_isa, + const std::shared_ptr& node) + : jit_emitter(host, host_isa, node, get_arithmetic_binary_exec_precision(node)) { + prepare_table(); + tanh_emitter = std::make_unique(h, host_isa, node); +} + +jit_gelu_tanh_emitter::jit_gelu_tanh_emitter(dnnl::impl::cpu::aarch64::jit_generator* host, + dnnl::impl::cpu::aarch64::cpu_isa_t host_isa, + const ov::element::Type exec_prc) : jit_emitter(host, host_isa, exec_prc) { + prepare_table(); + tanh_emitter = std::make_unique(h, host_isa, exec_prc); +} + +size_t jit_gelu_tanh_emitter::get_inputs_count() const { return 1; } + +size_t jit_gelu_tanh_emitter::get_aux_vecs_count() const { + return std::max(tanh_emitter->get_aux_vecs_count() + 2, 3); +} + +size_t jit_gelu_tanh_emitter::get_aux_gprs_count() const { + return tanh_emitter->get_aux_gprs_count() + 1; +} + +void jit_gelu_tanh_emitter::emit_impl(const std::vector &in_vec_idxs, const std::vector &out_vec_idxs) const { + if (host_isa_ == dnnl::impl::cpu::aarch64::asimd) { + emit_isa(in_vec_idxs, out_vec_idxs); + } else { + OV_CPU_JIT_EMITTER_THROW("Can't create jit eltwise kernel"); + } +} + +template +void jit_gelu_tanh_emitter::emit_isa(const std::vector &in_vec_idxs, const std::vector &out_vec_idxs) const { + OV_CPU_JIT_EMITTER_ASSERT(exec_prc_ == ov::element::f32, "unsupported precision: " + exec_prc_.to_string()); + + using TReg = typename dnnl::impl::cpu::aarch64::cpu_isa_traits::TReg; + const TReg vmm_src(in_vec_idxs[0]); + const TReg vmm_dst(out_vec_idxs[0]); + + const TReg vmm_aux1(aux_vec_idxs[0]); + const TReg vmm_aux0(aux_vec_idxs[std::max(tanh_emitter->get_aux_vecs_count(), 1)]); + const TReg vmm_aux2(aux_vec_idxs[std::max(tanh_emitter->get_aux_vecs_count() + 1, 2)]); + + // compute G(x) = sqrt_root_two_over_pi * x * (1 + fitting_const * x * x) + h->fmul(vmm_aux0.s, vmm_src.s, vmm_src.s); + h->ld1r(vmm_aux1.s, table_val2("gelu_tanh_fitting_const")); + h->ld1r(vmm_aux2.s, table_val2("one")); + h->fmla(vmm_aux2.s, vmm_aux1.s, vmm_aux0.s); + h->fmul(vmm_aux2.s, vmm_src.s, vmm_aux2.s); + h->ld1r(vmm_aux1.s, table_val2("gelu_tanh_sqrt_two_over_pi")); + h->fmul(vmm_aux0.s, vmm_aux1.s, vmm_aux2.s); + + const bool store_src = vmm_src.getIdx() == vmm_dst.getIdx(); + if (store_src) { + h->mov(vmm_aux2.b16, vmm_src.b16); + } + + tanh_emitter->emit_code( + { vmm_aux0.getIdx() }, + out_vec_idxs, + aux_vec_idxs, + aux_gpr_idxs); + + // compute 0.5 * x * (1 + tanh(G(x))) + h->ld1r(vmm_aux1.s, table_val2("one")); + h->fadd(vmm_dst.s, vmm_aux1.s, vmm_dst.s); + h->ld1r(vmm_aux1.s, table_val2("half")); + h->fmul(vmm_dst.s, vmm_aux1.s, vmm_dst.s); + h->fmul(vmm_dst.s, store_src ? vmm_aux2.s : vmm_src.s, vmm_dst.s); +} + +void jit_gelu_tanh_emitter::register_table_entries() { + push_arg_entry_of("one", 0x3f800000, true); + push_arg_entry_of("half", 0x3f000000, true); + push_arg_entry_of("gelu_tanh_fitting_const", 0x3d372713, true); + push_arg_entry_of("gelu_tanh_fitting_const_times_three", 0x3e095d4f, true); + push_arg_entry_of("gelu_tanh_sqrt_two_over_pi", 0x3f4c422a, true); +} + +void jit_gelu_tanh_emitter::emit_data() const { + jit_emitter::emit_data(); + tanh_emitter->emit_data(); +} + +std::set> jit_gelu_tanh_emitter::get_supported_precisions(const std::shared_ptr& node) { + return {{element::f32}}; +} + /// HARD_SWISH /// jit_hswish_emitter::jit_hswish_emitter(dnnl::impl::cpu::aarch64::jit_generator* host, dnnl::impl::cpu::aarch64::cpu_isa_t host_isa, diff --git a/src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_eltwise_emitters.hpp b/src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_eltwise_emitters.hpp index e4a4f548aff..1b46d5c2b76 100644 --- a/src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_eltwise_emitters.hpp +++ b/src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_eltwise_emitters.hpp @@ -224,6 +224,39 @@ private: void emit_isa(const std::vector &in_vec_idxs, const std::vector &out_vec_idxs) const; }; +class jit_tanh_emitter; + +class jit_gelu_tanh_emitter : public jit_emitter { +public: + jit_gelu_tanh_emitter(dnnl::impl::cpu::aarch64::jit_generator* host, + dnnl::impl::cpu::aarch64::cpu_isa_t host_isa, + const ov::element::Type exec_prc = ov::element::f32); + + jit_gelu_tanh_emitter(dnnl::impl::cpu::aarch64::jit_generator* host, + dnnl::impl::cpu::aarch64::cpu_isa_t host_isa, + const std::shared_ptr& node); + + size_t get_inputs_count() const override; + + size_t get_aux_vecs_count() const override; + + size_t get_aux_gprs_count() const override; + + void register_table_entries() override; + + void emit_data() const override; + + static std::set> get_supported_precisions(const std::shared_ptr& node = nullptr); + +private: + std::unique_ptr tanh_emitter; + + void emit_impl(const std::vector &in_vec_idxs, const std::vector &out_vec_idxs) const override; + + template + void emit_isa(const std::vector &in_vec_idxs, const std::vector &out_vec_idxs) const; +}; + class jit_hswish_emitter : public jit_emitter { public: jit_hswish_emitter(dnnl::impl::cpu::aarch64::jit_generator *host, diff --git a/src/plugins/intel_cpu/src/nodes/executors/aarch64/jit_eltwise.cpp b/src/plugins/intel_cpu/src/nodes/executors/aarch64/jit_eltwise.cpp index 055b87c299d..72458c3d292 100644 --- a/src/plugins/intel_cpu/src/nodes/executors/aarch64/jit_eltwise.cpp +++ b/src/plugins/intel_cpu/src/nodes/executors/aarch64/jit_eltwise.cpp @@ -26,6 +26,7 @@ bool JitEltwiseExecutor::isSupported( Algorithm::EltwiseEqual, Algorithm::EltwiseExp, Algorithm::EltwiseGeluErf, + Algorithm::EltwiseGeluTanh, Algorithm::EltwiseHswish, Algorithm::EltwiseMaximum, Algorithm::EltwiseMinimum, diff --git a/src/plugins/intel_cpu/src/nodes/kernels/aarch64/jit_uni_eltwise_generic.cpp b/src/plugins/intel_cpu/src/nodes/kernels/aarch64/jit_uni_eltwise_generic.cpp index 7a8f209b51d..8093f0617eb 100644 --- a/src/plugins/intel_cpu/src/nodes/kernels/aarch64/jit_uni_eltwise_generic.cpp +++ b/src/plugins/intel_cpu/src/nodes/kernels/aarch64/jit_uni_eltwise_generic.cpp @@ -641,6 +641,7 @@ std::shared_ptr jit_uni_eltwise_generic::create_eltwise_emitte OV_CASE(Algorithm::EltwiseMinimum, ov::intel_cpu::aarch64::jit_minimum_emitter), OV_CASE(Algorithm::EltwiseMish, ov::intel_cpu::aarch64::jit_mish_emitter), OV_CASE(Algorithm::EltwiseGeluErf, ov::intel_cpu::aarch64::jit_gelu_erf_emitter), + OV_CASE(Algorithm::EltwiseGeluTanh, ov::intel_cpu::aarch64::jit_gelu_tanh_emitter), OV_CASE(Algorithm::EltwiseMulAdd, ov::intel_cpu::aarch64::jit_mul_add_emitter), OV_CASE(Algorithm::EltwiseMod, ov::intel_cpu::aarch64::jit_mod_emitter), OV_CASE(Algorithm::EltwiseMultiply, ov::intel_cpu::aarch64::jit_multiply_emitter), @@ -805,6 +806,7 @@ std::set> eltwise_precision_helper::get_supported_pre OV_CASE(Algorithm::EltwiseEqual, jit_equal_emitter), OV_CASE(Algorithm::EltwiseExp, jit_exp_emitter), OV_CASE(Algorithm::EltwiseGeluErf, jit_gelu_erf_emitter), + OV_CASE(Algorithm::EltwiseGeluTanh, jit_gelu_tanh_emitter), OV_CASE(Algorithm::EltwiseHswish, jit_hswish_emitter), OV_CASE(Algorithm::EltwiseMaximum, jit_maximum_emitter), OV_CASE(Algorithm::EltwiseMinimum, jit_minimum_emitter), diff --git a/src/plugins/intel_cpu/src/nodes/kernels/aarch64/jit_uni_eltwise_generic.hpp b/src/plugins/intel_cpu/src/nodes/kernels/aarch64/jit_uni_eltwise_generic.hpp index fea311915cb..8f18a9815b4 100644 --- a/src/plugins/intel_cpu/src/nodes/kernels/aarch64/jit_uni_eltwise_generic.hpp +++ b/src/plugins/intel_cpu/src/nodes/kernels/aarch64/jit_uni_eltwise_generic.hpp @@ -118,7 +118,7 @@ private: // X5 | [not used] | RSI | d_bias // X6 | [not used] | RBP | d_weights // X7 | [not used] | RSP | - // X8 | [not used] | R8 | src ptr + // X8 | temporary | R8 | src ptr // X9 | work amount | R9 | src ptr // X10 | ker temporary| R10 | src ptr // X11 | ker temporary| R11 | src ptr @@ -163,7 +163,11 @@ private: OPENVINO_THROW("aux gpr register " + std::to_string(idx) + " is not supported"); } - return XReg(13 + idx); + if (idx == 0) { + return XReg(8); + } + + return XReg(13 + idx - 1); } // Vector registers mapping @@ -179,13 +183,14 @@ private: // 15 | aux // 16 | aux // 17 | aux - // 18 | src + // 18 | aux // 19 | src // 20 | src // 21 | src // 22 | src // 23 | src - // 24-31 | [not used] + // 24 | src + // 25-31 | [not used] TReg vmm_dst {9}; @@ -194,18 +199,18 @@ private: if (idx > MAX_ELTWISE_INPUTS) { OPENVINO_THROW("source vector register " + std::to_string(idx) + " is not supported"); } - return TReg(18 + idx); + return TReg(19 + idx); } inline SReg get_scl_reg(const uint32_t idx) { if (idx > MAX_ELTWISE_INPUTS) { OPENVINO_THROW("source scalar register " + std::to_string(idx) + " is not supported"); } - return SReg(18 + idx); + return SReg(19 + idx); } inline TReg get_aux_vmm(const uint32_t idx) { - if (idx > 7) { + if (idx > 8) { OPENVINO_THROW("aux vector register " + std::to_string(idx) + " is not supported"); } return TReg(10 + idx); diff --git a/src/plugins/intel_cpu/tests/functional/custom/single_layer_tests/classes/activation.cpp b/src/plugins/intel_cpu/tests/functional/custom/single_layer_tests/classes/activation.cpp index d5063655e0a..f5034d71a8e 100644 --- a/src/plugins/intel_cpu/tests/functional/custom/single_layer_tests/classes/activation.cpp +++ b/src/plugins/intel_cpu/tests/functional/custom/single_layer_tests/classes/activation.cpp @@ -115,7 +115,7 @@ void ActivationLayerCPUTest::SetUp() { inType = inPrecision; outType = outPrecision; const auto primitiveType = getPrimitiveType(activationType, inType, inputShapes); - selectedType = primitiveType.empty() ? "" : getPrimitiveType(activationType, inType, inputShapes) + "_" + netPrecision.to_string(); + selectedType = primitiveType.empty() ? "" : primitiveType + "_" + netPrecision.to_string(); #if defined(OPENVINO_ARCH_ARM) || defined(OPENVINO_ARCH_ARM64) # if defined(OPENVINO_ARCH_ARM) @@ -123,8 +123,10 @@ void ActivationLayerCPUTest::SetUp() { selectedType = std::string("ref_") + netPrecision.to_string(); # endif if ((primitiveType != "jit") && - (activationType == utils::ActivationTypes::GeluTanh || // @todo not supported by ACL, can be decomposed with transformation - activationType == utils::ActivationTypes::SoftSign || // @todo not supported by ACL, can be decomposed with transformation + (activationType == utils::ActivationTypes::SoftSign || // @todo not supported by ACL, can be decomposed with transformation +#if defined(OPENVINO_ARCH_ARM) + activationType == utils::ActivationTypes::GeluTanh || // @todo not supported by ACL, can be decomposed with transformation +#endif inputShapes.front().first.rank().get_length() > 5)) // @todo tmp fallback to ref, remove after 6D+ ranks are properly supported selectedType = std::string("ref_") + netPrecision.to_string(); #else @@ -161,6 +163,7 @@ std::string ActivationLayerCPUTest::getPrimitiveType(const utils::ActivationType (activation_type == utils::ActivationTypes::HardSigmoid) || (activation_type == utils::ActivationTypes::Mish) || (activation_type == utils::ActivationTypes::GeluErf) || + (activation_type == utils::ActivationTypes::GeluTanh) || (activation_type == utils::ActivationTypes::Relu) || (activation_type == utils::ActivationTypes::Sigmoid) || (activation_type == utils::ActivationTypes::Swish) ||