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
|
/*
* 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 <iomanip>
#include <stdexcept>
#include <type_traits>
#include "fbgemm/ConvUtils.h"
#include "fbgemm/Fbgemm.h"
namespace fbgemm {
template <typename PT, typename inpType, typename accType>
PackMatrix<PT, inpType, accType>::PackMatrix(
int32_t rows,
int32_t cols,
inpType* buf,
int groups)
: buf_(buf), nrows_(rows), ncols_(cols), G_(groups) {
bufAllocatedHere_ = false;
if (!cpuinfo_initialize()) {
throw std::runtime_error("Failed to initialize cpuinfo!");
}
}
template <typename PT, typename inpType, typename accType>
int PackMatrix<PT, inpType, accType>::packedBufferSize(int rows, int cols) {
if (cpuinfo_has_x86_avx512f()) {
if (isA()) {
return PackingTraits<inpType, accType, inst_set_t::avx512>::MCB *
PackingTraits<inpType, accType, inst_set_t::avx512>::KCB;
} else {
int rowBlock = PackingTraits<inpType, accType, inst_set_t::avx512>::KCB;
int colBlock = PackingTraits<inpType, accType, inst_set_t::avx512>::NCB;
return (((rows + rowBlock - 1) / rowBlock) * rowBlock) *
(((cols + colBlock - 1) / colBlock) * colBlock);
}
} else if (cpuinfo_has_x86_avx2()) {
if (isA()) {
return PackingTraits<inpType, accType, inst_set_t::avx2>::MCB *
PackingTraits<inpType, accType, inst_set_t::avx2>::KCB;
} else {
int rowBlock = PackingTraits<inpType, accType, inst_set_t::avx2>::KCB;
int colBlock = PackingTraits<inpType, accType, inst_set_t::avx2>::NCB;
return (((rows + rowBlock - 1) / rowBlock) * rowBlock) *
(((cols + colBlock - 1) / colBlock) * colBlock);
}
} else {
// TODO: Have default slower path
assert(0 && "unsupported architecure");
}
return -1;
}
// int32 accumulation
template class PackMatrix<PackAMatrix<uint8_t, int32_t>, uint8_t, int32_t>;
template class PackMatrix<
PackAWithRowOffset<uint8_t, int32_t>,
uint8_t,
int32_t>;
template class PackMatrix<PackAWithIm2Col<uint8_t, int32_t>, uint8_t, int32_t>;
template class PackMatrix<
PackAWithIm2Col<uint8_t, int32_t, 3>,
uint8_t,
int32_t>;
template class PackMatrix<
PackAWithQuantRowOffset<uint8_t, int32_t>,
uint8_t,
int32_t>;
template class PackMatrix<PackBMatrix<int8_t, int32_t>, int8_t, int32_t>;
// int16 accumulation
template class PackMatrix<PackAWithIm2Col<uint8_t, int16_t>, uint8_t, int16_t>;
template class PackMatrix<
PackAWithIm2Col<uint8_t, int16_t, 3>,
uint8_t,
int16_t>;
template class PackMatrix<
PackAWithRowOffset<uint8_t, int16_t>,
uint8_t,
int16_t>;
template class PackMatrix<PackAMatrix<uint8_t, int16_t>, uint8_t, int16_t>;
template class PackMatrix<PackBMatrix<int8_t, int16_t>, int8_t, int16_t>;
} // namespace fbgemm
|
