MCPcopy Create free account
hub / github.com/Oneflow-Inc/oneflow / PutBlocks

Method PutBlocks

oneflow/core/embedding/persistent_table.cpp:475–529  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

473
474template<typename Key, typename Engine>
475void PersistentTableImpl<Key, Engine>::PutBlocks(uint32_t num_keys, const void* keys,
476 const void* blocks) {
477 CHECK(!read_only_);
478 std::lock_guard<std::recursive_mutex> lock(mutex_);
479 const uint32_t num_blocks = RoundUp(num_keys, num_values_per_block_) / num_values_per_block_;
480 const uint32_t num_padded_keys = num_blocks * num_values_per_block_;
481 const uint64_t start_index = physical_table_size_;
482 physical_table_size_ += num_padded_keys;
483 CHECK_EQ(start_index % num_values_per_block_, 0);
484 const uint64_t start_block_id = start_index / num_values_per_block_;
485 uint64_t written_blocks = 0;
486 const uint64_t block_keys_size = num_values_per_block_ * sizeof(Key);
487 BlockingCounter bc(1);
488 workers_.at(0)->Schedule([&](Engine*) {
489 while (written_blocks < num_blocks) {
490 const uint64_t batch_start_block_id = start_block_id + written_blocks;
491 const uint64_t batch_chunk_id = batch_start_block_id / num_logical_blocks_per_chunk_;
492 if (batch_chunk_id == value_files_.size()) {
493 value_files_.emplace_back(ValueFilePath(batch_chunk_id), O_CREAT | O_RDWR | O_DIRECT, 0644);
494 } else {
495 CHECK_LE(batch_chunk_id, value_files_.size());
496 }
497 if ((!writable_key_file_.IsOpen()) || writable_key_file_chunk_id_ != batch_chunk_id) {
498 writable_key_file_ = PosixFile(KeyFilePath(batch_chunk_id), O_CREAT | O_RDWR, 0644);
499 }
500 PosixFile& value_file = value_files_.at(batch_chunk_id);
501 const uint64_t block_id_in_chunk =
502 batch_start_block_id - batch_chunk_id * num_logical_blocks_per_chunk_;
503 const uint64_t blocks_to_write =
504 std::min(num_blocks - written_blocks,
505 (batch_chunk_id + 1) * num_logical_blocks_per_chunk_ - batch_start_block_id);
506 const uint64_t values_bytes = blocks_to_write * logical_block_size_;
507 const uint64_t values_offset_in_file = block_id_in_chunk * logical_block_size_;
508 CHECK_LE(value_file.Size(), values_offset_in_file);
509 value_file.Truncate(values_offset_in_file + values_bytes);
510 PCHECK(pwrite(value_file.fd(), BytesOffset(blocks, written_blocks * logical_block_size_),
511 values_bytes, values_offset_in_file)
512 == values_bytes);
513 const uint64_t keys_offset_in_file = block_id_in_chunk * block_keys_size;
514 writable_key_file_.Truncate(keys_offset_in_file + blocks_to_write * block_keys_size);
515 const uint64_t keys_bytes = std::min(num_keys - written_blocks * num_values_per_block_,
516 blocks_to_write * num_values_per_block_)
517 * sizeof(Key);
518 PCHECK(pwrite(writable_key_file_.fd(), BytesOffset(keys, written_blocks * block_keys_size),
519 keys_bytes, keys_offset_in_file)
520 == keys_bytes);
521 written_blocks += blocks_to_write;
522 }
523 bc.Decrease();
524 });
525 for (uint64_t i = 0; i < num_keys; ++i) {
526 row_id_mapping_[static_cast<const Key*>(keys)[i]] = start_index + i;
527 }
528 bc.WaitForeverUntilCntEqualZero();
529}
530
531template<typename Key, typename Engine>
532void PersistentTableImpl<Key, Engine>::Put(uint32_t num_keys, const void* keys,

Callers

nothing calls this directly

Calls 13

RoundUpFunction · 0.85
PosixFileClass · 0.85
BytesOffsetFunction · 0.85
IsOpenMethod · 0.80
TruncateMethod · 0.80
fdMethod · 0.80
ScheduleMethod · 0.45
atMethod · 0.45
sizeMethod · 0.45
emplace_backMethod · 0.45
SizeMethod · 0.45

Tested by

no test coverage detected