Pass capacity and control bytes explicitly to SetCtrl, ResetCtrl, WasNeverFull, and find_first_non_full helpers. The compiler must assume that e.g. writing `ctrl[i] = h` in `SetCtrl*` can alias `CommonFields` and indirect calls (`transfer_n`, `hash_slot`) before `SetCtrl*` can clobber memory. Passing `cap` and `ctrl` explicitly to `SetCtrlNoSanitizeImpl`, `SetCtrl`, `BlockCtrl`, `SetCtrlInSingleGroupTable`, `BlockControlBytes`, `WasNeverFull`, `ResetCtrl`, `find_first_non_full_from_h1`, and `find_first_non_full` avoids redundantly reloading `c.capacity()` and `c.control()` from `CommonFields`. Also annotate `BlockControlBytes` with `ABSL_ASSUME(blocked_element_count <= HashtableInlineData::kMaxBlockedElementCount)` and write the cloned control byte first in `SetCtrlNoSanitizeImpl` for better codegen. PiperOrigin-RevId: 999623416
diff --git a/absl/container/internal/raw_hash_set.cc b/absl/container/internal/raw_hash_set.cc index 70d00cf..4847749 100644 --- a/absl/container/internal/raw_hash_set.cc +++ b/absl/container/internal/raw_hash_set.cc
@@ -178,9 +178,8 @@ } } -FindInfo find_first_non_full_from_h1(const ctrl_t* ctrl, size_t h1, - HashtableCapacity capacity) { - const size_t cap = capacity.capacity(); +FindInfo find_first_non_full_from_h1(size_t cap, const ctrl_t* ctrl, + size_t h1) { auto seq = probe_h1(ProbeCapacity{cap}, h1); if (IsEmptyOrDeleted(ctrl[seq.offset()])) { return {seq.offset(), /*probe_length=*/0}; @@ -196,9 +195,8 @@ // // NOTE: this function must work with tables having both empty and deleted // slots in the same group. Such tables appear during `erase()`. -FindInfo find_first_non_full(const CommonFields& common, size_t hash) { - return find_first_non_full_from_h1(common.control(), H1(hash), - common.capacity_impl()); +FindInfo find_first_non_full(size_t cap, const ctrl_t* ctrl, size_t hash) { + return find_first_non_full_from_h1(cap, ctrl, H1(hash)); } // Same as `find_first_non_full`, but returns the mask corresponding to the @@ -518,56 +516,61 @@ // // Unlike setting it directly, this function will perform bounds checks and // mirror the value to the cloned tail if necessary. -inline void SetCtrlNoSanitizeImpl(const CommonFields& c, size_t i, ctrl_t h) { - ABSL_SWISSTABLE_ASSERT(i < c.capacity()); - ctrl_t* ctrl = c.control(); - const size_t cap = c.capacity(); - ctrl[i] = h; +inline void SetCtrlNoSanitizeImpl(size_t cap, ctrl_t* ctrl, size_t i, + ctrl_t h) { + ABSL_SWISSTABLE_ASSERT(i < cap); + // Note: writing the cloned byte first results in better generated code. ctrl[i < NumClonedBytes() ? i + cap + 1 : i] = h; + ctrl[i] = h; } -inline void SetCtrl(const CommonFields& c, size_t i, ctrl_t h, - size_t slot_size) { +inline void SetCtrl(const CommonFields& c, size_t cap, ctrl_t* ctrl, size_t i, + ctrl_t h, size_t slot_size) { ABSL_SWISSTABLE_ASSERT(!c.is_small()); DoSanitizeOnSetCtrl(c, i, h, slot_size); - SetCtrlNoSanitizeImpl(c, i, h); + SetCtrlNoSanitizeImpl(cap, ctrl, i, h); } -// Overload for setting to an occupied `h2_t` rather than a special `ctrl_t`. -inline void SetCtrl(const CommonFields& c, size_t i, h2_t h, size_t slot_size) { - SetCtrl(c, i, static_cast<ctrl_t>(h), slot_size); + +inline void SetCtrl(const CommonFields& c, size_t cap, ctrl_t* ctrl, size_t i, + h2_t h, size_t slot_size) { + SetCtrl(c, cap, ctrl, i, static_cast<ctrl_t>(h), slot_size); } // Sets `ctrl[i]` to `ctrl_t::kSentinel`. // // Unlike setting it directly, this function will perform bounds checks and // mirror the value to the cloned tail if necessary. -inline void BlockCtrl(const CommonFields& c, size_t i) { - ABSL_SWISSTABLE_ASSERT(!c.is_small()); - SetCtrlNoSanitizeImpl(c, i, ctrl_t::kSentinel); +inline void BlockCtrl(size_t cap, ctrl_t* ctrl, size_t i) { + ABSL_SWISSTABLE_ASSERT(!IsSmallCapacity(cap)); + SetCtrlNoSanitizeImpl(cap, ctrl, i, ctrl_t::kSentinel); } // Like SetCtrl, but in a single group table, we can save some operations when // setting the cloned control byte. -inline void SetCtrlInSingleGroupTable(const CommonFields& c, size_t i, ctrl_t h, +inline void SetCtrlInSingleGroupTable(const CommonFields& c, size_t cap, + ctrl_t* ctrl, size_t i, ctrl_t h, size_t slot_size) { - const size_t cap = c.capacity(); ABSL_SWISSTABLE_ASSERT(!c.is_small()); + ABSL_SWISSTABLE_ASSERT(cap == c.capacity()); + ABSL_SWISSTABLE_ASSERT(ctrl == c.control()); ABSL_SWISSTABLE_ASSERT(is_single_group(cap)); DoSanitizeOnSetCtrl(c, i, h, slot_size); - ctrl_t* ctrl = c.control(); ctrl[i] = h; ctrl[i + cap + 1] = h; } // Overload for setting to an occupied `h2_t` rather than a special `ctrl_t`. -inline void SetCtrlInSingleGroupTable(const CommonFields& c, size_t i, h2_t h, +inline void SetCtrlInSingleGroupTable(const CommonFields& c, size_t cap, + ctrl_t* ctrl, size_t i, h2_t h, size_t slot_size) { - SetCtrlInSingleGroupTable(c, i, static_cast<ctrl_t>(h), slot_size); + SetCtrlInSingleGroupTable(c, cap, ctrl, i, static_cast<ctrl_t>(h), slot_size); } -void BlockControlBytes(CommonFields& common, size_t blocked_element_count) { - const size_t capacity = common.capacity(); +void BlockControlBytes(size_t capacity, ctrl_t* ctrl, + size_t blocked_element_count) { + ABSL_ASSUME(blocked_element_count <= + HashtableInlineData::kMaxBlockedElementCount); while (blocked_element_count > 0) { - BlockCtrl(common, capacity - blocked_element_count); + BlockCtrl(capacity, ctrl, capacity - blocked_element_count); --blocked_element_count; } } @@ -601,7 +604,7 @@ // repeat procedure for current slot with moved from element (target) const size_t blocked_element_count = common.blocked_element_count(); ConvertDeletedToEmptyAndFullToDeleted(ctrl, capacity); - BlockControlBytes(common, blocked_element_count); + BlockControlBytes(capacity, ctrl, blocked_element_count); const void* hash_fn = policy.hash_fn(common); auto hasher = policy.hash_slot; auto transfer_n = policy.transfer_n; @@ -624,14 +627,14 @@ } if (!IsDeleted(ctrl[i])) continue; const size_t hash = (*hasher)(hash_fn, slot_ptr, common.seed().seed()); - const FindInfo target = find_first_non_full(common, hash); + const FindInfo target = find_first_non_full(capacity, ctrl, hash); const size_t new_i = target.offset; total_probe_length += target.probe_length; // Verify if the old and new i fall within the same group wrt the hash. // If they do, we don't need to move the object as it falls already in the // best probe we can. - const size_t probe_offset = probe(common, hash).offset(); + const size_t probe_offset = probe(ProbeCapacity{capacity}, hash).offset(); const h2_t h2 = H2(hash); const auto probe_index = [probe_offset, capacity](size_t pos) { return ((pos - probe_offset) & capacity) / Group::kWidth; @@ -639,7 +642,7 @@ // Element doesn't move. if (ABSL_PREDICT_TRUE(probe_index(new_i) == probe_index(i))) { - SetCtrl(common, i, h2, slot_size); + SetCtrl(common, capacity, ctrl, i, h2, slot_size); continue; } @@ -648,14 +651,14 @@ // Transfer element to the empty spot. // SetCtrl poisons/unpoisons the slots so we have to call it at the // right time. - SetCtrl(common, new_i, h2, slot_size); + SetCtrl(common, capacity, ctrl, new_i, h2, slot_size); (*transfer_n)(set, new_slot_ptr, slot_ptr, 1); - SetCtrl(common, i, ctrl_t::kEmpty, slot_size); + SetCtrl(common, capacity, ctrl, i, ctrl_t::kEmpty, slot_size); // Initialize or change empty space id. tmp_space_id = i; } else { ABSL_SWISSTABLE_ASSERT(IsDeleted(ctrl[new_i])); - SetCtrl(common, new_i, h2, slot_size); + SetCtrl(common, capacity, ctrl, new_i, h2, slot_size); // Until we are done rehashing, DELETED marks previously FULL slots. if (tmp_space_id == kUnknownId) { @@ -679,20 +682,20 @@ // Prepare insert for the new element. PrepareInsertCommon(common); ResetGrowthLeft(capacity, common.size() + blocked_element_count, common); - FindInfo find_info = find_first_non_full(common, new_hash); - SetCtrl(common, find_info.offset, H2(new_hash), slot_size); + FindInfo find_info = find_first_non_full(capacity, ctrl, new_hash); + SetCtrl(common, capacity, ctrl, find_info.offset, H2(new_hash), slot_size); common.infoz().RecordInsertMiss(new_hash, find_info.probe_length); common.infoz().RecordRehash(total_probe_length); return SlotAddress(slot_array, find_info.offset, slot_size); } -bool WasNeverFull(CommonFields& c, size_t index) { - if (is_single_group(c.capacity())) { +bool WasNeverFull(size_t capacity, ctrl_t* ctrl, size_t index) { + if (is_single_group(capacity)) { return true; } - const size_t index_before = (index - Group::kWidth) & c.capacity(); - const auto empty_after = Group(c.control() + index).MaskEmpty(); - const auto empty_before = Group(c.control() + index_before).MaskEmpty(); + const size_t index_before = (index - Group::kWidth) & capacity; + const auto empty_after = Group(ctrl + index).MaskEmpty(); + const auto empty_before = Group(ctrl + index_before).MaskEmpty(); // We count how many consecutive non empties we have to the right and to the // left of `it`. If the sum is >= kWidth then there is at least one probe @@ -707,12 +710,10 @@ // control bytes are kEmpty except for the kSentinel bytes. // If the table has blocked elements, last `blocked_element_count` are set to // kSentinel. -void ResetCtrl(CommonFields& common, size_t slot_size, - size_t blocked_element_count) { - ABSL_SWISSTABLE_ASSERT(IsCapacityValidForBlockedElements(common.capacity()) || +void ResetCtrl(size_t capacity, ctrl_t* ctrl, void* slot_array, + size_t slot_size, size_t blocked_element_count) { + ABSL_SWISSTABLE_ASSERT(IsCapacityValidForBlockedElements(capacity) || blocked_element_count == 0); - const size_t capacity = common.capacity(); - ctrl_t* ctrl = common.control(); static constexpr size_t kTwoGroupCapacity = 2 * Group::kWidth - 1; if (ABSL_PREDICT_TRUE(capacity <= kTwoGroupCapacity)) { if (IsSmallCapacity(capacity)) return; @@ -728,9 +729,9 @@ capacity + 1 + NumClonedBytes()); } ctrl[capacity] = ctrl_t::kSentinel; - SanitizerPoisonMemoryRegion(common.slot_array(capacity), + SanitizerPoisonMemoryRegion(slot_array, slot_size * (capacity - blocked_element_count)); - BlockControlBytes(common, blocked_element_count); + BlockControlBytes(capacity, ctrl, blocked_element_count); } // Initializes control bytes for growing from capacity 1 to 3. @@ -881,14 +882,16 @@ c.decrement_size(); c.infoz().RecordErase(); - if (WasNeverFull(c, index)) { - SetCtrl(c, index, ctrl_t::kEmpty, slot_size); + const size_t cap = c.capacity(); + ctrl_t* ctrl = c.control(); + if (WasNeverFull(cap, ctrl, index)) { + SetCtrl(c, cap, ctrl, index, ctrl_t::kEmpty, slot_size); c.OverwriteFullAsEmpty(); return; } c.OverwriteFullAsDeleted(); - SetCtrl(c, index, ctrl_t::kDeleted, slot_size); + SetCtrl(c, cap, ctrl, index, ctrl_t::kDeleted, slot_size); } void ClearBackingArray(CommonFields& c, @@ -898,9 +901,12 @@ if (reuse) { const size_t blocked_element_count = c.blocked_element_count(); c.set_size_to_zero(); - ABSL_SWISSTABLE_ASSERT(c.capacity() > policy.soo_capacity()); - ResetCtrl(c, policy.slot_size, blocked_element_count); - ResetGrowthLeft(c.capacity(), blocked_element_count, c); + size_t cap = c.capacity(); + ctrl_t* ctrl = c.control(); + void* slot_array = c.slot_array(cap); + ABSL_SWISSTABLE_ASSERT(cap > policy.soo_capacity()); + ResetCtrl(cap, ctrl, slot_array, policy.slot_size, blocked_element_count); + ResetGrowthLeft(cap, blocked_element_count, c); ABSL_SWISSTABLE_ASSERT(c.blocked_element_count() == blocked_element_count); c.infoz().RecordStorageChanged(0, c.capacity()); } else { @@ -1032,8 +1038,10 @@ if (common.is_small()) { target = FindInfo{0, 0}; } else { - target = find_first_non_full(common, hash); - SetCtrl(common, target.offset, H2(hash), slot_size); + const size_t cap = common.capacity(); + ctrl_t* ctrl = common.control(); + target = find_first_non_full(cap, ctrl, hash); + SetCtrl(common, cap, ctrl, target.offset, H2(hash), slot_size); } policy.transfer_n(&common, SlotAddress(new_slots, target.offset, slot_size), slot, 1); @@ -1204,7 +1212,8 @@ common.set_control(new_ctrl); common.generate_new_seed(has_infoz); - ResetCtrl(common, slot_size, blocked_element_count); + ResetCtrl(new_capacity, new_ctrl, new_slots, slot_size, + blocked_element_count); ResetGrowthLeft(new_capacity, blocked_element_count, common); if (ABSL_PREDICT_FALSE(has_infoz)) { @@ -1234,8 +1243,10 @@ SanitizerUnpoisonMemoryRegion(target_slot, policy.slot_size); policy.transfer_n(&c, target_slot, c.soo_data(), 1); c.set_control(new_ctrl); - ResetCtrl(c, policy.slot_size, /*blocked_element_count=*/0); - SetCtrl(c, offset, H2(soo_slot_hash), policy.slot_size); + ResetCtrl(new_capacity, new_ctrl, new_slots, policy.slot_size, + /*blocked_element_count=*/0); + SetCtrl(c, new_capacity, new_ctrl, offset, H2(soo_slot_hash), + policy.slot_size); } enum class ResizeFullSooTableSamplingMode { @@ -1453,9 +1464,9 @@ CommonFields& c, const PolicyFunctions& __restrict policy, const ProbedItem* start, const ProbedItem* end, void* old_slots) { ABSL_SWISSTABLE_ASSERT(!c.is_small()); - const HashtableCapacity new_capacity = c.capacity_impl(); + const size_t new_capacity = c.capacity(); - void* new_slots = c.slot_array(new_capacity.capacity()); + void* new_slots = c.slot_array(new_capacity); ctrl_t* new_ctrl = c.control(); size_t total_probe_length = 0; @@ -1464,18 +1475,19 @@ for (; start < end; ++start) { const FindInfo target = find_first_non_full_from_h1( - new_ctrl, static_cast<size_t>(start->h1), new_capacity); + new_capacity, new_ctrl, static_cast<size_t>(start->h1)); total_probe_length += target.probe_length; const size_t old_index = static_cast<size_t>(start->source_offset); const size_t new_i = target.offset; - ABSL_SWISSTABLE_ASSERT(old_index < new_capacity.capacity() / 2); - ABSL_SWISSTABLE_ASSERT(new_i < new_capacity.capacity()); + ABSL_SWISSTABLE_ASSERT(old_index < new_capacity / 2); + ABSL_SWISSTABLE_ASSERT(new_i < new_capacity); ABSL_SWISSTABLE_ASSERT(IsEmpty(new_ctrl[new_i])); void* src_slot = SlotAddress(old_slots, old_index, slot_size); void* dst_slot = SlotAddress(new_slots, new_i, slot_size); SanitizerUnpoisonMemoryRegion(dst_slot, slot_size); transfer_n(&c, dst_slot, src_slot, 1); - SetCtrl(c, new_i, static_cast<h2_t>(start->h2), slot_size); + SetCtrl(c, new_capacity, new_ctrl, new_i, static_cast<h2_t>(start->h2), + slot_size); } return total_probe_length; } @@ -1491,10 +1503,9 @@ ABSL_ATTRIBUTE_NOINLINE size_t ProcessProbedMarkedElements( CommonFields& c, const PolicyFunctions& __restrict policy, ctrl_t* old_ctrl, void* old_slots, size_t start) { - const size_t new_capacity = c.capacity(); - const size_t old_capacity = PreviousCapacity(new_capacity); + const size_t old_capacity = PreviousCapacity(c.capacity()); const size_t slot_size = policy.slot_size; - void* new_slots = c.slot_array(new_capacity); + void* new_slots = c.slot_array(c.capacity()); size_t total_probe_length = 0; const void* hash_fn = policy.hash_fn(c); auto hash_slot = policy.hash_slot; @@ -1506,11 +1517,13 @@ } void* src_slot = SlotAddress(old_slots, old_index, slot_size); const size_t hash = hash_slot(hash_fn, src_slot, seed); - const FindInfo target = find_first_non_full(c, hash); + const size_t new_capacity = c.capacity(); + ctrl_t* new_ctrl = c.control(); + const FindInfo target = find_first_non_full(new_capacity, new_ctrl, hash); total_probe_length += target.probe_length; const size_t new_i = target.offset; void* dst_slot = SlotAddress(new_slots, new_i, slot_size); - SetCtrl(c, new_i, H2(hash), slot_size); + SetCtrl(c, new_capacity, new_ctrl, new_i, H2(hash), slot_size); transfer_n(&c, dst_slot, src_slot, 1); } return total_probe_length; @@ -1886,7 +1899,8 @@ offset = new_hash & 1 ? 0 : new_capacity - 1; ABSL_SWISSTABLE_ASSERT(IsEmpty(new_ctrl[offset])); - SetCtrlInSingleGroupTable(common, offset, new_h2, policy.slot_size); + SetCtrlInSingleGroupTable(common, new_capacity, new_ctrl, offset, new_h2, + policy.slot_size); find_info = FindInfo{offset, 0}; // Single group tables have all slots full on resize. So we can transfer // all slots without checking the control bytes. @@ -1898,8 +1912,9 @@ } else { total_probe_length = GrowToNextCapacityDispatch(common, policy, old_ctrl, old_slots); - find_info = find_first_non_full(common, new_hash); - SetCtrl(common, find_info.offset, new_h2, policy.slot_size); + find_info = find_first_non_full(new_capacity, new_ctrl, new_hash); + SetCtrl(common, new_capacity, new_ctrl, find_info.offset, new_h2, + policy.slot_size); } ABSL_SWISSTABLE_ASSERT(old_capacity > policy.soo_capacity()); @@ -2050,13 +2065,14 @@ // 1. Table with deleted slots that has space for the inserting element. // 2. Table without deleted slots that has space and GrowthInfoView was // rebalanced. - FindInfo target = find_first_non_full(common, hash); + ctrl_t* ctrl = common.control(); + FindInfo target = find_first_non_full(cap, ctrl, hash); PrepareInsertCommon(common); - common.OverwriteControlAsFull(common.control()[target.offset]); + common.OverwriteControlAsFull(ctrl[target.offset]); // Compute before `SetCtrl*` to avoid reloading `control_`. void* res = SlotAddress(common.slot_array(cap), target.offset, policy.slot_size); - SetCtrl(common, target.offset, H2(hash), policy.slot_size); + SetCtrl(common, cap, ctrl, target.offset, H2(hash), policy.slot_size); if (common.has_infoz()) { ReportInsertMissToInfoz(common, target.probe_length, hash); } @@ -2083,8 +2099,8 @@ // Compute before `SetCtrl*` to avoid reloading `control_`. void* res = SlotAddress(common.slot_array(kNewCapacity), SooSlotIndex(), policy.slot_size); - SetCtrlInSingleGroupTable(common, SooSlotIndex(), H2(new_hash), - policy.slot_size); + SetCtrlInSingleGroupTable(common, kNewCapacity, common.control(), + SooSlotIndex(), H2(new_hash), policy.slot_size); if (common.has_infoz()) { ReportInsertMissToInfoz(common, /*probe_length=*/0, new_hash); } @@ -2201,7 +2217,8 @@ common.generate_new_seed(has_infoz); size_t total_probe_length = 0; - ResetCtrl(common, slot_size, /*blocked_element_count=*/0); + ResetCtrl(new_capacity, new_ctrl, new_slots, slot_size, + /*blocked_element_count=*/0); ABSL_SWISSTABLE_ASSERT(old_capacity > 0); total_probe_length = FindNewPositionsAndTransferSlots( common, policy, old_ctrl, old_slots, old_capacity); @@ -2404,10 +2421,11 @@ // The table is guaranteed to be empty, so we can do faster than // a full `insert`. const size_t hash = (*hasher)(hash_fn, that_slot, seed); - FindInfo target = find_first_non_full(common, hash); - infoz.RecordInsertMiss(hash, target.probe_length); + ctrl_t* ctrl = common.control(); + FindInfo target = find_first_non_full(cap, ctrl, hash); offset = target.offset; - SetCtrl(common, offset, H2(hash), slot_size); + SetCtrl(common, cap, ctrl, offset, H2(hash), slot_size); + infoz.RecordInsertMiss(hash, target.probe_length); copy_fn(&common, SlotAddress(target_slot_array, offset, slot_size), that_slot); common.maybe_increment_generation_on_insert(); @@ -2463,7 +2481,8 @@ // Compute before `SetCtrl` to avoid reloading `control_`. void* res = SlotAddress(common.slot_array(cap), target_group_offset, policy.slot_size); - SetCtrl(common, target_group_offset, H2(hash), policy.slot_size); + SetCtrl(common, cap, common.control(), target_group_offset, H2(hash), + policy.slot_size); // Call NOINLINE function to move infoz instructions out of line. if (common.has_infoz()) { return ReportInsertMissToInfozAndGetSlot(common, policy, hash,