[AArch64] Extend vecreduce -> udot handling to v8i8

[dmgreen]

Summary

https://reviews.llvm.org/D88577 added v16i8 vecreduce to udot/sdot lowering. This extends that to v8i8 too, generalizing the pattern to handle the extra types.

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

11749

Looks like this comment needs updating now

[dmgreen]

Update comment.

[dmgreen]

ping

[fhahn]

LGTM, thanks!

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static SDValue foldVectorXorShiftIntoCmp(SDNode *N, SelectionDAG &DAG,
   auto *ShiftAmt = dyn_cast<ConstantSDNode>(Shift.getOperand(1));
   EVT ShiftEltTy = Shift.getValueType().getVectorElementType();
   if (!ShiftAmt || ShiftAmt->getZExtValue() != ShiftEltTy.getSizeInBits() - 1)
     return SDValue();
 
   return DAG.getNode(AArch64ISD::CMGEz, SDLoc(N), VT, Shift.getOperand(0));
 }
 
-// VECREDUCE_ADD( EXTEND(v16i8_type) ) to
-// VECREDUCE_ADD( DOTv16i8(v16i8_type) )
+// Turn a v8i8/v16i8 extended vecreduce into a udot/sdot and vecreduce
+//   vecreduce.add(ext(A)) to vecreduce.add(DOT(zero, A, one))
 static SDValue performVecReduceAddCombine(SDNode *N, SelectionDAG &DAG,
                                           const AArch64Subtarget *ST) {
   SDValue Op0 = N->getOperand(0);
   if (!ST->hasDotProd() || N->getValueType(0) != MVT::i32)
     return SDValue();
 
   if (Op0.getValueType().getVectorElementType() != MVT::i32)
     return SDValue();
 
   unsigned ExtOpcode = Op0.getOpcode();
   if (ExtOpcode != ISD::ZERO_EXTEND && ExtOpcode != ISD::SIGN_EXTEND)
     return SDValue();
 
   EVT Op0VT = Op0.getOperand(0).getValueType();
-  if (Op0VT != MVT::v16i8)
+  if (Op0VT != MVT::v8i8 && Op0VT != MVT::v16i8)
     return SDValue();
 
   SDLoc DL(Op0);
   SDValue Ones = DAG.getConstant(1, DL, Op0VT);
-  SDValue Zeros = DAG.getConstant(0, DL, MVT::v4i32);
+  SDValue Zeros =
+      DAG.getConstant(0, DL, Op0VT == MVT::v8i8 ? MVT::v2i32 : MVT::v4i32);
   auto DotOpcode =
       (ExtOpcode == ISD::ZERO_EXTEND) ? AArch64ISD::UDOT : AArch64ISD::SDOT;
   SDValue Dot = DAG.getNode(DotOpcode, DL, Zeros.getValueType(), Zeros,
                             Ones, Op0.getOperand(0));
   return DAG.getNode(ISD::VECREDUCE_ADD, DL, N->getValueType(0), Dot);
 }
 
 // Given a ABS node, detect the following pattern:

llvm/test/CodeGen/AArch64/neon-dotreduce.ll

entry:
   %7 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %6)
   ret i32 %7
 }
 
 define i32 @test_udot_v8i8_nomla(i8* nocapture readonly %a1) {
 ; CHECK-LABEL: test_udot_v8i8_nomla:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    ldr d0, [x0]
-; CHECK-NEXT:    ushll v0.8h, v0.8b, #0
-; CHECK-NEXT:    ushll v1.4s, v0.4h, #0
-; CHECK-NEXT:    uaddw2 v0.4s, v1.4s, v0.8h
-; CHECK-NEXT:    addv s0, v0.4s
+; CHECK-NEXT:    movi v1.2d, #0000000000000000
+; CHECK-NEXT:    movi v2.8b, #1
+; CHECK-NEXT:    udot v1.2s, v2.8b, v0.8b
+; CHECK-NEXT:    addp v0.2s, v1.2s, v1.2s
 ; CHECK-NEXT:    fmov w0, s0
 ; CHECK-NEXT:    ret
 entry:
   %0 = bitcast i8* %a1 to <8 x i8>*
   %1 = load <8 x i8>, <8 x i8>* %0
   %2 = zext <8 x i8> %1 to <8 x i32>
   %3 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %2)
   ret i32 %3
