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
|
/*
* Copyright (c) Facebook, Inc. and its affiliates.
* All rights reserved.
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#include <cpuinfo.h>
#include <cassert>
#include <iomanip>
#include "RefImplementations.h"
#include "fbgemm/Fbgemm.h"
namespace fbgemm {
template <typename T, typename accT, int SPATIAL_DIM>
PackWeightMatrixForGConv<T, accT, SPATIAL_DIM>::PackWeightMatrixForGConv(
matrix_op_t trans,
const conv_param_t<SPATIAL_DIM>& conv_param,
const T* sdata,
T* pdata)
: trans_(trans), conv_param_(conv_param), sdata_(sdata) {
assert(SPATIAL_DIM == 2 && "3D conv not supported yet");
if (!pdata) {
bufAllocatedHere_ = true;
pdata_ = static_cast<T*>(fbgemmAlignedAlloc(
64,
conv_param_.G * conv_param_.K[0] * conv_param_.K[1] *
(conv_param_.OC / conv_param_.G) *
(conv_param_.IC / conv_param_.G) * sizeof(T)));
} else {
bufAllocatedHere_ = false;
pdata_ = pdata;
}
pack();
}
/**
* @brief Get the index of the unpacked data for a given <r, s, k, g, c, tr>
*
* Non-transposed: G (R S C/G) K/G
* Transposed: G K/G (R S C/G)
* Using inline as this will be called frequently
*/
template <typename T, typename accT, int SPATIAL_DIM>
inline int PackWeightMatrixForGConv<T, accT, SPATIAL_DIM>::unpacked_index_(
int r, int s, int k, int g, int c, bool tr) {
// Get the full dimensions
int R = conv_param_.K[0];
int S = conv_param_.K[1];
int G = conv_param_.G;
int IC_per_G = conv_param_.IC / G;
int OC_per_G = conv_param_.OC / G;
int idx;
if (tr) {
idx = (((g * OC_per_G + k) * R + r) * S + s) * IC_per_G + c;
} else {
idx = (((g * R + r) * S + s) * IC_per_G + c) * OC_per_G + k;
}
return idx;
}
/**
* @brief Get the index of the packed data for a given <r, s, k, g, c>
*
* The index may differ depending on IC_per_G.
* Using inline as this will be called frequently
*/
template <typename T, typename accT, int SPATIAL_DIM>
inline int PackWeightMatrixForGConv<T, accT, SPATIAL_DIM>::packed_index_(
int r, int s, int k, int g, int c) {
// Get the full dimensions
int R = conv_param_.K[0];
int S = conv_param_.K[1];
int G = conv_param_.G;
int IC_per_G = conv_param_.IC / G;
int OC_per_G = conv_param_.OC / G;
int idx;
// For IC_per_G == 4, we need to work on 2 groups at a time
if (IC_per_G == 4) {
idx = (((((g / 2) * R + r) * S + s) * OC_per_G + k) * 2 + (g % 2))
* IC_per_G + c;
} else {
idx = ((((g * (IC_per_G / 4) + (c / 4)) * R + r) * S + s) * OC_per_G + k)
* 4 + (c % 4);
}
return idx;
}
/**
* @ brief Pack or unpack matrix
*
* Let IC_per_G be number of input channels per group and OC_per_G be number of
* output channels per group.
*
* For IC_per_G == 4 && OC_per_G == 4 optimized
* kernel works on 2 groups at a time hence input channels for g and g+1 group
* are laid out sequentially for each output channel, i.e., the layout is (G/2)
* R S K (2C) and K (2C) is in each 32B vector.
* We work on two groups at a time to fully utilize the avx2 SIMD width of
* 256-bits.
*
* For IC_per_G == 8, 16, 32 && OC_per_G == 8, 16, 32 there is no need to work
* on 2 groups at a time and full SIMD width can be efficiently utilized even
* while working on 1 group at a time.
* In this case, the layout is G (C/4) R S K 4
*/
template <typename T, typename accT, int SPATIAL_DIM>
void PackWeightMatrixForGConv<T, accT, SPATIAL_DIM>::pack_unpack_(
const T* src, T* dst, bool ispack) {
// filters are assumed to be in G RS C/G K/G format
int R = conv_param_.K[0];
int S = conv_param_.K[1];
int G = conv_param_.G;
int IC_per_G = conv_param_.IC / G;
int OC_per_G = conv_param_.OC / G;
// If transpose option is set, the weight matrix is in layout G K/G (R S C/G)
// instead of G (R S C/G) K/G
bool tr = (trans_ == matrix_op_t::Transpose);
if (fbgemmOptimizedGConv(conv_param_)) {
// currently only this case is supported
for (int r = 0; r < R; ++r) {
for (int s = 0; s < S; ++s) {
for (int k = 0; k < OC_per_G; ++k) {
for (int g = 0; g < G; ++g) {
for (int c = 0; c < IC_per_G; ++c) {
int p_idx = packed_index_(r, s, k, g, c);
int up_idx = unpacked_index_(r, s, k, g, c, tr);
// Pack: src (unpacked) -> dst (packed)
if (ispack) {
dst[p_idx] = src[up_idx];
} else {
dst[up_idx] = src[p_idx];
}
}
}
}
}
}
} else {
// For pack & transposed, call transposeConvWeights()
// G K/G (R S C/G) => G (R S C/G) K/G
if (tr) {
if (ispack) {
transposeConvWeights(conv_param_, src, dst);
} else {
// TODO: Wrap this as a inverseTransposeConvWeights()?
// For unpack & transposed, call transposeConvWeights()
// G (R S C/G) K/G => G K/G (R S C/G)
for (int r = 0; r < R; ++r) {
for (int s = 0; s < S; ++s) {
for (int k = 0; k < OC_per_G; ++k) {
for (int g = 0; g < G; ++g) {
for (int c = 0; c < IC_per_G; ++c) {
dst[(((g * OC_per_G + k) * R + r) * S + s)
* IC_per_G + c] =
src[(((g * R + r) * S + s) * IC_per_G + c)
* OC_per_G + k];
}
}
}
}
}
} // end if(ispack)
} else {
// just copy the data for not supported cases
memcpy(dst, src,
G * R * S * OC_per_G * IC_per_G * sizeof(inpType));
} //end if(tr)
} // end if(fbgemmOptimizedGConv(conv_param_)
}
/**
* @brief Pack weight tensor in a suitable format required for the optimized
* kernel.
*/
template <typename T, typename accT, int SPATIAL_DIM>
void PackWeightMatrixForGConv<T, accT, SPATIAL_DIM>::pack() {
pack_unpack_(sdata_, pdata_, true);
}
/**
* @brief Unpack the packed weight tensor (for the optimized kernel)
* to the original form.
*/
template <typename T, typename accT, int SPATIAL_DIM>
void PackWeightMatrixForGConv<T, accT, SPATIAL_DIM>::unpack(T* origin_buf) {
pack_unpack_(const_cast<const T*>(pdata_), origin_buf, false);
}
template class PackWeightMatrixForGConv<int8_t, int32_t, 2>;
template class PackWeightMatrixForGConv<int8_t, int16_t, 2>;
template class PackWeightMatrixForGConv<int8_t, int32_t, 3>;
template class PackWeightMatrixForGConv<int8_t, int16_t, 3>;
} // namespace fbgemm
|
