Allow to prefetch buckets

In Apache DataFusion we are using hashbrown extensively and I noticed in hash aggregation when the hash table (we use the raw API there) become very large we are memory bound so a natural solution would be to have some kind of prefetching (FYI we are using stable Rust), but there is no way to do that currently

We have something like this (simplified):

for (row, &target_hash) in batch_hashes.iter().enumerate() {
    let entry = self.map.entry(
        target_hash,
        // eq
        |(exist_hash, group_idx)| target_hash == *exist_hash && group_rows.row(row) == group_values.row(*group_idx),
        // hasher
        |(hash, _)|  *hash);

    let group_idx = match entry {
        // Existing group_index for this group value
        Entry::Occupied(o) => {
            let (_hash, group_idx) = o.get();  
            
            *group_idx
        }
        // Need to create new entry for the group
        Entry::Vacant(v) => {
            // Add new entry to aggr_state and save newly created index
            let group_idx = group_values.num_rows();
            group_values.push(group_rows.row(row));
            v.insert((target_hash, group_idx));

            group_idx
        }
    };
    groups.push(group_idx);
}

Allowing to prefetch from the map can help performance for our case.

Because we insert to the table on missing we would have to reserve the x amount of items beforehand for the prefetching to be valuable

So, I was a little confused what I meant but was able to read the code from your fork to understand. You're referring to actual cache prefectching (unstable feature), not just requesting all the buckets in ahead of time.

Although, I have to ask, would adding an entry version of get_disjoint_mut potentially cover your use case? Since this would load all the entries at once and potentially avoid the need for prefectching.

I don't think we can support disjoint entries, unless it's a new entry type that doesn't allow insertion or removal.

I was under the impression the existing entry types already accounted for this, but that's fair. I guess that the purpose here is to predefine every entry for the group, so, perhaps we could have a version that initializes them for you using a function, or Default?

My main thought process was this specific case, where get_disjoint_mut is insufficient since some of the entries may be missing.

So, I was a little confused what I meant but was able to read the code from your fork to understand. You're referring to actual cache prefectching (unstable feature), not just requesting all the buckets in ahead of time.

if requesting all buckets ahead of time is possible than yes, it would be great, but I don't think it is possible to do as it would disallow hash collision (either the same hash which is ok for me, or 2 different hashes but going to the same bucket) as it would return the same bucket.

You can get away with prefetching on stable like they do in the branches crate:

prefetch in stable

/// Prefetches data for reading into the cache.
///
/// This function hints to the CPU that the data at the given address
/// will be read soon, allowing the CPU to load the data into the cache
/// in advance. This can improve performance by reducing cache misses.
///
/// # Arguments
///
/// * `addr` - A pointer to the data to prefetch.
/// * `LOCALITY` - The cache locality to prefetch into.
#[inline(always)]
#[cfg(feature = "prefetch")]
pub fn prefetch_read_data<T, const LOCALITY: i32>(addr: *const T) {
    #[cfg(branches_stable)]
    {
        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
        unsafe {
            match LOCALITY {
                0 => core::arch::asm!(
                    "prefetcht0 [{}]",
                    in(reg) addr,
                    options(nostack, readonly, preserves_flags)
                ), // L1 cache
                1 => core::arch::asm!(
                    "prefetcht1 [{}]",
                    in(reg) addr,
                    options(nostack, readonly, preserves_flags)
                ), // L2 cache
                2 => core::arch::asm!(
                    "prefetcht2 [{}]",
                    in(reg) addr,
                    options(nostack, readonly, preserves_flags)
                ), // L3 cache
                _ => core::arch::asm!(
                    "prefetchnta [{}]",
                    in(reg) addr,
                    options(nostack, readonly, preserves_flags)
                ), // Non-temporal
            }
        }

        #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
        unsafe {
            match LOCALITY {
                0 => core::arch::asm!(
                    "prfm pldl1keep, [{}]",
                    in(reg) addr,
                    options(nostack, readonly, preserves_flags)
                ), // L1 cache
                1 => core::arch::asm!(
                    "prfm pldl2keep, [{}]",
                    in(reg) addr,
                    options(nostack, readonly, preserves_flags)
                ), // L2 cache
                _ => core::arch::asm!(
                    "prfm pldl3keep, [{}]",
                    in(reg) addr,
                    options(nostack, readonly, preserves_flags)
                ), // L3 or non-temporal
            }
        }

        #[cfg(target_arch = "riscv64")]
        unsafe {
            core::arch::asm!(
                "prefetch [{}]",
                in(reg) addr,
                options(nostack, readonly, preserves_flags)
            );
        }

        #[cfg(any(target_arch = "powerpc", target_arch = "powerpc64"))]
        unsafe {
            core::arch::asm!(
                "dcbt 0, {}",
                in(reg) addr,
                options(nostack, readonly, preserves_flags)
            );
        }
    }
    #[cfg(branches_nightly)]
    core::intrinsics::prefetch_read_data::<_, LOCALITY>(addr)
}

