[AArch64][SME] Add support for conditional select for svcount.

[sdesmalen]

Summary

This patch adds codegen support to match a conditional select with svcount
operands.

It also fixes an issue in InstCombineSelect.cpp where it tried to query the
(scalable) size of an aarch64_svcount type.

llvm/docs/LangRef.rst

 held in a AArch64 SVE predicate register. Such a value contains information
 about the number of active lanes, the element width and a bit that tells whether
 the generated mask should be inverted. Target-specific intrinsics should be used
 to move the aarch64_svcount value to/from a predicate vector.
 
 There are certain limitations on the type:
 * The type can be used for function parameters and return values.
 * The supported LLVM operations on this type are strictly limited to ``load``,
-  ``store``, ``phi`` and ``alloca`` instructions.
+  ``store``, ``phi``, ``select`` and ``alloca`` instructions.
 
 The predicate-as-counter type is a scalable type.
 
 :Syntax:
 
 ::
 
       aarch64_svcount

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Subtarget->useSVEForFixedLengthVectors()) {
       for (MVT VT : MVT::fp_fixedlen_vector_valuetypes())
         if (useSVEForFixedLengthVectorVT(VT))
           addRegisterClass(VT, &AArch64::ZPRRegClass);
     }
   }
 
   if (Subtarget->hasSVE2p1() || Subtarget->hasSME2()) {
     addRegisterClass(MVT::aarch64svcount, &AArch64::PPRRegClass);
+    setOperationAction(ISD::SELECT, MVT::aarch64svcount, Custom);
+    setOperationAction(ISD::SELECT_CC, MVT::aarch64svcount, Expand);
   }
 
   // Compute derived properties from the register classes
   computeRegisterProperties(Subtarget->getRegisterInfo());
 
   // Provide all sorts of operation actions
   setOperationAction(ISD::GlobalAddress, MVT::i64, Custom);
   setOperationAction(ISD::GlobalTLSAddress, MVT::i64, Custom);

 SDValue AArch64TargetLowering::LowerSELECT(SDValue Op,
                                            SelectionDAG &DAG) const {
   SDValue CCVal = Op->getOperand(0);
   SDValue TVal = Op->getOperand(1);
   SDValue FVal = Op->getOperand(2);
   SDLoc DL(Op);
 
   EVT Ty = Op.getValueType();
+  if (Ty == MVT::aarch64svcount) {
+    TVal = DAG.getNode(AArch64ISD::REINTERPRET_CAST, DL, MVT::nxv16i1, TVal);
+    FVal = DAG.getNode(AArch64ISD::REINTERPRET_CAST, DL, MVT::nxv16i1, FVal);
+    EVT CCVT = CCVal.getValueType();
+    SDValue ID =
+        DAG.getTargetConstant(Intrinsic::aarch64_sve_whilelo, DL, CCVT);
+    SDValue Zero = DAG.getConstant(0, DL, CCVT);
+    SDValue SplatVal = DAG.getNode(ISD::SIGN_EXTEND_INREG, DL, CCVT, CCVal,
+                                   DAG.getValueType(MVT::i1));
+    SDValue SplatPred =
+        DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL, MVT::nxv16i1, ID, Zero, SplatVal);
+    SDValue Sel =
+        DAG.getNode(ISD::VSELECT, DL, MVT::nxv16i1, SplatPred, TVal, FVal);
+    return DAG.getNode(AArch64ISD::REINTERPRET_CAST, DL, Ty, Sel);
+  }
+
   if (Ty.isScalableVector()) {
     SDValue TruncCC = DAG.getNode(ISD::TRUNCATE, DL, MVT::i1, CCVal);
     MVT PredVT = MVT::getVectorVT(MVT::i1, Ty.getVectorElementCount());
     SDValue SplatPred = DAG.getNode(ISD::SPLAT_VECTOR, DL, PredVT, TruncCC);
     return DAG.getNode(ISD::VSELECT, DL, Ty, SplatPred, TVal, FVal);
   }
 
   if (useSVEForFixedLengthVectorVT(Ty)) {
static SDValue performSelectCombine(SDNode *N,
                                     TargetLowering::DAGCombinerInfo &DCI) {
   SelectionDAG &DAG = DCI.DAG;
   SDValue N0 = N->getOperand(0);
   EVT ResVT = N->getValueType(0);
 
   if (N0.getOpcode() != ISD::SETCC)
     return SDValue();
 
-  if (ResVT.isScalableVector())
+  if (ResVT.isScalableVector() || ResVT == MVT::aarch64svcount)
     return SDValue();
 
   // Make sure the SETCC result is either i1 (initial DAG), or i32, the lowered
   // scalar SetCCResultType. We also don't expect vectors, because we assume
   // that selects fed by vector SETCCs are canonicalized to VSELECT.
   assert((N0.getValueType() == MVT::i1 || N0.getValueType() == MVT::i32) &&
          "Scalar-SETCC feeding SELECT has unexpected result type!");
 

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

let Predicates = [IsLE] in {
     def : Pat<(nxv8i16 (bitconvert (nxv8bf16 ZPR:$src))), (nxv8i16 ZPR:$src)>;
     def : Pat<(nxv4i32 (bitconvert (nxv8bf16 ZPR:$src))), (nxv4i32 ZPR:$src)>;
     def : Pat<(nxv2i64 (bitconvert (nxv8bf16 ZPR:$src))), (nxv2i64 ZPR:$src)>;
     def : Pat<(nxv8f16 (bitconvert (nxv8bf16 ZPR:$src))), (nxv8f16 ZPR:$src)>;
     def : Pat<(nxv4f32 (bitconvert (nxv8bf16 ZPR:$src))), (nxv4f32 ZPR:$src)>;
     def : Pat<(nxv2f64 (bitconvert (nxv8bf16 ZPR:$src))), (nxv2f64 ZPR:$src)>;
   }
 
+  // These allow casting from/to the opaque aarch64svcount type.
+  def : Pat<(aarch64svcount (reinterpret_cast (nxv16i1 PPR:$src))), (COPY_TO_REGCLASS PPR:$src, PPR)>;
+  def : Pat<(nxv16i1 (reinterpret_cast (aarch64svcount PPR:$src))), (COPY_TO_REGCLASS PPR:$src, PPR)>;
+
   // These allow casting from/to unpacked predicate types.
   def : Pat<(nxv16i1 (reinterpret_cast (nxv16i1 PPR:$src))), (COPY_TO_REGCLASS PPR:$src, PPR)>;
   def : Pat<(nxv16i1 (reinterpret_cast (nxv8i1 PPR:$src))), (COPY_TO_REGCLASS PPR:$src, PPR)>;
   def : Pat<(nxv16i1 (reinterpret_cast (nxv4i1 PPR:$src))), (COPY_TO_REGCLASS PPR:$src, PPR)>;
   def : Pat<(nxv16i1 (reinterpret_cast (nxv2i1 PPR:$src))), (COPY_TO_REGCLASS PPR:$src, PPR)>;
   def : Pat<(nxv16i1 (reinterpret_cast (nxv1i1 PPR:$src))), (COPY_TO_REGCLASS PPR:$src, PPR)>;
   def : Pat<(nxv8i1 (reinterpret_cast (nxv16i1 PPR:$src))), (COPY_TO_REGCLASS PPR:$src, PPR)>;
   def : Pat<(nxv8i1 (reinterpret_cast  (nxv4i1 PPR:$src))), (COPY_TO_REGCLASS PPR:$src, PPR)>;

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

if (Pred == ICmpInst::ICMP_SGT && match(CmpRHS, m_AllOnes()) &&
     Builder.SetInsertPoint(&SI);
     Value *IsNeg = Builder.CreateIsNeg(CmpLHS, ICI->getName());
     replaceOperand(SI, 0, IsNeg);
     SI.swapValues();
     SI.swapProfMetadata();
     return &SI;
   }
 
+  // The code below doesn't work on aarch64_svcount.
+  if (TrueVal->getType()->isAArch64SvcountTy())
+    return nullptr;
+
   // FIXME: This code is nearly duplicated in InstSimplify. Using/refactoring
   // decomposeBitTestICmp() might help.
   {
     unsigned BitWidth =
         DL.getTypeSizeInBits(TrueVal->getType()->getScalarType());
     APInt MinSignedValue = APInt::getSignedMinValue(BitWidth);
     Value *X;
     const APInt *Y, *C;

llvm/test/CodeGen/AArch64/sme-aarch64-svcount.ll

 ; CHECKO3-NEXT:    str p0, [sp, #7, mul vl]
 ; CHECKO3-NEXT:    bl take_svcount_5
 ; CHECKO3-NEXT:    addvl sp, sp, #1
 ; CHECKO3-NEXT:    ldp x29, x30, [sp], #16 // 16-byte Folded Reload
 ; CHECKO3-NEXT:    ret
   call void @take_svcount_5(aarch64_svcount %arg, aarch64_svcount %arg, aarch64_svcount %arg, aarch64_svcount %arg, aarch64_svcount %arg)
   ret void
 }
+
+; Test code-generation of aarch64svcount being used in a select to support e.g.
+;   return k < 42 ? y : z;
+; where y and z are of types svcount_t.
+define aarch64_svcount @foo(i32 %k, aarch64_svcount %x, aarch64_svcount %y) {
+; CHECKO0-LABEL: foo:
+; CHECKO0:       // %bb.0:
+; CHECKO0-NEXT:    mov p2.b, p1.b
+; CHECKO0-NEXT:    mov p1.b, p0.b
+; CHECKO0-NEXT:    subs w8, w0, #42
+; CHECKO0-NEXT:    cset w8, lt
+; CHECKO0-NEXT:    sbfx w9, w8, #0, #1
+; CHECKO0-NEXT:    mov w8, wzr
+; CHECKO0-NEXT:    whilelo p0.b, w8, w9
+; CHECKO0-NEXT:    sel p0.b, p0, p1.b, p2.b
+; CHECKO0-NEXT:    ret
+;
+; CHECKO3-LABEL: foo:
+; CHECKO3:       // %bb.0:
+; CHECKO3-NEXT:    cmp w0, #42
+; CHECKO3-NEXT:    cset w8, lt
+; CHECKO3-NEXT:    sbfx w8, w8, #0, #1
+; CHECKO3-NEXT:    whilelo p2.b, wzr, w8
+; CHECKO3-NEXT:    sel p0.b, p2, p0.b, p1.b
+; CHECKO3-NEXT:    ret
+  %cmp = icmp slt i32 %k, 42
+  %x.y = select i1 %cmp, aarch64_svcount %x, aarch64_svcount %y
+  ret aarch64_svcount %x.y
+}

llvm/test/Transforms/InstCombine/AArch64/sme-svcount.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -instcombine -S < %s | FileCheck %s
 
 define aarch64_svcount @test_alloca_store_reload(aarch64_svcount %val) nounwind {
 ; CHECK-LABEL: @test_alloca_store_reload(
 ; CHECK-NEXT:    ret aarch64_svcount [[VAL:%.*]]
 ;
   %ptr = alloca aarch64_svcount, align 1
   store aarch64_svcount %val, ptr %ptr
   %res = load aarch64_svcount, ptr %ptr
   ret aarch64_svcount %res
 }
+
+; Test that instcombine doesn't try to query the (scalable) size of aarch64_svcount
+; in foldSelectInstWithICmp.
+define aarch64_svcount @test_combine_on_select(aarch64_svcount %x, aarch64_svcount %y, i32 %k) {
+  %cmp = icmp sgt i32 %k, 42
+  %x.y = select i1 %cmp, aarch64_svcount %x, aarch64_svcount %y
+  ret aarch64_svcount %x.y
+}