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
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270

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

#ifndef CPU_X64_RNN_RNN_BRGEMM_UTILS_HPP
#define CPU_X64_RNN_RNN_BRGEMM_UTILS_HPP

#include <memory>
#include "common/c_types_map.hpp"
#include "common/memory_tracking.hpp"
#include "cpu/x64/brgemm/brgemm.hpp"
#include "cpu/x64/jit_brgemm_transpose_utils.hpp"
#include "cpu/x64/matmul/brgemm_matmul_copy_utils.hpp"
#include "cpu/x64/rnn/jit_brgemm_transpose_single_row.hpp"
#include "cpu/x64/rnn/jit_diff_weights_peephole.hpp"
#include "cpu/x64/rnn/jit_gates_reduction.hpp"

namespace dnnl {
namespace impl {
namespace cpu {

namespace rnn_utils {
struct rnn_conf_t;
}

namespace x64 {

struct jit_brgemm_trans_src_t;

namespace rnn_brgemm_utils {

using brgemm_ker_ptr_t = std::unique_ptr<brgemm_kernel_t>;
using brgemm_pallete_t = char[64];
using srcatch_gates_reorder_ker_ptr_t
        = std::unique_ptr<matmul::jit_brgemm_matmul_copy_b_t>;

struct rnn_brgemm_base_t {
    static void init_scratchpad(const cpu::rnn_utils::rnn_conf_t &rnn,
            memory_tracking::registrar_t &scratchpad, dim_t gemm_acc_type_size,
            dim_t gemm_acc_align);
    static constexpr dim_t num_base_kernels_ = 3;
    static constexpr dim_t num_proj_kernels_ = 4;
    static constexpr dim_t num_vanilla_gru_iter_part2_kernels_ = 4;
};

template <prop_kind_t aprop>
struct rnn_brgemm_t;

template <>
struct rnn_brgemm_t<prop_kind::forward> : public rnn_brgemm_base_t {
    using rnn_brgemm_base_t::init_scratchpad;

    static status_t configure_brgemm(cpu::rnn_utils::rnn_conf_t &rnn,
            alg_kind_t cell_kind, dim_t src_layer_type_size,
            dim_t scratch_type_size);
    status_t init_kernels(const cpu::rnn_utils::rnn_conf_t &rnn,
            data_type_t src_type, data_type_t weights_type);

    brgemm_desc_t desc_layer_b0_[num_base_kernels_];
    brgemm_desc_t desc_iter_b0_[num_base_kernels_];
    brgemm_desc_t desc_iter_b1_[num_base_kernels_];
    brgemm_desc_t desc_layer_N_tail_b0_[num_base_kernels_];
    brgemm_desc_t desc_iter_N_tail_b0_[num_base_kernels_];
    brgemm_desc_t desc_iter_N_tail_b1_[num_base_kernels_];

    brgemm_desc_t desc_layer_K1_tail_b1_[num_base_kernels_];
    brgemm_desc_t desc_layer_NK1_tail_b1_[num_base_kernels_];
    brgemm_desc_t desc_iter_K2_tail_b1_[num_base_kernels_];
    brgemm_desc_t desc_iter_NK2_tail_b1_[num_base_kernels_];

    brgemm_desc_t desc_layermerged_b0_[num_base_kernels_];
    brgemm_desc_t desc_layermerged_N_tail_b0_[num_base_kernels_];
    brgemm_desc_t desc_layermerged_K1_tail_b1_[num_base_kernels_];
    brgemm_desc_t desc_layermerged_NK1_tail_b1_[num_base_kernels_];

    brgemm_desc_t desc_proj_b0_[num_proj_kernels_];
    brgemm_desc_t desc_proj_N_tail_b0_[num_proj_kernels_];
    brgemm_desc_t desc_proj_N_tail_b1_[num_proj_kernels_];
    brgemm_desc_t desc_proj_K_tail_b1_[num_proj_kernels_];
    brgemm_desc_t desc_proj_NK_tail_b1_[num_proj_kernels_];

    // Set of brgemm descriptor for 2nd part of iteration gemm in vanulla GRU
    // cell
    brgemm_desc_t desc_iter_p2_b1_[num_vanilla_gru_iter_part2_kernels_];
    brgemm_desc_t desc_iter_p2_N_tail_b1_[num_vanilla_gru_iter_part2_kernels_];
    brgemm_desc_t desc_iter_p2_K2_tail_b1_[num_vanilla_gru_iter_part2_kernels_];
    brgemm_desc_t
            desc_iter_p2_NK2_tail_b1_[num_vanilla_gru_iter_part2_kernels_];