/// Prefetches data for writing into the cache.
///
/// This function hints to the CPU that the data at the given address
/// will be written soon, allowing the CPU to load the data into the cache
/// in advance. This can improve performance by reducing cache misses.
///
/// # Arguments
///
/// * `addr` - A pointer to the data to prefetch.
/// * `LOCALITY` - The cache locality to prefetch into.
#[inline(always)]
#[cfg(feature = "prefetch")]
pub fn prefetch_write_data<T, const LOCALITY: i32>(addr: *const T) {
    #[cfg(branches_stable)]
    {
        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
        unsafe {
            core::arch::asm!(
                "prefetchw [{}]",
                in(reg) addr,
                options(nostack, readonly, preserves_flags)
            ) // Write-prefetch for L1/L2/L3 cache
        }

        #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
        unsafe {
            core::arch::asm!(
                "prfm pstl1keep, [{}]",
                in(reg) addr,
                options(nostack, readonly, preserves_flags)
            ); // Write-prefetch for L1 cache
        }

        #[cfg(target_arch = "riscv64")]
        unsafe {
            core::arch::asm!(
                "prefetchw [{}]",
                in(reg) addr,
                options(nostack, readonly, preserves_flags)
            );
        }

        #[cfg(any(target_arch = "powerpc", target_arch = "powerpc64"))]
        unsafe {
            core::arch::asm!(
                "dcbtst 0, {}",
                in(reg) addr,
                options(nostack, readonly, preserves_flags)
            ); // Write-prefetch
        }
    }
    #[cfg(branches_nightly)]
    core::intrinsics::prefetch_write_data::<_, LOCALITY>(addr)
}

From:

https://github.com/fereidani/branches/blob/3904fcc53d152173b931e37c919f832bac67b81c/src/lib.rs#L155C1-L293C2

Although, I have to ask, would adding an entry version of get_disjoint_mut potentially cover your use case? Since this would load all the entries at once and potentially avoid the need for prefectching.

I was not familiar with get_disjoint_mut, but looking at the docs it would have the same issue if I need to insert all the values, due to hash collision (either the same hash which is ok for me, or 2 different hashes but going to the same bucket)

Also in get_disjoint_mut if a single hash was not found it will return None and also there is no prefetch there.

and both get_disjoint_mut and requesting all buckets ahead of time doesn't solve the problem that we now stall for all the buckets to be fetched from various places in memory, so with a prefetch we can at least allow us to do some work while the CPU prefetch the data

I'm generally in favor of adding a way to prefetch a hash table entry, and abseil has a similar method on its hash table implementation.

My main concern is what the API should look like, specifically with regards to the kind of prefetch that is being performed. That is, should we offer the full flexibility of locality and read/write or should we keep the interface simpler and only provide L1 read prefetching? I'm inclined to go with whatever ends up being decided for the std prefetch hints in rust-lang/rust#146941, but we would need to wait for that to stabilize before we could re-use the Locality enum.

we can start by having a function that returns optional raw pointer to the bucket (mark it unsafe of course) and then the user could implement prefetch however they like, and if the prefetch hints are getting evolved we could add an explicit API for it.

returning raw API with enough docs that says that it should be used for any way of getting data or modifying data and this pointer may be invalid after modification to the hash table and should not be held in memory.

You can already get a raw pointer to a bucket from get which returns a reference. What you actually want here is a "guess" at where the bucket probably is without doing any of the equality checks. I don't think that is reasonable to expose as an API because it's really only useful for prefetching, so I would rather just expose prefetching directly.

