src: cpu: x64: augru: fwd, bwd: fix sse4.1

oneapi-src · Sep 15, 2022 · d8ffd44 · d8ffd44
1 parent 65e270b
commit d8ffd44
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 8 deletions.
diff --git a/src/cpu/x64/rnn/jit_uni_gru_cell_postgemm_1_bwd.hpp b/src/cpu/x64/rnn/jit_uni_gru_cell_postgemm_1_bwd.hpp
@@ -167,7 +167,7 @@ struct jit_uni_gru_cell_postgemm_part1_bwd : public jit_uni_rnn_postgemm {
                 // 1. compute dAttention -= dG0 * G
                 uni_vfnmadd231ps(diff_attn_acc, dG0, G0, tmp2);
                 // 2. Compute dG0 *= 1 - Attention
-                uni_vsubps(tmp1, one_vmm, attn);
+                uni_vsubps(tmp1, one_vmm, attn, tmp2);
                 uni_vmulps(dG0, dG0, tmp1);
             }
 

diff --git a/src/cpu/x64/rnn/jit_uni_gru_cell_postgemm_2_fwd.hpp b/src/cpu/x64/rnn/jit_uni_gru_cell_postgemm_2_fwd.hpp
@@ -81,6 +81,7 @@ struct jit_uni_gru_cell_postgemm_part2_fwd : public jit_uni_rnn_postgemm {
     }
     const Vmm tmp1_vmm = Vmm(9);
     const Vmm tmp2_vmm = Vmm(10);
