[llvm] [LoongArch] Fold shifted vector extract comparisons (PR #214204)

[via llvm-commits]

https://github.com/heiher updated https://github.com/llvm/llvm-project/pull/214204

>From 1183eac050110a2c60bd33e4d59e395eb5ed6241 Mon Sep 17 00:00:00 2001
From: WANG Rui <wangrui at loongson.cn>
Date: Wed, 5 Aug 2026 19:06:23 +0800
Subject: [PATCH] [LoongArch] Fold shifted vector extract comparisons

Fold comparisons of the form:

  (shl (extract_vector_elt X, I), GRLen - EleBits) ==/!= 0

into:

  (extract_vector_elt X, I) ==/!= 0

When the shift amount equals `GRLen - EleBits`, the left shift only moves
the extracted element into the most significant bits without affecting
whether the value is zero. This canonicalization exposes
EXTRACT_VECTOR_ELT to later combines and enables selecting
VPICKVE2GR_* instructions directly.
---
 .../LoongArch/LoongArchISelLowering.cpp       | 17 ++++++++
 .../LoongArch/lasx/vec-extract-brcond.ll      | 39 ++++++-------------
 .../LoongArch/lsx/vec-extract-brcond.ll       | 39 ++++++-------------
 3 files changed, 41 insertions(+), 54 deletions(-)

diff --git a/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp b/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
index e4a369c096545..3b57a338b86f2 100644
--- a/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
+++ b/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
@@ -7297,6 +7297,23 @@ static bool combine_CC(SDValue &LHS, SDValue &RHS, SDValue &CC, const SDLoc &DL,
     return true;
   }
 
+  // Fold ((shl (extract_vector_elt X, I), GRLen - EleBits)), 0, eq/ne) ->
+  //      ((extract_vector_elt X, I), 0, eq/ne)
+  if (isNullConstant(RHS) && (CCVal == ISD::SETEQ || CCVal == ISD::SETNE) &&
+      LHS.getOpcode() == ISD::SHL && LHS.hasOneUse() &&
+      isa<ConstantSDNode>(LHS.getOperand(1))) {
+    SDValue Ext = LHS.getOperand(0);
+    unsigned Sht = LHS.getConstantOperandVal(1);
+    if (Ext.getOpcode() == ISD::EXTRACT_VECTOR_ELT) {
+      SDValue Vec = Ext.getOperand(0);
+      unsigned EleBits = Vec.getScalarValueSizeInBits();
+      if ((EleBits + Sht) == Subtarget.getGRLen()) {
+        LHS = Ext;
+        return true;
+      }
+    }
+  }
+
   return false;
 }
 
