[CPU] [ARM] JIT GELU Tanh (#23969)
### Details: - *[CPU] [ARM] jit gelu tanh* ### Tickets: - *CVS-138293*
This commit is contained in:
Edward Shogulin
GitHub
parent
38ce1fff0f
commit
31fccc801f
|
|
@ -613,6 +613,97 @@ std::set<std::vector<element::Type>> 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<ov::Node>& node)
|
||||
: jit_emitter(host, host_isa, node, get_arithmetic_binary_exec_precision(node)) {
|
||||
prepare_table();
|
||||
tanh_emitter = std::make_unique<jit_tanh_emitter>(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<jit_tanh_emitter>(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<size_t>(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<size_t> &in_vec_idxs, const std::vector<size_t> &out_vec_idxs) const {
|
||||
if (host_isa_ == dnnl::impl::cpu::aarch64::asimd) {
|
||||
emit_isa<dnnl::impl::cpu::aarch64::asimd>(in_vec_idxs, out_vec_idxs);
|
||||
} else {
|
||||
OV_CPU_JIT_EMITTER_THROW("Can't create jit eltwise kernel");
|
||||
}
|
||||
}
|
||||
|
||||
template <dnnl::impl::cpu::aarch64::cpu_isa_t isa>
|
||||
void jit_gelu_tanh_emitter::emit_isa(const std::vector<size_t> &in_vec_idxs, const std::vector<size_t> &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<isa>::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<size_t>(tanh_emitter->get_aux_vecs_count(), 1)]);
|
||||
const TReg vmm_aux2(aux_vec_idxs[std::max<size_t>(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<std::vector<element::Type>> jit_gelu_tanh_emitter::get_supported_precisions(const std::shared_ptr<ov::Node>& 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,
|
||||
|
|
|
|||
|
|
@ -224,6 +224,39 @@ private:
|
|||
void emit_isa(const std::vector<size_t> &in_vec_idxs, const std::vector<size_t> &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<ov::Node>& 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<std::vector<element::Type>> get_supported_precisions(const std::shared_ptr<ov::Node>& node = nullptr);
|
||||
|
||||
private:
|
||||
std::unique_ptr<jit_tanh_emitter> tanh_emitter;
|
||||
|
||||
void emit_impl(const std::vector<size_t> &in_vec_idxs, const std::vector<size_t> &out_vec_idxs) const override;
|
||||
|
||||
template <dnnl::impl::cpu::aarch64::cpu_isa_t isa>
|
||||
void emit_isa(const std::vector<size_t> &in_vec_idxs, const std::vector<size_t> &out_vec_idxs) const;
|
||||
};
|
||||
|
||||
class jit_hswish_emitter : public jit_emitter {
|
||||
public:
|
||||
jit_hswish_emitter(dnnl::impl::cpu::aarch64::jit_generator *host,
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ bool JitEltwiseExecutor::isSupported(
|
|||
Algorithm::EltwiseEqual,
|
||||
Algorithm::EltwiseExp,
|
||||
Algorithm::EltwiseGeluErf,
|
||||
Algorithm::EltwiseGeluTanh,
|
||||
Algorithm::EltwiseHswish,
|
||||
Algorithm::EltwiseMaximum,
|
||||
Algorithm::EltwiseMinimum,
|
||||
|
|
|
|||
|
|
@ -641,6 +641,7 @@ std::shared_ptr<jit_emitter> jit_uni_eltwise_generic<isa>::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<std::vector<element::Type>> 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),
|
||||
|
|
|
|||
|
|
@ -118,7 +118,7 @@ private:
|
|||
// X5 | [not used] | RSI | d_bias
|
||||
// X6 | [not used] | RBP | d_weights
|
||||
// X7 | [not used] | RSP | <stack pointer>
|
||||
// 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);
|
||||
|
|
|
|||
|
|
@ -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) ||
|
||||
|
|
|
|||
Loading…
Reference in New Issue