path: root/thirdparty/oidn/mkl-dnn/src/common/rnn.cpp

/*******************************************************************************
* Copyright 2018 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/

#include "mkldnn.h"

#include "c_types_map.hpp"
#include "type_helpers.hpp"
#include "utils.hpp"
#include "cpu/gemm/os_blas.hpp"

using namespace mkldnn::impl;
using namespace mkldnn::impl::status;
using namespace mkldnn::impl::types;
using namespace mkldnn::impl::utils;

namespace {
memory_desc_t copy_maybe_null(const memory_desc_t *md) {
    return md ? *md : zero_md();
}

rnn_desc_t zero_rnn_desc() {
    auto rd = rnn_desc_t();
    rd.src_layer_desc = zero_md();
    rd.src_iter_desc = zero_md();
    rd.weights_layer_desc = zero_md();
    rd.weights_iter_desc = zero_md();
    rd.bias_desc = zero_md();
    rd.dst_layer_desc = zero_md();
    rd.dst_iter_desc = zero_md();
    rd.diff_src_layer_desc = zero_md();
    rd.diff_src_iter_desc = zero_md();
    rd.diff_weights_layer_desc = zero_md();
    rd.diff_weights_iter_desc = zero_md();
    rd.diff_bias_desc = zero_md();
    rd.diff_dst_layer_desc = zero_md();
    rd.diff_dst_iter_desc = zero_md();
    return rd;
}
}

/* Public C Api */

status_t mkldnn_rnn_cell_desc_init(rnn_cell_desc_t *rnn_cell_desc,
        mkldnn_alg_kind_t cell_kind, mkldnn_alg_kind_t act_f,
        unsigned int flags, float alpha, float clipping) {
    using namespace mkldnn::impl::alg_kind;

    bool args_ok = true
            && one_of(cell_kind, vanilla_rnn, vanilla_lstm, vanilla_gru,
                    gru_linear_before_reset)
            && IMPLICATION(cell_kind == vanilla_rnn,
                    one_of(act_f, eltwise_relu, eltwise_tanh, eltwise_logistic));
    if (!args_ok)
        return invalid_arguments;

    auto rcd = mkldnn_rnn_cell_desc_t();

    rcd.cell_kind = cell_kind;
    rcd.activation_kind = act_f;
    rcd.flags = flags;
    rcd.alpha = rcd.flags & mkldnn_rnn_cell_with_relu ? alpha : 0;
    rcd.clipping = rcd.flags & mkldnn_rnn_cell_with_clipping ? clipping : 0;

    *rnn_cell_desc = rcd;

    return success;
}

int mkldnn_rnn_cell_get_gates_count(const rnn_cell_desc_t *rnn_cell_desc) {
    switch (rnn_cell_desc->cell_kind) {
    case mkldnn::impl::alg_kind::vanilla_rnn: return 1;
    case mkldnn::impl::alg_kind::vanilla_gru: return 3;
    case mkldnn::impl::alg_kind::gru_linear_before_reset: return 3;
    case mkldnn::impl::alg_kind::vanilla_lstm: return 4;
    default: assert(!"unknown cell kind"); return 0;
    }
    return 0;
}

int mkldnn_rnn_cell_get_states_count(const rnn_cell_desc_t *rnn_cell_desc) {
    switch (rnn_cell_desc->cell_kind) {
    case mkldnn::impl::alg_kind::vanilla_rnn: return 1;
    case mkldnn::impl::alg_kind::vanilla_gru: return 1;
    case mkldnn::impl::alg_kind::gru_linear_before_reset: return 1;
    case mkldnn::impl::alg_kind::vanilla_lstm: return 2;
    default: assert(!"unknown cell kind"); return 0;
    }
    return 0;
}