entry:
   %7 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %6)
   ret i32 %7
 }
 
 define i32 @test_sdot_v8i8_nomla(i8* nocapture readonly %a1) {
 ; CHECK-LABEL: test_sdot_v8i8_nomla:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    ldr d0, [x0]
-; CHECK-NEXT:    sshll v0.8h, v0.8b, #0
-; CHECK-NEXT:    sshll v1.4s, v0.4h, #0
-; CHECK-NEXT:    saddw2 v0.4s, v1.4s, v0.8h
-; CHECK-NEXT:    addv s0, v0.4s
+; CHECK-NEXT:    movi v1.2d, #0000000000000000
+; CHECK-NEXT:    movi v2.8b, #1
+; CHECK-NEXT:    sdot v1.2s, v2.8b, v0.8b
+; CHECK-NEXT:    addp v0.2s, v1.2s, v1.2s
 ; CHECK-NEXT:    fmov w0, s0
 ; CHECK-NEXT:    ret
 entry:
   %0 = bitcast i8* %a1 to <8 x i8>*
   %1 = load <8 x i8>, <8 x i8>* %0
   %2 = sext <8 x i8> %1 to <8 x i32>
   %3 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %2)
   ret i32 %3
entry:
   %r2 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %m2)
   %x = add i32 %r1, %r2
   ret i32 %x
 }
 
 define i32 @test_udot_v8i8_double_nomla(<8 x i8> %a, <8 x i8> %b, <8 x i8> %c, <8 x i8> %d) {
 ; CHECK-LABEL: test_udot_v8i8_double_nomla:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    ushll v0.8h, v0.8b, #0
-; CHECK-NEXT:    ushll v1.8h, v2.8b, #0
-; CHECK-NEXT:    ushll v2.4s, v0.4h, #0
-; CHECK-NEXT:    ushll v3.4s, v1.4h, #0
-; CHECK-NEXT:    uaddw2 v0.4s, v2.4s, v0.8h
-; CHECK-NEXT:    uaddw2 v1.4s, v3.4s, v1.8h
-; CHECK-NEXT:    add v0.4s, v0.4s, v1.4s
-; CHECK-NEXT:    addv s0, v0.4s
+; CHECK-NEXT:    movi v1.2d, #0000000000000000
+; CHECK-NEXT:    movi v3.8b, #1
+; CHECK-NEXT:    udot v1.2s, v3.8b, v2.8b
+; CHECK-NEXT:    udot v1.2s, v3.8b, v0.8b
+; CHECK-NEXT:    addp v0.2s, v1.2s, v1.2s
 ; CHECK-NEXT:    fmov w0, s0
 ; CHECK-NEXT:    ret
 entry:
   %az = zext <8 x i8> %a to <8 x i32>
   %r1 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %az)
   %cz = zext <8 x i8> %c to <8 x i32>
   %r2 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %cz)
   %x = add i32 %r1, %r2
entry:
   %r2 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %m2)
   %x = add i32 %r1, %r2
   ret i32 %x
 }
 
 define i32 @test_sdot_v8i8_double_nomla(<8 x i8> %a, <8 x i8> %b, <8 x i8> %c, <8 x i8> %d) {
 ; CHECK-LABEL: test_sdot_v8i8_double_nomla:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    sshll v0.8h, v0.8b, #0
-; CHECK-NEXT:    sshll v1.8h, v2.8b, #0
-; CHECK-NEXT:    sshll v2.4s, v0.4h, #0
-; CHECK-NEXT:    sshll v3.4s, v1.4h, #0
-; CHECK-NEXT:    saddw2 v0.4s, v2.4s, v0.8h
-; CHECK-NEXT:    saddw2 v1.4s, v3.4s, v1.8h
-; CHECK-NEXT:    add v0.4s, v0.4s, v1.4s
-; CHECK-NEXT:    addv s0, v0.4s
+; CHECK-NEXT:    movi v1.2d, #0000000000000000
+; CHECK-NEXT:    movi v3.8b, #1
+; CHECK-NEXT:    sdot v1.2s, v3.8b, v2.8b
+; CHECK-NEXT:    sdot v1.2s, v3.8b, v0.8b
+; CHECK-NEXT:    addp v0.2s, v1.2s, v1.2s
 ; CHECK-NEXT:    fmov w0, s0
 ; CHECK-NEXT:    ret
 entry:
   %az = sext <8 x i8> %a to <8 x i32>
   %r1 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %az)
   %cz = sext <8 x i8> %c to <8 x i32>
   %r2 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %cz)
   %x = add i32 %r1, %r2