| 1 | #include <list> |
| 2 | #include <optional> |
| 3 | |
| 4 | #include "common/memory_desc.hpp" |
| 5 | #include "common/memory.hpp" |
| 6 | |
| 7 | #include "cpu/utils.hpp" |
| 8 | #include "cpu/dnnl_helper.h" |
| 9 | |
| 10 | static dnnl::engine& default_engine() { |
| 11 | static dnnl::engine engine(dnnl::engine::kind::cpu, 0); |
| 12 | return engine; |
| 13 | } |
| 14 | |
| 15 | static dnnl::stream& default_stream() { |
| 16 | static dnnl::stream stream(default_engine()); |
| 17 | return stream; |
| 18 | } |
| 19 | |
| 20 | void release_dnnl_matmul_handler(int64_t handler) { |
| 21 | DNNLMatMulPrimitiveHandler* ptr = |
| 22 | reinterpret_cast<DNNLMatMulPrimitiveHandler*>(handler); |
| 23 | delete ptr; |
| 24 | } |
| 25 | |
| 26 | template <typename KT, typename VT> |
| 27 | class DNNLPrimitiveCache { |
| 28 | public: |
| 29 | using cache_value_t = std::pair<KT, VT>; |
| 30 | using result_value_t = VT; |
| 31 | using container_t = std::list<cache_value_t>; |
| 32 | using value_iterator_t = typename container_t::iterator; |
| 33 | using map_t = std::unordered_map<KT, value_iterator_t>; |
| 34 | using creator_t = VT (*)(); |
| 35 | |
| 36 | public: |
| 37 | DNNLPrimitiveCache(size_t capacity) |
| 38 | : capacity_(capacity), |
| 39 | values_(), |
| 40 | key_to_value_(std::min(256lu, capacity)) { |
| 41 | assert(capacity > 0); |
| 42 | } |
| 43 | |
| 44 | template <typename F> |
| 45 | result_value_t get_or_create(const KT& key, F&& creator) { |
| 46 | std::optional<value_iterator_t> value = get_value(key); |
| 47 | if (value.has_value()) { |
| 48 | return value.value()->second; |
| 49 | } else { |
| 50 | return add_value({key, creator()})->second; |
| 51 | } |
| 52 | } |
| 53 | |
| 54 | size_t size() const { return values_.size(); } |
| 55 | |
| 56 | private: |
| 57 | void dump_data() { |
| 58 | std::stringstream ss; |
| 59 | ss << "table_id: " << std::hex << reinterpret_cast<size_t>(this) << std::dec |
| 60 | << "\n"; |
| 61 | ss << "container: ["; |
| 62 | for (auto&& iter : values_) { |
| 63 | ss << "(" << iter.first << ", " << std::hex |
| 64 | << reinterpret_cast<size_t>(iter.second.get()) << "), " << std::dec; |
| 65 | } |
| 66 | ss << "]\n"; |
| 67 | |
| 68 | ss << "map: ["; |
| 69 | for (auto&& iter : key_to_value_) { |
| 70 | ss << "(" << iter.first << ", " << iter.second->first << ", " << std::hex |
| 71 | << reinterpret_cast<size_t>(iter.second->second.get()) << std::dec |
| 72 | << "), "; |
| 73 | } |
| 74 | ss << "]\n"; |
| 75 | std::printf("%s\n", ss.str().c_str()); |
| 76 | } |
| 77 | |
| 78 | value_iterator_t add_value(cache_value_t&& new_value) { |
| 79 | if (size() == capacity_) { |
| 80 | cache_value_t& last_item = values_.back(); |
| 81 | key_to_value_.erase(last_item.first); |
| 82 | values_.pop_back(); |
| 83 | } |
| 84 | |
| 85 | auto& added_value_ = values_.emplace_front(std::move(new_value)); |
| 86 | key_to_value_.emplace(added_value_.first, values_.begin()); |
| 87 | return values_.begin(); |
| 88 | } |
| 89 | |
| 90 | std::optional<value_iterator_t> get_value(const KT& key) { |
| 91 | if (key_to_value_.size() > 0 && key == values_.begin()->first) { |
| 92 | return values_.begin(); |
| 93 | } |
| 94 | |
| 95 | auto value_map_iterator = key_to_value_.find(key); |
| 96 | if (value_map_iterator != key_to_value_.end()) { |
| 97 | values_.splice(values_.begin(), values_, value_map_iterator->second); |
| 98 | return value_map_iterator->second; |
| 99 | } else { |
| 100 | return {}; |
| 101 | } |
| 102 | } |
| 103 | |
| 104 | private: |
| 105 | const size_t capacity_; |
| 106 | container_t values_; |
| 107 | map_t key_to_value_; |
| 108 | }; |
| 109 | |
| 110 | DNNLMatMulPrimitiveHandler::DNNLMatMulPrimitiveHandler( |
| 111 | const Args& args, dnnl::memory::data_type b_type) |
| 112 | : b_n_size_(args.b_n_size), |
| 113 | b_n_stride_(args.b_n_stride), |
| 114 | b_k_size_(args.b_k_size), |
| 115 | b_k_stride_(args.b_k_stride), |
| 116 | b_type_(b_type), |
| 117 | c_type_(args.c_type), |
| 118 | runtime_memory_ptrs_(8), |
| 119 | primitive_cache_size_(args.primitive_cache_size) { |
| 120 | assert(primitive_cache_size_ > 0); |
| 121 | } |
| 122 | |
| 123 | void DNNLMatMulPrimitiveHandler::prepack_weight( |
| 124 | void* original_b_ptr, dnnl::memory::desc original_b_md, |
| 125 | dnnl::memory::desc b_target_mem_desc) { |
| 126 | dnnl::memory original_weight(original_b_md, default_engine(), original_b_ptr); |
| 127 | dnnl::memory packed_weight(b_target_mem_desc, default_engine()); |
| 128 | { |
| 129 | dnnl::reorder(original_weight, packed_weight) |
| 130 | .execute(default_stream(), original_weight, packed_weight); |
| 131 | default_stream().wait(); |
| 132 | } |
| 133 | memory_cache_[DNNL_ARG_WEIGHTS] = packed_weight; |
| 134 | b_target_mem_desc_ = b_target_mem_desc; |
| 135 | } |
| 136 | |
| 137 | void DNNLMatMulPrimitiveHandler::set_runtime_memory_ptr( |
| 138 | size_t index, dnnl_memory* memory_ptr) { |
| 139 | dnnl::impl::memory_storage_t* mem_storage_ptr = memory_ptr->memory_storage(); |
| 140 | dnnl_memory_desc* mem_desc = const_cast<dnnl_memory_desc*>(memory_ptr->md()); |
| 141 | runtime_memory_ptrs_[index] = {mem_storage_ptr, mem_desc}; |
| 142 | } |
| 143 | |
| 144 | std::pair<dnnl::impl::memory_storage_t*, dnnl_memory_desc*> |
| 145 | DNNLMatMulPrimitiveHandler::get_runtime_memory_ptr(size_t index) { |
| 146 | return runtime_memory_ptrs_[index]; |
| 147 | } |
| 148 | |
| 149 | namespace std { |
| 150 | template <> |
| 151 | struct hash<W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey> { |
| 152 | size_t operator()( |
| 153 | const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& val) const { |
| 154 | return hash<dnnl_dim_t>()(val.b_n_size) ^ hash<dnnl_dim_t>()(val.b_k_size) ^ |
| 155 | hash<int>()(static_cast<int>(val.a_qs)) ^ |
| 156 | hash<int>()(static_cast<int>(val.b_qs)) ^ hash<bool>()(val.use_azp) ^ |
| 157 | hash<int>()(static_cast<int>(val.c_type)); |
| 158 | } |
| 159 | }; |
| 160 | |
| 161 | template <> |
| 162 | struct hash<W8A8MatMulPrimitiveHandler::MSizeCacheKey> { |
| 163 | size_t operator()( |
| 164 | const W8A8MatMulPrimitiveHandler::MSizeCacheKey& val) const { |
| 165 | return hash<dnnl_dim_t>()(val.a_m_size) ^ hash<bool>()(val.use_bias) ^ |
| 166 | hash<int>()(static_cast<int>(val.bias_type)); |
| 167 | } |
| 168 | }; |
| 169 | |
| 170 | template <> |
| 171 | struct hash<MatMulPrimitiveHandler::ClassMatmulCacheKey> { |
| 172 | size_t operator()( |
| 173 | const MatMulPrimitiveHandler::ClassMatmulCacheKey& val) const { |
| 174 | return hash<dnnl_dim_t>()(val.b_n_size) ^ hash<dnnl_dim_t>()(val.b_k_size) ^ |
| 175 | hash<int>()(static_cast<int>(val.b_type)) ^ |
| 176 | hash<int>()(static_cast<int>(val.c_type)); |
| 177 | } |
| 178 | }; |
| 179 | |
| 180 | template <> |
| 181 | struct hash<MatMulPrimitiveHandler::MSizeCacheKey> { |
| 182 | size_t operator()(const MatMulPrimitiveHandler::MSizeCacheKey& val) const { |
| 183 | return hash<dnnl_dim_t>()(val.a_m_size) ^ |
| 184 | hash<dnnl_dim_t>()(val.a_m_stride) ^ hash<bool>()(val.use_bias) ^ |
| 185 | hash<int>()(static_cast<int>(val.bias_type)); |
| 186 | } |
| 187 | }; |
| 188 | } // namespace std |
| 189 | |
| 190 | bool operator==(const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& l, |
| 191 | const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& r) { |
| 192 | return l.b_n_size == r.b_n_size && l.b_k_size == r.b_k_size && |
| 193 | l.a_qs == r.a_qs && l.b_qs == r.b_qs && l.use_azp == r.use_azp && |
| 194 | l.c_type == r.c_type; |
| 195 | } |
| 196 | |
| 197 | bool operator==(const W8A8MatMulPrimitiveHandler::MSizeCacheKey& l, |
| 198 | const W8A8MatMulPrimitiveHandler::MSizeCacheKey& r) { |
| 199 | return l.use_bias == r.use_bias && l.a_m_size == r.a_m_size && |
| 200 | l.bias_type == r.bias_type; |
| 201 | } |
| 202 | |
| 203 | bool operator==(const MatMulPrimitiveHandler::ClassMatmulCacheKey& l, |
| 204 | const MatMulPrimitiveHandler::ClassMatmulCacheKey& r) { |
| 205 | return l.b_n_size == r.b_n_size && l.b_k_size == r.b_k_size && |
| 206 | l.b_type == r.b_type && l.c_type == r.c_type; |
| 207 | } |
| 208 | |
| 209 | bool operator==(const MatMulPrimitiveHandler::MSizeCacheKey& l, |
| 210 | const MatMulPrimitiveHandler::MSizeCacheKey& r) { |
| 211 | return l.a_m_size == r.a_m_size && l.a_m_stride == r.a_m_stride && |
| 212 | l.use_bias == r.use_bias && l.bias_type == r.bias_type; |
| 213 | } |
| 214 | |
| 215 | static std::shared_ptr<W8A8MatMulPrimitiveHandler::MSizeCache> |
| 216 | get_w8a8_class_primitive_cache( |
| 217 | const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& key, |
| 218 | int64_t cache_size) { |
| 219 | static W8A8MatMulPrimitiveHandler::ClassMatmulCache cache(128); |
| 220 | assert(cache_size > 0); |
| 221 | return cache.get_or_create(key, [&]() { |
| 222 | return std::make_shared<W8A8MatMulPrimitiveHandler::MSizeCache>(cache_size); |
| 223 | }); |
| 224 | } |
| 225 | |
| 226 | W8A8MatMulPrimitiveHandler::W8A8MatMulPrimitiveHandler(const Args& args) |
| 227 | : DNNLMatMulPrimitiveHandler( |
| 228 | static_cast<const DNNLMatMulPrimitiveHandler::Args&>(args), |
| 229 | dnnl::memory::data_type::s8), |
| 230 | use_azp_(args.use_a_zero_point), |
| 231 | a_qs_(args.a_quantization_strategy), |
| 232 | b_qs_(args.b_quantization_strategy), |
| 233 | m_size_cache_(nullptr) { |
| 234 | assert(a_qs_ != QuantizationStrategy::PER_OUTPUT_CHANNEL); |
| 235 | assert(b_qs_ != QuantizationStrategy::PER_TOKEN); |
| 236 | if (a_qs_ == QuantizationStrategy::PER_TOKEN) { |
| 237 | assert(!use_azp_); |
| 238 | }; |
| 239 | dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_, |
| 240 | {b_k_stride_, b_n_stride_}); |
| 241 | |
| 242 | // dummy M size for prepacking weights |
| 243 | // Prepacking weights improves performance and avoid runtime reorders |
| 244 | constexpr dnnl_dim_t kProbeM = 128; |
| 245 | |
| 246 | prepack_weight(args.b_ptr, original_b_md, |
| 247 | create_primitive_desc( |
| 248 | MSizeCacheKey{.a_m_size = kProbeM, |
| 249 | .use_bias = false, |
| 250 | .bias_type = dnnl::memory::data_type::undef}, |
| 251 | /*first_time=*/true) |
| 252 | .weights_desc()); |
| 253 | init_runtime_memory_cache(args); |
| 254 | } |
| 255 | |
| 256 | void W8A8MatMulPrimitiveHandler::execute(ExecArgs& args) { |
| 257 | auto&& [a_storage, a_mem_desc] = get_runtime_memory_ptr(0); |
| 258 | auto&& [c_storage, c_mem_desc] = get_runtime_memory_ptr(1); |
| 259 | a_storage->set_data_handle((void*)args.a_ptr); |
| 260 | a_mem_desc->dims[0] = args.a_m_size; |
| 261 | c_storage->set_data_handle((void*)args.c_ptr); |
| 262 | c_mem_desc->dims[0] = args.a_m_size; |
| 263 | |
| 264 | if (a_qs_ == QuantizationStrategy::PER_TENSOR) { |
| 265 | auto&& [a_scale_storage, a_scale_mem_desc] = get_runtime_memory_ptr(2); |
| 266 | a_scale_storage->set_data_handle((void*)args.a_scales_ptr); |
| 267 | } |
| 268 | if (use_azp_) { |
| 269 | auto&& [a_zero_point_storage, a_zero_point_mem_desc] = |
| 270 | get_runtime_memory_ptr(3); |
| 271 | a_zero_point_storage->set_data_handle((void*)args.a_zero_points_ptr); |
| 272 | } |
| 273 | |
| 274 | if (args.use_bias) { |
| 275 | auto&& [bias_storage, bias_mem_desc] = get_runtime_memory_ptr(4); |
| 276 | bias_storage->set_data_handle((void*)args.bias_ptr); |
| 277 | } |
| 278 | |
| 279 | dnnl::matmul matmul = get_matmul_cache(args); |
| 280 | |
| 281 | auto&& [scratchpad_storage, scratchpad_mem_desc] = get_runtime_memory_ptr(5); |
| 282 | scratchpad_storage->set_data_handle( |
| 283 | cpu_utils::ScratchPadManager::get_scratchpad_manager()->get_data<void>()); |
| 284 | |
| 285 | matmul.execute(default_stream(), memory_cache_); |
| 286 | default_stream().wait(); |
| 287 | } |
| 288 | |
| 289 | dnnl::matmul W8A8MatMulPrimitiveHandler::get_matmul_cache( |
| 290 | const MSizeCacheKey& key) { |
| 291 | if (m_size_cache_.get() == nullptr) { |
| 292 | ClassMatmulCacheKey key = {.b_n_size = b_n_size_, |
| 293 | .b_k_size = b_k_size_, |
| 294 | .a_qs = a_qs_, |
| 295 | .b_qs = b_qs_, |
| 296 | .use_azp = use_azp_, |
| 297 | .c_type = c_type_}; |
| 298 | m_size_cache_ = get_w8a8_class_primitive_cache(key, primitive_cache_size_); |
| 299 | } |
| 300 | |
| 301 | return m_size_cache_->get_or_create(key, [&]() { |
| 302 | dnnl::matmul::primitive_desc desc = this->create_primitive_desc(key, false); |
| 303 | auto manager = cpu_utils::ScratchPadManager::get_scratchpad_manager(); |
| 304 | manager->realloc(desc.scratchpad_desc().get_size()); |
| 305 | return dnnl::matmul(desc); |
| 306 | }); |
| 307 | } |
| 308 | |
| 309 | void W8A8MatMulPrimitiveHandler::init_runtime_memory_cache(const Args& args) { |
| 310 | memory_cache_[DNNL_ARG_SRC] = dnnl::memory({{1, b_k_size_}, |
| 311 | dnnl::memory::data_type::s8, |
| 312 | dnnl::memory::format_tag::ab}, |
| 313 | default_engine(), nullptr); |
| 314 | set_runtime_memory_ptr(0, memory_cache_[DNNL_ARG_SRC].get()); |
| 315 | memory_cache_[DNNL_ARG_DST] = |
| 316 | dnnl::memory({{1, b_n_size_}, c_type_, dnnl::memory::format_tag::ab}, |
| 317 | default_engine(), nullptr); |
| 318 | set_runtime_memory_ptr(1, memory_cache_[DNNL_ARG_DST].get()); |
| 319 | |
| 320 | // For PER_TOKEN, scales will be applied in outside epilogue |
| 321 | if (a_qs_ == QuantizationStrategy::PER_TENSOR) { |
| 322 | memory_cache_[DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC] = dnnl::memory( |
| 323 | {{1}, dnnl::memory::data_type::f32, {1}}, default_engine(), nullptr); |
| 324 | set_runtime_memory_ptr( |
| 325 | 2, memory_cache_[DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC].get()); |
| 326 | if (use_azp_) { |
| 327 | memory_cache_[DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_SRC] = dnnl::memory( |
| 328 | {{1}, dnnl::memory::data_type::s32, {1}}, default_engine(), nullptr); |
| 329 | set_runtime_memory_ptr( |
| 330 | 3, memory_cache_[DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_SRC].get()); |
| 331 | } |
| 332 | } |
| 333 | |
| 334 | if (b_qs_ == QuantizationStrategy::PER_TENSOR) { |
| 335 | memory_cache_[DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS] = |
| 336 | dnnl::memory({{1}, dnnl::memory::data_type::f32, {1}}, default_engine(), |
| 337 | (void*)args.b_scales_ptr); |
| 338 | } else if (b_qs_ == QuantizationStrategy::PER_OUTPUT_CHANNEL) { |
| 339 | memory_cache_[DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS] = |
| 340 | dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}}, |
| 341 | default_engine(), (void*)args.b_scales_ptr); |
| 342 | } |
| 343 | |
| 344 | memory_cache_[DNNL_ARG_BIAS] = |
| 345 | dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}}, |
| 346 | default_engine(), nullptr); |
| 347 | set_runtime_memory_ptr(4, memory_cache_[DNNL_ARG_BIAS].get()); |
| 348 | |
| 349 | memory_cache_[DNNL_ARG_SCRATCHPAD] = |
| 350 | dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}}, |
| 351 | default_engine(), nullptr); |
| 352 | set_runtime_memory_ptr(5, memory_cache_[DNNL_ARG_SCRATCHPAD].get()); |
| 353 | } |
| 354 | |
| 355 | dnnl::matmul::primitive_desc W8A8MatMulPrimitiveHandler::create_primitive_desc( |
| 356 | const MSizeCacheKey& key, bool first_time) { |
| 357 | dnnl::memory::desc a_md({key.a_m_size, b_k_size_}, |
| 358 | dnnl::memory::data_type::s8, |
| 359 | dnnl::memory::format_tag::ab); |
| 360 | dnnl::memory::desc b_md; |
| 361 | if (first_time) { |
| 362 | b_md = |
| 363 | dnnl::memory::desc({b_k_size_, b_n_size_}, dnnl::memory::data_type::s8, |
| 364 | dnnl::memory::format_tag::any); |
| 365 | } else { |
| 366 | b_md = b_target_mem_desc_; |
| 367 | } |
| 368 | dnnl::memory::desc c_md({key.a_m_size, b_n_size_}, c_type_, |
| 369 | dnnl::memory::format_tag::ab); |
| 370 | |
| 371 | dnnl::primitive_attr attr; |
| 372 | |
| 373 | attr.set_scratchpad_mode(dnnl::scratchpad_mode::user); |
| 374 | |
| 375 | // For PER_TOKEN, scales will be applied in outside epilogue |
| 376 | if (a_qs_ == QuantizationStrategy::PER_TENSOR) { |
| 377 | attr.set_scales_mask(DNNL_ARG_SRC, 0); |
| 378 | if (use_azp_) { |
| 379 | attr.set_zero_points_mask(DNNL_ARG_SRC, 0); |
| 380 | } |
| 381 | } |
| 382 | |
| 383 | if (b_qs_ == QuantizationStrategy::PER_TENSOR) { |
| 384 | attr.set_scales_mask(DNNL_ARG_WEIGHTS, 0); |
| 385 | } else if (b_qs_ == QuantizationStrategy::PER_OUTPUT_CHANNEL) { |
| 386 | attr.set_scales_mask(DNNL_ARG_WEIGHTS, 2); |
| 387 | } |
| 388 | |
| 389 | if (key.use_bias) { |
| 390 | // For PER_TOKEN, bias will be applied in epilogue |
| 391 | assert(a_qs_ == QuantizationStrategy::PER_TENSOR); |
| 392 | dnnl::memory::desc bias_md({1, b_n_size_}, key.bias_type, {b_n_size_, 1}); |
| 393 | return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, bias_md, |
| 394 | c_md, attr); |
| 395 | } else { |
| 396 | return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, c_md, |
| 397 | attr); |
| 398 | } |
| 399 | } |
| 400 | |
| 401 | MatMulPrimitiveHandler::MatMulPrimitiveHandler(const Args& args) |
| 402 | : DNNLMatMulPrimitiveHandler( |
| 403 | static_cast<DNNLMatMulPrimitiveHandler::Args>(args), args.ab_type), |
| 404 | m_size_cache_(nullptr) { |
| 405 | assert(b_type_ == dnnl::memory::data_type::f32 || |
| 406 | b_type_ == dnnl::memory::data_type::bf16 || |
| 407 | b_type_ == dnnl::memory::data_type::f16); |
| 408 | |
| 409 | dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_, |
| 410 | {b_k_stride_, b_n_stride_}); |
| 411 | |
| 412 | // dummy M size for prepacking weights |
| 413 | // Prepacking weights improves performance and avoid runtime reorders |
| 414 | constexpr dnnl_dim_t kProbeM = 128; |
| 415 | |
| 416 | prepack_weight(args.b_ptr, original_b_md, |
| 417 | create_primitive_desc( |
| 418 | MSizeCacheKey{// Use a concrete M so oneDNN's kernel |
| 419 | // selector can choose an optimally blocked |
| 420 | // weight layout. |
| 421 | .a_m_size = kProbeM, |
| 422 | .a_m_stride = b_k_size_, |
| 423 | .use_bias = false, |
| 424 | .bias_type = dnnl::memory::data_type::undef}, |
| 425 | true) |
| 426 | .weights_desc()); |
| 427 | init_runtime_memory_cache(args); |
| 428 | } |
| 429 | |
| 430 | static std::shared_ptr<MatMulPrimitiveHandler::MSizeCache> |
| 431 | get_matul_class_primitive_cache( |
| 432 | const MatMulPrimitiveHandler::ClassMatmulCacheKey& key, |
| 433 | int64_t cache_size) { |
| 434 | static MatMulPrimitiveHandler::ClassMatmulCache cache(128); |
| 435 | assert(cache_size > 0); |
| 436 | return cache.get_or_create(key, [&]() { |
| 437 | return std::make_shared<MatMulPrimitiveHandler::MSizeCache>(cache_size); |
| 438 | }); |
| 439 | } |
| 440 | |
| 441 | void MatMulPrimitiveHandler::execute(ExecArgs& args) { |
| 442 | auto&& [a_storage, a_mem_desc] = get_runtime_memory_ptr(0); |
| 443 | auto&& [c_storage, c_mem_desc] = get_runtime_memory_ptr(1); |
| 444 | a_storage->set_data_handle((void*)args.a_ptr); |
| 445 | a_mem_desc->dims[0] = args.a_m_size; |
| 446 | a_mem_desc->format_desc.blocking.strides[0] = args.a_m_stride; |
| 447 | c_storage->set_data_handle((void*)args.c_ptr); |
| 448 | c_mem_desc->dims[0] = args.a_m_size; |
| 449 | |
| 450 | #ifndef VLLM_USE_ACL |
| 451 | // We do not support in ACL backend of oneDNN, we handle bias by: |
| 452 | // 1. copying it into the result tensor |
| 453 | // 2. attaching a fused-sum post-op to the matmul primitive |
| 454 | if (args.use_bias) { |
| 455 | auto&& [bias_storage, bias_mem_desc] = get_runtime_memory_ptr(2); |
| 456 | bias_storage->set_data_handle((void*)args.bias_ptr); |
| 457 | } |
| 458 | #endif |
| 459 | dnnl::matmul matmul = get_matmul_cache(args); |
| 460 | |
| 461 | // With ACL backend of oneDNN, the required memory format might change when the |
| 462 | // source tensor dims change. This does not really happen in practice, so isn't |
| 463 | // a performance hit, but we need to support it because the API allows for it. |
| 464 | #ifdef VLLM_USE_ACL |
| 465 | auto new_expected_wei_desc = |
| 466 | dnnl::matmul::primitive_desc( |
| 467 | const_cast<dnnl_primitive_desc_t>(matmul.get_primitive_desc())) |
| 468 | .weights_desc(); |
| 469 | if (new_expected_wei_desc != b_target_mem_desc_) { |
| 470 | prepack_weight(memory_cache_[DNNL_ARG_WEIGHTS].get_data_handle(), |
| 471 | b_target_mem_desc_, new_expected_wei_desc); |
| 472 | } |
| 473 | #endif |
| 474 | |
| 475 | auto&& [scratchpad_storage, scratchpad_mem_desc] = get_runtime_memory_ptr(3); |
| 476 | scratchpad_storage->set_data_handle( |
| 477 | cpu_utils::ScratchPadManager::get_scratchpad_manager()->get_data<void>()); |
| 478 | |
| 479 | matmul.execute(default_stream(), memory_cache_); |
| 480 | default_stream().wait(); |
| 481 | } |
| 482 | |
| 483 | dnnl::matmul MatMulPrimitiveHandler::get_matmul_cache( |
| 484 | const MSizeCacheKey& key) { |
| 485 | if (m_size_cache_.get() == nullptr) { |
| 486 | ClassMatmulCacheKey class_key = {.b_n_size = b_n_size_, |
| 487 | .b_k_size = b_k_size_, |
| 488 | .b_type = b_type_, |
| 489 | .c_type = c_type_}; |
| 490 | m_size_cache_ = |
| 491 | get_matul_class_primitive_cache(class_key, primitive_cache_size_); |
| 492 | } |
| 493 | return m_size_cache_->get_or_create(key, [&]() { |
| 494 | dnnl::matmul::primitive_desc desc = this->create_primitive_desc(key, false); |
| 495 | auto manager = cpu_utils::ScratchPadManager::get_scratchpad_manager(); |
| 496 | manager->realloc(desc.scratchpad_desc().get_size()); |
| 497 | return dnnl::matmul(desc); |
| 498 | }); |
| 499 | } |
| 500 | |
| 501 | dnnl::matmul::primitive_desc MatMulPrimitiveHandler::create_primitive_desc( |
| 502 | const MSizeCacheKey& key, bool first_time) { |
| 503 | dnnl::memory::desc a_md; |
| 504 | dnnl::memory::desc b_md; |
| 505 | if (first_time) { |
| 506 | a_md = dnnl::memory::desc({key.a_m_size, b_k_size_}, b_type_, |
| 507 | dnnl::memory::format_tag::ab); |
| 508 | b_md = dnnl::memory::desc({b_k_size_, b_n_size_}, b_type_, |
| 509 | dnnl::memory::format_tag::any); |
| 510 | } else { |
| 511 | a_md = dnnl::memory::desc({key.a_m_size, b_k_size_}, b_type_, |
| 512 | {key.a_m_stride, 1}); |
| 513 | #ifdef VLLM_USE_ACL |
| 514 | // ACL's backend of oneDNN always expects the weight format to be "any" |
| 515 | b_md = dnnl::memory::desc({b_k_size_, b_n_size_}, b_type_, |
| 516 | dnnl::memory::format_tag::any); |
| 517 | #else |
| 518 | b_md = b_target_mem_desc_; |
| 519 | #endif |
| 520 | } |
| 521 | dnnl::memory::desc c_md({key.a_m_size, b_n_size_}, c_type_, |
| 522 | dnnl::memory::format_tag::ab); |
| 523 | |
| 524 | dnnl::primitive_attr attr; |
| 525 | attr.set_scratchpad_mode(dnnl::scratchpad_mode::user); |
| 526 | |
| 527 | if (key.use_bias) { |
| 528 | dnnl::memory::desc bias_md({1, b_n_size_}, key.bias_type, {b_n_size_, 1}); |
| 529 | // Since ACL's matmuls don't support passing a bias_md, we apply the bias |
| 530 | // through a fused-sum post-op |
| 531 | #ifdef VLLM_USE_ACL |
| 532 | dnnl::post_ops post_ops; |
| 533 | post_ops.append_sum(); |
| 534 | attr.set_post_ops(post_ops); |
| 535 | return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, c_md, |
| 536 | attr); |
| 537 | #else |
| 538 | return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, bias_md, |
| 539 | c_md, attr); |
| 540 | #endif |
| 541 | } else { |
| 542 | return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, c_md, |
| 543 | attr); |
| 544 | } |
| 545 | } |
| 546 | |
| 547 | void MatMulPrimitiveHandler::init_runtime_memory_cache(const Args& args) { |
| 548 | memory_cache_[DNNL_ARG_SRC] = dnnl::memory( |
| 549 | {{1, b_k_size_}, b_type_, {b_k_size_, 1}}, default_engine(), nullptr); |
| 550 | set_runtime_memory_ptr(0, memory_cache_[DNNL_ARG_SRC].get()); |
| 551 | memory_cache_[DNNL_ARG_DST] = |
| 552 | dnnl::memory({{1, b_n_size_}, c_type_, dnnl::memory::format_tag::ab}, |
| 553 | default_engine(), nullptr); |
| 554 | set_runtime_memory_ptr(1, memory_cache_[DNNL_ARG_DST].get()); |
| 555 | |
| 556 | // ACL matmuls don't support bias_md, so we don't need these |
| 557 | #ifndef VLLM_USE_ACL |
| 558 | memory_cache_[DNNL_ARG_BIAS] = |
| 559 | dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}}, |
| 560 | default_engine(), nullptr); |
| 561 | set_runtime_memory_ptr(2, memory_cache_[DNNL_ARG_BIAS].get()); |
| 562 | #endif |
| 563 | memory_cache_[DNNL_ARG_SCRATCHPAD] = |
| 564 | dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}}, |
| 565 | default_engine(), nullptr); |
| 566 | set_runtime_memory_ptr(3, memory_cache_[DNNL_ARG_SCRATCHPAD].get()); |
| 567 | } |
| 568 | |
| 569 | bool is_onednn_acl_supported() { |
| 570 | #ifdef VLLM_USE_ACL |
| 571 | return true; |
| 572 | #else |
| 573 | return false; |
| 574 | #endif |
| 575 | } |