status_t check_data_type_consistency_fwd(const rnn_cell_desc_t *rnn_cell_desc,
        prop_kind_t prop_kind, const memory_desc_t *src_layer_desc,
        const memory_desc_t *src_iter_desc,
        const memory_desc_t *weights_layer_desc,
        const memory_desc_t *weights_iter_desc, const memory_desc_t *bias_desc,
        const memory_desc_t *dst_layer_desc,
        const memory_desc_t *dst_iter_desc) {
    using namespace data_type;
    data_type_t src_layer_dt = src_layer_desc->data_type;
    data_type_t dst_layer_dt = dst_layer_desc->data_type;
    data_type_t weights_iter_dt = weights_iter_desc->data_type;
    data_type_t weights_layer_dt = weights_layer_desc->data_type;

    bool is_f32 = everyone_is(f32, src_layer_dt, dst_layer_dt, weights_iter_dt,
                          weights_layer_dt)
            && IMPLICATION(!is_zero_md(src_iter_desc),
                          src_iter_desc->data_type == f32)
            && IMPLICATION(!is_zero_md(dst_iter_desc),
                          dst_iter_desc->data_type == f32)
            && IMPLICATION(!is_zero_md(bias_desc), bias_desc->data_type == f32);

#if USE_MKL_PACKED_GEMM
    bool is_u8u8u8 = src_layer_dt == u8
            && IMPLICATION(!is_zero_md(src_iter_desc),
                             src_iter_desc->data_type == u8)
            && IMPLICATION(!is_zero_md(dst_iter_desc),
                             dst_iter_desc->data_type == u8)
            && one_of(dst_layer_dt, u8, f32)
            && everyone_is(s8, weights_iter_dt, weights_layer_dt)
            && IMPLICATION(!is_zero_md(bias_desc), bias_desc->data_type == f32);

    bool is_f32u8f32 = src_layer_dt == u8
            && IMPLICATION(!is_zero_md(src_iter_desc),
                               src_iter_desc->data_type == f32)
            && IMPLICATION(!is_zero_md(dst_iter_desc),
                               dst_iter_desc->data_type == f32)
            && one_of(dst_layer_dt, u8, f32)
            && everyone_is(s8, weights_iter_dt, weights_layer_dt)
            && IMPLICATION(!is_zero_md(bias_desc), bias_desc->data_type == f32);

    bool is_inference = prop_kind == prop_kind::forward_inference;
    bool is_lstm = rnn_cell_desc->cell_kind == mkldnn_vanilla_lstm;

    return (is_f32 || ((is_u8u8u8 || is_f32u8f32) && is_lstm && is_inference))
            ? success
            : unimplemented;
#else
    return is_f32 ? success : unimplemented;
#endif
}

status_t check_dim_consistency(const rnn_cell_desc_t *rnn_cell_desc,
        rnn_direction_t direction, int L, int D, int T, int N, int S, int G,
        int SLC, int SIC, int DLC, int DIC, const memory_desc_t *src_layer_desc,
        const memory_desc_t *src_iter_desc,
        const memory_desc_t *weights_layer_desc,
        const memory_desc_t *weights_iter_desc, const memory_desc_t *bias_desc,
        const memory_desc_t *dst_layer_desc,
        const memory_desc_t *dst_iter_desc) {
    bool args_ok;

    // * algorithm specific
    args_ok = true
        && IMPLICATION(rnn_cell_desc->cell_kind == alg_kind::vanilla_gru,
                       DIC == SIC);
    if (!args_ok) return invalid_arguments;
    int extra_bias =
            rnn_cell_desc->cell_kind == alg_kind::gru_linear_before_reset;

    // * on num layers
    args_ok = true
        && L == weights_layer_desc->dims[0]
        && L == weights_iter_desc->dims[0]
        && IMPLICATION(!is_zero_md(bias_desc), L == bias_desc->dims[0])
        && IMPLICATION(!is_zero_md(src_iter_desc), L == src_iter_desc->dims[0])
        && IMPLICATION(!is_zero_md(dst_iter_desc), L == dst_iter_desc->dims[0]);
    if (!args_ok) return invalid_arguments;

    // * on num directions
    args_ok = true
        && D == weights_layer_desc->dims[1]
        && D == weights_iter_desc->dims[1]
        && IMPLICATION(!is_zero_md(bias_desc), D == bias_desc->dims[1])
        && IMPLICATION(!is_zero_md(src_iter_desc), D == src_iter_desc->dims[1])
        && IMPLICATION(!is_zero_md(dst_iter_desc), D == dst_iter_desc->dims[1]);
    if (!args_ok) return invalid_arguments;

    // * on num iterations
    args_ok = true
        && T == src_layer_desc->dims[0]
        && T == dst_layer_desc->dims[0];
    if (!args_ok) return invalid_arguments;

    // * on mb
    args_ok = true
        && N == src_layer_desc->dims[1]
        && N == dst_layer_desc->dims[1]
        && IMPLICATION(!is_zero_md(src_iter_desc), N == src_iter_desc->dims[3])
        && IMPLICATION(!is_zero_md(dst_iter_desc), N == dst_iter_desc->dims[3]);
    if (!args_ok) return invalid_arguments;

    // * on num gates
    args_ok = true
        && G == mkldnn_rnn_cell_get_gates_count(rnn_cell_desc)
        && G == weights_layer_desc->dims[3]
        && G == weights_iter_desc->dims[3]
        && IMPLICATION(!is_zero_md(bias_desc),
                G + extra_bias == bias_desc->dims[2]);
    if (!args_ok) return invalid_arguments;

    // * on num states
    args_ok = true
        && S == mkldnn_rnn_cell_get_states_count(rnn_cell_desc)
        && IMPLICATION(!is_zero_md(src_iter_desc), S == src_iter_desc->dims[2])
        && IMPLICATION(!is_zero_md(dst_iter_desc), S == dst_iter_desc->dims[2]);
    if (!args_ok) return invalid_arguments;

    // * on slc
    args_ok = true
        && SLC == weights_layer_desc->dims[2]
        && SLC == src_layer_desc->dims[2];
    if (!args_ok) return invalid_arguments;

    // * on sic
    args_ok = true
        && SIC == weights_iter_desc->dims[2]
        && IMPLICATION(!is_zero_md(src_iter_desc),
                SIC == src_iter_desc->dims[4]);
    if (!args_ok) return invalid_arguments;

    // * on dlc
    int dlc_multiplier = (direction == mkldnn_bidirectional_concat) ? 2 : 1;
    args_ok = true
        && DLC == dlc_multiplier * DIC
        && DLC == dst_layer_desc->dims[2];
    if (!args_ok) return invalid_arguments;

    // * on dic
    args_ok = true
        && DIC == weights_layer_desc->dims[4]
        && DIC == weights_iter_desc->dims[4]
        && IMPLICATION(!is_zero_md(bias_desc), DIC == bias_desc->dims[3])
        && IMPLICATION(!is_zero_md(dst_iter_desc),
                DIC == dst_iter_desc->dims[4]);
    if (!args_ok) return invalid_arguments;

    // * unrolling/fusion conditions
    args_ok = true
        && IMPLICATION(L > 1, (dlc_multiplier * SLC) == DLC)
        && IMPLICATION(T > 1, SIC == DIC);
    if (!args_ok) return invalid_arguments;

    return success;
}

