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

[Mykyta Yatsenko <mykyta.yatsenko5@gmail.com>]

From: Mykyta Yatsenko <yatsenko@meta.com>

Four- and eight-byte keys are common enough to warrant specialized htab
lookups. Use jhash_1word() and jhash_2words() for these key sizes instead
of the generic jhash2() path.

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 hashing 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.

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   197.77   168.65    90.88
    medium       124.97   147.48   137.76    77.51
    large        131.06    78.51    76.04    47.25

  Change:
                         key size (bytes)
                      1        4        8       10
    small         -2.79% +116.83%  +63.11%   -0.94%
    medium        -0.79%  +83.78%  +49.83%   -1.28%
    large         -6.84%  +46.58%  +29.58%   -0.12%

The u32 specialization improves lookup throughput by 47-117%, 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 v3:
- Use __allways_inline everywhere compile-time key_size is
used (Alexei)
- Remove hashrnd adjustment in htab init, call jhash_1word()/jhash_2word()
- Link to v2: https://patch.msgid.link/20260921-hashtab_fast_hashfn-v2-1-79aa962ec56b@meta.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 | 63 +++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 47 insertions(+), 16 deletions(-)

diff --git a/kernel/bpf/hashtab.c b/kernel/bpf/hashtab.c
index 6f331c80130d..f744a42bb813 100644
--- a/kernel/bpf/hashtab.c
+++ b/kernel/bpf/hashtab.c
@@ -677,11 +677,18 @@ static struct bpf_map *htab_map_alloc(union bpf_attr *attr)
 	return ERR_PTR(err);
 }
 
-static inline u32 htab_map_hash(const void *key, u32 key_len, u32 hashrnd)
+static __always_inline u32 htab_map_hash(const void *key, u32 key_len, u32 hashrnd)
 {
-	if (likely(key_len % 4 == 0))
+	const u32 *k = key;
+
+	if (key_len == sizeof(u32))
+		return jhash_1word(k[0], hashrnd);
+	else if (key_len == sizeof(u64))
+		return jhash_2words(k[0], k[1], hashrnd);
+	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);
 }
 
 static inline struct bucket *__select_bucket(struct bpf_htab *htab, u32 hash)
@@ -712,9 +719,9 @@ static struct htab_elem *lookup_elem_raw(struct hlist_nulls_head *head, u32 hash
  * the unlikely event when elements moved from one bucket into another
  * while link list is being walked
  */
-static struct htab_elem *lookup_nulls_elem_raw(struct hlist_nulls_head *head,
-					       u32 hash, void *key,
-					       u32 key_size, u32 n_buckets)
+static __always_inline struct htab_elem *lookup_nulls_elem_raw(struct hlist_nulls_head *head,
+							       u32 hash, void *key,
+							       u32 key_size, u32 n_buckets)
 {
 	struct hlist_nulls_node *n;
 	struct htab_elem *l;
@@ -730,22 +737,15 @@ static struct htab_elem *lookup_nulls_elem_raw(struct hlist_nulls_head *head,
 	return NULL;
 }
 
-/* Called from syscall or from eBPF program directly, so
- * arguments have to match bpf_map_lookup_elem() exactly.
- * 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 +755,28 @@ static void *__htab_map_lookup_elem(struct bpf_map *map, void *key)
 	return l;
 }
 
+/*
+ * Called from syscall or from eBPF program directly, so
+ * arguments have to match bpf_map_lookup_elem() exactly.
+ * The return value is adjusted by BPF instructions
+ * in htab_map_gen_lookup(). u32/u64 specializations
+ * allow utilizing compile-time key_size for optimization.
+ */
+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 +805,16 @@ 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);
+	BUILD_BUG_ON(!__same_type(&__htab_map_lookup_elem_u32,
+				  (void *(*)(struct bpf_map *map, void *key))NULL));
+	BUILD_BUG_ON(!__same_type(&__htab_map_lookup_elem_u64,
+				  (void *(*)(struct bpf_map *map, void *key))NULL));
+	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);
 	*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>