[AArch64][SVE] Lowering sve.dot to DOT node

[junparser]

Summary

As is the title, this patch also takes care about dup node in performAddDotCombine which tries to support scalable vector type.

TestPlan: check-llvm

[junparser]

address pre-check comments.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13365

Can we add a isZerosVector function that encapsulates isBuildVectorAllZeros and the splats/dups logic? It sounds generally useful to have a function that checks for the various ways a vector can be all zeros.

isConstantSplatVectorAllZeros may also be better, as it may already handle the ISD::SPLAT_VECTOR part.

[junparser]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13365

sound good to me, I'll update this.

[junparser]

address the comments.

[joechrisellis]

Thanks for the patch @junparser! Just a few minor comments from me. 🙂

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

2161–2163

Might be better to do this _before_ the above if(ISD::isConstantSplatVectorAllZeroes(N)), in my opinion, to avoid running this loop twice (once in this function, once in ISD::isConstantSplatVectorAllZeros).

2165

A cursory glance at ISD::isConstantSplatVectorAllZeros suggests that the SPLAT_VECTOR case is entirely captured earlier on, I think. I reckon you can get away with just:

if (N->getOpcode() != AArch64ISD::DUP)
  return false;

2171–2173

nit:

return (CINT && CINT->isNullValue()) || (CFP && CFP->isZero());

13361

nit: superfluous newline.

13365

Following on from this, it might be a good idea to write isZerosVector as:

static bool isZerosVector(const SDNode *N) {
  return ISD::isConstantSplatVectorAllZeros(N) || isConstantDupVectorAllZeros(N);
}

Might be a nicer separation of the logic here.

Your call, though!

llvm/test/CodeGen/AArch64/sve-intrinsics-int-arith.ll

41

The new CHECK: lines introduced look simple enough to write out by hand, so I would prefer to not use utils/update_llc_checks.py in this case because it has made the diff a lot noisier (e.g. all of the tests have been changed syntactically in this patch, but AFAICT none of them have changed functionally).

In this case, I think it should be sufficient to write your new tests like this:

define <vscale x 2 x i64> @test_sdot_i64_zero(<vscale x 2 x i64> %a, <vscale x 8 x i16> %b, <vscale x 8 x i16> %c) {
; CHECK-LABEL: test_sdot_i64_zero:
; CHECK:         sdot z0.d, z1.h, z2.h
; CHECK-NEXT:    ret
entry:
  %vdot1.i = call <vscale x 2 x i64> @llvm.aarch64.sve.sdot.nxv2i64(<vscale x 2 x i64> zeroinitializer, <vscale x 8 x i16> %b, <vscale x 8 x i16> %c)
  %ret = add <vscale x 2 x i64> %vdot1.i, %a
  ret <vscale x 2 x i64> %ret
}

define <vscale x 4 x i32> @test_sdot_i32_zero(<vscale x 4 x i32> %a, <vscale x 16 x i8> %b, <vscale x 16 x i8> %c) {
; CHECK-LABEL: test_sdot_i32_zero:
; CHECK:         sdot z0.s, z1.b, z2.b
; CHECK-NEXT:    ret
entry:
  %vdot1.i = call <vscale x 4 x i32> @llvm.aarch64.sve.sdot.nxv4i32(<vscale x 4 x i32> zeroinitializer, <vscale x 16 x i8> %b, <vscale x 16 x i8> %c)
  %ret = add <vscale x 4 x i32> %vdot1.i, %a
  ret <vscale x 4 x i32> %ret
}

... and:

define <vscale x 2 x i64> @test_udot_i64_zero(<vscale x 2 x i64> %a, <vscale x 8 x i16> %b, <vscale x 8 x i16> %c) {
; CHECK-LABEL: test_udot_i64_zero:
; CHECK:         udot z0.d, z1.h, z2.h
; CHECK-NEXT:    ret
entry:
  %vdot1.i = call <vscale x 2 x i64> @llvm.aarch64.sve.udot.nxv2i64(<vscale x 2 x i64> zeroinitializer, <vscale x 8 x i16> %b, <vscale x 8 x i16> %c)
  %ret = add <vscale x 2 x i64> %vdot1.i, %a
  ret <vscale x 2 x i64> %ret
}

