onednn-src 0.1.13

Source of oneAPI Deep Neural Network Library (oneDNN)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196

/*******************************************************************************
* 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.
*******************************************************************************/

/*
 * Cell execution of Vanilla RNN
 */

#include "common/dnnl_thread.hpp"
#include "common/math_utils.hpp"

#include "cpu/rnn/postgemm_dispatcher.hpp"

namespace dnnl {
namespace impl {
namespace cpu {

using namespace dnnl::impl::utils;
using namespace dnnl::impl::math;
using namespace rnn_utils;

float activation(alg_kind_t alg_kind, prop_kind_t prop_kind, float s,
        float alpha, float cliping) {
    using namespace dnnl::impl::alg_kind;

    if (prop_kind == prop_kind::forward
            || prop_kind == prop_kind::forward_inference) {
        switch (alg_kind) {
            case eltwise_relu: return relu_fwd<float>(s, alpha);
            case eltwise_tanh: return tanh_fwd<float>(s);
            case eltwise_logistic: return logistic_fwd<float>(s);
            default: assert(!"unsupported algorithm");
        }
    } else if (prop_kind == prop_kind::backward) {
        switch (alg_kind) {
            case eltwise_relu: return relu_bwd<float>(s, alpha);
            case eltwise_tanh: return one_m_square<float>(s);
            case eltwise_logistic: return x_m_square<float>(s);
            default: assert(!"unsupported algorithm");
        }
    } else {
        assert(!"unsupported propagation kind");
    }
    return NAN;
}

constexpr float linear(float s, float alpha, float clipping) {
    return alpha * s;
}

template <typename T, typename src_data_t, typename scratch_data_t>
void rnn_fwd_postgemm_template(T func1, const float *scales, float alpha,
        const rnn_utils::rnn_conf_t &rnn,
        rnn_utils::cell_position_t cell_position, src_data_t *ws_gates_,
        scratch_data_t *scratch_gates_, src_data_t *dst_layer_,
        src_data_t *dst_iter_, const src_data_t *src_iter_, const void *bias_,
        int block_step) {

    const ws_gates_aoc_t<src_data_t> ws_gates(rnn, ws_gates_);
    const scratch_gates_aoc_t<scratch_data_t> scratch_gates(
            rnn, scratch_gates_);
    const auto bias_aoc = rnn_utils::make_raw_aoc(
            bias_, types::data_type_size(rnn.bias_dt), rnn.n_bias, rnn.dhc);
    const auto bias = [&](int gate_id, int dhc_id) {
        return to_float(bias_aoc(gate_id, dhc_id), rnn.bias_dt);
    };
    const auto dst_layer_ld = rnn.dst_layer_ld(cell_position);
    const auto dst_iter_ld = rnn.dst_iter_ld(cell_position);
    const ws_states_layer_aoc_t<src_data_t> dst_layer(
            rnn, dst_layer_, dst_layer_ld);
    const ws_states_iter_aoc_t<src_data_t> dst_iter(
            rnn, dst_iter_, dst_iter_ld);

    if (scales != nullptr) alpha = scales[0];

    const int n_elem = block_step / sizeof(scratch_data_t);

    const auto postgemm_call = [&](dim_t i) {
        for (int j = 0; j < n_elem; j++) {
            const float h
                    = func1(scratch_gates(i, 0, j) + bias(0, j), alpha, 0);
            if (dst_layer_ != nullptr) dst_layer(i, j) = h;
            if (dst_iter_ != nullptr) dst_iter(i, j) = h;
            if (rnn.is_training) ws_gates(i, 0, j) = h;
        }
    };

    if (rnn.is_brgemm && !rnn.unfused_post_gemm) {
        for (int i = 0; i < rnn.m_block; i++)
            postgemm_call(i);
    } else
        parallel_nd(rnn.mb, postgemm_call);
}

template <data_type_t src_type, data_type_t scratch_type, data_type_t acc_type>
rnn_postgemm_sig(
        (rnn_postgemm_fwd_t<src_type, scratch_type, acc_type>::rnn_postgemm)) {
    const float *scales = this->pd_->attr()->rnn_tparams_.scales_;
    const auto act_f = [this](float a, float alpha, float clipping) {
        return gates_t(activation(this->pd_->activation_kind(),
                this->pd_->get_prop_kind(), a, alpha, clipping));
    };
    const auto linear_f = [](float a, float alpha, float clipping) {
        return gates_t(linear(a, alpha, clipping));
    };
    const auto alpha = this->pd_->desc()->alpha;
    if (!this->pd_->attr()->rnn_tparams_.test_mode_)
        rnn_fwd_postgemm_template(act_f, nullptr, alpha, rnn, cell_position,
                ws_gates_, scratch_gates_, dst_layer_, dst_iter_, src_iter_,
                bias_, block_step);
    else
        rnn_fwd_postgemm_template(linear_f, scales, alpha, rnn, cell_position,
                ws_gates_, scratch_gates_, dst_layer_, dst_iter_, src_iter_,
                bias_, block_step);
}

template rnn_postgemm_sig(rnn_postgemm_fwd_f32_t::rnn_postgemm);
template rnn_postgemm_sig(rnn_postgemm_fwd_bf16_t::rnn_postgemm);
template rnn_postgemm_sig(rnn_postgemm_fwd_f16_t::rnn_postgemm);

template <>
rnn_postgemm_sig(rnn_postgemm_fwd_u8_t::rnn_postgemm) {
    assert(!"VANILLA RNN int8 is not supported");
}

template <>
rnn_postgemm_sig(rnn_postgemm_fwd_s8_t::rnn_postgemm) {
    assert(!"VANILLA RNN int8 is not supported");
}

template <typename T1, typename T2, typename src_data_t, typename acc_data_t,
        typename scratch_data_t>
void rnn_bwd_postgemm_template(T1 func1, T2 to_src, const float *scales,
        float alpha, const rnn_utils::rnn_conf_t &rnn, src_data_t *ws_gates_,
        scratch_data_t *scratch_gates_, acc_data_t *diff_dst_iter_,
        acc_data_t *diff_dst_layer_) {
    const ws_gates_aoc_t<src_data_t> ws_gates(rnn, ws_gates_);
    const scratch_gates_aoc_t<scratch_data_t> scratch_gates(
            rnn, scratch_gates_);
    const ws_diff_states_iter_aoc_t<acc_data_t> diff_dst_iter(
            rnn, diff_dst_iter_);
    const ws_diff_states_layer_aoc_t<acc_data_t> diff_dst_layer(
            rnn, diff_dst_layer_);
    if (scales != nullptr) alpha = scales[0];

    parallel_nd(rnn.mb, [&](dim_t i) {
        for (int j = 0; j < rnn.dhc; ++j) {
            const float dH = diff_dst_layer(i, j) + diff_dst_iter(i, j);
            const auto g = (float)ws_gates(i, 0, j);
            const float res = dH * func1(g, alpha, 0);
            src_data_t res_converted = to_src(res);
            scratch_gates(i, 0, j) = res_converted;
        }
    });
}

template <data_type_t src_type, data_type_t scratch_type, data_type_t acc_type>
rnn_postgemm_sig(
        (rnn_postgemm_bwd_t<src_type, scratch_type, acc_type>::rnn_postgemm)) {
    const float *scales = this->pd_->attr()->rnn_tparams_.scales_;
    const auto act_f = [this](float a, float alpha, float clipping) {
        return activation(this->pd_->activation_kind(),
                this->pd_->get_prop_kind(), a, alpha, 0);
    };
    const auto linear_f = [](float a, float alpha, float clipping) {
        return linear(a, alpha, 0);
    };
    const auto to_src = [&](float a) { return gates_t(a); };
    const auto alpha = this->pd_->desc()->alpha;
    if (!this->pd_->attr()->rnn_tparams_.test_mode_)
        rnn_bwd_postgemm_template(act_f, to_src, nullptr, alpha, rnn, ws_gates_,
                scratch_gates_, diff_dst_iter_, diff_dst_layer_);
    else
        rnn_bwd_postgemm_template(linear_f, to_src, scales, alpha, rnn,
                ws_gates_, scratch_gates_, diff_dst_iter_, diff_dst_layer_);
}

template rnn_postgemm_sig(rnn_postgemm_bwd_f32_t::rnn_postgemm);
template rnn_postgemm_sig(rnn_postgemm_bwd_bf16_t::rnn_postgemm);
template rnn_postgemm_sig(rnn_postgemm_bwd_f16_t::rnn_postgemm);

} // namespace cpu
} // namespace impl
} // namespace dnnl