[llvm] [X86] SimplifyDemandedBitsForTargetNode - add X86ISD::BZHI handling (PR #199828)

[Qiongsi Wu via llvm-commits]

https://github.com/qiongsiwu updated https://github.com/llvm/llvm-project/pull/199828

>From b79e5b5453e27ded4942a378999bc47e6fd27fcc Mon Sep 17 00:00:00 2001
From: Qiongsi Wu <qiongsiwu at gmail.com>
Date: Tue, 26 May 2026 21:05:49 -0700
Subject: [PATCH] Implementing case X86ISD::BZHI for
 X86TargetLowering::SimplifyDemandedBitsForTargetNode.

---
 llvm/lib/Target/X86/X86ISelLowering.cpp | 30 +++++++++
 llvm/test/CodeGen/X86/combine-bzhi.ll   | 90 +++++++++++++++++++++++++
 2 files changed, 120 insertions(+)

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 69eab53dca965..04ed1dda21b6d 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -46091,6 +46091,36 @@ bool X86TargetLowering::SimplifyDemandedBitsForTargetNode(
 
     break;
   }
+  case X86ISD::BZHI: {
+    SDValue Op0 = Op.getOperand(0);
+    SDValue Op1 = Op.getOperand(1);
+
+    if (auto *Cst1 = dyn_cast<ConstantSDNode>(Op1)) {
+      uint64_t Val = Cst1->getZExtValue();
+      uint64_t Masked = Val & 0xFF;
+      if (Masked != Val) {
+        SDLoc DL(Op);
+        return TLO.CombineTo(
+            Op, TLO.DAG.getNode(X86ISD::BZHI, DL, VT, Op0,
+                                TLO.DAG.getConstant(Masked, DL, VT)));
+      }
+    }
+
+    KnownBits Known1;
+    APInt DemandedMask(APInt::getLowBitsSet(BitWidth, 8));
+    if (SimplifyDemandedBits(Op1, DemandedMask, Known1, TLO, Depth + 1))
+      return true;
+
+    uint64_t MaxIdx = Known1.trunc(8).getMaxValue().getZExtValue();
+    if (MaxIdx < BitWidth) {
+      APInt DemandedSrc =
+          OriginalDemandedBits & APInt::getLowBitsSet(BitWidth, MaxIdx);
+      if (SimplifyDemandedBits(Op0, DemandedSrc, Known, TLO, Depth + 1))
+        return true;
+    }
+
+    break;
+  }
   case X86ISD::VPMADD52L:
   case X86ISD::VPMADD52H: {
     KnownBits KnownOp0, KnownOp1, KnownOp2;
diff --git a/llvm/test/CodeGen/X86/combine-bzhi.ll b/llvm/test/CodeGen/X86/combine-bzhi.ll
index 37b0c78347373..0aa6cf9c9510d 100644
--- a/llvm/test/CodeGen/X86/combine-bzhi.ll
+++ b/llvm/test/CodeGen/X86/combine-bzhi.ll
@@ -81,3 +81,93 @@ define i64 @test_bzhi64_range(i64 %arg) nounwind readnone {
   %4 = tail call i64 @llvm.x86.bmi.bzhi.64(i64 30064771240, i64 %3)
   ret i64 %4
 }
+
+define i32 @test_bzhi32_control_upper_bits(i32 %a, i32 %b) nounwind {
+; CHECK-LABEL: test_bzhi32_control_upper_bits:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bzhil %esi, %edi, %eax
+; CHECK-NEXT:    retq
+  %control = or i32 %b, 256
+  %result = tail call i32 @llvm.x86.bmi.bzhi.32(i32 %a, i32 %control)
+  ret i32 %result
+}
+
+define i64 @test_bzhi64_control_upper_bits(i64 %a, i64 %b) nounwind {
+; CHECK-LABEL: test_bzhi64_control_upper_bits:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bzhiq %rsi, %rdi, %rax
+; CHECK-NEXT:    retq
+  %control = or i64 %b, 256
+  %result = tail call i64 @llvm.x86.bmi.bzhi.64(i64 %a, i64 %control)
+  ret i64 %result
+}
+
+; Constant control 0x110 (272) — low 8 bits are 0x10 (16), upper bits are junk.
+; Should simplify to bzhi(a, 16), which keeps low 16 bits.
+define i32 @test_bzhi32_constant_control_upper_bits(i32 %a) nounwind {
+; CHECK-LABEL: test_bzhi32_constant_control_upper_bits:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    movl $16, %eax
+; CHECK-NEXT:    bzhil %eax, %edi, %eax
+; CHECK-NEXT:    retq
+  %result = tail call i32 @llvm.x86.bmi.bzhi.32(i32 %a, i32 272)
+  ret i32 %result
+}
+
+define i64 @test_bzhi64_constant_control_upper_bits(i64 %a) nounwind {
+; CHECK-LABEL: test_bzhi64_constant_control_upper_bits:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    movl $16, %eax
+; CHECK-NEXT:    bzhiq %rax, %rdi, %rax
+; CHECK-NEXT:    retq
+  %result = tail call i64 @llvm.x86.bmi.bzhi.64(i64 %a, i64 272)
+  ret i64 %result
+}
+
+define i32 @test_bzhi32_src_upper_bits(i32 %a, i32 %b) nounwind {
+; CHECK-LABEL: test_bzhi32_src_upper_bits:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    andl $15, %esi
+; CHECK-NEXT:    bzhil %esi, %edi, %eax
+; CHECK-NEXT:    retq
+  %src = or i32 %a, 65536
+  %control = and i32 %b, 15
+  %result = tail call i32 @llvm.x86.bmi.bzhi.32(i32 %src, i32 %control)
+  ret i32 %result
+}
+
+define i64 @test_bzhi64_src_upper_bits(i64 %a, i64 %b) nounwind {
+; CHECK-LABEL: test_bzhi64_src_upper_bits:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    andl $15, %esi
+; CHECK-NEXT:    bzhiq %rsi, %rdi, %rax
+; CHECK-NEXT:    retq
+  %src = or i64 %a, 65536
+  %control = and i64 %b, 15
+  %result = tail call i64 @llvm.x86.bmi.bzhi.64(i64 %src, i64 %control)
+  ret i64 %result
+}
+
+define i32 @test_bzhi32_known_bits(i32 %a, i32 %b) nounwind {
+; CHECK-LABEL: test_bzhi32_known_bits:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    andl $7, %esi
+; CHECK-NEXT:    bzhil %esi, %edi, %eax
+; CHECK-NEXT:    retq
+  %control = and i32 %b, 7
+  %bzhi = tail call i32 @llvm.x86.bmi.bzhi.32(i32 %a, i32 %control)
+  %result = and i32 %bzhi, 127
+  ret i32 %result
+}
+
+define i64 @test_bzhi64_known_bits(i64 %a, i64 %b) nounwind {
+; CHECK-LABEL: test_bzhi64_known_bits:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    andl $7, %esi
+; CHECK-NEXT:    bzhiq %rsi, %rdi, %rax
+; CHECK-NEXT:    retq
+  %control = and i64 %b, 7
+  %bzhi = tail call i64 @llvm.x86.bmi.bzhi.64(i64 %a, i64 %control)
+  %result = and i64 %bzhi, 127
+  ret i64 %result
+}