[AArch64] Select SMULL for zero extended vectors when top bit is zero

[zjaffal]

Summary

we can safely replace a zext instruction with sext if the top bit is zero. This is useful because we can select smull when both operands are sign extended.

[HsiangKai]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4466

NewOpc has no value if LastBitValue is not zero in the previous if block. It will lost some opportunities to use UMULL to do unsigned multiplication. It is why some test cases become worse.

Do you think it is a good way to check NewOpc before checking isN0ZExt && isN1ZExt? That is

if (!NewOpc) {
  if (isN0ZExt && isN1ZExt)
    NewOpc = AArch64ISD::UMULL;
  else if (isN1SExt || isN1ZExt) {
    // ...
  }
}

[fhahn]

Thanks for the patch @zjaffal! Please update the code as suggested by @HsiangKai to fix the regressions in the test case.

[zjaffal]

Fix regression caused by not selecting UMULL when we can

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4448

The logic with the nested ifs is a bit hard to follow now.

Would it be possible to simplify the code by moving all those checks to a static helper function which returns the new opcode (and updates the other variables as required)? This would allow to use early exits to reduce the need for nested of conditions.

[spatel]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4455

Just a couple of drive-by comments:

1. Why is this passing depth 4?
2. Use DAG.SignBitIsZero() to make this easier to read? ("last bit" usually means the LSB, not the MSB/signbit).

[zjaffal]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4455

The depth is 4 following a similar code block in the same file

// Check if the value is zero-extended from i1 to i8
static bool checkZExtBool(SDValue Arg, const SelectionDAG &DAG) {
  unsigned SizeInBits = Arg.getValueType().getSizeInBits();
  if (SizeInBits < 8)
    return false;

  APInt RequredZero(SizeInBits, 0xFE);
  KnownBits Bits = DAG.computeKnownBits(Arg, 4);
  bool ZExtBool = (Bits.Zero & RequredZero) == RequredZero;
  return ZExtBool;
}

I tested the implementation using DAG.SignBitIsZero() and it works with depth=0

[zjaffal]

Refactor common code for selecting multiplication instruction out of the main code for LowerMul, rewrite checking for the top bit to use DAG.SignBitIsZero and refactor on top of the added new tests

[fhahn]

Thanks for the update! Could you split out the code that adds selectUmullSmull into a separate NFC patch and then update this patch to just add the new handling for converting zext to sext?

[zjaffal]

Rebase patch on top of https://reviews.llvm.org/D134875

[zjaffal]

rebase after lint of parent patch

[zjaffal]

Rebase on top of main

[zjaffal]

ping

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4381

Is there some other changes needed here, to pass a DL through to the function from the callsites of selectUmullSmull?

[zjaffal]

Rebase on top of main and edit the function call to include DL

[dmgreen]

Thanks. A small nit, but otherwise this looks like a good patch. The newly created node will never be used, but I don't think that would be a problem. LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4407

Small nit but can you move this above the new code? This condition and the new one are mutually exclusive, and it would seem a little better to have the check for two zext next to the bool IsN0ZExt = ..., the same as the base smull case.

[zjaffal]

Move code blocko

[fhahn]

LGTM, thanks!

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

