| 1324 | |
| 1325 | template <typename TL, typename TR> |
| 1326 | inline void SparseMatMul<TL, TR>::ComputeBlockSizes( |
| 1327 | const typename SparseMatMul<TL, TR>::ConstMatrixMapL& left, |
| 1328 | const typename SparseMatMul<TL, TR>::ConstMatrixMapR& right, |
| 1329 | bool transpose_left, int num_threads, int* KR, int* NR, int* KL, int* JB, |
| 1330 | int* IB) { |
| 1331 | // Heuristics for calculating block sizes |
| 1332 | // Assume two hyperthreads per core. |
| 1333 | const int est_num_cores = std::max(1, (num_threads + 1) / 2); |
| 1334 | // Use block of rhs with at most 128K floats per core. |
| 1335 | const int mem = est_num_cores * 128 * 1024; |
| 1336 | *KR = std::min(static_cast<int>(right.dimension(0)), mem / 256); |
| 1337 | *NR = right.dimension(1); |
| 1338 | if (*KR * *NR > mem) { |
| 1339 | // 4096 may be enough to amortize the cost of writes. |
| 1340 | *KR = std::min<int>(*KR, 4096); |
| 1341 | } |
| 1342 | // Use sizes that are multiples of K and 256. |
| 1343 | *KR = std::max(1, *KR / K) * K; |
| 1344 | *NR = std::max(1, *NR / 256) * 256; |
| 1345 | if (*KR * *NR > mem) { |
| 1346 | *NR = mem / *KR; |
| 1347 | } |
| 1348 | *NR = std::max(1, *NR / 256) * 256; |
| 1349 | |
| 1350 | const int left_dim0 = transpose_left ? left.dimension(1) : left.dimension(0); |
| 1351 | const int left_dim1 = transpose_left ? left.dimension(0) : left.dimension(1); |
| 1352 | for (*KL = 1024; *KL > K; *KL /= 2) { |
| 1353 | if (*KR % *KL == 0 && |
| 1354 | std::max<int>(1, left_dim0 / 64) * (left_dim1 / *KL) > est_num_cores) { |
| 1355 | break; |
| 1356 | } |
| 1357 | } |
| 1358 | DCHECK_EQ(*KL % K, 0); |
| 1359 | DCHECK_GE(*KR, *KL); |
| 1360 | if (*KR < right.dimension(0)) { |
| 1361 | CHECK_EQ(*KR % *KL, 0); |
| 1362 | } |
| 1363 | |
| 1364 | *JB = std::max(1, static_cast<int>(sqrt(num_threads) / 2.0)); |
| 1365 | *IB = 8 * *JB; |
| 1366 | DCHECK_EQ(N * sizeof(float) % 64, size_t{0}); |
| 1367 | } |
| 1368 | |
| 1369 | #ifdef TENSORFLOW_USE_LIBXSMM |
| 1370 | |