bpf, x86: Use BMI2 shifts also when the count is in rcx

Commit 77d8f5d47bfb ("bpf,x64: use shrx/sarx/shlx when available") kept
the legacy shift by cl when the count already is in rcx, since that
encoding is two bytes shorter. However, shl/shr/sar r,cl are 3 uops on
Intel from Haswell through Skylake and still 2 uops on Ice Lake, as a
zero count must leave the flags untouched, whereas shlx/shrx/sarx are a
single uop with a latency of 1. On AMD both forms are a single uop, so
the two extra bytes are the only cost there. Use the BMI2 form whenever
it is available. LLVM keeps the count in whichever register it was
allocated to, r4 included.

Before/after for r1 <<= r4:

  48 d3 e7        shlq %cl,%rdi
  c4 e2 f1 f7 ff  shlxq %rcx,%rdi,%rdi

Signed-off-by: Daniel Borkmann <daniel@iogearbox.net>
Link: https://lore.kernel.org/bpf/20261013171230.523677-11-daniel@iogearbox.net
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>

1 files changed, 6 insertions, 2 deletions

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index 50158d4b85b2..de49d7f51a54 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -2539,8 +2539,12 @@ static int do_jit(struct bpf_verifier_env *env, struct bpf_prog *bpf_prog, int *
 		case BPF_ALU64 | BPF_LSH | BPF_X:
 		case BPF_ALU64 | BPF_RSH | BPF_X:
 		case BPF_ALU64 | BPF_ARSH | BPF_X:
-			/* BMI2 shifts aren't better when shift count is already in rcx */
-			if (boot_cpu_has(X86_FEATURE_BMI2) && src_reg != BPF_REG_4) {
+			/*
+			 * BMI2 shifts take the count from any register and,
+			 * unlike the cl forms, are a single uop on Intel, so
+			 * prefer them also when the count already is in rcx.
+			 */
+			if (boot_cpu_has(X86_FEATURE_BMI2)) {
 				/* shrx/sarx/shlx dst_reg, dst_reg, src_reg */
 				bool w = (BPF_CLASS(insn->code) == BPF_ALU64);
 				u8 op;