@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 223 lines 5.69 KB
Raw
1 #ifndef DNNL_HELPER_H
2 #define DNNL_HELPER_H
3
4 #include <optional>
5 #include <cassert>
6
7 #include "oneapi/dnnl/dnnl.hpp"
8
9 namespace c10 {
10 struct BFloat16;
11 struct Half;
12 } // namespace c10
13
14 namespace dnnl {
15 namespace impl {
16 struct memory_storage_t;
17 struct matmul_pd_t;
18 struct matmul_desc_t;
19 } // namespace impl
20 } // namespace dnnl
21 struct dnnl_memory_desc;
22
23 template <typename KT, typename VT>
24 class DNNLPrimitiveCache;
25
26 template <typename T>
27 struct DNNLType {
28 static constexpr dnnl::memory::data_type type =
29 dnnl::memory::data_type::undef;
30 };
31
32 template <>
33 struct DNNLType<int8_t> {
34 static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::s8;
35 };
36
37 template <>
38 struct DNNLType<int32_t> {
39 static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::s32;
40 };
41
42 template <>
43 struct DNNLType<float> {
44 static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::f32;
45 };
46
47 template <>
48 struct DNNLType<c10::BFloat16> {
49 static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::bf16;
50 };
51
52 template <>
53 struct DNNLType<c10::Half> {
54 static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::f16;
55 };
56
57 template <typename T>
58 constexpr inline dnnl::memory::data_type get_dnnl_type() {
59 return DNNLType<std::decay_t<T>>::type;
60 }
61
62 class DNNLMatMulPrimitiveHandler {
63 public:
64 virtual ~DNNLMatMulPrimitiveHandler() = default;
65
66 protected:
67 struct Args {
68 dnnl_dim_t b_n_size;
69 dnnl_dim_t b_n_stride;
70 dnnl_dim_t b_k_size;
71 dnnl_dim_t b_k_stride;
72 void* b_ptr;
73 dnnl::memory::data_type c_type;
74 size_t primitive_cache_size;
75 };
76
77 protected:
78 DNNLMatMulPrimitiveHandler(const Args& args, dnnl::memory::data_type b_type);
79
80 void prepack_weight(void* original_b_ptr, dnnl::memory::desc original_b_md,
81 dnnl::memory::desc b_target_mem_desc);
82
83 void set_runtime_memory_ptr(size_t index, dnnl_memory* memory_ptr);
84
85 std::pair<dnnl::impl::memory_storage_t*, dnnl_memory_desc*>
86 get_runtime_memory_ptr(size_t index);
87
88 protected:
89 const dnnl_dim_t b_n_size_;
90 const dnnl_dim_t b_n_stride_;
91 const dnnl_dim_t b_k_size_;
92 const dnnl_dim_t b_k_stride_;
93 dnnl::memory::data_type b_type_;
94 dnnl::memory::data_type c_type_;
95 std::unordered_map<int, dnnl::memory> memory_cache_;
96 std::vector<std::pair<dnnl::impl::memory_storage_t*, dnnl_memory_desc*>>
97 runtime_memory_ptrs_;
98 dnnl::memory::desc b_target_mem_desc_;
99 int64_t primitive_cache_size_;
100 };
101
102 class W8A8MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler {
103 public:
104 enum class QuantizationStrategy { PER_TOKEN, PER_TENSOR, PER_OUTPUT_CHANNEL };
105
106 struct Args : public DNNLMatMulPrimitiveHandler::Args {
107 bool use_a_zero_point;
108 QuantizationStrategy a_quantization_strategy;
109 QuantizationStrategy b_quantization_strategy;
110 float* b_scales_ptr;
111 };
112
113 struct ClassMatmulCacheKey {
114 dnnl_dim_t b_n_size;
115 dnnl_dim_t b_k_size;
116 QuantizationStrategy a_qs;
117 QuantizationStrategy b_qs;
118 bool use_azp;
119 dnnl::memory::data_type c_type;
120
121 friend bool operator==(const ClassMatmulCacheKey& l,
122 const ClassMatmulCacheKey& r);
123 };
124
125 struct MSizeCacheKey {
126 dnnl_dim_t a_m_size;
127 bool use_bias;
128 dnnl::memory::data_type bias_type;
129
130 friend bool operator==(const MSizeCacheKey& l, const MSizeCacheKey& r);
131 };
132
133 using MSizeCache = DNNLPrimitiveCache<MSizeCacheKey, dnnl::matmul>;
134 using ClassMatmulCache =
135 DNNLPrimitiveCache<ClassMatmulCacheKey, std::shared_ptr<MSizeCache>>;
136
137 struct ExecArgs : public MSizeCacheKey {
138 const int8_t* a_ptr;
139 const float* a_scales_ptr;
140 const int32_t* a_zero_points_ptr;
141 const void* bias_ptr;
142 void* c_ptr;
143 };
144
145 public:
146 W8A8MatMulPrimitiveHandler(const Args& args);
147
148 QuantizationStrategy get_input_scale_strategy() const { return a_qs_; }
149
150 bool get_input_use_zero_point() const { return use_azp_; }
151
152 void execute(ExecArgs& args);
153
154 private:
155 dnnl::matmul::primitive_desc create_primitive_desc(const MSizeCacheKey& key,
156 bool first_time);
157
158 void init_runtime_memory_cache(const Args& args);
159
160 dnnl::matmul get_matmul_cache(const MSizeCacheKey& key);
161
162 private:
163 const bool use_azp_;
164 const QuantizationStrategy a_qs_;
165 const QuantizationStrategy b_qs_;
166 std::shared_ptr<MSizeCache> m_size_cache_;
167 };
168
169 class MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler {
170 public:
171 struct Args : public DNNLMatMulPrimitiveHandler::Args {
172 dnnl::memory::data_type ab_type;
173 };
174
175 struct ClassMatmulCacheKey {
176 dnnl_dim_t b_n_size;
177 dnnl_dim_t b_k_size;
178 dnnl::memory::data_type b_type;
179 dnnl::memory::data_type c_type;
180
181 friend bool operator==(const ClassMatmulCacheKey& l,
182 const ClassMatmulCacheKey& r);
183 };
184
185 struct MSizeCacheKey {
186 dnnl_dim_t a_m_size;
187 dnnl_dim_t a_m_stride;
188 bool use_bias;
189 dnnl::memory::data_type bias_type;
190
191 friend bool operator==(const MSizeCacheKey& l, const MSizeCacheKey& r);
192 };
193
194 using MSizeCache = DNNLPrimitiveCache<MSizeCacheKey, dnnl::matmul>;
195 using ClassMatmulCache =
196 DNNLPrimitiveCache<ClassMatmulCacheKey, std::shared_ptr<MSizeCache>>;
197
198 struct ExecArgs : public MSizeCacheKey {
199 const void* a_ptr;
200 const void* bias_ptr;
201 void* c_ptr;
202 };
203
204 public:
205 MatMulPrimitiveHandler(const Args& args);
206
207 dnnl::memory::data_type get_output_type() const { return c_type_; }
208
209 void execute(ExecArgs& args);
210
211 private:
212 dnnl::matmul::primitive_desc create_primitive_desc(const MSizeCacheKey& key,
213 bool first_time);
214
215 void init_runtime_memory_cache(const Args& args);
216
217 dnnl::matmul get_matmul_cache(const MSizeCacheKey& key);
218
219 private:
220 std::shared_ptr<MSizeCache> m_size_cache_;
221 };
222
223 #endif