[AArch64] Use neon instructions for i64/i128 ISD::PARITY calculation

[RKSimon]

Summary

As noticed on D129765 and reported on Issue #56531 - aarch64 targets can use the neon ctpop + add-reduce instructions to speed up scalar ctpop instructions, but we fail to do this for parity calculations.

I'm not sure where the cutoff should be, but i64 (+ i128 special case) shows a definite reduction in instruction count. i32 is about the same (not sure if scalar <-> neon transfers are particularly costly?), and sub-i32 promotion looks to be a definite regression compared to parity expansion optimized for those widths.

[deadalnix]

It took me a bit a grave digging to figure out the motivation behind PARITY, but this seems to do the right thing.

That being said, shouldn't this be the default strategy to lower PARITY, rather than special case it for AArch64?

[dmgreen]

I've been trying to add up latencies to see which is better between then two sequences. I think you are right about i32 case - it is better to avoid the fpr register moves.

The code changes looks good to me. I was just not sure which is better between the i64 eor's and moving to float regs to use a cnt. It will depend on the cpu - but an eor is either a quick 1 cycle instruction, which is hard to beat with neon instructions, or it is a 2 cycle instruction and the cnt; addv and fmov's will have longer latencies.

I ended up having to get a simulator out to measure the differences. Whilst it is slower on some cpus, it seems to be quicker in more cases and by more of a margin. So looks OK to me.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

7790

Formatting.

[RKSimon]

Thanks!

llvm/lib/Target/AArch64/AArch64ISelLowering.h

