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

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

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/



>  	*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>
>