define <vscale x 4 x i32> @test_udot_i32_zero(<vscale x 4 x i32> %a, <vscale x 16 x i8> %b, <vscale x 16 x i8> %c) {
; CHECK-LABEL: test_udot_i32_zero:
; CHECK:         udot z0.s, z1.b, z2.b
; CHECK-NEXT:    ret
entry:
  %vdot1.i = call <vscale x 4 x i32> @llvm.aarch64.sve.udot.nxv4i32(<vscale x 4 x i32> zeroinitializer, <vscale x 16 x i8> %b, <vscale x 16 x i8> %c)
  %ret = add <vscale x 4 x i32> %vdot1.i, %a
  ret <vscale x 4 x i32> %ret
}

... and then just revert all other changes in the file. 😄

[paulwalker-arm]

llvm/test/CodeGen/AArch64/sve-intrinsics-int-arith.ll

41

I think not using update_llc_checks.py was a mistake when landing the original file. Although using it now makes the patch look bigger it will make it easier to update in the future, which is a good thing.

[david-arm]

llvm/test/CodeGen/AArch64/sve-intrinsics-int-arith.ll

41

I think the LLVM community prefers using utils/update_llc_test_checks.py where possible, in particular for small functions like those here.

[junparser]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

2165

the ISD::isConstantSplatVectorAllZeros does not check ConstantFPSDNode for SPLAT_VECTOR, that's I added it here.

2165

I will add another patch to check check ConstantFPSDNode for SPLAT_VECTOR in ISD::isConstantSplatVectorAllZeros and then remove the code here? is that ok?

llvm/test/CodeGen/AArch64/sve-intrinsics-int-arith.ll

41

I also prefer to use update_llc_checks.py here.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

2165

I had not realized that isConstantSplatVectorAllZeros didn't already handle SPLAT_VECTOR with FP constants. Adding it there sounds like a good idea, either here or in a separate patch.

llvm/test/CodeGen/AArch64/sve-intrinsics-int-arith.ll

41

Yeah, using the update script has benefits in terms of maintainability and consistency between test. As well as not missing things in the generated assembly.

Feel free regenerate the existing test and commit that separately, so just the changes from the patch are shown here. I often to so far as to commit the new tests with the current trunk codegen, so that in the review it is obvious what is changing in the codegen.

[junparser]

address comments.

[joechrisellis]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

2165

ACK -- I would have expected it to already handle SPLAT_VECTOR with float constants. @dmgreen's suggestion sounds good to me, hopefully that's a straightforward change! 🙂

llvm/test/CodeGen/AArch64/sve-intrinsics-int-arith.ll

41

ACK -- any which way if fine by me. I like @dmgreen's idea of regenerating the test as a separate commit, but no complaints from me if you just leave it as-is. 🙂

[junparser]

address comments.

@joechrisellis @david-arm Add fp constants handle in ISD::isConstantSplatVector , also split the testcase.

[dmgreen]

Thanks. LGTM

