Deduplicate backing array allocation in PrepareInsertSmallNonSoo to reduce binary size. In `PrepareInsertSmallNonSoo`, growing from capacity 0 to 1 and growing from capacity 1 to 3 via `Grow1To3AndPrepareInsert` previously maintained two separate inlined copies of `policy.get_char_alloc`, `RawHashSetLayout`, and `AllocBackingArray`. Naively sharing the general `RawHashSetLayout` constructor between both paths introduces dynamic branching for `MetadataBeforeControlSize` and `NumControlBytes`, as well as extra `PHI` nodes across `policy.get_char_alloc` that spill registers to the stack. Introduce `RawHashSetLayout::CreateForSmallNonSooGrowth` and an `AllocBackingArray` overload accepting `const RawHashSetLayout&` so both small non-SOO growth paths share a single allocation sequence in `PrepareInsertSmallNonSoo`. Parameterizing the layout by `size = common.size()` (`0` for 0 → 1, `1` for 1 → 3) computes `num_control_bytes = size * NumControlBytes(3)` and `new_capacity = NextCapacity(size)` directly from `size`, and computes `metadata_size` branchlessly without `GrowthInfoSizeForCapacity` (which is 0 for capacities 1 and 3). Because `Grow1To3AndPrepareInsert`, `AllocBackingArray`, and `RawHashSetLayout::CreateForSmallNonSooGrowth` are all inlined into `PrepareInsertSmallNonSoo`, this reduces the total compiled size of `PrepareInsertSmallNonSoo` (including all functions inlined into it) from 781 B to 716 B (-65 B / -8.3%) and shrinks its stack frame by 16 bytes. PiperOrigin-RevId: 1001834810
diff --git a/absl/container/internal/raw_hash_set.cc b/absl/container/internal/raw_hash_set.cc index 374ab9b..0837bf9 100644 --- a/absl/container/internal/raw_hash_set.cc +++ b/absl/container/internal/raw_hash_set.cc
@@ -1170,10 +1170,8 @@ BackingArrayPtrs AllocBackingArray(CommonFields& common, const PolicyFunctions& __restrict policy, - size_t new_capacity, bool has_infoz, - void* alloc, size_t blocked_element_count) { - RawHashSetLayout layout(new_capacity, policy.slot_size, policy.slot_align, - has_infoz, blocked_element_count); + const RawHashSetLayout& layout, + void* alloc) { // Perform a direct call in the common case to allow for profile-guided // heap optimization (PGHO) to understand which allocation function is used. constexpr size_t kDefaultAlignment = BackingArrayAlignment(alignof(size_t)); @@ -1193,6 +1191,17 @@ mem + layout.slot_offset()}; } +BackingArrayPtrs AllocBackingArray(CommonFields& common, + const PolicyFunctions& __restrict policy, + size_t new_capacity, bool has_infoz, + void* alloc, size_t blocked_element_count) { + return AllocBackingArray( + common, policy, + RawHashSetLayout(new_capacity, policy.slot_size, policy.slot_align, + has_infoz, blocked_element_count), + alloc); +} + void ResizeEmptyNonAllocatedTableImpl(CommonFields& common, const PolicyFunctions& __restrict policy, size_t new_capacity, @@ -1706,10 +1715,10 @@ ABSL_SWISSTABLE_ASSERT(new_capacity <= ProbedItem::kMaxNewCapacity); Encoder encoder(old_ctrl); policy.transfer_unprobed_elements_to_next_capacity( - common, old_ctrl, old_slots, &encoder, - [](void* probed_storage, h2_t h2, size_t source_offset, size_t h1) { + common, old_ctrl, &encoder, old_slots, PreviousCapacity(new_capacity), + [](void* probed_storage, size_t hash, size_t source_offset) { auto encoder_ptr = static_cast<Encoder*>(probed_storage); - encoder_ptr->EncodeItem(ProbedItem(h2, source_offset, h1)); + encoder_ptr->EncodeItem(ProbedItem(H2(hash), source_offset, H1(hash))); }); InitializeMirroredControlBytes(common.control(), new_capacity); return encoder.DecodeAndInsertToTable(common, policy, old_slots); @@ -1807,7 +1816,8 @@ void* Grow1To3AndPrepareInsert(CommonFields& common, const PolicyFunctions& __restrict policy, - void* hasher, HashForKeyFn get_hash) { + void* hasher, HashForKeyFn get_hash, void* alloc, + bool has_infoz, void* new_slots) { // TODO(b/413062340): Refactor to reuse more code with // GrowSooTableToNextCapacityAndPrepareInsert. ABSL_SWISSTABLE_ASSERT(common.capacity() == 1); @@ -1817,21 +1827,8 @@ ABSL_SWISSTABLE_ASSERT(common.blocked_element_count() == 0); constexpr size_t kOldCapacity = 1; constexpr size_t kNewCapacity = NextCapacity(kOldCapacity); - void* old_slots = common.single_non_soo_slot(); - // old_slots == old_ctrl in case of capacity == 1. - ctrl_t* old_ctrl = static_cast<ctrl_t*>(old_slots); - - const size_t slot_size = policy.slot_size; - const size_t slot_align = policy.slot_align; - void* alloc = policy.get_char_alloc(common); - const bool has_infoz = common.has_infoz(); - common.set_capacity(kNewCapacity); - - const auto [new_ctrl, new_slots] = - AllocBackingArray(common, policy, kNewCapacity, has_infoz, alloc, - /*blocked_element_count=*/0); - common.set_control(new_ctrl); - SanitizerPoisonMemoryRegion(new_slots, kNewCapacity * slot_size); + ctrl_t* new_ctrl = + static_cast<ctrl_t*>(new_slots) - NumControlBytes(kNewCapacity); if (ABSL_PREDICT_TRUE(!has_infoz)) { // When we're sampled, we already have a seed. @@ -1839,29 +1836,45 @@ } const size_t new_hash = get_hash(hasher, common.seed().seed()); h2_t new_h2 = H2(new_hash); - size_t orig_hash = - policy.hash_slot(policy.hash_fn(common), old_slots, common.seed().seed()); + // `hash_fn` is in separate statement, so the hash_slot function pointer is + // not loaded before the hash_fn call. + const void* hash_fn = policy.hash_fn(common); + size_t orig_hash = policy.hash_slot( + hash_fn, common.single_non_soo_slot(), common.seed().seed()); size_t offset = Resize1To3NewOffset(new_hash); InitializeThreeElementsControlBytes(H2(orig_hash), new_h2, offset, new_ctrl); + // Read `old_slots` and `slot_size` after the hash calls to avoid keeping + // them live across indirect calls. + void* old_slots = common.single_non_soo_slot(); + // old_slots == old_ctrl in case of capacity == 1. + ctrl_t* old_ctrl = static_cast<ctrl_t*>(old_slots); + common.set_control(new_ctrl); + const size_t slot_size = policy.slot_size; + SanitizerPoisonMemoryRegion(new_slots, kNewCapacity * slot_size); + void* old_element_target = NextSlot(new_slots, slot_size); + // Compute `new_element_target_slot` before `transfer_n` to free up registers + // of `new_slots` and `offset`. + void* new_element_target_slot = SlotAddress(new_slots, offset, slot_size); SanitizerUnpoisonMemoryRegion(old_element_target, slot_size); policy.transfer_n(&common, old_element_target, old_slots, 1); - - void* new_element_target_slot = SlotAddress(new_slots, offset, slot_size); SanitizerUnpoisonMemoryRegion(new_element_target_slot, slot_size); + // Group `common` updates together so `inline_data_` is modified in a single + // load-store sequence. + common.set_capacity(kNewCapacity); + common.InitGrowthLeftNoDeleted(kNewCapacity - 2, kNewCapacity); PrepareInsertCommon(common); ABSL_SWISSTABLE_ASSERT(common.size() == 2); - common.InitGrowthLeftNoDeleted(kNewCapacity - 2, kNewCapacity); if (ABSL_PREDICT_FALSE(has_infoz)) { ReportSingleGroupTableGrowthToInfoz(common, old_ctrl, new_hash); } // Deallocate after `ReportSingleGroupTableGrowthToInfoz`, because it reads // infoz from old_ctrl. - DeallocBackingArrayImpl(alloc, kOldCapacity, old_ctrl, slot_size, slot_align, - has_infoz, + DeallocBackingArrayImpl(alloc, kOldCapacity, old_ctrl, slot_size, + policy.slot_align, has_infoz, /*blocked_element_count=*/0, policy.dealloc); return new_element_target_slot; } @@ -1951,49 +1964,86 @@ } // namespace +RawHashSetLayout RawHashSetLayout::CreateForSmallNonSooGrowth(size_t size, + size_t slot_size, + size_t slot_align, + bool has_infoz) { + ABSL_SWISSTABLE_ASSERT(size <= 1); + static_assert(NumControlBytes(1) == 0); + const size_t num_control_bytes = size * NumControlBytes(3); + static_assert(NextCapacity(0) == 1); + static_assert(NextCapacity(1) == 3); + const size_t new_capacity = NextCapacity(size); + const size_t slots_size = new_capacity * slot_size; + + // Since `GrowthInfoSizeForCapacity` is 0 for capacities 1 and 3, compute + // `metadata_size` directly without `MetadataBeforeControlSize` to avoid its + // `ABSL_PREDICT_FALSE(has_infoz)` branch and allow a branchless shift. + static_assert(GrowthInfoSizeForCapacity(1) == 0); + static_assert(GrowthInfoSizeForCapacity(3) == 0); + const size_t metadata_size = + (has_infoz ? sizeof(HashtablezInfoHandle) + sizeof(uint64_t) : 0) + + NumGenerationBytes(); + ABSL_SWISSTABLE_ASSERT(metadata_size == + MetadataBeforeControlSize(has_infoz, new_capacity)); + + RawHashSetLayout layout(uninitialized_tag_t{}); + layout.slot_offset_ = + AlignUpTo(metadata_size + num_control_bytes, slot_align); + layout.control_offset_ = layout.slot_offset_ - num_control_bytes; + layout.generation_offset_ = layout.control_offset_ - NumGenerationBytes(); + layout.alloc_size_ = layout.slot_offset_ + slots_size; + return layout; +} + void* PrepareInsertSmallNonSoo(CommonFields& common, const PolicyFunctions& __restrict policy, void* hasher, HashForKeyFn get_hash) { ABSL_SWISSTABLE_ASSERT(common.is_small()); ABSL_SWISSTABLE_ASSERT(!policy.soo_enabled); - if (common.capacity() == 1) { - if (common.empty()) { - IncrementSmallSizeNonSoo(common, policy); - // Compute before `RecordInsertMissCold` to avoid reloading `control_`. - void* res = common.single_non_soo_slot(); - // Call NOINLINE function to move infoz instructions out of line. - if (common.has_infoz()) { - ReportInsertMissToInfozAndComputeHash(common, /*probe_length=*/0, - hasher, get_hash); - } - return res; - } else { - return Grow1To3AndPrepareInsert(common, policy, hasher, get_hash); + const bool is_empty = common.empty(); + if (common.capacity() == 1 && is_empty) { + IncrementSmallSizeNonSoo(common, policy); + // Compute before `RecordInsertMissCold` to avoid reloading `control_`. + void* res = common.single_non_soo_slot(); + // Call NOINLINE function to move infoz instructions out of line. + if (common.has_infoz()) { + ReportInsertMissToInfozAndComputeHash(common, /*probe_length=*/0, hasher, + get_hash); } + return res; } - // Growing from 0 to 1 capacity. - ABSL_SWISSTABLE_ASSERT(common.capacity() == 0); - constexpr size_t kNewCapacity = 1; - - common.set_capacity(kNewCapacity); + const size_t size = common.size(); const bool has_infoz = - policy.is_hashtablez_enabled && ShouldSampleNextTable(); + is_empty ? (policy.is_hashtablez_enabled && ShouldSampleNextTable()) + : common.has_infoz(); void* alloc = policy.get_char_alloc(common); + void* new_slots = + AllocBackingArray( + common, policy, + RawHashSetLayout::CreateForSmallNonSooGrowth( + size, policy.slot_size, policy.slot_align, has_infoz), + alloc) + .slots; + if (is_empty) { + // Growing from 0 to 1 capacity. + ABSL_SWISSTABLE_ASSERT(common.capacity() == 0); + constexpr size_t kNewCapacity = 1; + common.set_capacity(kNewCapacity); + common.set_control(static_cast<ctrl_t*>(new_slots)); - const auto [new_ctrl, new_slots] = - AllocBackingArray(common, policy, kNewCapacity, has_infoz, alloc, - /*blocked_element_count=*/0); - common.set_control(new_ctrl); + static_assert(NextCapacity(0) == 1); + PrepareInsertCommon(common); - static_assert(NextCapacity(0) == 1); - PrepareInsertCommon(common); - - if (ABSL_PREDICT_FALSE(has_infoz)) { - ReportSingleElementTableGrowthToInfozAndForceSample(common, policy, hasher, - get_hash); + if (ABSL_PREDICT_FALSE(has_infoz)) { + ReportSingleElementTableGrowthToInfozAndForceSample( + common, policy, hasher, get_hash); + } + return new_slots; } - return new_slots; + return Grow1To3AndPrepareInsert(common, policy, hasher, get_hash, alloc, + has_infoz, new_slots); } namespace { @@ -2264,10 +2314,16 @@ } ABSL_SWISSTABLE_ASSERT(common.size() == policy.soo_capacity()); static constexpr size_t kNewCapacity = NextCapacity(SooCapacity()); - const size_t slot_size = policy.slot_size; - void* alloc = policy.get_char_alloc(common); - common.set_capacity(kNewCapacity); + // Generate the seed and compute hashes before `AllocBackingArray` so + // `hasher`, `get_hash`, and `new_ctrl`/`new_slots` do not need to be spilled + // across calls. + common.generate_new_seed(/*has_infoz=*/false); + const size_t new_hash = get_hash(hasher, common.seed().seed()); + const void* hash_fn = policy.hash_fn(common); + const h2_t soo_slot_h2 = + H2(policy.hash_slot(hash_fn, common.soo_data(), common.seed().seed())); + void* alloc = policy.get_char_alloc(common); // Since the table is not empty, it will not be sampled. // The decision to sample was already made during the first insertion. // @@ -2276,21 +2332,21 @@ const auto [new_ctrl, new_slots] = AllocBackingArray( common, policy, kNewCapacity, /*has_infoz=*/false, alloc, /*blocked_element_count=*/0); - + // Group `common` updates together after `AllocBackingArray` so `inline_data_` + // is modified in a single load-store sequence. + common.set_capacity(kNewCapacity); + common.InitGrowthLeftNoDeleted(kNewCapacity - 2, kNewCapacity); PrepareInsertCommon(common); ABSL_SWISSTABLE_ASSERT(common.size() == 2); - common.InitGrowthLeftNoDeleted(kNewCapacity - 2, kNewCapacity); - common.generate_new_seed(/*has_infoz=*/false); - const h2_t soo_slot_h2 = H2(policy.hash_slot( - policy.hash_fn(common), common.soo_data(), common.seed().seed())); - const size_t new_hash = get_hash(hasher, common.seed().seed()); const size_t offset = Resize1To3NewOffset(new_hash); InitializeThreeElementsControlBytes(soo_slot_h2, H2(new_hash), offset, new_ctrl); + const size_t slot_size = policy.slot_size; SanitizerPoisonMemoryRegion(new_slots, slot_size * kNewCapacity); void* target_slot = SlotAddress(new_slots, SooSlotIndex(), slot_size); + void* new_slot = SlotAddress(new_slots, offset, slot_size); SanitizerUnpoisonMemoryRegion(target_slot, slot_size); if constexpr (TransferUsesMemcpy) { // Target slot is placed at index 1, but capacity is at @@ -2314,7 +2370,6 @@ // Full SOO table couldn't be sampled. If SOO table is sampled, it would // have been resized to the next capacity. ABSL_SWISSTABLE_ASSERT(!common.infoz().IsSampled()); - void* new_slot = SlotAddress(new_slots, offset, slot_size); SanitizerUnpoisonMemoryRegion(new_slot, slot_size); return new_slot; } @@ -2634,7 +2689,7 @@ const DtorPolicy& policy); // Extern template instantiation for inline functions. -template size_t TryFindNewIndexWithoutProbing(size_t h1, size_t old_index, +template size_t TryFindNewIndexWithoutProbing(size_t hash, size_t old_index, size_t old_capacity, ctrl_t* new_ctrl, size_t new_capacity);
diff --git a/absl/container/internal/raw_hash_set.h b/absl/container/internal/raw_hash_set.h index ad79652..2d43ec4 100644 --- a/absl/container/internal/raw_hash_set.h +++ b/absl/container/internal/raw_hash_set.h
@@ -1273,6 +1273,13 @@ alloc_size_ = slot_offset_ + (capacity - blocked_element_count) * slot_size; } + // Creates a layout for a small non-SOO table growing from capacity + // 0 -> 1 (`size == 0`) or 1 -> 3 (`size == 1`). + static RawHashSetLayout CreateForSmallNonSooGrowth(size_t size, + size_t slot_size, + size_t slot_align, + bool has_infoz); + // Returns precomputed offset from the start of the backing allocation of // control. size_t control_offset() const { return control_offset_; } @@ -1290,6 +1297,8 @@ size_t alloc_size() const { return alloc_size_; } private: + explicit RawHashSetLayout(uninitialized_tag_t) {} + size_t control_offset_; size_t generation_offset_; size_t slot_offset_; @@ -1952,8 +1961,8 @@ constexpr size_t kProbedElementIndexSentinel = ~size_t{}; // Implementation detail of transfer_unprobed_elements_to_next_capacity_fn. -// Tries to find the new index for an element whose hash corresponds to -// `h1` for growth to the next capacity. +// Tries to find the new index for an element with `hash` for growth to the next +// capacity. // Returns kProbedElementIndexSentinel if full probing is required. // // If element is located in the first probing group in the table before growth, @@ -1963,23 +1972,18 @@ // table. We only attempt to do so if the first probe group is already // initialized. template <typename = void> -inline size_t TryFindNewIndexWithoutProbing(size_t h1, size_t old_index, +inline size_t TryFindNewIndexWithoutProbing(size_t hash, size_t old_index, size_t old_capacity, ctrl_t* new_ctrl, size_t new_capacity) { - size_t index_diff = old_index - h1; + size_t h1 = H1(hash); // The first probe group starts with h1 & capacity. // All following groups start at (h1 + Group::kWidth * K) & capacity. - // We can find an index within the floating group as index_diff modulo - // Group::kWidth. - // Both old and new capacity are larger than Group::kWidth so we can avoid - // computing `& capacity`. - size_t in_floating_group_index = index_diff & (Group::kWidth - 1); - // By subtracting we will get the difference between the first probe group - // and the probe group corresponding to old_index. - index_diff -= in_floating_group_index; - if (ABSL_PREDICT_TRUE((index_diff & old_capacity) == 0)) { - size_t new_index = (h1 + in_floating_group_index) & new_capacity; + // If old_index is in the first probe group, (old_index - h1) & old_capacity + // is the index within the floating group (< Group::kWidth). + size_t index_diff = (old_index - h1) & old_capacity; + if (ABSL_PREDICT_TRUE(index_diff < Group::kWidth)) { + size_t new_index = (h1 + index_diff) & new_capacity; ABSL_ASSUME(new_index != kProbedElementIndexSentinel); return new_index; } @@ -2090,11 +2094,11 @@ // Different encoding is used depending on the capacity of the table. // See ProbedItem*Bytes classes for details. void (*transfer_unprobed_elements_to_next_capacity)( - CommonFields& common, const ctrl_t* old_ctrl, void* old_slots, + CommonFields& common, const ctrl_t* old_ctrl, // TODO(b/382423690): Try to use absl::FunctionRef here. - void* probed_storage, - void (*encode_probed_element)(void* probed_storage, h2_t h2, - size_t source_offset, size_t h1)); + void* probed_storage, void* old_slots, size_t old_capacity, + void (*encode_probed_element)(void* probed_storage, size_t hash, + size_t source_offset)); uint8_t soo_capacity() const { return static_cast<uint8_t>(soo_enabled ? SooCapacity() : 0); @@ -4249,15 +4253,15 @@ // by GetKey + Hash for memcpyable types. // TODO(b/382423690): Try to type erase for big slots: sizeof(slot_type) > 16. static void transfer_unprobed_elements_to_next_capacity_fn( - CommonFields& common, const ctrl_t* old_ctrl, void* old_slots, - void* probed_storage, - void (*encode_probed_element)(void* probed_storage, h2_t h2, - size_t source_offset, size_t h1)) { - const size_t new_capacity = common.capacity(); - ABSL_ASSUME(new_capacity > kMaxSmallCapacity); - const size_t old_capacity = PreviousCapacity(new_capacity); + CommonFields& common, const ctrl_t* old_ctrl, void* probed_storage, + void* old_slots, size_t old_capacity, + void (*encode_probed_element)(void* probed_storage, size_t hash, + size_t source_offset)) { ABSL_ASSUME(old_capacity + 1 >= Group::kWidth); ABSL_ASSUME((old_capacity + 1) % Group::kWidth == 0); + const size_t new_capacity = NextCapacity(old_capacity); + ABSL_ASSUME(new_capacity > kMaxSmallCapacity); + ABSL_SWISSTABLE_ASSERT(common.capacity() == new_capacity); auto* set = reinterpret_cast<raw_hash_set*>(&common); slot_type* old_slots_ptr = to_slot(old_slots); @@ -4279,19 +4283,17 @@ // TODO(b/382423690): try to avoid entire hash calculation since we need // only one new bit of h1. size_t hash = set->hash_of(old_slot); - size_t h1 = H1(hash); - h2_t h2 = H2(hash); size_t new_index = TryFindNewIndexWithoutProbing( - h1, old_index, old_capacity, new_ctrl, new_capacity); + hash, old_index, old_capacity, new_ctrl, new_capacity); // Note that encode_probed_element is allowed to use old_ctrl buffer // till and included the old_index. if (ABSL_PREDICT_FALSE(new_index == kProbedElementIndexSentinel)) { - encode_probed_element(probed_storage, h2, old_index, h1); + encode_probed_element(probed_storage, hash, old_index); continue; } ABSL_SWISSTABLE_ASSERT((new_index & old_capacity) <= old_index); ABSL_SWISSTABLE_ASSERT(IsEmpty(new_ctrl[new_index])); - new_ctrl[new_index] = static_cast<ctrl_t>(h2); + new_ctrl[new_index] = static_cast<ctrl_t>(H2(hash)); auto* new_slot = new_slots + new_index; SanitizerUnpoisonMemoryRegion(new_slot, sizeof(slot_type)); set->transfer(new_slot, old_slot); @@ -4541,7 +4543,7 @@ // Extern template for inline function keeps possibility of inlining. // When compiler decided to not inline, no symbols will be added to the // corresponding translation unit. -extern template size_t TryFindNewIndexWithoutProbing(size_t h1, +extern template size_t TryFindNewIndexWithoutProbing(size_t hash, size_t old_index, size_t old_capacity, ctrl_t* new_ctrl,
diff --git a/absl/container/internal/raw_hash_set_test.cc b/absl/container/internal/raw_hash_set_test.cc index 429d7cf..ed25c05 100644 --- a/absl/container/internal/raw_hash_set_test.cc +++ b/absl/container/internal/raw_hash_set_test.cc
@@ -147,6 +147,26 @@ } } +TEST(RawHashSetLayout, CreateForSmallNonSooGrowth) { + for (size_t size : {0u, 1u}) { + for (size_t slot_align : {1u, 2u, 4u, 8u, 16u}) { + for (size_t slot_size : {slot_align, slot_align * 2, slot_align * 3}) { + for (bool has_infoz : {false, true}) { + RawHashSetLayout expected(NextCapacity(size), slot_size, slot_align, + has_infoz, /*blocked_element_count=*/0); + RawHashSetLayout actual = + RawHashSetLayout::CreateForSmallNonSooGrowth( + size, slot_size, slot_align, has_infoz); + EXPECT_EQ(actual.control_offset(), expected.control_offset()); + EXPECT_EQ(actual.generation_offset(), expected.generation_offset()); + EXPECT_EQ(actual.slot_offset(), expected.slot_offset()); + EXPECT_EQ(actual.alloc_size(), expected.alloc_size()); + } + } + } + } +} + void VerifyMiddleSizeTableLayout(size_t capacity, size_t slot_size, size_t slot_align, bool has_infoz, size_t blocked_element_count,