+    const Vmm tmp3_vmm = Vmm(11);
 
     void generate() override {
         using namespace Xbyak;
@@ -219,8 +220,8 @@ struct jit_uni_gru_cell_postgemm_part2_fwd : public jit_uni_rnn_postgemm {
                                 scratch_dt_size);
                         uni_vbroadcastss(tmp2_vmm, tmp2s_vmm);
                         // G01 = (1 - a) * G0
-                        compute_vsubps(
-                                tmp2_vmm, tmp1_vmm, tmp2_vmm, current_vlen);
+                        compute_vsubps(tmp2_vmm, tmp1_vmm, tmp2_vmm, tmp3_vmm,
+                                current_vlen);
                         compute_vmulps(G0(loop_ur_idx), G0(loop_ur_idx),
                                 tmp2_vmm, current_vlen);
                         to_float(tmp2_vmm,
@@ -232,7 +233,7 @@ struct jit_uni_gru_cell_postgemm_part2_fwd : public jit_uni_rnn_postgemm {
                                 current_vlen);
                         // tmp1 = G2 * tmp1
                         compute_vmulps(tmp1_vmm, G2(loop_ur_idx), tmp1_vmm,
-                                current_vlen);
+                                tmp3_vmm, current_vlen);
                         // states_t_l = G01 * states_tm1_l + tmp1
                         compute_vfmadd213ps(G0(loop_ur_idx), tmp2_vmm, tmp1_vmm,
                                 current_vlen);

diff --git a/src/cpu/x64/rnn/jit_uni_gru_lbr_cell_postgemm_bwd.hpp b/src/cpu/x64/rnn/jit_uni_gru_lbr_cell_postgemm_bwd.hpp
@@ -172,7 +172,7 @@ struct jit_uni_gru_lbr_cell_postgemm_bwd : public jit_uni_rnn_postgemm {
                 // 1. compute dAttention = -dG0 * G
                 uni_vfnmadd231ps(diff_attn_acc, dG0, G0, tmp2);
                 // 2. Compute dG0 *= 1 - Attention
-                uni_vsubps(tmp1, one_vmm, attn);
+                uni_vsubps(tmp1, one_vmm, attn, tmp2);
                 uni_vmulps(dG0, dG0, tmp1);
             }
             // compute dG2

diff --git a/src/cpu/x64/rnn/jit_uni_gru_lbr_cell_postgemm_fwd.hpp b/src/cpu/x64/rnn/jit_uni_gru_lbr_cell_postgemm_fwd.hpp
@@ -82,7 +82,7 @@ struct jit_uni_gru_lbr_cell_postgemm_fwd : public jit_uni_rnn_postgemm {
         const Reg64 table_reg(rbx); // table is used for data scale and shifts
 
         // We skip vmm0 as it can be used by the injector for masks on sse4.1
-        const Vmm G0(1), G1(2), G2(3), tmp1_vmm(5), tmp2_vmm(6);
+        const Vmm G0(1), G1(2), G2(3), tmp1_vmm(5), tmp2_vmm(6), tmp3_vmm(7);
 
         // constant table map
         const Address one_addr = ptr[table_reg];
@@ -189,12 +189,14 @@ struct jit_uni_gru_lbr_cell_postgemm_fwd : public jit_uni_rnn_postgemm {
                             scratch_dt_size);
                     uni_vbroadcastss(tmp2_vmm, tmp2s_vmm);
                     // G01 = (1 - a) * G0
-                    compute_vsubps(tmp2_vmm, tmp1_vmm, tmp2_vmm, current_vlen);
+                    compute_vsubps(tmp2_vmm, tmp1_vmm, tmp2_vmm, tmp3_vmm,
+                            current_vlen);
                     compute_vmulps(G0, G0, tmp2_vmm, current_vlen);
                     // tmp1 = 1 - G01
                     compute_vsubps(tmp1_vmm, tmp1_vmm, G0, current_vlen);
                     // tmp1 = G2 * tmp1
-                    compute_vmulps(tmp1_vmm, G2, tmp1_vmm, current_vlen);
+                    compute_vmulps(
+                            tmp1_vmm, G2, tmp1_vmm, tmp3_vmm, current_vlen);
                     // states_t_l = G01 * states_tm1_l + tmp2
                     to_float(tmp2_vmm, ptr[addr_states_tm1_l_reg], src_data_t,
                             current_vlen);

diff --git a/src/cpu/x64/rnn/jit_uni_rnn_common_postgemm.hpp b/src/cpu/x64/rnn/jit_uni_rnn_common_postgemm.hpp
@@ -750,6 +750,17 @@ struct jit_uni_rnn_postgemm : public jit_generator {
             uni_vsubps(v1, v2, v3);
     }
 
+    template <typename Vmm>
+    void compute_vsubps(const Vmm &v1, const Vmm &v2, const Vmm &v3,
+            const Vmm &buf, int vlen_bytes) {
+        if (vlen_bytes == 4)
+            // special case for scalar-based tail processing
+            uni_vsubss(Xbyak::Xmm(v1.getIdx()), Xbyak::Xmm(v2.getIdx()),
+                    Xbyak::Xmm(v3.getIdx()), Xbyak::Xmm(buf.getIdx()));
+        else
+            uni_vsubps(v1, v2, v3, buf);
+    }
+
     template <typename Vmm>
     void compute_vmulps(
             const Vmm &v1, const Vmm &v2, const Vmm &v3, int vlen_bytes) {
@@ -761,6 +772,17 @@ struct jit_uni_rnn_postgemm : public jit_generator {
             uni_vmulps(v1, v2, v3);
     }
 
+    template <typename Vmm>
+    void compute_vmulps(const Vmm &v1, const Vmm &v2, const Vmm &v3,
+            const Vmm &buf, int vlen_bytes) {
+        if (vlen_bytes == 4)
+            // special case for scalar-based tail processing
+            uni_vmulss(Xbyak::Xmm(v1.getIdx()), Xbyak::Xmm(v2.getIdx()),
+                    Xbyak::Xmm(v3.getIdx()), Xbyak::Xmm(buf.getIdx()));
+        else
+            uni_vmulps(v1, v2, v3, buf);
+    }
+
     template <typename Vmm>
     void compute_vfmadd231ps(
             const Vmm &v1, const Vmm &v2, const Vmm &v3, int vlen_bytes) {