[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

9290–9296

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

9290–9296

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

9290–9296

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

3323

typo "implemented"

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9280–9289

You can use getPromotedVTForPredicate here.

9292

Can ANY_EXTEND be used here?

9296

As above can this be getAnyExtOrTrunc?

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

490

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

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);
+      setOperationAction(ISD::EXTRACT_VECTOR_ELT, VT, Custom);
 
       // 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));
         setOperationAction(ISD::UINT_TO_FP, VT, Promote);
         AddPromotedToType(ISD::UINT_TO_FP, VT, getPromotedVTForPredicate(VT));
       }
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/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 // =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) #0 {
+; 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) #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 // =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) #0 {
+; 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
+}
+
 attributes #0 = { "target-features"="+sve" }