Re: [PATCH bpf-next v2] bpf: Speed up htab lookups for u32/u64 keys

[Anton Protopopov <a.s.protopopov@gmail.com>]

On 26/09/22 12:45PM, Mykyta Yatsenko wrote:
> 
> 
> On 9/22/26 8:31 AM, Anton Protopopov wrote:
> > On 26/09/21 02:28PM, Mykyta Yatsenko wrote:
> >> From: Mykyta Yatsenko <yatsenko@meta.com>
> >>
> >> Four- and eight-byte keys are common enough to warrant specialized htab
> >> lookups. Fold JHASH_INITVAL and the key length into hashrnd once when
> >> allocating maps with these key sizes.
> >>
> >> For JITed BPF_MAP_TYPE_HASH lookups, htab_map_gen_lookup() also knows the
> >> map's fixed key size. Select u32- and u64-specific lookup entry points so
> >> the compiler can specialize both hash setup and key comparison. Other key
> >> sizes continue to use the generic lookup entry point.
> >>
> >> This preserves the jhash2() result and lookup semantics while avoiding
> >> generic length handling and repeated state setup.
> >>
> >> Measure this with six runs of:
> >>
> >>   ./bench -w3 -d10 -a bpf-hashmap-lookup \
> >>       --key_size KEY_SIZE --max_entries MAX_ENTRIES \
> >>       --nr_entries NR_ENTRIES --nr_loops NR_LOOPS \
> >>       --map_flags 0x40
> >>
> >> inside a vng guest with two vCPUs and 4 GiB of RAM. Run the baseline
> >> first and the optimized kernel second. Use these workloads:
> >>
> >>              max_entries   nr_entries    nr_loops
> >>   small              512          256   8,388,608
> >>   medium          10,000        5,000   8,000,000
> >>   large          100,000       50,000   8,000,000
> >>
> >> Mean throughput in million lookups per second is:
> >>
> >>   Baseline:
> >>                          key size (bytes)
> >>                       1        4        8       10
> >>     small         89.45    91.21   103.40    91.74
> >>     medium       125.97    80.25    91.94    78.52
> >>     large        140.67    53.56    58.68    47.31
> >>
> >>   Optimized:
> >>                          key size (bytes)
> >>                       1        4        8       10
> >>     small         86.95   233.63   168.65    90.88
> >>     medium       124.97   150.20   137.76    77.51
> >>     large        131.06    80.51    76.04    47.25
> >>
> >>   Change:
> >>                          key size (bytes)
> >>                       1        4        8       10
> >>     small         -2.79% +156.15%  +63.11%   -0.94%
> >>     medium        -0.79%  +87.17%  +49.83%   -1.28%
> >>     large         -6.84%  +50.32%  +29.58%   -0.12%
> >>
> >> The u32 specialization improves lookup throughput by 50-156%, and the
> >> u64 specialization improves it by 30-63%. Ten-byte controls remain
> >> within a 1.3% slowdown. One-byte controls range from -6.8% to -0.8%.
> >>
> >> Signed-off-by: Mykyta Yatsenko <yatsenko@meta.com>
> >> Acked-by: Anton Protopopov <a.s.protopopov@gmail.com>
> >> ---
> >> Changes in v2:
> >> - Added variants of __htab_map_lookup_elem() functions with inlined
> >> u32/u64 key sizes, use them in htab_map_gen_lookup();
> >> improved lookup perf significantly (Alexei)
> >> - Link to v1: https://patch.msgid.link/20260921-hashtab_fast_hashfn-v1-1-f92ae72cefcd@meta.com
> >> ---
> >>  kernel/bpf/hashtab.c | 49 ++++++++++++++++++++++++++++++++++++++++++-------
> >>  1 file changed, 42 insertions(+), 7 deletions(-)
> >>
> >> diff --git a/kernel/bpf/hashtab.c b/kernel/bpf/hashtab.c
> >> index 6f331c80130d..77c96105eda8 100644
> >> --- a/kernel/bpf/hashtab.c
> >> +++ b/kernel/bpf/hashtab.c
> >> @@ -612,6 +612,13 @@ static struct bpf_map *htab_map_alloc(union bpf_attr *attr)
> >>  		htab->hashrnd = 0;
> >>  	else
> >>  		htab->hashrnd = get_random_u32();
> >> +	/*
> >> +	 * Fold the jhash constant and key length into hashrnd once for the
> >> +	 * fixed-size fast paths instead of doing it on every lookup.
> >> +	 */
> >> +	if (htab->map.key_size == sizeof(u32) ||
> >> +	    htab->map.key_size == sizeof(u64))
> >> +		htab->hashrnd += JHASH_INITVAL + htab->map.key_size;
> >>  
> >>  	htab_init_buckets(htab);
> >>  
> >> @@ -679,9 +686,19 @@ static struct bpf_map *htab_map_alloc(union bpf_attr *attr)
> >>  
> >>  static inline u32 htab_map_hash(const void *key, u32 key_len, u32 hashrnd)
> >>  {
> >> -	if (likely(key_len % 4 == 0))
> >> +	const u32 *k = key;
> >> +	u32 b;
> >> +
> >> +	if (key_len == sizeof(u32))
> >> +		b = 0;
> >> +	else if (key_len == sizeof(u64))
> >> +		b = k[1];
> >> +	else if (likely(key_len % 4 == 0))
> >>  		return jhash2(key, key_len / 4, hashrnd);
> >> -	return jhash(key, key_len, hashrnd);
> >> +	else
> >> +		return jhash(key, key_len, hashrnd);
> >> +
> >> +	return __jhash_nwords(k[0], b, 0, hashrnd);
> >>  
> >>  static inline struct bucket *__select_bucket(struct bpf_htab *htab, u32 hash)
> >> @@ -735,17 +752,15 @@ static struct htab_elem *lookup_nulls_elem_raw(struct hlist_nulls_head *head,
> >>   * The return value is adjusted by BPF instructions
> >>   * in htab_map_gen_lookup().
> >>   */
> >> -static void *__htab_map_lookup_elem(struct bpf_map *map, void *key)
> >> +static __always_inline void *__htab_lookup(struct bpf_map *map, void *key, u32 key_size)
> >>  {
> >>  	struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
> >>  	struct hlist_nulls_head *head;
> >>  	struct htab_elem *l;
> >> -	u32 hash, key_size;
> >> +	u32 hash;
> >>  
> >>  	WARN_ON_ONCE(!bpf_rcu_lock_held());
> >>  
> >> -	key_size = map->key_size;
> >> -
> >>  	hash = htab_map_hash(key, key_size, htab->hashrnd);
> >>  
> >>  	head = select_bucket(htab, hash);
> >> @@ -755,6 +770,21 @@ static void *__htab_map_lookup_elem(struct bpf_map *map, void *key)
> >>  	return l;
> >>  }
> >>  
> >> +static void *__htab_map_lookup_elem(struct bpf_map *map, void *key)
> >> +{
> >> +	return __htab_lookup(map, key, map->key_size);
> >> +}
> >> +
> >> +static void *__htab_map_lookup_elem_u32(struct bpf_map *map, void *key)
> >> +{
> >> +	return __htab_lookup(map, key, sizeof(u32));
> >> +}
> >> +
> >> +static void *__htab_map_lookup_elem_u64(struct bpf_map *map, void *key)
> >> +{
> >> +	return __htab_lookup(map, key, sizeof(u64));
> >> +}
> >> +
> >>  static void *htab_map_lookup_elem(struct bpf_map *map, void *key)
> >>  {
> >>  	struct htab_elem *l = __htab_map_lookup_elem(map, key);
> >> @@ -783,7 +813,12 @@ static int htab_map_gen_lookup(struct bpf_map *map, struct bpf_insn *insn_buf)
> >>  
> >>  	BUILD_BUG_ON(!__same_type(&__htab_map_lookup_elem,
> >>  		     (void *(*)(struct bpf_map *map, void *key))NULL));
> >> -	*insn++ = BPF_EMIT_CALL(__htab_map_lookup_elem);
> >> +	if (map->key_size == sizeof(u32))
> >> +		*insn++ = BPF_EMIT_CALL(__htab_map_lookup_elem_u32);
> >> +	else if (map->key_size == sizeof(u64))
> >> +		*insn++ = BPF_EMIT_CALL(__htab_map_lookup_elem_u64);
> >> +	else
> >> +		*insn++ = BPF_EMIT_CALL(__htab_map_lookup_elem);
> > 
> > I am trying to remember why I didn't do this when I was touching this
> > code, I was definitely thinking about it... Given your changes,
> > it is worth to approach xxh3 again, it works way better for bigger key
> > sizes, see [1].
> > 
> >   [1] https://archive.fosdem.org/2023/schedule/event/bpf_hashing/
> > 
> > 
> Let me try to run the benchmark. Did you use an actual xxh3 (not
> available in kernel) or xxh32/xxh64?> 

The actual xxh3. For small keys xxh3 was really fragile (the hash was
outperforming in hash-only benchmarks, but was not that good when used in the
hashmap), so I ended up only patching code to use jhash2. For keys >=20,
I think, xxh3 becomes definitely better for all benchmarks. One further note:
the original xxh3 uses vectorized instructions for keys >=240 and for such keys
the fastest ist xxh64. (Well, for keys >5-6K the spooky hash was the best,
but are there any such use cases?)

For the code, I've used my own kernel port of xxh3. The XXH3 also uses a
key-size-based switch internally, so for generated lookups something like
xxh3_17_to_128 can be used specifically. Unfortanetely, I am not sure
what was the latest version I had... You can find one here:

  https://github.com/aspsk/bpf-next/commit/e86be3c926051ee87375da4f0df4b873d2f7e9c2

(notably see xxh3_17_to_128() and xxh3_129_to_240()). The code should work,
but definitely needs some love, maybe it's easier to re-import it using
modern tools.

> >>  	*insn++ = BPF_JMP_IMM(BPF_JEQ, ret, 0, 1);
> >>  	*insn++ = BPF_ALU64_IMM(BPF_ADD, ret,
> >>  				offsetof(struct htab_elem, key) +
> >>
> >> ---
> >> base-commit: 6e1ac8758fcd07bee12ccaff0087cc66930a4187
> >> change-id: 20260916-hashtab_fast_hashfn-9e52b3c2276e
> >>
> >> Best regards,
> >> --  
> >> Mykyta Yatsenko <yatsenko@meta.com>
> >>
>