[DAGCombiner] Fix issue #121372

PR #118823 added a DAG combine for extracting elements of a
vector returned from SETCC, however it doesn't correctly deal
with the case where the vector element type is not i1. In
this case we have to take account of the boolean contents,
which are represent differently between vectors and scalars.
The code now explicitly performs an inreg sign extend in
order to get the same result.

Fixes #121372

Author: david-arm

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -22808,9 +22808,26 @@ static SDValue scalarizeExtractedBinOp(SDNode *ExtElt, SelectionDAG &DAG,
     return SDValue();
 
   EVT ResVT = ExtElt->getValueType(0);
-  if (Opc == ISD::SETCC &&
-      (ResVT != Vec.getValueType().getVectorElementType() || LegalTypes))
-    return SDValue();
+  bool SetCCNeedsSignExt = false;
+  if (Opc == ISD::SETCC) {
+    EVT VecVT = Vec.getValueType();
+    if (ResVT != VecVT.getVectorElementType() || LegalTypes)
+      return SDValue();
+
+    if (ResVT != MVT::i1) {
+      bool VecRequiresSignExt = TLI.getBooleanContents(VecVT) ==
+                                TargetLowering::ZeroOrNegativeOneBooleanContent;
+      bool ScalarRequiresSignExt =
+          TLI.getBooleanContents(ResVT) ==
+          TargetLowering::ZeroOrNegativeOneBooleanContent;
+      if (VecRequiresSignExt && !ScalarRequiresSignExt)
+        SetCCNeedsSignExt = true;
+      else if (!VecRequiresSignExt && ScalarRequiresSignExt) {
+        // There are currently no targets with this behaviour.
+        return SDValue();
+      }
+    }
+  }
 
   // Targets may want to avoid this to prevent an expensive register transfer.
   if (!TLI.shouldScalarizeBinop(Vec))
@@ -22834,8 +22851,14 @@ static SDValue scalarizeExtractedBinOp(SDNode *ExtElt, SelectionDAG &DAG,
     EVT OpVT = Op0.getValueType().getVectorElementType();
     Op0 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, OpVT, Op0, Index);
     Op1 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, OpVT, Op1, Index);
-    return DAG.getSetCC(DL, ResVT, Op0, Op1,
-                        cast<CondCodeSDNode>(Vec->getOperand(2))->get());
+    SDValue NewVal = DAG.getSetCC(
+        DL, ResVT, Op0, Op1, cast<CondCodeSDNode>(Vec->getOperand(2))->get());
+    // We may need to sign-extend the result to match the same behaviour as the
+    // vector version of SETCC.
+    if (SetCCNeedsSignExt)
+      NewVal = DAG.getNode(ISD::SIGN_EXTEND_INREG, DL, ResVT, NewVal,
+                           DAG.getValueType(MVT::i1));
+    return NewVal;
   }
   Op0 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, ResVT, Op0, Index);
   Op1 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, ResVT, Op1, Index);

llvm/test/CodeGen/AArch64/extract-vector-cmp.ll

@@ -58,10 +58,11 @@ define i128 @extract_icmp_v1i128(ptr %p) {
 ; CHECK-LABEL: extract_icmp_v1i128:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldp x9, x8, [x0]
-; CHECK-NEXT:    mov x1, xzr
 ; CHECK-NEXT:    orr x8, x9, x8
 ; CHECK-NEXT:    cmp x8, #0
-; CHECK-NEXT:    cset w0, eq
+; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx x0, x8, #0, #1
+; CHECK-NEXT:    mov x1, x0
 ; CHECK-NEXT:    ret
   %load = load <1 x i128>, ptr %p, align 16
   %cmp = icmp eq <1 x i128> %load, zeroinitializer
@@ -141,12 +142,14 @@ for.cond.cleanup:
 }
 
 
+; TODO: Combine the sbfx(cset) into a csetm
 define i32 @issue_121372(<4 x i32> %v) {
 ; CHECK-LABEL: issue_121372:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    fmov w8, s0
 ; CHECK-NEXT:    cmp w8, #0
 ; CHECK-NEXT:    cset w8, eq
+; CHECK-NEXT:    sbfx w8, w8, #0, #1
 ; CHECK-NEXT:    cmp w8, #1
 ; CHECK-NEXT:    csetm w0, lt
 ; CHECK-NEXT:    ret