One possibility you may want to consider while the std prefetch hint is still unstable is to manually calculate the bucket index from the hash and then use the new bucket API (#657) to get the address of the first bucket for this hash. Since hashbrown doesn't guarantee that its hashing scheme is going to be stable, you'll most likely have to pin your hashbrown version.

For prefetch purposes, wouldn't you want both the control word and data pointers? e.g. the previously-linked abseil prefetch does both. You're trying to hint the memory system to load the pieces you'll need soon for a proper lookup.

How exactly do the existing prefetch intrinsics work in Rust's memory model specifically? Since the control and data pointers have the same provenance, would the memory system would see them as being part of the same portion to be prefetched?

Or does it work based upon whatever memory page a pointer is in, irrespective of its provenance? Does it have a heuristic based upon the size of the allocation?

One other big thing to consider is how relatively large the actual data is in the hash table, compared to the control pointers. Even the best case scenario of a single pointer for the data means that the data is around 8x larger than the control. For a simple vector, which is three pointers wide, you would have 24x larger data compared to control. It seems very likely that, even if your data is large, the entire control segment might fit in a handful of cache pages, if not a single page, and so maybe prefetching the data segment is more important since the control will just stay in cache.

Perhaps it would be more useful for the sake of precaching to explicitly optimise for cases where the central pointer is aligned to a page boundary to support this behaviour, at least if the table is larger than that size.

The prefetch intrinsic is a hint which is a no-op in the rust abstract machine. Even passing an invalid pointer is not UB. This means that it doesn't care about provenance. Instead it will prefetch the cache line that contains the given address. Cache lines are always aligned and typically have a size of 32/64/128 bytes depending on the specific CPU. This means that we need to emit 2 separate prefetches: one for the control bytes and one for the data.

@rluvaton

Two comments;

I wonder if a variant could be added for get_disjoint_mut for read-only cases (DataFusion join probing)?
I guess that could be made faster than unrolling manually (some micro benchmark gives a ~2.4x speedup using this approach)?

let r0 = map.find(chunk[0], |(h, _)| chunk[0] == *h);                                                                                            
let r1 = map.find(chunk[1], |(h, _)| chunk[1] == *h);                                                                                            
let r2 = map.find(chunk[2], |(h, _)| chunk[2] == *h);                                                                                            
let r3 = map.find(chunk[3], |(h, _)| chunk[3] == *h);

@rluvaton Perhaps we could change some algorithms in DataFusion to handle hash-collisions later (for aggregates: lookup first, handle missing entries later).

@Dandandan

I ran some cross-architecture benchmarks that might be useful for the API discussion here. Vendored hashbrown with RawTable::prefetch(hash) added (prefetching both ctrl and data), tested at sizes from 16K to 4M slots on an i5-13600K. (48KB L1d) and an Apple M4 (192KB L1d). 16-byte keys, ahash, 50% load, randomized lookup order, median of 5 runs.

Single-lookup prefetch in find() is mostly noise on M4 and only ~8% on x86 at 1M entries. The prefetch-to-use distance in the probe loop is just too short. It only starts to matter at 4M where ctrl blows past L2 - then ~15% on both architectures.

Batched lookups are a different story. Hashing 4 keys, issuing all 4 prefetches, then probing gives 15-20% with ahash on M4 and ~28% on x86 at 1M. The 4 hash computations create enough distance for prefetches to resolve. Slower hashers benefit more - siphash sees ~50% because there's more compute time between prefetch and use.

The underlying issue is the ctrl byte array. At 1M entries it's 1MB, way past L1 on both machines. Stock hashbrown runs at IPC 0.62 on x86 / 0.71 on M4 — clearly memory-stalled. Batching pushes M4 to 0.96.

For the API: RawTable::prefetch(hash) feels like the right primitive, but callers won't get much from it unless they can batch.

The 4 way find pattern @Dandandan mentioned is where the real gains are.

Navigation Menu

Repository navigation

Allow to prefetch buckets #677

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Allow to prefetch buckets #677

Description

Activity

clarfonthey commented on Dec 28, 2025

cuviper commented on Dec 28, 2025

clarfonthey commented on Dec 28, 2025

rluvaton commented on Dec 28, 2025

rluvaton commented on Dec 28, 2025

Amanieu commented on Dec 31, 2025

rluvaton commented on Jan 1, 2026

Amanieu commented on Jan 1, 2026

cuviper commented on Jan 1, 2026

clarfonthey commented on Jan 9, 2026

Amanieu commented on Jan 10, 2026

Dandandan commented on Mar 15, 2026

joshuaisaact commented on Mar 28, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions