1 | /******************************************************************************* |
2 | * Copyright 2022 Intel Corporation |
3 | * |
4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
5 | * you may not use this file except in compliance with the License. |
6 | * You may obtain a copy of the License at |
7 | * |
8 | * http://www.apache.org/licenses/LICENSE-2.0 |
9 | * |
10 | * Unless required by applicable law or agreed to in writing, software |
11 | * distributed under the License is distributed on an "AS IS" BASIS, |
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
13 | * See the License for the specific language governing permissions and |
14 | * limitations under the License. |
15 | *******************************************************************************/ |
16 | |
17 | /// @example augru.cpp |
18 | /// > Annotated version: @ref augru_example_cpp |
19 | /// |
20 | /// @page augru_example_cpp_short |
21 | /// |
22 | /// This C++ API example demonstrates how to create and execute an |
23 | /// [AUGRU RNN](@ref dev_guide_rnn) primitive in forward training propagation |
24 | /// mode. |
25 | /// |
26 | /// Key optimizations included in this example: |
27 | /// - Creation of optimized memory format from the primitive descriptor. |
28 | /// |
29 | /// @page augru_example_cpp AUGRU RNN Primitive Example |
30 | /// @copydetails augru_example_cpp_short |
31 | /// |
32 | /// @include augru.cpp |
33 | |
34 | #include <algorithm> |
35 | #include <cmath> |
36 | #include <iostream> |
37 | #include <string> |
38 | #include <vector> |
39 | |
40 | #include "example_utils.hpp" |
41 | #include "oneapi/dnnl/dnnl.hpp" |
42 | |
43 | using namespace dnnl; |
44 | |
45 | using tag = memory::format_tag; |
46 | using dt = memory::data_type; |
47 | |
48 | void augru_example(dnnl::engine::kind engine_kind) { |
49 | |
50 | if (engine_kind == engine::kind::gpu) |
51 | throw example_allows_unimplemented { |
52 | "No AUGRU implementation is available for GPU.\n" }; |
53 | |
54 | // Create execution dnnl::engine. |
55 | dnnl::engine engine(engine_kind, 0); |
56 | |
57 | // Create dnnl::stream. |
58 | dnnl::stream engine_stream(engine); |
59 | |
60 | // Tensor dimensions. |
61 | const memory::dim N = 26, // batch size |
62 | T = 6, // time steps |
63 | C = 12, // channels |
64 | G = 3, // gates |
65 | L = 1, // layers |
66 | D = 1; // directions |
67 | |
68 | // Source (src), weights, bias, attention, and destination (dst) tensors |
69 | // dimensions. |
70 | memory::dims src_dims = {T, N, C}; |
71 | memory::dims attention_dims = {T, N, 1}; |
72 | memory::dims weights_dims = {L, D, C, G, C}; |
73 | memory::dims bias_dims = {L, D, G, C}; |
74 | memory::dims dst_dims = {T, N, C}; |
75 | |
76 | // Allocate buffers. |
77 | std::vector<float> src_layer_data(product(src_dims)); |
78 | std::vector<float> attention_data(product(attention_dims)); |
79 | std::vector<float> weights_layer_data(product(weights_dims)); |
80 | std::vector<float> weights_iter_data(product(weights_dims)); |
81 | std::vector<float> bias_data(product(bias_dims)); |
82 | std::vector<float> dst_layer_data(product(dst_dims)); |
83 | |
84 | // Initialize src, weights, and bias tensors. |
85 | std::generate(src_layer_data.begin(), src_layer_data.end(), []() { |
86 | static int i = 0; |
87 | return std::cos(i++ / 10.f); |
88 | }); |
89 | std::generate(attention_data.begin(), attention_data.end(), []() { |
90 | static int i = 0; |
91 | return std::sin(i++ * 2.f); |
92 | }); |
93 | std::generate(weights_layer_data.begin(), weights_layer_data.end(), []() { |
94 | static int i = 0; |
95 | return std::sin(i++ * 2.f); |
96 | }); |
97 | std::generate(bias_data.begin(), bias_data.end(), []() { |
98 | static int i = 0; |
99 | return std::tanh(float(i++)); |
100 | }); |
101 | |
102 | // Create memory descriptors and memory objects for src, bias, and dst. |
103 | auto src_layer_md = memory::desc(src_dims, dt::f32, tag::tnc); |
104 | auto attention_md = memory::desc(attention_dims, dt::f32, tag::tnc); |
105 | auto bias_md = memory::desc(bias_dims, dt::f32, tag::ldgo); |
106 | auto dst_layer_md = memory::desc(dst_dims, dt::f32, tag::tnc); |
107 | |
108 | auto src_layer_mem = memory(src_layer_md, engine); |
109 | auto attention_mem = memory(attention_md, engine); |
110 | auto bias_mem = memory(bias_md, engine); |
111 | auto dst_layer_mem = memory(dst_layer_md, engine); |
112 | |
113 | // Create memory objects for weights using user's memory layout. In this |
114 | // example, LDIGO is assumed. |
115 | auto user_weights_layer_mem |
116 | = memory({weights_dims, dt::f32, tag::ldigo}, engine); |
117 | auto user_weights_iter_mem |
118 | = memory({weights_dims, dt::f32, tag::ldigo}, engine); |
119 | |
120 | // Write data to memory object's handle. |
121 | write_to_dnnl_memory(src_layer_data.data(), src_layer_mem); |
122 | write_to_dnnl_memory(attention_data.data(), attention_mem); |
123 | write_to_dnnl_memory(bias_data.data(), bias_mem); |
124 | write_to_dnnl_memory(weights_layer_data.data(), user_weights_layer_mem); |
125 | write_to_dnnl_memory(weights_iter_data.data(), user_weights_iter_mem); |
126 | |
127 | // Create memory descriptors for weights with format_tag::any. This enables |
128 | // the AUGRU primitive to choose the optimized memory layout. |
129 | auto augru_weights_layer_md = memory::desc(weights_dims, dt::f32, tag::any); |
130 | auto augru_weights_iter_md = memory::desc(weights_dims, dt::f32, tag::any); |
131 | |
132 | // Optional memory descriptors for recurrent data. |
133 | auto src_iter_md = memory::desc(); |
134 | auto dst_iter_md = memory::desc(); |
135 | |
136 | // Create primitive descriptor. |
137 | auto augru_pd |
138 | = augru_forward::primitive_desc(engine, prop_kind::forward_training, |
139 | rnn_direction::unidirectional_left2right, src_layer_md, |
140 | src_iter_md, attention_md, augru_weights_layer_md, |
141 | augru_weights_iter_md, bias_md, dst_layer_md, dst_iter_md); |
142 | |
143 | // For now, assume that the weights memory layout generated by the primitive |
144 | // and the ones provided by the user are identical. |
145 | auto augru_weights_layer_mem = user_weights_layer_mem; |
146 | auto augru_weights_iter_mem = user_weights_iter_mem; |
147 | |
148 | // Reorder the data in case the weights memory layout generated by the |
149 | // primitive and the one provided by the user are different. In this case, |
150 | // we create additional memory objects with internal buffers that will |
151 | // contain the reordered data. |
152 | if (augru_pd.weights_desc() != user_weights_layer_mem.get_desc()) { |
153 | augru_weights_layer_mem = memory(augru_pd.weights_desc(), engine); |
154 | reorder(user_weights_layer_mem, augru_weights_layer_mem) |
155 | .execute(engine_stream, user_weights_layer_mem, |
156 | augru_weights_layer_mem); |
157 | } |
158 | |
159 | if (augru_pd.weights_iter_desc() != user_weights_iter_mem.get_desc()) { |
160 | augru_weights_iter_mem = memory(augru_pd.weights_iter_desc(), engine); |
161 | reorder(user_weights_iter_mem, augru_weights_iter_mem) |
162 | .execute(engine_stream, user_weights_iter_mem, |
163 | augru_weights_iter_mem); |
164 | } |
165 | |
166 | // Create the memory objects from the primitive descriptor. A workspace is |
167 | // also required for AUGRU. |
168 | // NOTE: Here, the workspace is required for later usage in backward |
169 | // propagation mode. |
170 | auto src_iter_mem = memory(augru_pd.src_iter_desc(), engine); |
171 | auto weights_iter_mem = memory(augru_pd.weights_iter_desc(), engine); |
172 | auto dst_iter_mem = memory(augru_pd.dst_iter_desc(), engine); |
173 | auto workspace_mem = memory(augru_pd.workspace_desc(), engine); |
174 | |
175 | // Create the primitive. |
176 | auto augru_prim = augru_forward(augru_pd); |
177 | |
178 | // Primitive arguments |
179 | std::unordered_map<int, memory> augru_args; |
180 | augru_args.insert({DNNL_ARG_SRC_LAYER, src_layer_mem}); |
181 | augru_args.insert({DNNL_ARG_AUGRU_ATTENTION, attention_mem}); |
182 | augru_args.insert({DNNL_ARG_WEIGHTS_LAYER, augru_weights_layer_mem}); |
183 | augru_args.insert({DNNL_ARG_WEIGHTS_ITER, augru_weights_iter_mem}); |
184 | augru_args.insert({DNNL_ARG_BIAS, bias_mem}); |
185 | augru_args.insert({DNNL_ARG_DST_LAYER, dst_layer_mem}); |
186 | augru_args.insert({DNNL_ARG_SRC_ITER, src_iter_mem}); |
187 | augru_args.insert({DNNL_ARG_DST_ITER, dst_iter_mem}); |
188 | augru_args.insert({DNNL_ARG_WORKSPACE, workspace_mem}); |
189 | |
190 | // Primitive execution: AUGRU. |
191 | augru_prim.execute(engine_stream, augru_args); |
192 | |
193 | // Wait for the computation to finalize. |
194 | engine_stream.wait(); |
195 | |
196 | // Read data from memory object's handle. |
197 | read_from_dnnl_memory(dst_layer_data.data(), dst_layer_mem); |
198 | } |
199 | |
200 | int main(int argc, char **argv) { |
201 | return handle_example_errors(augru_example, parse_engine_kind(argc, argv)); |
202 | } |
203 | |