[AArch64][SVE] Implement extractelement of i1 vectors.

[efriedma]

Summary

The implementation just extends the vector to a larger element type, and extracts from that. Not fancy, but generates reasonable code.

While I'm here, fix warning from computeKnown bits.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9050–9056

This is a common route when expanding i1 based operations so downstream we did it in common code, for example:

setOperationAction(ISD::EXTRACT_VECTOR_ELT, MVT::nxv2i1, Promote);
AddPromotedToType(ISD::EXTRACT_VECTOR_ELT, MVT::nxv2i1, MVT::nxv2i64);

and updated SelectionDAGLegalize::PromoteNode accordingly. Is this route something we can do upstream? instead of having custom lowering.

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9050–9056

The problem is that "promote" is currently interpreted to mean "bitcast the vector operands" by vector legalization, not integer promotion, so it would be confusing to use it like this. Maybe worth messing with LegalizeAction to distingush between the two kinds of "promotion", though.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9050–9056

Within AArch64ISelLowering.cpp I can see:

setOperationAction(ISD::FADD, MVT::v4f16, Promote);
AddPromotedToType(ISD::FADD, MVT::v4f16, MVT::v4f32);

With SelectionDAGLegalize::PromoteNode applying the appropriate extension. I can see similar usages like this for other targets. To me it looks like "Promote" can mean promotion or bitcast, with nodes having support for whichever has been required in the past. Based on this it seems safe to extend PromoteNode's handling of EXTRACT_VECTOR_ELT to include real promotion.

[efriedma]

Rebase and clang-format.

I looked into handling this in LegalizeVectorOps, but that doesn't really work out: we actually legalize EXTRACT_VECTOR_ELT nodes in LegalizeDAG. And I don't want to do this in LegalizeDAG: LegalizeDAG can't handle arbitrary illegal vector operations.

There might be some way to sort out this mess, but it isn't trivial.

[paulwalker-arm]

Perhaps I'm missing something but I've created D89950 to show what I mentioned in my previous comment. To me it's preferable to have this as target independent code as it seems a common enough solution. I know it results in "Pomote" having multiple meanings but that boat has sailed because different nodes have already chosen different meanings and I think it's pretty clear from the old and new VTs what's being asked for.

[efriedma]

I'm worried the usage of ANY_EXTEND and TRUNCATE in D89950 is going to cause a mess in the future. LegalizeDAG doesn't know how to legalize most vector operations; LegalizeVectorOps handles them. If someone tries to use "Promote" on a target where those operations aren't legal, it'll break. And then someone will try to "fix" LegalizeDAG, and cause other problems.

Note this is specific to the vector operations we legalize in LegalizeDAG; so essentially, INSERT_VECTOR_ELT, EXTRACT_VECTOR_ELT, and VECTOR_SHUFFLE.

[paulwalker-arm]

I just don't buy this argument. You're saying that we must force all target's to perform custom lowering for promotable illegal operations on i1 vectors because somebody in the future might try to "fix" LegalizeDAG when they should implement custom lowering for extend/truncate operations. This is the situation Target/AArch64 is in and we never considered "fixing" LegalizeDAG so I don't see why others wouldn't just follow our example.

I've already pointed out:

setOperationAction(ISD::FADD, MVT::v4f16, Promote);
AddPromotedToType(ISD::FADD, MVT::v4f16, MVT::v4f32);

that exists today and is happy to emit ISD::FP_EXTEND. We mirrored this to support i1 based int->fp operations, although perhaps that's a case of two wrongs not making a right :)

Sure we can teach LegalizeVectorOps LegalizeDAG's PromoteNode abilities but that feels like wasted effort considering it just works. My understand of the two legalisers is that LegalizeVectorOps might emit scalar operations that require further legalisation, which is a situation that doesn't apply here considering it's vector in vector out.

So ultimately I don't know what to do here. There's nothing functionally wrong with the patch (other than a couple of minor suggestions) so I guess I'll begrudgingly accept the patch but I just think you're causing artificial restrictions that others have not felt the need to apply and thus hope you'll reconsider.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3311

typo "implemented"

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9040–9049

You can use getPromotedVTForPredicate here.

9052

Can ANY_EXTEND be used here?

9056

As above can this be getAnyExtOrTrunc?

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

252

Not a problem for this patch but I'm surprised to see the and here.

[paulwalker-arm]

Just making a note here that I've push D90093 so that we follow the same principle shown here for the lowering of i1 based int_to_fp conversions. This also means getPromotedVTForPredicate is now useable for this patch also.

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

for (MVT VT : MVT::integer_scalable_vector_valuetypes()) {
         setOperationAction(ISD::UMAX, VT, Custom);
         setOperationAction(ISD::SHL, VT, Custom);
         setOperationAction(ISD::SRL, VT, Custom);
         setOperationAction(ISD::SRA, VT, Custom);
         if (VT.getScalarType() == MVT::i1) {
           setOperationAction(ISD::SETCC, VT, Custom);
           setOperationAction(ISD::TRUNCATE, VT, Custom);
           setOperationAction(ISD::CONCAT_VECTORS, VT, Legal);
+          setOperationAction(ISD::EXTRACT_VECTOR_ELT, VT, Custom);
         }
       }
     }
 
     for (auto VT : {MVT::nxv8i8, MVT::nxv4i16, MVT::nxv2i32}) {
       setOperationAction(ISD::EXTRACT_SUBVECTOR, VT, Custom);
       setOperationAction(ISD::INSERT_SUBVECTOR, VT, Custom);
     }
