[SVE] Implement extractelement of i1 from scalable vectors.

llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp

case ISD::BUILD_VECTOR: {
     SDValue CvtVec = DAG.getNode(ISD::BITCAST, SL, OVT, Concat);
     Results.push_back(CvtVec);
     break;
   }
   case ISD::EXTRACT_VECTOR_ELT: {
     MVT EltVT = OVT.getVectorElementType();
     MVT NewEltVT = NVT.getVectorElementType();
 
+    SDValue Vec = Node->getOperand(0);
+    SDValue Idx = Node->getOperand(1);
+    SDLoc SL(Node);
+
+    if (OVT.getSizeInBits() != NVT.getSizeInBits()) {
+      assert(NVT.isVector() &&
+             OVT.getVectorElementCount() == NVT.getVectorElementCount() &&
+             "Invalid promote type for extract_vector_elt!");
+      assert(NewEltVT.bitsGT(EltVT) && "Cannot promote to a smaller type!");
+
+      // Ensure the extract VT is not smaller than its operand element type.
+      EVT ResVT = Node->getValueType(0);
+      EVT ExtactVT = ResVT.bitsGT(NewEltVT) ? ResVT : NewEltVT;
+
+      // Perform the extract using a promoted input vector.
+      SDValue V = DAG.getNode(ISD::ANY_EXTEND, SL, NVT, Vec);
+      SDValue Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, ExtactVT, V, Idx);
+
+      Results.push_back(DAG.getNode(ISD::TRUNCATE, SL, ResVT, Elt));
+      break;
+    }
+
     // Handle bitcasts to a different vector type with the same total bit size.
     //
     // e.g. v2i64 = extract_vector_elt x:v2i64, y:i32
     //  =>
     //  v4i32:castx = bitcast x:v2i64
     //
     // i64 = bitcast
     //   (v2i32 build_vector (i32 (extract_vector_elt castx, (2 * y))),
     //                       (i32 (extract_vector_elt castx, (2 * y + 1)))
     //
 
     assert(NVT.isVector() && OVT.getSizeInBits() == NVT.getSizeInBits() &&
            "Invalid promote type for extract_vector_elt");
     assert(NewEltVT.bitsLT(EltVT) && "not handled");
 
     MVT MidVT = getPromotedVectorElementType(TLI, EltVT, NewEltVT);
     unsigned NewEltsPerOldElt = MidVT.getVectorNumElements();
 
-    SDValue Idx = Node->getOperand(1);
     EVT IdxVT = Idx.getValueType();
-    SDLoc SL(Node);
     SDValue Factor = DAG.getConstant(NewEltsPerOldElt, SL, IdxVT);
     SDValue NewBaseIdx = DAG.getNode(ISD::MUL, SL, IdxVT, Idx, Factor);
-
-    SDValue CastVec = DAG.getNode(ISD::BITCAST, SL, NVT, Node->getOperand(0));
+    SDValue CastVec = DAG.getNode(ISD::BITCAST, SL, NVT, Vec);
 
     SmallVector<SDValue, 8> NewOps;
     for (unsigned I = 0; I < NewEltsPerOldElt; ++I) {
       SDValue IdxOffset = DAG.getConstant(I, SL, IdxVT);
       SDValue TmpIdx = DAG.getNode(ISD::ADD, SL, IdxVT, NewBaseIdx, IdxOffset);
 
       SDValue Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, SL, NewEltVT,
                                 CastVec, TmpIdx);

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