status_t MKLDNN_API mkldnn_rnn_forward_desc_init(mkldnn_rnn_desc_t *rnn_desc,
        prop_kind_t prop_kind, const rnn_cell_desc_t *rnn_cell_desc,
        const rnn_direction_t direction, const memory_desc_t *src_layer_desc,
        const memory_desc_t *src_iter_desc,
        const memory_desc_t *weights_layer_desc,
        const memory_desc_t *weights_iter_desc, const memory_desc_t *bias_desc,
        const memory_desc_t *dst_layer_desc,
        const memory_desc_t *dst_iter_desc) {
    bool args_ok = true && rnn_cell_desc != nullptr
            && !any_null(src_layer_desc, weights_layer_desc, weights_iter_desc,
                       dst_layer_desc);
    if (!args_ok) return invalid_arguments;

    //check dimensions consistency
    int L = weights_layer_desc->dims[0];
    int T = src_layer_desc->dims[0];
    int N = src_layer_desc->dims[1];
    const int D = one_of(direction, mkldnn_unidirectional_left2right,
                          mkldnn_unidirectional_right2left) ?
            1 :
            2;
    int G = mkldnn_rnn_cell_get_gates_count(rnn_cell_desc);
    int S = mkldnn_rnn_cell_get_states_count(rnn_cell_desc);
    int SLC = src_layer_desc->dims[2];
    int SIC = weights_iter_desc->dims[2];
    int DLC = dst_layer_desc->dims[2];
    int DIC = weights_layer_desc->dims[4];

    CHECK(check_dim_consistency(rnn_cell_desc, direction, L, D, T, N, S,
            G, SLC, SIC, DLC, DIC, src_layer_desc, src_iter_desc,
            weights_layer_desc, weights_iter_desc, bias_desc, dst_layer_desc,
            dst_iter_desc));

    CHECK(check_data_type_consistency_fwd(rnn_cell_desc, prop_kind,
            src_layer_desc, src_iter_desc, weights_layer_desc,
            weights_iter_desc, bias_desc, dst_layer_desc, dst_iter_desc));

    // Create the descriptor
    mkldnn_rnn_desc_t rd = zero_rnn_desc();

    rd.primitive_kind = primitive_kind::rnn;
    rd.prop_kind = prop_kind;
    rd.cell_desc = *rnn_cell_desc;
    rd.direction = direction;
    rd.src_layer_desc = copy_maybe_null(src_layer_desc);
    rd.src_iter_desc = copy_maybe_null(src_iter_desc);
    rd.weights_layer_desc = copy_maybe_null(weights_layer_desc);
    rd.weights_iter_desc = copy_maybe_null(weights_iter_desc);
    rd.bias_desc = copy_maybe_null(bias_desc);
    rd.dst_layer_desc = copy_maybe_null(dst_layer_desc);
    rd.dst_iter_desc = copy_maybe_null(dst_iter_desc);

    *rnn_desc = rd;

    return success;
}