SDValue AArch64TargetLowering::LowerSET_ROUNDING(SDValue Op,
   FPCR = DAG.getNode(ISD::OR, DL, MVT::i64, FPCR, RMValue);
   SDValue Ops2[] = {
       Chain, DAG.getTargetConstant(Intrinsic::aarch64_set_fpcr, DL, MVT::i64),
       FPCR};
   return DAG.getNode(ISD::INTRINSIC_VOID, DL, MVT::Other, Ops2);
 }
 
 static unsigned selectUmullSmull(SDNode *&N0, SDNode *&N1, SelectionDAG &DAG,
-                                 bool &IsMLA) {
+                                 SDLoc DL, bool &IsMLA) {
   bool IsN0SExt = isSignExtended(N0, DAG);
   bool IsN1SExt = isSignExtended(N1, DAG);
   if (IsN0SExt && IsN1SExt)
     return AArch64ISD::SMULL;
 
   bool IsN0ZExt = isZeroExtended(N0, DAG);
   bool IsN1ZExt = isZeroExtended(N1, DAG);
 
+  // Select SMULL if we can replace zext with sext.
+  if ((IsN0SExt && IsN1ZExt) || (IsN0ZExt && IsN1SExt)) {
+    SDValue ZextOperand;
+    if (IsN0ZExt)
+      ZextOperand = N0->getOperand(0);
+    else
+      ZextOperand = N1->getOperand(0);
+    if (DAG.SignBitIsZero(ZextOperand)) {
+      SDNode *NewSext =
+          DAG.getSExtOrTrunc(ZextOperand, DL, N0->getValueType(0)).getNode();
+      if (IsN0ZExt)
+        N0 = NewSext;
+      else
+        N1 = NewSext;
+      return AArch64ISD::SMULL;
+    }
+  }
   if (IsN0ZExt && IsN1ZExt)
     return AArch64ISD::UMULL;
 
   if (!IsN1SExt && !IsN1ZExt)
     return 0;
   // Look for (s/zext A + s/zext B) * (s/zext C). We want to turn these
   // into (s/zext A * s/zext C) + (s/zext B * s/zext C)
   if (IsN1SExt && isAddSubSExt(N0, DAG)) {
     IsMLA = true;
SDValue AArch64TargetLowering::LowerMUL(SDValue Op, SelectionDAG &DAG) const {
   // VMULL can be detected.  Otherwise v2i64 multiplications are not legal.
   assert(VT.is128BitVector() && VT.isInteger() &&
          "unexpected type for custom-lowering ISD::MUL");
   SDNode *N0 = Op.getOperand(0).getNode();
   SDNode *N1 = Op.getOperand(1).getNode();
   bool isMLA = false;
   unsigned NewOpc = selectUmullSmull(N0, N1, DAG, isMLA);
 
   if (!NewOpc) {
     if (VT == MVT::v2i64)
       // Fall through to expand this.  It is not legal.
       return SDValue();
     else
       // Other vector multiplications are legal.
       return Op;
   }
 
   // Legalize to a S/UMULL instruction
   SDLoc DL(Op);
   SDValue Op0;
   SDValue Op1 = skipExtensionForVectorMULL(N1, DAG);
   if (!isMLA) {
     Op0 = skipExtensionForVectorMULL(N0, DAG);
     assert(Op0.getValueType().is64BitVector() &&
            Op1.getValueType().is64BitVector() &&
            "unexpected types for extended operands to VMULL");
     return DAG.getNode(NewOpc, DL, VT, Op0, Op1);
   }
   // Optimizing (zext A + zext B) * C, to (S/UMULL A, C) + (S/UMULL B, C) during
   // isel lowering to take advantage of no-stall back to back s/umul + s/umla.
   // This is true for CPUs with accumulate forwarding such as Cortex-A53/A57
   SDValue N00 = skipExtensionForVectorMULL(N0->getOperand(0).getNode(), DAG);
   SDValue N01 = skipExtensionForVectorMULL(N0->getOperand(1).getNode(), DAG);
   EVT Op1VT = Op1.getValueType();
   return DAG.getNode(N0->getOpcode(), DL, VT,

llvm/test/CodeGen/AArch64/aarch64-smull.ll

; CHECK-NEXT:    ret
   %tmp5 = mul <2 x i64> %tmp3, %tmp4
   ret <2 x i64> %tmp5
 }
 
 define <8 x i32> @smull_zext_v8i8_v8i32(<8 x i8>* %A, <8 x i16>* %B) nounwind {
 ; CHECK-LABEL: smull_zext_v8i8_v8i32:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldr d0, [x0]
-; CHECK-NEXT:    ldr q1, [x1]
+; CHECK-NEXT:    ldr q2, [x1]
 ; CHECK-NEXT:    ushll v0.8h, v0.8b, #0
-; CHECK-NEXT:    sshll v2.4s, v1.4h, #0
-; CHECK-NEXT:    sshll2 v1.4s, v1.8h, #0
-; CHECK-NEXT:    ushll2 v3.4s, v0.8h, #0
-; CHECK-NEXT:    ushll v0.4s, v0.4h, #0
-; CHECK-NEXT:    mul v1.4s, v3.4s, v1.4s
-; CHECK-NEXT:    mul v0.4s, v0.4s, v2.4s
+; CHECK-NEXT:    smull2 v1.4s, v0.8h, v2.8h
+; CHECK-NEXT:    smull v0.4s, v0.4h, v2.4h
 ; CHECK-NEXT:    ret
   %load.A = load <8 x i8>, <8 x i8>* %A
   %load.B = load <8 x i16>, <8 x i16>* %B
   %zext.A = zext <8 x i8> %load.A to <8 x i32>
   %sext.B = sext <8 x i16> %load.B to <8 x i32>
   %res = mul <8 x i32> %zext.A, %sext.B
   ret <8 x i32> %res
 }
 
 define <8 x i32> @smull_zext_v8i8_v8i32_sext_first_operand(<8 x i16>* %A, <8 x i8>* %B) nounwind {
 ; CHECK-LABEL: smull_zext_v8i8_v8i32_sext_first_operand:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldr d1, [x1]
-; CHECK-NEXT:    ldr q0, [x0]
-; CHECK-NEXT:    ushll v1.8h, v1.8b, #0
-; CHECK-NEXT:    sshll v2.4s, v0.4h, #0
-; CHECK-NEXT:    sshll2 v0.4s, v0.8h, #0
-; CHECK-NEXT:    ushll2 v3.4s, v1.8h, #0
-; CHECK-NEXT:    ushll v4.4s, v1.4h, #0
-; CHECK-NEXT:    mul v1.4s, v0.4s, v3.4s
-; CHECK-NEXT:    mul v0.4s, v2.4s, v4.4s
+; CHECK-NEXT:    ldr d0, [x1]
+; CHECK-NEXT:    ldr q2, [x0]
+; CHECK-NEXT:    ushll v0.8h, v0.8b, #0
+; CHECK-NEXT:    smull2 v1.4s, v2.8h, v0.8h
+; CHECK-NEXT:    smull v0.4s, v2.4h, v0.4h
 ; CHECK-NEXT:    ret
   %load.A = load <8 x i16>, <8 x i16>* %A
   %load.B = load <8 x i8>, <8 x i8>* %B
   %sext.A = sext <8 x i16> %load.A to <8 x i32>
   %zext.B = zext <8 x i8> %load.B to <8 x i32>
   %res = mul <8 x i32> %sext.A, %zext.B
   ret <8 x i32> %res
 }

 }
 
 define <4 x i32> @smull_zext_v4i16_v4i32(<4 x i8>* %A, <4 x i16>* %B) nounwind {
 ; CHECK-LABEL: smull_zext_v4i16_v4i32:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldr s0, [x0]
 ; CHECK-NEXT:    ldr d1, [x1]
 ; CHECK-NEXT:    ushll v0.8h, v0.8b, #0
-; CHECK-NEXT:    sshll v1.4s, v1.4h, #0
-; CHECK-NEXT:    ushll v0.4s, v0.4h, #0
-; CHECK-NEXT:    mul v0.4s, v0.4s, v1.4s
+; CHECK-NEXT:    smull v0.4s, v0.4h, v1.4h
 ; CHECK-NEXT:    ret
   %load.A = load <4 x i8>, <4 x i8>* %A
   %load.B = load <4 x i16>, <4 x i16>* %B
   %zext.A = zext <4 x i8> %load.A to <4 x i32>
   %sext.B = sext <4 x i16> %load.B to <4 x i32>
   %res = mul <4 x i32> %zext.A, %sext.B
   ret <4 x i32> %res
 }

 }
 
 define <2 x i64> @smull_zext_and_v2i32_v2i64(<2 x i32>* %A, <2 x i32>* %B) nounwind {
 ; CHECK-LABEL: smull_zext_and_v2i32_v2i64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ldr d0, [x0]
 ; CHECK-NEXT:    ldr d1, [x1]
 ; CHECK-NEXT:    bic v0.2s, #128, lsl #24
-; CHECK-NEXT:    sshll v1.2d, v1.2s, #0
-; CHECK-NEXT:    ushll v0.2d, v0.2s, #0
-; CHECK-NEXT:    fmov x9, d1
-; CHECK-NEXT:    fmov x10, d0
-; CHECK-NEXT:    mov x8, v1.d[1]
-; CHECK-NEXT:    mov x11, v0.d[1]
-; CHECK-NEXT:    mul x9, x10, x9
-; CHECK-NEXT:    mul x8, x11, x8
-; CHECK-NEXT:    fmov d0, x9
-; CHECK-NEXT:    mov v0.d[1], x8
+; CHECK-NEXT:    smull v0.2d, v0.2s, v1.2s
 ; CHECK-NEXT:    ret
   %load.A = load <2 x i32>, <2 x i32>* %A
   %and.A = and <2 x i32> %load.A, <i32 u0x7FFFFFFF, i32 u0x7FFFFFFF>
   %load.B = load <2 x i32>, <2 x i32>* %B
   %zext.A = zext <2 x i32> %and.A to <2 x i64>
   %sext.B = sext <2 x i32> %load.B to <2 x i64>
   %res = mul <2 x i64> %zext.A, %sext.B
   ret <2 x i64> %res