generic: sycl: Adding support for RNN FWD r2l, sum & concat

Author: ShanoToni

src/gpu/generic/sycl/README.md

@@ -193,4 +193,4 @@ The implementation supports forward propagation and vanilla RNN cell kind.
 
 * Supported formats: `ldigo`, `ldgoi`
 * Supported data types: `f32`, `bf16`, `f16`, `s8`, `u8`
-* Supported direction: `left2right`
+* Supported direction: `left2right`, `right2left`, `concat`, `sum`

src/gpu/generic/sycl/rnn/cell_common.cpp

@@ -29,13 +29,17 @@ using namespace rnn_utils;
 
 status_t _ref_rnn_common_t::cell_execution(const cell_ctx_t &cell_struct) {
 
-    auto cell_layer = cell_struct.workspace.states_range(cell_struct.lay - 1,
-            cell_struct.lay - 1, cell_struct.dir, cell_struct.dir,
-            cell_struct.iter - 1, cell_struct.iter);
+    auto cell_layer = cell_struct.workspace.states_range(cell_struct.lay,
+            cell_struct.lay, cell_struct.dir, cell_struct.dir, cell_struct.iter,
+            cell_struct.iter);
 
-    auto cell_iter = cell_struct.workspace.states_range(cell_struct.lay,
-            cell_struct.lay, cell_struct.dir, cell_struct.dir,
-            cell_struct.iter - 2, cell_struct.iter - 1);
+    auto iter_off = cell_struct.iter == 0
+            ? (-1 * (cell_struct.rnn.n_dir - 1) * (cell_struct.rnn.n_iter + 1))
+                    - 1
+            : cell_struct.iter - 1;
+    auto cell_iter = cell_struct.workspace.states_range(cell_struct.lay + 1,
+            cell_struct.lay + 1, cell_struct.dir, cell_struct.dir, iter_off,
+            iter_off);
 
     auto scratch_gates = cell_struct.scratch.gates(0);
 

src/gpu/generic/sycl/rnn/ref_rnn.cpp

@@ -127,9 +127,6 @@ status_t _ref_rnn_common_t::pd_t::init(impl::engine_t *engine) {
     VDISPATCH_RNN_SC(set_weights_desc(this->weights_iter_md_, rnn_conf),
             "unsupported weights iter memory descriptor");
 
-    // Currently only run L2R
-    VDISPATCH_RNN(this->direction() == dnnl_unidirectional_left2right,
-            VERBOSE_BAD_ALGORITHM);
     // Check dimensions consistency
     VDISPATCH_RNN((this->SIC() == this->DHC() || (this->T() == 1)),
             VERBOSE_INCONSISTENT_DIM, "SIC", (int)this->SIC(), "DHC",
@@ -154,10 +151,19 @@ status_t _ref_rnn_common_t::pd_t::init(impl::engine_t *engine) {
     CHECK(memory_desc_init_by_tag(state_md, 5, state_dims,
             rnn_conf.src_data_type, format_tag::abcde));
 
-    copy_init_layer_conf_ = sycl_rnn_copy_conf_t {
-            xpu::sycl::md_t(this->src_md(0)), xpu::sycl::md_t(&state_md),
-            rnn_conf.slc, rnn_conf.n_dir, rnn_conf.n_layer, rnn_conf.n_iter,
-            rnn_conf.mb, rnn_conf.states_ws_ld, true, true};
+    // using is_l2r/r2l to account for bidirectional as well
+    // if both l2r and r2l are true, case is bidirectional concat
+    // is_sum changes behaviour from concat to sum for bidirectional case
+
+    bool is_l2r = !(this->desc()->direction == dnnl_unidirectional_right2left);
+    bool is_r2l = !(this->desc()->direction == dnnl_unidirectional_left2right);
+    bool is_sum = this->desc()->direction == dnnl_bidirectional_sum;
+
+    copy_init_layer_conf_
+            = sycl_rnn_copy_conf_t {xpu::sycl::md_t(this->src_md(0)),
+                    xpu::sycl::md_t(&state_md), rnn_conf.slc, rnn_conf.n_dir,
+                    rnn_conf.n_layer, rnn_conf.n_iter, rnn_conf.mb,
+                    rnn_conf.states_ws_ld, true, true, is_l2r, is_r2l, false};
 
     xpu::sycl::md_t src_iter_md = this->src_md(1)->data_type == data_type::undef
             ? xpu::sycl::md_t()
@@ -166,20 +172,21 @@ status_t _ref_rnn_common_t::pd_t::init(impl::engine_t *engine) {
     copy_init_iter_conf_ = sycl_rnn_copy_conf_t {src_iter_md,
             xpu::sycl::md_t(&state_md), rnn_conf.sic, rnn_conf.n_dir,
             rnn_conf.n_layer, rnn_conf.n_iter, rnn_conf.mb,
-            rnn_conf.states_ws_ld, false, true};
+            rnn_conf.states_ws_ld, false, true, is_l2r, is_r2l, false};
 
     copy_res_layer_conf_ = sycl_rnn_copy_conf_t {xpu::sycl::md_t(&state_md),
             xpu::sycl::md_t(this->dst_md(0)), rnn_conf.dhc, rnn_conf.n_dir,
             rnn_conf.n_layer, rnn_conf.n_iter, rnn_conf.mb,
-            rnn_conf.states_ws_ld, true, false};
+            rnn_conf.states_ws_ld, true, false, is_l2r, is_r2l, is_sum};
 
     xpu::sycl::md_t dst_iter_md = this->dst_md(1)->data_type == data_type::undef
             ? xpu::sycl::md_t()
             : xpu::sycl::md_t(this->dst_md(1));
 
     copy_res_iter_conf_ = sycl_rnn_copy_conf_t {xpu::sycl::md_t(&state_md),
             dst_iter_md, rnn_conf.dhc, rnn_conf.n_dir, rnn_conf.n_layer,
-            rnn_conf.n_iter, rnn_conf.mb, rnn_conf.states_ws_ld, false, false};
+            rnn_conf.n_iter, rnn_conf.mb, rnn_conf.states_ws_ld, false, false,
+            is_l2r, is_r2l, false};
 
     sycl_rnn_bias_conf_t_ = sycl_rnn_bias_conf_t();
     sycl_rnn_bias_conf_t_.dst_md = xpu::sycl::md_t(this->dst_md(0));
@@ -580,7 +587,7 @@ status_t _ref_rnn_common_t::rnn_bias(const exec_ctx_t &ctx, dim_t batch,
 
         auto dst_mem_arg
                 = utils::downcast<const xpu::sycl::memory_storage_base_t *>(
-                        ws.states(lay, dir, iter - 1).get())
+                        ws.states(lay + 1, dir, iter).get())
                           ->get_out_memory_arg(ctx.stream(), cgh);
         ref_rnn_bias bias_kernel(pd()->sycl_rnn_bias_conf_t_, src_mem_arg,
                 bias_mem_arg, dst_mem_arg);

src/gpu/generic/sycl/rnn/rnn_kernels.hpp

@@ -49,35 +49,83 @@ struct ref_rnn_copy_t {
         : src_ {src}, dst_ {dst}, conf_ {conf} {}
 
     void operator()(::sycl::nd_item<3> item) const {
-        const dim_t tl = item.get_global_id(0) / conf_.n_dir; // timestep/layer
-        const dim_t dir = item.get_global_id(0) % conf_.n_dir; // direction
+        const dim_t tl = item.get_global_id(0) // timestep/layer
+                / (conf_.layer ? 1 : conf_.n_dir);
+        dim_t dir = conf_.layer
+                ? 0
+                : item.get_global_id(0) % conf_.n_dir; // direction
         const dim_t n = item.get_global_id(1); // batch
         const dim_t c = item.get_global_id(2); // channel
 
         if (dir >= conf_.n_dir || n >= conf_.batch || c >= conf_.range) return;
 
         dim_t src_offset = 0;
         dim_t dst_offset = 0;
+
         if (conf_.layer) { // layer
             if (tl >= conf_.n_iter) return;
             if (conf_.to_state) { // init
-                src_offset = conf_.src_md.off(tl, n, c);
-                dst_offset = conf_.dst_md.off(0, dir, tl, n, c);
+                if (conf_.l2r) { // l2r
+                    src_offset = conf_.src_md.off(tl, n, c);
+                    dst_offset = conf_.dst_md.off(0, dir, tl, n, c);
+                    do_copy(src_offset, dst_offset, src_ptr(), dst_ptr());
+                    dir = 1;
+                }
+                if (conf_.r2l) { // r2l
+                    src_offset = conf_.src_md.off(tl, n, c);
+                    dst_offset = conf_.dst_md.off(
+                            0, conf_.n_dir - 1, conf_.n_iter - tl - 1, n, c);
+                    do_copy(src_offset, dst_offset, src_ptr(), dst_ptr());
+                }
             } else { // res
-                src_offset = conf_.src_md.off(conf_.n_layer, dir, tl, n, c);
-                dst_offset = conf_.dst_md.off(tl, n, dir * conf_.range + c);
+                if (conf_.l2r) {
+                    dst_offset = conf_.dst_md.off(tl, n, dir * conf_.range + c);
+                    src_offset = conf_.src_md.off(conf_.n_layer, dir, tl, n, c);
+                    do_copy(src_offset, dst_offset, src_ptr(), dst_ptr());
+                    dir = 1;
+                }
+                if (conf_.r2l) {
+                    dst_offset = conf_.dst_md.off(tl, n, dir * conf_.range + c);
+                    src_offset = conf_.src_md.off(
+                            conf_.n_layer, dir, conf_.n_iter - tl - 1, n, c);
+                    if (conf_.sum) {
+                        dst_offset = conf_.dst_md.off(tl, n, c);
+                        auto src = load_float_value(
+                                src_md().data_type(), src_ptr(), src_offset);
+                        auto dst = load_float_value(conf_.dst_md.data_type(),
+                                dst_ptr(), dst_offset);
+                        store_float_value(src_md().data_type(), src + dst,
+                                dst_ptr(), dst_offset);
+                    } else {
+                        do_copy(src_offset, dst_offset, src_ptr(), dst_ptr());
+                    }
+                }
             }
         } else { // iter
             if (tl >= conf_.n_layer) return;
             if (conf_.to_state) { // init
                 src_offset = conf_.src_md.off(tl, dir, n, c);
                 dst_offset = conf_.dst_md.off(tl, dir, conf_.n_iter, n, c);
+                do_copy(src_offset, dst_offset, src_ptr(), dst_ptr());
             } else { // res
                 src_offset
                         = conf_.src_md.off(tl + 1, dir, conf_.n_iter - 1, n, c);
                 dst_offset = conf_.dst_md.off(tl, dir, n, c);
+                do_copy(src_offset, dst_offset, src_ptr(), dst_ptr());
             }
         }
+    }
+
+    xpu::sycl::in_memory_arg_t src_;
+    xpu::sycl::out_memory_arg_t dst_;
+    sycl_rnn_copy_conf_t conf_;
+
+    const xpu::sycl::md_t &src_md() const { return conf_.src_md; }
+    void *src_ptr() const { return src_.get_pointer(); }
+    void *dst_ptr() const { return dst_.get_pointer(); }
+
+    void do_copy(
+            dim_t src_offset, dim_t dst_offset, void *from, void *to) const {
         if (src_ptr()) {
             auto src = load_float_value(
                     src_md().data_type(), src_ptr(), src_offset);
@@ -92,14 +140,6 @@ struct ref_rnn_copy_t {
             }
         }
     }
-
-    xpu::sycl::in_memory_arg_t src_;
-    xpu::sycl::out_memory_arg_t dst_;
-    sycl_rnn_copy_conf_t conf_;
-
-    const xpu::sycl::md_t &src_md() const { return conf_.src_md; }
-    void *src_ptr() const { return src_.get_pointer(); }
-    void *dst_ptr() const { return dst_.get_pointer(); }
 };
 
 struct ref_rnn_bias {

src/gpu/generic/sycl/rnn/rnn_utils.hpp

@@ -205,14 +205,9 @@ struct workspace_t : public data_helper_t {
     }
 
     dim_t calc_off_ws_state(
-            dim_t i0_, dim_t i1, dim_t i2_, dim_t i3, dim_t i4) const {
-        //lay,dir,time
-        // Logical index into workspace grid
-        auto i0 = i0_ + 1;
-        auto i2 = i2_ + 1;
-
+            dim_t i0, dim_t i1, dim_t i2, dim_t i3, dim_t i4) const {
         assert(i0 >= 0);
-
+        //lay,dir,time
         return calc_4d_off(i0, i1, conf_.n_dir, i2, conf_.n_iter + 1, i3,
                 conf_.mb, i4, conf_.states_ws_ld);
     }
@@ -241,10 +236,7 @@ struct workspace_t : public data_helper_t {
 
     std::unique_ptr<mst> states(dim_t layer, dim_t dir, dim_t time) const {
         if (!states_) return {};
-
-        auto i0 = layer + 1;
-        auto i2 = time + 1;
-        auto off_ = get_offset(states_strides(), {i0, dir, i2, 0})
+        auto off_ = get_offset(states_strides(), {layer, dir, time, 0})
                 * conf_.ws_states_elsz;
         return states().clone_ptr_off(off_);
     }

src/gpu/generic/sycl/sycl_primitive_conf.hpp

@@ -471,6 +471,9 @@ struct sycl_rnn_copy_conf_t {
     dim_t states_ws_ld;
     bool layer;
     bool to_state;
+    bool l2r;
+    bool r2l;
+    bool sum;
 };
 
 struct sycl_rnn_bias_conf_t {