* @brief Load a partition of the sample pool. Update the cache automatically. * @param _head_partition_id id of head partition * @param _tail_partition_id id of tail partition */
| 1433 | * @param _tail_partition_id id of tail partition |
| 1434 | */ |
| 1435 | void load_partition(int _head_partition_id, int _tail_partition_id) { |
| 1436 | // check hit/miss |
| 1437 | bool cold_cache = head_partition_id == -1 || tail_partition_id == -1; |
| 1438 | std::vector<bool> hit(num_embedding, false); |
| 1439 | if (!cold_cache) { |
| 1440 | for (int i = 0; i < num_embedding; i++) { |
| 1441 | Protocol protocol = protocols[i]; |
| 1442 | if ((protocol & kSharedWithPredecessor) && !(protocol & kGlobal) && !hit[i - 1]) { |
| 1443 | // check swap hit |
| 1444 | if (head_partition_id == _tail_partition_id && tail_partition_id == _head_partition_id) { |
| 1445 | embeddings[i].swap(embeddings[i - 1]); |
| 1446 | moments[i].swap(moments[i - 1]); |
| 1447 | hit[i] = true; |
| 1448 | hit[i - 1] = true; |
| 1449 | } |
| 1450 | } |
| 1451 | // if the current embedding is shared with predecessor |
| 1452 | // then it is necessary to have hit on the predecessor |
| 1453 | if (!(i > 0 && embeddings[i] == embeddings[i - 1] && !hit[i - 1])){ |
| 1454 | hit[i] = hit[i] || ((protocol & kHeadPartition) && head_partition_id == _head_partition_id); |
| 1455 | hit[i] = hit[i] || ((protocol & kTailPartition) && tail_partition_id == _tail_partition_id); |
| 1456 | } |
| 1457 | } |
| 1458 | // we don't need to write back during prediction |
| 1459 | if (solver->is_train) |
| 1460 | for (int i = 0; i < num_embedding; i++) |
| 1461 | if (!hit[i]) |
| 1462 | write_embedding(i); |
| 1463 | } |
| 1464 | bool sampler_hit = (sampler_protocol & kGlobal) && !cold_cache; |
| 1465 | sampler_hit = sampler_hit || ((sampler_protocol & kHeadPartition) && (sampler_protocol & kTailPartition) |
| 1466 | && head_partition_id == _tail_partition_id |
| 1467 | && _head_partition_id == tail_partition_id); |
| 1468 | sampler_hit = sampler_hit || ((sampler_protocol & (kHeadPartition | kTailPartition)) == kHeadPartition |
| 1469 | && head_partition_id == _head_partition_id); |
| 1470 | sampler_hit = sampler_hit || ((sampler_protocol & (kHeadPartition | kTailPartition)) == kTailPartition |
| 1471 | && tail_partition_id == _tail_partition_id); |
| 1472 | |
| 1473 | // load partition mappings |
| 1474 | if (head_partition_id != _head_partition_id) { |
| 1475 | head_partition_id = _head_partition_id; |
| 1476 | if (!solver->naive_parallel) { |
| 1477 | head_global_ids = solver->head_partitions[head_partition_id]; |
| 1478 | head_partition_size = head_global_ids.size(); |
| 1479 | } |
| 1480 | } |
| 1481 | if (tail_partition_id != _tail_partition_id) { |
| 1482 | tail_partition_id = _tail_partition_id; |
| 1483 | if (!solver->naive_parallel) { |
| 1484 | tail_global_ids = solver->tail_partitions[tail_partition_id]; |
| 1485 | tail_partition_size = tail_global_ids.size(); |
| 1486 | } |
| 1487 | } |
| 1488 | if (!sampler_hit) { |
| 1489 | build_negative_sampler(); |
| 1490 | negative_sampler.to_device_async(); |
| 1491 | } |
| 1492 | for (int i = 0; i < num_embedding; i++) |
nothing calls this directly
no test coverage detected