| 1433 | |
| 1434 | template <typename TL, typename TR> |
| 1435 | inline void LibxsmmSparseMatMul<TL, TR>::Compute( |
| 1436 | typename LibxsmmSparseMatMul<TL, TR>::TensorInfoCache* cache, |
| 1437 | const typename LibxsmmSparseMatMul<TL, TR>::ConstMatrixMapL& left, |
| 1438 | const typename LibxsmmSparseMatMul<TL, TR>::ConstMatrixMapR& right, |
| 1439 | bool transpose_left, const DeviceBase::CpuWorkerThreads* thread_pool, |
| 1440 | bool transpose_output, MatrixMap* output) { |
| 1441 | if (false) { |
| 1442 | // Not handled by libxsmm currently |
| 1443 | SparseMatMul<TL, TR>::Compute( |
| 1444 | nullptr /* Assumes no cached data for fallback */, left, right, |
| 1445 | transpose_left, thread_pool, transpose_output, output); |
| 1446 | return; |
| 1447 | } |
| 1448 | const int num_threads = thread_pool->num_threads; |
| 1449 | const int left_dim0 = transpose_left ? left.dimension(1) : left.dimension(0); |
| 1450 | const int left_dim1 = transpose_left ? left.dimension(0) : left.dimension(1); |
| 1451 | const int right_dim0 = right.dimension(0); |
| 1452 | const int right_dim1 = right.dimension(1); |
| 1453 | CHECK_EQ(left_dim1, right_dim0); |
| 1454 | CHECK_EQ(left_dim0, |
| 1455 | (transpose_output ? output->dimension(1) : output->dimension(0))); |
| 1456 | CHECK_EQ(right_dim1, |
| 1457 | (transpose_output ? output->dimension(0) : output->dimension(1))); |
| 1458 | if (left_dim0 < 32 || left_dim1 < 32 || right_dim1 < 32) { |
| 1459 | // Causes problems in libxsmm |
| 1460 | SparseMatMul<TL, TR>::Compute( |
| 1461 | nullptr /* Assumes no cached data for fallback */, left, right, |
| 1462 | transpose_left, thread_pool, transpose_output, output); |
| 1463 | return; |
| 1464 | } |
| 1465 | auto left_data = left.data(); |
| 1466 | auto right_data = right.data(); |
| 1467 | auto output_data = output->data(); |
| 1468 | // Initialize libxsmm for this matrix; make sure another thread doesn't use |
| 1469 | // this handle |
| 1470 | auto entry = |
| 1471 | cache->take_cache_entry(left_dim0, right_dim0, right_dim1, num_threads); |
| 1472 | // Convert the left matrix to compressed sparse row (CSR) format |
| 1473 | ptrdiff_t total_num_creation_blocks = |
| 1474 | libxsmm_spmdm_get_num_createSparseSlice_blocks(&entry->handle); |
| 1475 | std::atomic<int> cur_create_block_number; |
| 1476 | cur_create_block_number.store(0); |
| 1477 | do_on_all_threads(thread_pool, [&](int i) { |
| 1478 | while (true) { |
| 1479 | int work_item = cur_create_block_number.fetch_add(1); |
| 1480 | if (work_item >= total_num_creation_blocks) break; |
| 1481 | wrapper_libxsmm_spmdm_createSparseSlice_generic_thread( |
| 1482 | empty_type_wrapper<TL>{}, &entry->handle, |
| 1483 | (transpose_left ? 'T' : 'N'), left_data, entry->output_csr, work_item, |
| 1484 | i, num_threads); |
| 1485 | } |
| 1486 | }); |
| 1487 | // Do matrix-matrix multiplication |
| 1488 | ptrdiff_t total_num_mult_blocks = |
| 1489 | libxsmm_spmdm_get_num_compute_blocks(&entry->handle); |
| 1490 | std::atomic<int> cur_mult_block_number; |
| 1491 | cur_mult_block_number.store(0); |
| 1492 | do_on_all_threads(thread_pool, [&](int i) { |
nothing calls this directly
no test coverage detected