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
|
#include "callbacks/configs.h"
#include "callbacks/output_buffer_info.h"
#include "intrinsics.h"
#include "kernels.h"
#include "types.h"
#include "vec_traits.h"
#if defined(THIS_IS_SSE2)
#define CPU_NAME SSE2
#define CPU_ATTR INTGEMM_SSE2
#elif defined(THIS_IS_AVX2)
#define CPU_NAME AVX2
#define CPU_ATTR INTGEMM_AVX2
#elif defined(THIS_IS_AVX512BW)
#define CPU_NAME AVX512BW
#define CPU_ATTR INTGEMM_AVX512BW
#else
#error "Only SSE2, AVX2 and AVX512BW are supported"
#endif
#if defined(THIS_IS_SSE2)
#define vi vector_t<CPUType::SSE2, int>
#define vf vector_t<CPUType::SSE2, float>
#define vd vector_t<CPUType::SSE2, double>
#else
#define vi vector_t<CPUType::AVX2, int>
#define vf vector_t<CPUType::AVX2, float>
#define vd vector_t<CPUType::AVX2, double>
#endif
namespace intgemm {
namespace callbacks {
template <CPUType CpuType, typename CallbackConfig>
class CallbackImpl;
}}
/*
* Callbacks implementations....
*/
namespace intgemm {
namespace callbacks {
/*
* Dummy
*/
template <> class CallbackImpl<CPUType::CPU_NAME, Dummy> {
public:
CPU_ATTR CallbackImpl(const Dummy&) {}
CPU_ATTR void operator()(vi, const OutputBufferInfo&) {}
};
/*
* Write
*/
template <typename OutputBufferType>
class CallbackImpl<CPUType::CPU_NAME, Write<OutputBufferType>> {
public:
CPU_ATTR CallbackImpl(const Write<OutputBufferType>& config) : config(config) {}
CPU_ATTR void operator()(vi input, const OutputBufferInfo& info) {
kernels::write(input, config.addr, info.row_idx * info.cols + info.col_idx);
}
private:
Write<OutputBufferType> config;
};
/*
* UnquantizeAndWrite
*/
template <> class CallbackImpl<CPUType::CPU_NAME, UnquantizeAndWrite> {
public:
CPU_ATTR CallbackImpl(const UnquantizeAndWrite& config) : config(config) {
unquant_mult = set1_ps<vf>(config.unquant_mult);
}
CPU_ATTR void operator()(vi input, const OutputBufferInfo& info) {
auto result = kernels::unquantize(input, unquant_mult);
kernels::write(result, config.addr, info.row_idx * info.cols + info.col_idx);
}
private:
UnquantizeAndWrite config;
vf unquant_mult;
};
/*
* AddBiasAndWrite
*/
template <> class CallbackImpl<CPUType::CPU_NAME, AddBiasAndWrite> {
public:
CPU_ATTR CallbackImpl(const AddBiasAndWrite& config) : config(config) {}
CPU_ATTR void operator()(vi input, const OutputBufferInfo& info) {
auto result = kernels::add_bias(input, config.bias_addr, info.col_idx);
kernels::write(result, config.output_addr, info.row_idx * info.cols + info.col_idx);
}
private:
AddBiasAndWrite config;
};
/*
* UnquantizeAndAddBiasAndWrite
*/
template <> class CallbackImpl<CPUType::CPU_NAME, UnquantizeAndAddBiasAndWrite> {
public:
CPU_ATTR CallbackImpl(const UnquantizeAndAddBiasAndWrite& config) : config(config) {
unquant_mult = set1_ps<vf>(config.unquant_mult);
}
CPU_ATTR void operator()(vi input, const OutputBufferInfo& info) {
auto result = kernels::unquantize(input, unquant_mult);
result = kernels::add_bias(result, config.bias_addr, info.col_idx);
kernels::write(result, config.output_addr, info.row_idx * info.cols + info.col_idx);
}
private:
UnquantizeAndAddBiasAndWrite config;
vf unquant_mult;
};
}
}
#undef CPU_NAME
#undef CPU_ATTR
#undef vi
#undef vf
#undef vd
|
