min & flags

openvinotoolkit · Feb 22, 2024 · 25987d1 · 25987d1
1 parent b9b8744
commit 25987d1
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 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
@@ -127,7 +127,7 @@ jit_exp_emitter::jit_exp_emitter(dnnl::impl::cpu::aarch64::jit_generator* host,
 
 size_t jit_exp_emitter::get_inputs_count() const { return 1; }
 
-size_t jit_exp_emitter::get_aux_vecs_count() const { return 3; }
+size_t jit_exp_emitter::get_aux_vecs_count() const { return 4; }
 
 size_t jit_exp_emitter::get_aux_gprs_count() const { return 1; }
 
@@ -152,9 +152,15 @@ void jit_exp_emitter::emit_isa(const std::vector<size_t> &in_vec_idxs, const std
     const TReg vmm_aux2(aux_vec_idxs[1]);
     const TReg vmm_aux0(aux_vec_idxs[2]);
 
+    const TReg vmm_mask(aux_vec_idxs[3]);
+
     h->ld1r(vmm_aux0.s, table_val2("exp_ln_flt_max_f"));
     h->fmin(vmm_dst.s, vmm_src.s, vmm_aux0.s);
     h->ld1r(vmm_aux0.s, table_val2("exp_ln_flt_min_f"));
+
+    // get mask of values lower than log(FLT_MIN) to zero them in the output
+    h->facgt(vmm_mask.s, vmm_aux0.s, vmm_src.s);
+
     h->fmax(vmm_dst.s, vmm_dst.s, vmm_aux0.s);
     h->mov(vmm_aux1.b16, vmm_dst.b16);
 
@@ -189,6 +195,9 @@ void jit_exp_emitter::emit_isa(const std::vector<size_t> &in_vec_idxs, const std
 
     h->sqshl(vmm_aux2.s, vmm_aux2.s, 23);
 
+    // set zeroes at those points which were < log(FLT_MIN)
+    h->and_(vmm_aux2.b16, vmm_mask.b16, vmm_aux2.b16);
+
     // compute polynomial
     h->ld1r(vmm_aux0.s, table_val2("exp_pol5"));
     h->ld1r(vmm_dst.s, table_val2("exp_pol4"));

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
@@ -174,7 +174,8 @@ struct jit_uni_eltwise_generic : public jit_uni_eltwise_kernel, jit_generator {
     // 11      | aux
     // 12      | aux
     // 13      | aux
-    // 14-15   | [not used]
+    // 14      | aux
+    // 15      | [not used]
     // 16      | src
     // 17      | src
     // 18      | src
@@ -202,7 +203,7 @@ struct jit_uni_eltwise_generic : public jit_uni_eltwise_kernel, jit_generator {
     }
 
     inline TReg get_aux_vmm(const uint32_t idx) {
-        if (idx > 2) {
+        if (idx > 3) {
             OPENVINO_THROW("aux vector register " + std::to_string(idx) + " is not supported");
         }
         return TReg(10 + idx);