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
20namespace dnnl {
21namespace impl {
22namespace gpu {
23namespace jit {
24
25// Adds size one spatial dimensions according to input parameters. Spatial
26// dimensions are assumed to be the last dimensions.
27layout_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
54layout_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
64layout_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
75layout_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
111layout_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
127layout_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
141std::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
150void 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