[llvm] [AArch64] Improve mull generation (PR #114997)

[David Sherwood via llvm-commits]

================
@@ -5452,34 +5398,27 @@ static unsigned selectUmullSmull(SDValue &N0, SDValue &N1, SelectionDAG &DAG,
   if (IsN0ZExt && IsN1ZExt)
     return AArch64ISD::UMULL;
 
-  // Select SMULL if we can replace zext with sext.
-  if (((IsN0SExt && IsN1ZExt) || (IsN0ZExt && IsN1SExt)) &&
-      !isExtendedBUILD_VECTOR(N0, DAG, false) &&
-      !isExtendedBUILD_VECTOR(N1, DAG, false)) {
-    SDValue ZextOperand;
-    if (IsN0ZExt)
-      ZextOperand = N0.getOperand(0);
-    else
-      ZextOperand = N1.getOperand(0);
-    if (DAG.SignBitIsZero(ZextOperand)) {
-      SDValue NewSext =
-          DAG.getSExtOrTrunc(ZextOperand, DL, N0.getValueType());
-      if (IsN0ZExt)
-        N0 = NewSext;
-      else
-        N1 = NewSext;
-      return AArch64ISD::SMULL;
-    }
-  }
-
   // Select UMULL if we can replace the other operand with an extend.
-  if (IsN0ZExt || IsN1ZExt) {
-    EVT VT = N0.getValueType();
-    APInt Mask = APInt::getHighBitsSet(VT.getScalarSizeInBits(),
-                                       VT.getScalarSizeInBits() / 2);
+  EVT VT = N0.getValueType();
+  APInt Mask = APInt::getHighBitsSet(VT.getScalarSizeInBits(),
+                                     VT.getScalarSizeInBits() / 2);
+  if (IsN0ZExt || IsN1ZExt)
     if (DAG.MaskedValueIsZero(IsN0ZExt ? N1 : N0, Mask))
       return AArch64ISD::UMULL;
-  }
+  // For v2i64 we look more aggresively at both operands being zero, to avoid
+  // scalarization.
+  if (VT == MVT::v2i64 && DAG.MaskedValueIsZero(N0, Mask) &&
+      DAG.MaskedValueIsZero(N1, Mask))
+    return AArch64ISD::UMULL;
+
+  if (IsN0SExt || IsN1SExt)
+    if (DAG.ComputeNumSignBits(IsN0SExt ? N1 : N0) >
+        VT.getScalarSizeInBits() / 2)
----------------
david-arm wrote:

nit: Is it worth pulling `VT.getScalarSizeInBits()` into it's own variable given it's used 3 times?

https://github.com/llvm/llvm-project/pull/114997