case ISD::EXTRACT_ELEMENT: {
     // Remove high part of known bit mask
     Known = Known.trunc(EltBitWidth);
     break;
   }
   case ISD::EXTRACT_VECTOR_ELT: {
     SDValue InVec = Op.getOperand(0);
     SDValue EltNo = Op.getOperand(1);
     EVT VecVT = InVec.getValueType();
+    // computeKnownBits not yet impemented for scalable vectors.
+    if (VecVT.isScalableVector())
+      break;
     const unsigned EltBitWidth = VecVT.getScalarSizeInBits();
     const unsigned NumSrcElts = VecVT.getVectorNumElements();
 
     // If BitWidth > EltBitWidth the value is anyext:ed. So we do not know
     // anything about the extended bits.
     if (BitWidth > EltBitWidth)
       Known = Known.trunc(EltBitWidth);
 

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Subtarget->hasSVE()) {
 
     for (auto VT : {MVT::nxv16i1, MVT::nxv8i1, MVT::nxv4i1, MVT::nxv2i1}) {
       setOperationAction(ISD::CONCAT_VECTORS, VT, Custom);
       setOperationAction(ISD::SELECT, VT, Custom);
       setOperationAction(ISD::SETCC, VT, Custom);
       setOperationAction(ISD::SPLAT_VECTOR, VT, Custom);
       setOperationAction(ISD::TRUNCATE, VT, Custom);
 
+      MVT PromotedVT = getPromotedVTForPredicate(VT);
+      setOperationAction(ISD::EXTRACT_VECTOR_ELT, VT, Promote);
+      AddPromotedToType(ISD::EXTRACT_VECTOR_ELT, VT, PromotedVT);
+
       // There are no legal MVT::nxv16f## based types.
       if (VT != MVT::nxv16i1) {
         setOperationAction(ISD::SINT_TO_FP, VT, Promote);
-        AddPromotedToType(ISD::SINT_TO_FP, VT, getPromotedVTForPredicate(VT));
+        AddPromotedToType(ISD::SINT_TO_FP, VT, PromotedVT);
         setOperationAction(ISD::UINT_TO_FP, VT, Promote);
-        AddPromotedToType(ISD::UINT_TO_FP, VT, getPromotedVTForPredicate(VT));
+        AddPromotedToType(ISD::UINT_TO_FP, VT, PromotedVT);
       }
     }
 
     for (auto VT : {MVT::nxv2f16, MVT::nxv4f16, MVT::nxv8f16, MVT::nxv2f32,
                     MVT::nxv4f32, MVT::nxv2f64}) {
       setOperationAction(ISD::CONCAT_VECTORS, VT, Custom);
       setOperationAction(ISD::INSERT_SUBVECTOR, VT, Custom);
       setOperationAction(ISD::SPLAT_VECTOR, VT, Custom);

llvm/test/CodeGen/AArch64/sve-extract-element.ll

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ret
   %a = insertelement <vscale x 2 x i64> undef, i64 %x, i32 0
   %b = shufflevector <vscale x 2 x i64> %a, <vscale x 2 x i64> undef, <vscale x 2 x i32> zeroinitializer
   %c = extractelement <vscale x 2 x i64> %b, i32 %y
   ret i64 %c
 }
 
+define i1 @test_lane0_16xi1(<vscale x 16 x i1> %a) #0 {
+; CHECK-LABEL: test_lane0_16xi1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.b, p0/z, #1 // =0x1
+; CHECK-NEXT:    fmov w8, s0
+; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    ret
+  %b = extractelement <vscale x 16 x i1> %a, i32 0
+  ret i1 %b
+}
+
+define i1 @test_lane9_8xi1(<vscale x 8 x i1> %a) #0 {
+; CHECK-LABEL: test_lane9_8xi1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.h, p0/z, #1 // =0x1
+; CHECK-NEXT:    mov z0.h, z0.h[9]
+; CHECK-NEXT:    fmov w8, s0
+; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    ret
+  %b = extractelement <vscale x 8 x i1> %a, i32 9
+  ret i1 %b
+}
+
+define i1 @test_lanex_4xi1(<vscale x 4 x i1> %a, i32 %x) #0 {
+; CHECK-LABEL: test_lanex_4xi1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    // kill: def $w0 killed $w0 def $x0
+; CHECK-NEXT:    sxtw x8, w0
+; CHECK-NEXT:    whilels p1.s, xzr, x8
+; CHECK-NEXT:    mov z0.s, p0/z, #1 // =0x1
+; CHECK-NEXT:    lastb w8, p1, z0.s
+; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    ret
+  %b = extractelement <vscale x 4 x i1> %a, i32 %x
+  ret i1 %b
+}
+
+define i1 @test_lane4_2xi1(<vscale x 2 x i1> %a) #0 {
+; CHECK-LABEL: test_lane4_2xi1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.d, p0/z, #1 // =0x1
+; CHECK-NEXT:    mov z0.d, z0.d[4]
+; CHECK-NEXT:    fmov x8, d0
+; CHECK-NEXT:    and w0, w8, #0x1
+; CHECK-NEXT:    ret
+  %b = extractelement <vscale x 2 x i1> %a, i32 4
+  ret i1 %b
+}
+
 attributes #0 = { "target-features"="+sve" }