SDValue AArch64TargetLowering::LowerINSERT_VECTOR_ELT(SDValue Op,
   // Re-narrow the resultant vector.
   return NarrowVector(Node, DAG);
 }
 
 SDValue
 AArch64TargetLowering::LowerEXTRACT_VECTOR_ELT(SDValue Op,
                                                SelectionDAG &DAG) const {
   assert(Op.getOpcode() == ISD::EXTRACT_VECTOR_ELT && "Unknown opcode!");
+  EVT VT = Op.getOperand(0).getValueType();
+
+  if (VT.getScalarType() == MVT::i1) {
+    // We can't directly extract from an SVE predicate; extend it first.
+    // (This isn't the only possible lowering, but it's straightforward.)
+    MVT ScalarVT;
+    auto Count = VT.getVectorElementCount();
+    if (Count.getKnownMinValue() == 2)
+      ScalarVT = MVT::i64;
+    else if (Count.getKnownMinValue() == 4)
+      ScalarVT = MVT::i32;
+    else if (Count.getKnownMinValue() == 8)
+      ScalarVT = MVT::i16;
+    else
+      ScalarVT = MVT::i8;
+    EVT VectorVT = EVT::getVectorVT(*DAG.getContext(), ScalarVT, Count);
+    SDLoc DL(Op);
+    SDValue Extend = DAG.getNode(ISD::SIGN_EXTEND, DL, VectorVT,
+                                 Op.getOperand(0));
+    MVT ExtractTy = ScalarVT == MVT::i64 ? MVT::i64 : MVT::i32;
+    SDValue Extract = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, ExtractTy,
+                                  Extend, Op.getOperand(1));
+    return DAG.getSExtOrTrunc(Extract, DL, Op.getValueType());
+  }
 
   // Check for non-constant or out of range lane.
-  EVT VT = Op.getOperand(0).getValueType();
   ConstantSDNode *CI = dyn_cast<ConstantSDNode>(Op.getOperand(1));
   if (!CI || CI->getZExtValue() >= VT.getVectorNumElements())
     return SDValue();
 
-
   // Insertion/extraction are legal for V128 types.
   if (VT == MVT::v16i8 || VT == MVT::v8i16 || VT == MVT::v4i32 ||
       VT == MVT::v2i64 || VT == MVT::v4f32 || VT == MVT::v2f64 ||
       VT == MVT::v8f16 || VT == MVT::v8bf16)
     return Op;
 
   if (VT != MVT::v8i8 && VT != MVT::v4i16 && VT != MVT::v2i32 &&
       VT != MVT::v1i64 && VT != MVT::v2f32 && VT != MVT::v4f16 &&

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

def : Pat<(nxv4f32 (vector_insert (nxv4f32 ZPR:$vec), (f32 FPR32:$src), GPR64:$index)),
                         $src)>;
   def : Pat<(nxv2f64 (vector_insert (nxv2f64 ZPR:$vec), (f64 FPR64:$src), GPR64:$index)),
             (CPY_ZPmV_D ZPR:$vec,
                         (CMPEQ_PPzZZ_D (PTRUE_D 31),
                                        (INDEX_II_D 0, 1),
                                        (DUP_ZR_D $index)),
                         $src)>;
 
