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 | #include "gpu/jit/conv/normalization.hpp" |
18 | #include "gpu/jit/ir/tensor.hpp" |
19 | |
20 | namespace dnnl { |
21 | namespace impl { |
22 | namespace gpu { |
23 | namespace jit { |
24 | |
25 | // Adds size one spatial dimensions according to input parameters. Spatial |
26 | // dimensions are assumed to be the last dimensions. |
27 | layout_t normalize_conv_spatial(const layout_t &layout, int old_sp_ndims, |
28 | int reduced_dim, bool fuse_spatial) { |
29 | int old_ndims = layout.ndims(); |
30 | int new_ndims = old_ndims - old_sp_ndims + (fuse_spatial ? 1 : 3); |
31 | |
32 | dim_assignment_t to_3d(old_ndims, new_ndims); |
33 | for (int i = 0; i < old_ndims; i++) { |
34 | if (i < old_ndims - old_sp_ndims) { |
35 | // Non-spatial dimensions. |
36 | to_3d.assign(i, i); |
37 | } else if (fuse_spatial) { |
38 | // Spatial dimensions, fuse to one. |
39 | to_3d.assign(i, new_ndims - 1); |
40 | } else { |
41 | // Spatial dimensions. |
42 | int sp_idx = 3 - (old_ndims - i); |
43 | if (reduced_dim == 3) { |
44 | sp_idx = 2; |
45 | } else if (sp_idx < reduced_dim) { |
46 | sp_idx += 1; |
47 | } |
48 | to_3d.assign(i, new_ndims - (3 - sp_idx)); |
49 | } |
50 | } |
51 | return to_3d.map(layout); |
52 | } |
53 | |
54 | layout_t insert_dimension(const layout_t &layout, int dim_idx) { |
55 | auto new_blocks = layout.blocks(); |
56 | for (auto &b : new_blocks) { |
57 | if (b.dim_idx >= dim_idx) b.dim_idx++; |
58 | } |
59 | return layout_t(layout.type(), layout.ndims() + 1, layout.offset(), |
60 | new_blocks, |
61 | /*do_normalize=*/false); |
62 | } |
63 | |
64 | layout_t remove_size_1_dimension(const layout_t &layout, int dim_idx) { |
65 | ir_assert(0 <= dim_idx && dim_idx < layout.ndims()); |
66 | ir_assert(layout.dim(dim_idx) == 1); |
67 | dim_assignment_t a(layout.ndims(), layout.ndims() - 1); |
68 | for (int i = 0; i < layout.ndims(); i++) { |
69 | if (i == dim_idx) continue; |
70 | a.assign(i, i < dim_idx ? i : i - 1); |
71 | } |
72 | return a.map(layout); |
73 | } |
74 | |
75 | layout_t split_dimension( |
76 | const layout_t &_layout, int dim_idx, int outer_block) { |
77 | int rem_inner_block |
78 | = ir_utils::safe_divide(_layout.dim(dim_idx), outer_block); |
79 | auto layout = insert_dimension(_layout, dim_idx); |
80 | std::vector<block_t> new_blocks; |
81 | for (auto &eb : layout.enumerated_blocks()) { |
82 | auto &b = eb.second; |
83 | if (b.dim_idx != dim_idx + 1) { |
84 | new_blocks.push_back(b); |
85 | continue; |
86 | } |
87 | if (b.block % rem_inner_block == 0) { |
88 | new_blocks.emplace_back(dim_idx + 1, rem_inner_block, b.stride); |
89 | new_blocks.emplace_back(dim_idx, b.block / rem_inner_block, |
90 | dim_t(b.stride) * rem_inner_block); |
91 | rem_inner_block = 1; |
92 | } else { |
93 | new_blocks.push_back(b); |
94 | rem_inner_block = ir_utils::safe_divide(rem_inner_block, b.block); |
95 | } |
96 | } |
97 | |
98 | // Remove inner blocks with size one. |
99 | std::vector<block_t> _new_blocks; |
100 | std::vector<bool> seen(layout.ndims()); |
101 | for (auto it = new_blocks.rbegin(); it != new_blocks.rend(); ++it) { |
102 | if (it->block == 1 && seen[it->dim_idx]) continue; |
103 | _new_blocks.push_back(*it); |
104 | seen[it->dim_idx] = true; |
105 | } |
106 | std::reverse(_new_blocks.begin(), _new_blocks.end()); |
107 | return layout_t(layout.type(), layout.ndims(), layout.offset(), _new_blocks, |
108 | /*do_normalize=*/false); |
109 | } |
110 | |
111 | layout_t normalize_conv_groups(const layout_t &layout, bool with_groups, |
112 | int groups, bool is_dw, bool add_groups, bool is_wei) { |
113 | if (with_groups == add_groups) return layout; |
114 | if (is_wei) { |
115 | ir_assert(groups == 1) |
116 | << "Adding/removing groups can be done only for single group." ; |
117 | if (add_groups) return insert_dimension(layout, 0); |
118 | return remove_size_1_dimension(layout, 0); |
119 | } |
120 | |
121 | ir_assert(!with_groups) << "Unexpected groups in source/destination." ; |
122 | if (is_dw) groups = layout.dim(1); |
123 | if (layout.dim(1) == 1) groups = 1; |
124 | return split_dimension(layout, /*dim_idx=*/1, groups); |
125 | } |
126 | |
127 | layout_t normalize_conv_layout(const layout_t &_layout, bool with_groups, |
128 | int groups, bool is_dw, int reduced_dim, bool fuse_spatial, |
129 | bool add_groups, bool is_wei) { |
130 | int old_sp_ndims = _layout.ndims() - (with_groups ? 3 : 2); |
131 | |
132 | layout_t layout = _layout; |
133 | layout = normalize_conv_spatial( |
134 | layout, old_sp_ndims, reduced_dim, fuse_spatial); |
135 | layout = normalize_conv_groups( |
136 | layout, with_groups, groups, is_dw, add_groups, is_wei); |
137 | |
138 | return layout; |
139 | } |
140 | |
141 | std::vector<dim_t> normalize_conv_dims(std::vector<dim_t> &dims, |
142 | bool with_groups, int groups, bool is_dw, int reduced_dim, |
143 | bool fuse_spatial, bool add_groups, bool is_wei) { |
144 | layout_t dummy_layout(type_t::u8(), 0, dims); |
145 | return normalize_conv_layout(dummy_layout, with_groups, groups, is_dw, |
146 | reduced_dim, fuse_spatial, add_groups, is_wei) |
147 | .dims(); |
148 | } |
149 | |
150 | void normalize_conv_layouts(layout_t &src_layout, layout_t &wei_layout, |
151 | layout_t &dst_layout, layout_t &bia_layout, bool with_groups, int g, |
152 | int ic, int oc, bool is_dw, int reduced_dim, bool fuse_spatial, |
153 | bool add_groups) { |
154 | src_layout = normalize_conv_layout(src_layout, /*with_groups=*/false, |
155 | g > 1 ? src_layout.dim(1) / ic : 1, is_dw, reduced_dim, |
156 | fuse_spatial, add_groups, |
157 | /*is_wei=*/false); |
158 | wei_layout = normalize_conv_layout(wei_layout, with_groups, g, is_dw, |
159 | reduced_dim, |
160 | /*fuse_spatial=*/false, add_groups, /*is_wei=*/true); |
161 | dst_layout = normalize_conv_layout(dst_layout, /*with_groups=*/false, |
162 | g > 1 ? dst_layout.dim(1) / oc : 1, is_dw, reduced_dim, |
163 | fuse_spatial, add_groups, |
164 | /*is_wei=*/false); |
165 | if (add_groups && !bia_layout.is_empty()) { |
166 | ir_assert(bia_layout.ndims() == 1) << bia_layout; |
167 | bia_layout = split_dimension(bia_layout, 0, g); |
168 | } |
169 | } |
170 | |
171 | } // namespace jit |
172 | } // namespace gpu |
173 | } // namespace impl |
174 | } // namespace dnnl |
175 | |