[paulwalker-arm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

149

@junparser This looks a little wired to me. Specifically the truncOrSelf part which to my mind is never something that is likely safe to do with a floating point constant (even after bitcasting to an APInt). Looking at the equivalent code in BuildVectorSDNode::isConstantSplat I can see that it stops after the bitcastToAPInt() stage. To be consistent, can this function follow the same pattern?

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

 
 bool ISD::isConstantSplatVector(const SDNode *N, APInt &SplatVal) {
   if (N->getOpcode() == ISD::SPLAT_VECTOR) {
     unsigned EltSize =
         N->getValueType(0).getVectorElementType().getSizeInBits();
     if (auto *Op0 = dyn_cast<ConstantSDNode>(N->getOperand(0))) {
       SplatVal = Op0->getAPIntValue().truncOrSelf(EltSize);
       return true;
+    } else if (auto *Op0 = dyn_cast<ConstantFPSDNode>(N->getOperand(0))) {
+      SplatVal = Op0->getValueAPF().bitcastToAPInt().truncOrSelf(EltSize);
+      return true;
     }
   }
 
   auto *BV = dyn_cast<BuildVectorSDNode>(N);
   if (!BV)
     return false;
 
   APInt SplatUndef;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 //===----------------------------------------------------------------------===//
 // AArch64 Lowering private implementation.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Lowering Code
 //===----------------------------------------------------------------------===//
 
+/// isZerosVector - Check whether SDNode N is a zero-filled vector.
+static bool isZerosVector(const SDNode *N) {
+  // Look through a bit convert.
+  while (N->getOpcode() == ISD::BITCAST)
+    N = N->getOperand(0).getNode();
+
+  if (ISD::isConstantSplatVectorAllZeros(N))
+    return true;
+
+  if (N->getOpcode() != AArch64ISD::DUP)
+    return false;
+
+  auto Opnd0 = N->getOperand(0);
+  auto *CINT = dyn_cast<ConstantSDNode>(Opnd0);
+  auto *CFP = dyn_cast<ConstantFPSDNode>(Opnd0);
+  return (CINT && CINT->isNullValue()) || (CFP && CFP->isZero());
+}
+
 /// changeIntCCToAArch64CC - Convert a DAG integer condition code to an AArch64
 /// CC
 static AArch64CC::CondCode changeIntCCToAArch64CC(ISD::CondCode CC) {
   switch (CC) {
   default:
     llvm_unreachable("Unknown condition code!");
   case ISD::SETNE:
     return AArch64CC::NE;
SDValue AArch64TargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
   case Intrinsic::aarch64_neon_sabd:
   case Intrinsic::aarch64_neon_uabd: {
     unsigned Opcode = IntNo == Intrinsic::aarch64_neon_uabd ? AArch64ISD::UABD
                                                             : AArch64ISD::SABD;
     return DAG.getNode(Opcode, dl, Op.getValueType(), Op.getOperand(1),
                        Op.getOperand(2));
   }
   case Intrinsic::aarch64_neon_sdot:
-  case Intrinsic::aarch64_neon_udot: {
-    unsigned Opcode = IntNo == Intrinsic::aarch64_neon_udot ? AArch64ISD::UDOT
-                                                            : AArch64ISD::SDOT;
+  case Intrinsic::aarch64_neon_udot:
+  case Intrinsic::aarch64_sve_sdot:
+  case Intrinsic::aarch64_sve_udot: {
+    unsigned Opcode = (IntNo == Intrinsic::aarch64_neon_udot ||
+                       IntNo == Intrinsic::aarch64_sve_udot)
+                          ? AArch64ISD::UDOT
     return DAG.getNode(Opcode, dl, Op.getValueType(), Op.getOperand(1),
                        Op.getOperand(2), Op.getOperand(3));
   }
   }
 }
 
 bool AArch64TargetLowering::shouldExtendGSIndex(EVT VT, EVT &EltTy) const {
   if (VT.getVectorElementType() == MVT::i8 ||

 // ADD(UDOT(zero, x, y), A) -->  UDOT(A, x, y)
 static SDValue performAddDotCombine(SDNode *N, SelectionDAG &DAG) {
   EVT VT = N->getValueType(0);
   if (N->getOpcode() != ISD::ADD)
     return SDValue();
 
   SDValue Dot = N->getOperand(0);
   SDValue A = N->getOperand(1);
   // Handle commutivity
   auto isZeroDot = [](SDValue Dot) {
     return (Dot.getOpcode() == AArch64ISD::UDOT ||
             Dot.getOpcode() == AArch64ISD::SDOT) &&
-           ISD::isBuildVectorAllZeros(Dot.getOperand(0).getNode());
+           isZerosVector(Dot.getOperand(0).getNode());
   };
   if (!isZeroDot(Dot))
     std::swap(Dot, A);
   if (!isZeroDot(Dot))
     return SDValue();
 
   return DAG.getNode(Dot.getOpcode(), SDLoc(N), VT, A, Dot.getOperand(1),
                      Dot.getOperand(2));

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

let Predicates = [HasSVE] in {
   defm SDIV_ZPmZ  : sve_int_bin_pred_arit_2_div<0b100, "sdiv",  "SDIV_ZPZZ", int_aarch64_sve_sdiv, DestructiveBinaryCommWithRev, "SDIVR_ZPmZ">;
   defm UDIV_ZPmZ  : sve_int_bin_pred_arit_2_div<0b101, "udiv",  "UDIV_ZPZZ", int_aarch64_sve_udiv, DestructiveBinaryCommWithRev, "UDIVR_ZPmZ">;
   defm SDIVR_ZPmZ : sve_int_bin_pred_arit_2_div<0b110, "sdivr", "SDIVR_ZPZZ", int_aarch64_sve_sdivr, DestructiveBinaryCommWithRev, "SDIV_ZPmZ", /*isReverseInstr*/ 1>;
   defm UDIVR_ZPmZ : sve_int_bin_pred_arit_2_div<0b111, "udivr", "UDIVR_ZPZZ", int_aarch64_sve_udivr, DestructiveBinaryCommWithRev, "UDIV_ZPmZ", /*isReverseInstr*/ 1>;
 
   defm SDIV_ZPZZ  : sve_int_bin_pred_sd<AArch64sdiv_p>;
   defm UDIV_ZPZZ  : sve_int_bin_pred_sd<AArch64udiv_p>;
 
-  defm SDOT_ZZZ : sve_intx_dot<0b0, "sdot", int_aarch64_sve_sdot>;
-  defm UDOT_ZZZ : sve_intx_dot<0b1, "udot", int_aarch64_sve_udot>;
+  defm SDOT_ZZZ : sve_intx_dot<0b0, "sdot", AArch64sdot>;
+  defm UDOT_ZZZ : sve_intx_dot<0b1, "udot", AArch64udot>;
 
   defm SDOT_ZZZI : sve_intx_dot_by_indexed_elem<0b0, "sdot", int_aarch64_sve_sdot_lane>;
   defm UDOT_ZZZI : sve_intx_dot_by_indexed_elem<0b1, "udot", int_aarch64_sve_udot_lane>;
 
   defm SXTB_ZPmZ : sve_int_un_pred_arit_0_h<0b000, "sxtb", AArch64sxt_mt>;
   defm UXTB_ZPmZ : sve_int_un_pred_arit_0_h<0b001, "uxtb", AArch64uxt_mt>;
   defm SXTH_ZPmZ : sve_int_un_pred_arit_0_w<0b010, "sxth", AArch64sxt_mt>;
   defm UXTH_ZPmZ : sve_int_un_pred_arit_0_w<0b011, "uxth", AArch64uxt_mt>;

llvm/test/CodeGen/AArch64/sve-intrinsics-int-arith.ll

 ; CHECK: abs z0.s, p0/m, z1.s
 ; CHECK-NEXT: ret
   %out = call <vscale x 4 x i32> @llvm.aarch64.sve.abs.nxv4i32(<vscale x 4 x i32> %a,
                                                                <vscale x 4 x i1> %pg,
                                                                <vscale x 4 x i32> %b)
   ret <vscale x 4 x i32> %out
 }
 
 define <vscale x 2 x i64> @abs_i64(<vscale x 2 x i64> %a, <vscale x 2 x i1> %pg, <vscale x 2 x i64> %b) {
 ; CHECK-LABEL: abs_i64:
 ; CHECK: abs z0.d, p0/m, z1.d
 ; CHECK-NEXT: ret
   %out = call <vscale x 2 x i64> @llvm.aarch64.sve.abs.nxv2i64(<vscale x 2 x i64> %a,
                                                                <vscale x 2 x i1> %pg,
                                                                <vscale x 2 x i64> %b)
   ret <vscale x 2 x i64> %out
 }

 ; CHECK: sdot z0.d, z1.h, z2.h
 ; CHECK-NEXT: ret
   %out = call <vscale x 2 x i64> @llvm.aarch64.sve.sdot.nxv2i64(<vscale x 2 x i64> %a,
                                                                 <vscale x 8 x i16> %b,
                                                                 <vscale x 8 x i16> %c)
   ret <vscale x 2 x i64> %out
 }
 
+define <vscale x 2 x i64> @test_sdot_i64_zero(<vscale x 2 x i64> %a, <vscale x 8 x i16> %b, <vscale x 8 x i16> %c) {
+; CHECK-LABEL: test_sdot_i64_zero:
+; CHECK:         sdot z0.d, z1.h, z2.h
+; CHECK-NEXT:    ret
+entry:
+  %vdot1.i = call <vscale x 2 x i64> @llvm.aarch64.sve.sdot.nxv2i64(<vscale x 2 x i64> zeroinitializer, <vscale x 8 x i16> %b, <vscale x 8 x i16> %c)
+  %ret = add <vscale x 2 x i64> %vdot1.i, %a
+  ret <vscale x 2 x i64> %ret
+}
+
+define <vscale x 4 x i32> @test_sdot_i32_zero(<vscale x 4 x i32> %a, <vscale x 16 x i8> %b, <vscale x 16 x i8> %c) {
+; CHECK-LABEL: test_sdot_i32_zero:
+; CHECK:         sdot z0.s, z1.b, z2.b
+; CHECK-NEXT:    ret
+entry:
+  %vdot1.i = call <vscale x 4 x i32> @llvm.aarch64.sve.sdot.nxv4i32(<vscale x 4 x i32> zeroinitializer, <vscale x 16 x i8> %b, <vscale x 16 x i8> %c)
+  %ret = add <vscale x 4 x i32> %vdot1.i, %a
+  ret <vscale x 4 x i32> %ret
+}
+
 ; SDOT (Indexed)
 
 define <vscale x 4 x i32> @sdot_lane_i32(<vscale x 4 x i32> %a, <vscale x 16 x i8> %b, <vscale x 16 x i8> %c) {
 ; CHECK-LABEL: sdot_lane_i32:
 ; CHECK: sdot z0.s, z1.b, z2.b[2]
 ; CHECK-NEXT: ret
   %out = call <vscale x 4 x i32> @llvm.aarch64.sve.sdot.lane.nxv4i32(<vscale x 4 x i32> %a,
                                                                      <vscale x 16 x i8> %b,

 ; CHECK: udot z0.d, z1.h, z2.h
 ; CHECK-NEXT: ret
   %out = call <vscale x 2 x i64> @llvm.aarch64.sve.udot.nxv2i64(<vscale x 2 x i64> %a,
                                                                 <vscale x 8 x i16> %b,
                                                                 <vscale x 8 x i16> %c)
   ret <vscale x 2 x i64> %out
 }
 
+define <vscale x 2 x i64> @test_udot_i64_zero(<vscale x 2 x i64> %a, <vscale x 8 x i16> %b, <vscale x 8 x i16> %c) {
+; CHECK-LABEL: test_udot_i64_zero:
+; CHECK:         udot z0.d, z1.h, z2.h
+; CHECK-NEXT:    ret
+entry:
+  %vdot1.i = call <vscale x 2 x i64> @llvm.aarch64.sve.udot.nxv2i64(<vscale x 2 x i64> zeroinitializer, <vscale x 8 x i16> %b, <vscale x 8 x i16> %c)
+  %ret = add <vscale x 2 x i64> %vdot1.i, %a
+  ret <vscale x 2 x i64> %ret
+}
+
+define <vscale x 4 x i32> @test_udot_i32_zero(<vscale x 4 x i32> %a, <vscale x 16 x i8> %b, <vscale x 16 x i8> %c) {
+; CHECK-LABEL: test_udot_i32_zero:
+; CHECK:         udot z0.s, z1.b, z2.b
+; CHECK-NEXT:    ret
+entry:
+  %vdot1.i = call <vscale x 4 x i32> @llvm.aarch64.sve.udot.nxv4i32(<vscale x 4 x i32> zeroinitializer, <vscale x 16 x i8> %b, <vscale x 16 x i8> %c)
+  %ret = add <vscale x 4 x i32> %vdot1.i, %a
+  ret <vscale x 4 x i32> %ret
+}
+
 ; UDOT (Indexed)
 
 define <vscale x 4 x i32> @udot_lane_i32(<vscale x 4 x i32> %a, <vscale x 16 x i8> %b, <vscale x 16 x i8> %c) {
 ; CHECK-LABEL: udot_lane_i32:
 ; CHECK: udot z0.s, z1.b, z2.b[2]
 ; CHECK-NEXT: ret
   %out = call <vscale x 4 x i32> @llvm.aarch64.sve.udot.lane.nxv4i32(<vscale x 4 x i32> %a,
                                                                      <vscale x 16 x i8> %b,