+  // Extract element from vector with zero index
+  def : Pat<(i32 (vector_extract (nxv16i8 ZPR:$vec), 0)),
+            (EXTRACT_SUBREG ZPR:$vec, ssub)>;
+  def : Pat<(i32 (vector_extract (nxv8i16 ZPR:$vec), 0)),
+            (EXTRACT_SUBREG ZPR:$vec, ssub)>;
+  def : Pat<(i32 (vector_extract (nxv4i32 ZPR:$vec), 0)),
+            (EXTRACT_SUBREG ZPR:$vec, ssub)>;
+  def : Pat<(i64 (vector_extract (nxv2i64 ZPR:$vec), 0)),
+            (EXTRACT_SUBREG ZPR:$vec, dsub)>;
+  def : Pat<(f16 (vector_extract (nxv2i64 ZPR:$vec), 0)),
+            (EXTRACT_SUBREG ZPR:$vec, hsub)>;
+  def : Pat<(f32 (vector_extract (nxv4f32 ZPR:$vec), 0)),
+            (EXTRACT_SUBREG ZPR:$vec, ssub)>;
+  def : Pat<(f64 (vector_extract (nxv2f64 ZPR:$vec), 0)),
+            (EXTRACT_SUBREG ZPR:$vec, dsub)>;
+
   // Extract element from vector with immediate index
   def : Pat<(i32 (vector_extract (nxv16i8 ZPR:$vec), sve_elm_idx_extdup_b:$index)),
             (EXTRACT_SUBREG (DUP_ZZI_B ZPR:$vec, sve_elm_idx_extdup_b:$index), ssub)>;
   def : Pat<(i32 (vector_extract (nxv8i16 ZPR:$vec), sve_elm_idx_extdup_h:$index)),
             (EXTRACT_SUBREG (DUP_ZZI_H ZPR:$vec, sve_elm_idx_extdup_h:$index), ssub)>;
   def : Pat<(i32 (vector_extract (nxv4i32 ZPR:$vec), sve_elm_idx_extdup_s:$index)),
             (EXTRACT_SUBREG (DUP_ZZI_S ZPR:$vec, sve_elm_idx_extdup_s:$index), ssub)>;
   def : Pat<(i64 (vector_extract (nxv2i64 ZPR:$vec), sve_elm_idx_extdup_d:$index)),

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

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s 2>%t | FileCheck %s
 ; RUN: FileCheck --check-prefix=WARN --allow-empty %s <%t
 
 ; If this check fails please read test/CodeGen/AArch64/README for instructions on how to resolve it.
 ; WARN-NOT: warning
 
 define i8 @test_lane0_16xi8(<vscale x 16 x i8> %a) {
 ; CHECK-LABEL: test_lane0_16xi8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov z0.b, b0
 ; CHECK-NEXT:    fmov w0, s0
 ; CHECK-NEXT:    ret
   %b = extractelement <vscale x 16 x i8> %a, i32 0
   ret i8 %b
 }
 
 define i16 @test_lane0_8xi16(<vscale x 8 x i16> %a) {
 ; CHECK-LABEL: test_lane0_8xi16:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov z0.h, h0
 ; CHECK-NEXT:    fmov w0, s0
 ; CHECK-NEXT:    ret
   %b = extractelement <vscale x 8 x i16> %a, i32 0
   ret i16 %b
 }
 
 define i32 @test_lane0_4xi32(<vscale x 4 x i32> %a) {
 ; CHECK-LABEL: test_lane0_4xi32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov z0.s, s0
 ; CHECK-NEXT:    fmov w0, s0
 ; CHECK-NEXT:    ret
   %b = extractelement <vscale x 4 x i32> %a, i32 0
   ret i32 %b
 }
 
 define i64 @test_lane0_2xi64(<vscale x 2 x i64> %a) {
 ; CHECK-LABEL: test_lane0_2xi64:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov z0.d, d0
 ; CHECK-NEXT:    fmov x0, d0
 ; CHECK-NEXT:    ret
   %b = extractelement <vscale x 2 x i64> %a, i32 0
   ret i64 %b
 }
 
 define double @test_lane0_2xf64(<vscale x 2 x double> %a) {
 ; CHECK-LABEL: test_lane0_2xf64:
; CHECK-NEXT:    ret
   %b = extractelement <vscale x 2 x double> %a, i32 9
   ret double %b
 }
 
 ; Deliberately choose an index that is undefined
 define i32 @test_lane64_4xi32(<vscale x 4 x i32> %a) {
 ; CHECK-LABEL: test_lane64_4xi32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov z0.s, s0
 ; CHECK-NEXT:    fmov w0, s0
 ; CHECK-NEXT:    ret
   %b = extractelement <vscale x 4 x i32> %a, i32 undef
   ret i32 %b
 }
 
 define i8 @extract_of_insert_undef_16xi8(i8 %a) {
 ; CHECK-LABEL: extract_of_insert_undef_16xi8:

 ; CHECK-LABEL: test_lanex_splat_2xi64:
 ; 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) {
+; CHECK-LABEL: test_lane0_16xi1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.b, p0/z, #-1 // =0xffffffffffffffff
+; 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) {
+; CHECK-LABEL: test_lane9_8xi1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.h, p0/z, #-1 // =0xffffffffffffffff
+; 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) {
+; 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 // =0xffffffffffffffff
+; 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) {
+; CHECK-LABEL: test_lane4_2xi1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.d, p0/z, #-1 // =0xffffffffffffffff
+; 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
+}

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

 ; CHECK-NEXT:    ret
   %b = insertelement <vscale x 16 x i8> undef, i8 %a, i32 0
   ret <vscale x 16 x i8> %b
 }
 
 define <vscale x 16 x i8> @test_insert0_of_extract0_16xi8(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b) {
 ; CHECK-LABEL: test_insert0_of_extract0_16xi8:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov z1.b, b1
-; CHECK-NEXT:    ptrue p0.b, vl1
 ; CHECK-NEXT:    fmov w8, s1
+; CHECK-NEXT:    ptrue p0.b, vl1
 ; CHECK-NEXT:    mov z0.b, p0/m, w8
 ; CHECK-NEXT:    ret
   %c = extractelement <vscale x 16 x i8> %b, i32 0
   %d = insertelement <vscale x 16 x i8> %a, i8 %c, i32 0
   ret <vscale x 16 x i8> %d
 }
 
 define <vscale x 16 x i8> @test_insert64_of_extract64_16xi8(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b) {