diff --git a/llvm/test/CodeGen/LoongArch/lasx/vec-extract-brcond.ll b/llvm/test/CodeGen/LoongArch/lasx/vec-extract-brcond.ll
index e16d9be1a95d9..3f5ddcee5236a 100644
--- a/llvm/test/CodeGen/LoongArch/lasx/vec-extract-brcond.ll
+++ b/llvm/test/CodeGen/LoongArch/lasx/vec-extract-brcond.ll
@@ -29,33 +29,18 @@ bb2:
 }
 
 define void @extract_16xi16_br_ne_0(ptr %src, ptr %dst) nounwind {
-; LA32-LABEL: extract_16xi16_br_ne_0:
-; LA32:       # %bb.0: # %entry
-; LA32-NEXT:    xvld $xr0, $a0, 0
-; LA32-NEXT:    vpickve2gr.h $a0, $vr0, 0
-; LA32-NEXT:    slli.w $a0, $a0, 16
-; LA32-NEXT:    beqz $a0, .LBB1_2
-; LA32-NEXT:  # %bb.1: # %bb1
-; LA32-NEXT:    st.b $zero, $a1, 0
-; LA32-NEXT:    ret
-; LA32-NEXT:  .LBB1_2: # %bb2
-; LA32-NEXT:    ori $a0, $zero, 1
-; LA32-NEXT:    st.b $a0, $a1, 0
-; LA32-NEXT:    ret
-;
-; LA64-LABEL: extract_16xi16_br_ne_0:
-; LA64:       # %bb.0: # %entry
-; LA64-NEXT:    xvld $xr0, $a0, 0
-; LA64-NEXT:    vpickve2gr.h $a0, $vr0, 0
-; LA64-NEXT:    slli.d $a0, $a0, 48
-; LA64-NEXT:    beqz $a0, .LBB1_2
-; LA64-NEXT:  # %bb.1: # %bb1
-; LA64-NEXT:    st.b $zero, $a1, 0
-; LA64-NEXT:    ret
-; LA64-NEXT:  .LBB1_2: # %bb2
-; LA64-NEXT:    ori $a0, $zero, 1
-; LA64-NEXT:    st.b $a0, $a1, 0
-; LA64-NEXT:    ret
+; CHECK-LABEL: extract_16xi16_br_ne_0:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    xvld $xr0, $a0, 0
+; CHECK-NEXT:    vpickve2gr.h $a0, $vr0, 0
+; CHECK-NEXT:    beqz $a0, .LBB1_2
+; CHECK-NEXT:  # %bb.1: # %bb1
+; CHECK-NEXT:    st.b $zero, $a1, 0
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB1_2: # %bb2
+; CHECK-NEXT:    ori $a0, $zero, 1
+; CHECK-NEXT:    st.b $a0, $a1, 0
+; CHECK-NEXT:    ret
 entry:
   %0 = load volatile <16 x i16>, ptr %src
   %1 = extractelement <16 x i16> %0, i64 0
diff --git a/llvm/test/CodeGen/LoongArch/lsx/vec-extract-brcond.ll b/llvm/test/CodeGen/LoongArch/lsx/vec-extract-brcond.ll
index e78590b3a62db..b748d48e3d341 100644
--- a/llvm/test/CodeGen/LoongArch/lsx/vec-extract-brcond.ll
+++ b/llvm/test/CodeGen/LoongArch/lsx/vec-extract-brcond.ll
@@ -29,33 +29,18 @@ bb2:
 }
 
 define void @extract_8xi16_br_ne_0(ptr %src, ptr %dst) nounwind {
-; LA32-LABEL: extract_8xi16_br_ne_0:
-; LA32:       # %bb.0: # %entry
-; LA32-NEXT:    vld $vr0, $a0, 0
-; LA32-NEXT:    vpickve2gr.h $a0, $vr0, 0
-; LA32-NEXT:    slli.w $a0, $a0, 16
-; LA32-NEXT:    beqz $a0, .LBB1_2
-; LA32-NEXT:  # %bb.1: # %bb1
-; LA32-NEXT:    st.b $zero, $a1, 0
-; LA32-NEXT:    ret
-; LA32-NEXT:  .LBB1_2: # %bb2
-; LA32-NEXT:    ori $a0, $zero, 1
-; LA32-NEXT:    st.b $a0, $a1, 0
-; LA32-NEXT:    ret
-;
-; LA64-LABEL: extract_8xi16_br_ne_0:
-; LA64:       # %bb.0: # %entry
-; LA64-NEXT:    vld $vr0, $a0, 0
-; LA64-NEXT:    vpickve2gr.h $a0, $vr0, 0
-; LA64-NEXT:    slli.d $a0, $a0, 48
-; LA64-NEXT:    beqz $a0, .LBB1_2
-; LA64-NEXT:  # %bb.1: # %bb1
-; LA64-NEXT:    st.b $zero, $a1, 0
-; LA64-NEXT:    ret
-; LA64-NEXT:  .LBB1_2: # %bb2
-; LA64-NEXT:    ori $a0, $zero, 1
-; LA64-NEXT:    st.b $a0, $a1, 0
-; LA64-NEXT:    ret
+; CHECK-LABEL: extract_8xi16_br_ne_0:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vld $vr0, $a0, 0
+; CHECK-NEXT:    vpickve2gr.h $a0, $vr0, 0
+; CHECK-NEXT:    beqz $a0, .LBB1_2
+; CHECK-NEXT:  # %bb.1: # %bb1
+; CHECK-NEXT:    st.b $zero, $a1, 0
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB1_2: # %bb2
+; CHECK-NEXT:    ori $a0, $zero, 1
+; CHECK-NEXT:    st.b $a0, $a1, 0
+; CHECK-NEXT:    ret
 entry:
   %0 = load volatile <8 x i16>, ptr %src
   %1 = extractelement <8 x i16> %0, i64 0