Allow to prefetch buckets

[rluvaton]

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

[clarfonthey]

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.

[cuviper]

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

[clarfonthey]

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.

[rluvaton]

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)

[rluvaton]

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

[Amanieu]

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.

[rluvaton]

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.

[Amanieu]

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.

[cuviper]

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.