    brgemm_ker_ptr_t kernel_layer_b0_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_layer_b1_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_iter_b0_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_iter_b1_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_layer_N_tail_b0_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_layer_N_tail_b1_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_iter_N_tail_b0_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_iter_N_tail_b1_[num_base_kernels_];

    brgemm_ker_ptr_t kernel_layer_K1_tail_b1_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_layer_NK1_tail_b1_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_iter_K2_tail_b1_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_iter_NK2_tail_b1_[num_base_kernels_];

    brgemm_ker_ptr_t kernel_layermerged_b0_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_layermerged_b1_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_layermerged_N_tail_b0_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_layermerged_N_tail_b1_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_layermerged_K1_tail_b1_[num_base_kernels_];
    brgemm_ker_ptr_t kernel_layermerged_NK1_tail_b1_[num_base_kernels_];

    brgemm_ker_ptr_t kernel_proj_b0_[num_proj_kernels_];
    brgemm_ker_ptr_t kernel_proj_N_tail_b0_[num_proj_kernels_];
    brgemm_ker_ptr_t kernel_proj_N_tail_b1_[num_proj_kernels_];
    brgemm_ker_ptr_t kernel_proj_K_tail_b1_[num_proj_kernels_];
    brgemm_ker_ptr_t kernel_proj_NK_tail_b1_[num_proj_kernels_];

    // Set of brgemm kernels for 2nd part of iteration gemm in vanulla GRU cell
    brgemm_ker_ptr_t kernel_iter_p2_b1_[num_vanilla_gru_iter_part2_kernels_];
    brgemm_ker_ptr_t
            kernel_iter_p2_N_tail_b1_[num_vanilla_gru_iter_part2_kernels_];
    brgemm_ker_ptr_t
            kernel_iter_p2_K2_tail_b1_[num_vanilla_gru_iter_part2_kernels_];
    brgemm_ker_ptr_t
            kernel_iter_p2_NK2_tail_b1_[num_vanilla_gru_iter_part2_kernels_];

    brgemm_pallete_t pallete_buff_iter_;
    brgemm_pallete_t pallete_buff_iter_n_tail_;
    brgemm_pallete_t pallete_buff_layer_;
    brgemm_pallete_t pallete_buff_layer_n_tail_;

    brgemm_pallete_t pallete_buff_k1_tail_;
    brgemm_pallete_t pallete_buff_k2_tail_;
    brgemm_pallete_t pallete_buff_nk1_tail_;
    brgemm_pallete_t pallete_buff_nk2_tail_;
    brgemm_pallete_t pallete_buff_proj_;
    brgemm_pallete_t pallete_buff_nproj_tail_;
    brgemm_pallete_t pallete_buff_kproj_tail_;
    brgemm_pallete_t pallete_buff_nkproj_tail_;

    brgemm_pallete_t pallete_buff_layermerged_;
    brgemm_pallete_t pallete_buff_layermerged_n_tail_;
    brgemm_pallete_t pallete_buff_layermerged_k1_tail_;
    brgemm_pallete_t pallete_buff_layermerged_nk1_tail_;

private:
    status_t brgemm_rnn_init_tiles(
            brgemm_desc_t *desc, brgemm_pallete_t pallete);
    status_t brgemm_rnn_init_tiles_proj(
            brgemm_desc_t *desc, brgemm_pallete_t pallete);
    status_t brgemm_rnn_init_tiles(
            brgemm_desc_t *desc, dim_t size, brgemm_pallete_t pallete);
};

struct rnn_diff_src_brgemm_t {
    brgemm_desc_t desc_iter_layer_beta0_;
    brgemm_desc_t desc_iter_layer_beta1_;
    brgemm_desc_t desc_layer_N_tail_beta0_;
    brgemm_desc_t desc_layer_N_tail_beta1_;
    brgemm_desc_t desc_iter_N_tail_beta0_;
    brgemm_desc_t desc_iter_N_tail_beta1_;
    brgemm_desc_t desc_iter_layer_K_tail_beta1_;
    brgemm_desc_t desc_layer_NK_tail_beta1_;
    brgemm_desc_t desc_iter_NK_tail_beta1_;

    brgemm_ker_ptr_t kernel_iter_layer_beta0_ = nullptr;
    brgemm_ker_ptr_t kernel_iter_layer_beta1_ = nullptr;
    brgemm_ker_ptr_t kernel_layer_N_tail_beta0_ = nullptr;
    brgemm_ker_ptr_t kernel_layer_N_tail_beta1_ = nullptr;
    brgemm_ker_ptr_t kernel_iter_N_tail_beta0_ = nullptr;
    brgemm_ker_ptr_t kernel_iter_N_tail_beta1_ = nullptr;
    brgemm_ker_ptr_t kernel_iter_layer_K_tail_beta1_ = nullptr;
    brgemm_ker_ptr_t kernel_layer_NK_tail_beta1_ = nullptr;
    brgemm_ker_ptr_t kernel_iter_NK_tail_beta1_ = nullptr;

    brgemm_pallete_t pallete_buff_iter_layer_ = {};
    brgemm_pallete_t pallete_buff_iter_layer_k_tail_ = {};
    brgemm_pallete_t pallete_buff_iter_n_tail_ = {};
    brgemm_pallete_t pallete_buff_layer_n_tail_ = {};
    brgemm_pallete_t pallete_buff_iter_nk_tail_ = {};
    brgemm_pallete_t pallete_buff_layer_nk_tail_ = {};
};

struct rnn_diff_wei_brgemm_t {
    brgemm_desc_t desc_iter_beta1_;
    brgemm_desc_t desc_layer_beta1_;
    brgemm_desc_t desc_iter_N_tail_beta1_;
    brgemm_desc_t desc_layer_N_tail_beta1_;
    brgemm_desc_t desc_iter_NK_tail_beta1_;
    brgemm_desc_t desc_layer_NK_tail_beta1_;
    brgemm_desc_t desc_iter_K_tail_beta1_;
    brgemm_desc_t desc_layer_K_tail_beta1_;

    brgemm_ker_ptr_t kernel_iter_beta1_ = nullptr;
    brgemm_ker_ptr_t kernel_layer_beta1_ = nullptr;
    brgemm_ker_ptr_t kernel_iter_N_tail_beta1_ = nullptr;
    brgemm_ker_ptr_t kernel_layer_N_tail_beta1_ = nullptr;
    brgemm_ker_ptr_t kernel_iter_NK_tail_beta1_ = nullptr;
    brgemm_ker_ptr_t kernel_layer_NK_tail_beta1_ = nullptr;
    brgemm_ker_ptr_t kernel_iter_K_tail_beta1_ = nullptr;
    brgemm_ker_ptr_t kernel_layer_K_tail_beta1_ = nullptr;

    brgemm_pallete_t pallete_buff_iter_ = {};
    brgemm_pallete_t pallete_buff_layer_ = {};
    brgemm_pallete_t pallete_buff_iter_n_tail_ = {};
    brgemm_pallete_t pallete_buff_layer_n_tail_ = {};
    brgemm_pallete_t pallete_buff_iter_nk_tail_ = {};
    brgemm_pallete_t pallete_buff_layer_nk_tail_ = {};
    brgemm_pallete_t pallete_buff_iter_k_tail_ = {};
    brgemm_pallete_t pallete_buff_layer_k_tail_ = {};

    srcatch_gates_reorder_ker_ptr_t srcatch_gates_reorder_kernel_;
};

template <>
struct rnn_brgemm_t<prop_kind::backward> : public rnn_brgemm_base_t {
public:
    static void init_scratchpad(const cpu::rnn_utils::rnn_conf_t &rnn,
            memory_tracking::registrar_t &scratchpad, dim_t gemm_acc_type_size,
            dim_t gemm_acc_align);
    static status_t configure_brgemm(cpu::rnn_utils::rnn_conf_t &rnn,
            alg_kind_t cell_kind, dim_t src_layer_type_size,
            dim_t scratch_type_size);

    status_t init_kernels(const cpu::rnn_utils::rnn_conf_t &rnn,
            data_type_t src_type, data_type_t weights_type);

    rnn_diff_src_brgemm_t diff_src_;
    rnn_diff_wei_brgemm_t diff_wei_;

    std::unique_ptr<jit_gates_reduction_t> kernel_gates_reduction_;
    std::unique_ptr<jit_gates_reduction_t> kernel_gates_reduction_tail_;

    std::unique_ptr<jit_brgemm_transpose_single_row_t>
            kernel_transpose_single_row_iter_;
    std::unique_ptr<jit_brgemm_transpose_single_row_t>
            kernel_transpose_single_row_layer_;

    std::unique_ptr<jit_brgemm_trans_src_t>
            kernel_transpose_iter_[num_base_kernels_];
    std::unique_ptr<jit_brgemm_trans_src_t>
            kernel_transpose_layer_[num_base_kernels_];

    std::unique_ptr<jit_diff_weights_peephole_t> kernel_peephole_;
    std::unique_ptr<jit_diff_weights_peephole_t> kernel_peephole_tail_;

private:
    static void configure_brgemm_peephole(cpu::rnn_utils::rnn_conf_t &rnn);

    status_t init_peephole_kernels(const cpu::rnn_utils::rnn_conf_t &rnn);
};

} // namespace rnn_brgemm_utils
} // namespace x64
} // namespace cpu
} // namespace impl
} // namespace dnnl

#endif