private:
   SDValue LowerVECTOR_SPLICE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerEXTRACT_SUBVECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerINSERT_SUBVECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerDIV(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerMUL(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVectorSRA_SRL_SHL(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerShiftParts(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVSETCC(SDValue Op, SelectionDAG &DAG) const;
-  SDValue LowerCTPOP(SDValue Op, SelectionDAG &DAG) const;
+  SDValue LowerCTPOP_PARITY(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerCTTZ(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerBitreverse(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerMinMax(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerFCOPYSIGN(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerFP_EXTEND(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerFP_ROUND(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVectorFP_TO_INT(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVectorFP_TO_INT_SAT(SDValue Op, SelectionDAG &DAG) const;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
   // AArch64 doesn't have {U|S}MUL_LOHI.
   setOperationAction(ISD::UMUL_LOHI, MVT::i64, Expand);
   setOperationAction(ISD::SMUL_LOHI, MVT::i64, Expand);
 
   setOperationAction(ISD::CTPOP, MVT::i32, Custom);
   setOperationAction(ISD::CTPOP, MVT::i64, Custom);
   setOperationAction(ISD::CTPOP, MVT::i128, Custom);
 
+  setOperationAction(ISD::PARITY, MVT::i64, Custom);
+  setOperationAction(ISD::PARITY, MVT::i128, Custom);
+
   setOperationAction(ISD::ABS, MVT::i32, Custom);
   setOperationAction(ISD::ABS, MVT::i64, Custom);
 
   setOperationAction(ISD::SDIVREM, MVT::i32, Expand);
   setOperationAction(ISD::SDIVREM, MVT::i64, Expand);
   for (MVT VT : MVT::fixedlen_vector_valuetypes()) {
     setOperationAction(ISD::SDIVREM, VT, Expand);
     setOperationAction(ISD::UDIVREM, VT, Expand);
SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
   case ISD::SRL:
   case ISD::SHL:
     return LowerVectorSRA_SRL_SHL(Op, DAG);
   case ISD::SHL_PARTS:
   case ISD::SRL_PARTS:
   case ISD::SRA_PARTS:
     return LowerShiftParts(Op, DAG);
   case ISD::CTPOP:
-    return LowerCTPOP(Op, DAG);
+  case ISD::PARITY:
+    return LowerCTPOP_PARITY(Op, DAG);
   case ISD::FCOPYSIGN:
     return LowerFCOPYSIGN(Op, DAG);
   case ISD::OR:
     return LowerVectorOR(Op, DAG);
   case ISD::XOR:
     return LowerXOR(Op, DAG);
   case ISD::PREFETCH:
     return LowerPREFETCH(Op, DAG);
SDValue AArch64TargetLowering::LowerFCOPYSIGN(SDValue Op,
   if (VT == MVT::f32)
     return DAG.getTargetExtractSubreg(AArch64::ssub, DL, VT, BSP);
   if (VT == MVT::f64)
     return DAG.getTargetExtractSubreg(AArch64::dsub, DL, VT, BSP);
 
   return BitCast(VT, BSP, DAG);
 }
 
-SDValue AArch64TargetLowering::LowerCTPOP(SDValue Op, SelectionDAG &DAG) const {
+SDValue AArch64TargetLowering::LowerCTPOP_PARITY(SDValue Op, SelectionDAG &DAG) const {
   if (DAG.getMachineFunction().getFunction().hasFnAttribute(
           Attribute::NoImplicitFloat))
     return SDValue();
 
   if (!Subtarget->hasNEON())
     return SDValue();
 
+  bool IsParity = Op.getOpcode() == ISD::PARITY;
+
   // While there is no integer popcount instruction, it can
   // be more efficiently lowered to the following sequence that uses
   // AdvSIMD registers/instructions as long as the copies to/from
   // the AdvSIMD registers are cheap.
   //  FMOV    D0, X0        // copy 64-bit int to vector, high bits zero'd
   //  CNT     V0.8B, V0.8B  // 8xbyte pop-counts
   //  ADDV    B0, V0.8B     // sum 8xbyte pop-counts
   //  UMOV    X0, V0.B[0]   // copy byte result back to integer reg
   SDValue Val = Op.getOperand(0);
   SDLoc DL(Op);
   EVT VT = Op.getValueType();
 
   if (VT == MVT::i32 || VT == MVT::i64) {
     if (VT == MVT::i32)
       Val = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i64, Val);
     Val = DAG.getNode(ISD::BITCAST, DL, MVT::v8i8, Val);
 
     SDValue CtPop = DAG.getNode(ISD::CTPOP, DL, MVT::v8i8, Val);
     SDValue UaddLV = DAG.getNode(
         ISD::INTRINSIC_WO_CHAIN, DL, MVT::i32,
         DAG.getConstant(Intrinsic::aarch64_neon_uaddlv, DL, MVT::i32), CtPop);
 
+    if (IsParity)
+      UaddLV = DAG.getNode(ISD::AND, DL, MVT::i32, UaddLV,
+                           DAG.getConstant(1, DL, MVT::i32));
+
     if (VT == MVT::i64)
       UaddLV = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i64, UaddLV);
     return UaddLV;
   } else if (VT == MVT::i128) {
     Val = DAG.getNode(ISD::BITCAST, DL, MVT::v16i8, Val);
 
     SDValue CtPop = DAG.getNode(ISD::CTPOP, DL, MVT::v16i8, Val);
     SDValue UaddLV = DAG.getNode(
         ISD::INTRINSIC_WO_CHAIN, DL, MVT::i32,
         DAG.getConstant(Intrinsic::aarch64_neon_uaddlv, DL, MVT::i32), CtPop);
 
+    if (IsParity)
+      UaddLV = DAG.getNode(ISD::AND, DL, MVT::i32, UaddLV,
+                           DAG.getConstant(1, DL, MVT::i32));
+
     return DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i128, UaddLV);
   }
 
+  assert(!IsParity && "ISD::PARITY of vector types not supported");
+
   if (VT.isScalableVector() || useSVEForFixedLengthVectorVT(VT))
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::CTPOP_MERGE_PASSTHRU);
 
   assert((VT == MVT::v1i64 || VT == MVT::v2i64 || VT == MVT::v2i32 ||
           VT == MVT::v4i32 || VT == MVT::v4i16 || VT == MVT::v8i16) &&
          "Unexpected type for custom ctpop lowering");
 
   EVT VT8Bit = VT.is64BitVector() ? MVT::v8i8 : MVT::v16i8;
case ISD::VECREDUCE_UMIN:
     Results.push_back(LowerVECREDUCE(SDValue(N, 0), DAG));
     return;
   case ISD::ADD:
   case ISD::FADD:
     ReplaceAddWithADDP(N, Results, DAG, Subtarget);
     return;
 
   case ISD::CTPOP:
-    if (SDValue Result = LowerCTPOP(SDValue(N, 0), DAG))
+  case ISD::PARITY:
+    if (SDValue Result = LowerCTPOP_PARITY(SDValue(N, 0), DAG))
       Results.push_back(Result);
     return;
   case AArch64ISD::SADDV:
     ReplaceReductionResults(N, Results, DAG, ISD::ADD, AArch64ISD::SADDV);
     return;
   case AArch64ISD::UADDV:
     ReplaceReductionResults(N, Results, DAG, ISD::ADD, AArch64ISD::UADDV);
     return;

llvm/test/CodeGen/AArch64/parity.ll

; CHECK-NEXT:    ret
   %1 = tail call i32 @llvm.ctpop.i32(i32 %x)
   %2 = and i32 %1, 1
   ret i32 %2
 }
 
 define i64 @parity_64(i64 %x) {
 ; CHECK-LABEL: parity_64:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    eor x8, x0, x0, lsr #32
-; CHECK-NEXT:    eor x8, x8, x8, lsr #16
-; CHECK-NEXT:    eor x8, x8, x8, lsr #8
-; CHECK-NEXT:    eor x8, x8, x8, lsr #4
-; CHECK-NEXT:    eor x8, x8, x8, lsr #2
-; CHECK-NEXT:    eor w8, w8, w8, lsr #1
-; CHECK-NEXT:    and x0, x8, #0x1
+; CHECK-NEXT:    fmov d0, x0
+; CHECK-NEXT:    cnt v0.8b, v0.8b
+; CHECK-NEXT:    uaddlv h0, v0.8b
+; CHECK-NEXT:    fmov w8, s0
+; CHECK-NEXT:    and w0, w8, #0x1
 ; CHECK-NEXT:    ret
   %1 = tail call i64 @llvm.ctpop.i64(i64 %x)
   %2 = and i64 %1, 1
   ret i64 %2
 }
 
 define i128 @parity_128(i128 %x) {
 ; CHECK-LABEL: parity_128:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    eor x8, x0, x1
+; CHECK-NEXT:    fmov d0, x0
+; CHECK-NEXT:    mov v0.d[1], x1
 ; CHECK-NEXT:    mov x1, xzr
-; CHECK-NEXT:    eor x8, x8, x8, lsr #32
-; CHECK-NEXT:    eor x8, x8, x8, lsr #16
-; CHECK-NEXT:    eor x8, x8, x8, lsr #8
-; CHECK-NEXT:    eor x8, x8, x8, lsr #4
-; CHECK-NEXT:    eor x8, x8, x8, lsr #2
-; CHECK-NEXT:    eor w8, w8, w8, lsr #1
-; CHECK-NEXT:    and x0, x8, #0x1
+; CHECK-NEXT:    cnt v0.16b, v0.16b
+; CHECK-NEXT:    uaddlv h0, v0.16b
+; CHECK-NEXT:    fmov w8, s0
+; CHECK-NEXT:    and w0, w8, #0x1
 ; CHECK-NEXT:    ret
   %1 = tail call i128 @llvm.ctpop.i128(i128 %x)
   %2 = and i128 %1, 1
   ret i128 %2
 }
 
 define i32 @parity_64_trunc(i64 %x) {
 ; CHECK-LABEL: parity_64_trunc:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    eor x8, x0, x0, lsr #32
-; CHECK-NEXT:    eor x8, x8, x8, lsr #16
-; CHECK-NEXT:    eor x8, x8, x8, lsr #8
-; CHECK-NEXT:    eor x8, x8, x8, lsr #4
-; CHECK-NEXT:    eor x8, x8, x8, lsr #2
-; CHECK-NEXT:    eor w8, w8, w8, lsr #1
+; CHECK-NEXT:    fmov d0, x0
+; CHECK-NEXT:    cnt v0.8b, v0.8b
+; CHECK-NEXT:    uaddlv h0, v0.8b
+; CHECK-NEXT:    fmov w8, s0
 ; CHECK-NEXT:    and w0, w8, #0x1
 ; CHECK-NEXT:    ret
   %1 = tail call i64 @llvm.ctpop.i64(i64 %x)
   %2 = trunc i64 %1 to i32
   %3 = and i32 %2, 1
   ret i32 %3
 }