status_t MKLDNN_API mkldnn_rnn_backward_desc_init(mkldnn_rnn_desc_t *rnn_desc,
        prop_kind_t prop_kind, const rnn_cell_desc_t *rnn_cell_desc,
        const rnn_direction_t direction, const memory_desc_t *src_layer_desc,
        const memory_desc_t *src_iter_desc,
        const memory_desc_t *weights_layer_desc,
        const memory_desc_t *weights_iter_desc, const memory_desc_t *bias_desc,
        const memory_desc_t *dst_layer_desc, const memory_desc_t *dst_iter_desc,
        const memory_desc_t *diff_src_layer_desc,
        const memory_desc_t *diff_src_iter_desc,
        const memory_desc_t *diff_weights_layer_desc,
        const memory_desc_t *diff_weights_iter_desc,
        const memory_desc_t *diff_bias_desc,
        const memory_desc_t *diff_dst_layer_desc,
        const memory_desc_t *diff_dst_iter_desc) {
    bool args_ok = true
            && !any_null(src_layer_desc, weights_layer_desc, weights_iter_desc,
                       dst_layer_desc, diff_src_layer_desc,
                       diff_weights_layer_desc, diff_weights_iter_desc,
                       diff_dst_layer_desc);
    if (!args_ok)
        return invalid_arguments;

    auto xnor_md = [=](const memory_desc_t *a_md, const memory_desc_t *b_md) {
        return is_zero_md(a_md) == is_zero_md(b_md);
    };

    args_ok = args_ok && xnor_md(bias_desc, diff_bias_desc)
            && xnor_md(dst_iter_desc, diff_dst_iter_desc)
            && xnor_md(src_iter_desc, diff_src_iter_desc);
    if (!args_ok)
        return invalid_arguments;

    //check dimensions consistency
    int L = weights_layer_desc->dims[0];
    int T = src_layer_desc->dims[0];
    int N = src_layer_desc->dims[1];
    const int D = one_of(direction, mkldnn_unidirectional_left2right,
                          mkldnn_unidirectional_right2left) ?
            1 :
            2;
    int G = mkldnn_rnn_cell_get_gates_count(rnn_cell_desc);
    int S = mkldnn_rnn_cell_get_states_count(rnn_cell_desc);
    int SLC = src_layer_desc->dims[2];
    int SIC = weights_iter_desc->dims[2];
    int DLC = dst_layer_desc->dims[2];
    int DIC = weights_layer_desc->dims[4];

    status_t st = check_dim_consistency(rnn_cell_desc, direction, L, D, T, N, S,
            G, SLC, SIC, DLC, DIC, src_layer_desc, src_iter_desc,
            weights_layer_desc, weights_iter_desc, bias_desc, dst_layer_desc,
            dst_iter_desc);
    if (st != success) return st;

    st = check_dim_consistency(rnn_cell_desc, direction, L, D, T, N, S,
            G, SLC, SIC, DLC, DIC, diff_src_layer_desc, diff_src_iter_desc,
            diff_weights_layer_desc, diff_weights_iter_desc, diff_bias_desc,
            diff_dst_layer_desc, diff_dst_iter_desc);
    if (st != success) return st;

    mkldnn_rnn_desc_t rd = zero_rnn_desc();

    rd.primitive_kind = primitive_kind::rnn;
    rd.prop_kind = prop_kind;
    rd.cell_desc = *rnn_cell_desc;
    rd.direction = direction;

    rd.src_layer_desc = copy_maybe_null(src_layer_desc);
    rd.src_iter_desc = copy_maybe_null(src_iter_desc);
    rd.weights_layer_desc = copy_maybe_null(weights_layer_desc);
    rd.weights_iter_desc = copy_maybe_null(weights_iter_desc);
    rd.bias_desc = copy_maybe_null(bias_desc);
    rd.dst_layer_desc = copy_maybe_null(dst_layer_desc);
    rd.dst_iter_desc = copy_maybe_null(dst_iter_desc);
    rd.diff_src_layer_desc = copy_maybe_null(diff_src_layer_desc);
    rd.diff_src_iter_desc = copy_maybe_null(diff_src_iter_desc);
    rd.diff_weights_layer_desc = copy_maybe_null(diff_weights_layer_desc);
    rd.diff_weights_iter_desc = copy_maybe_null(diff_weights_iter_desc);
    rd.diff_bias_desc = copy_maybe_null(diff_bias_desc);
    rd.diff_dst_layer_desc = copy_maybe_null(diff_dst_layer_desc);
    rd.diff_dst_iter_desc = copy_maybe_null(diff_dst_iter_desc);

    *rnn_desc = rd;

    return success;
}