[AArch64] Don't rely on (zext (trunc x)) pattern to detect zext_inreg MULL patterns - use value tracking directly

[RKSimon]

Summary

As explained on D159533, I'm trying to generalize the "(zext (trunc x)) -> x iff the upper bits are known zero" fold in getNode() and I was seeing assertions in the aarch64 mull matching code as it was assuming these 'zero-extend-inreg' patterns will remain from earlier in LowerMUL.

Instead I've updated selectUmullSmull/skipExtensionForVectorMULL to just use value tracking to detect when the upper bits are known zero, and to insert the truncation nodes later if necessary.

I really don't like creating SDValue(N, 0) on the fly from SDNode value as technically we could be using any result index from these nodes - so I've ended up cleaning up a lot of mul code to use SDValue directly instead of peeking through to the SDNode. I'm happy to undo this and just rely use SDValue(N, 0) if there's resistance, but this is much cleaner imo. I'd push this change as pre-commit NFC.

(Sorry for still using Phab but I'm frantically trying to get my local backlog dealt with before moving over to using github branches).

[dmgreen]

I think this looks OK. Thanks

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4446–4447

I believe this should always be a 128bit vector at this point.

[RKSimon]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

4446–4447

Thanks, I'll make it an assert

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

for (const SDValue &Elt : N->op_values()) {
     }
     return false;
   }
 
   return true;
 }
 
 static SDValue skipExtensionForVectorMULL(SDValue N, SelectionDAG &DAG) {
+  EVT VT = N.getValueType();
+  assert(VT.is128BitVector() && "Unexpected vector MULL size");
+
+  unsigned NumElts = VT.getVectorNumElements();
+  unsigned OrigEltSize = VT.getScalarSizeInBits();
+  unsigned EltSize = OrigEltSize / 2;
+  MVT TruncVT = MVT::getVectorVT(MVT::getIntegerVT(EltSize), NumElts);
+
+  APInt HiBits = APInt::getHighBitsSet(OrigEltSize, EltSize);
+  if (DAG.MaskedValueIsZero(N, HiBits))
+    return DAG.getNode(ISD::TRUNCATE, SDLoc(N), TruncVT, N);
+
   if (ISD::isExtOpcode(N.getOpcode()))
     return addRequiredExtensionForVectorMULL(N.getOperand(0), DAG,
-                                             N.getOperand(0).getValueType(),
-                                             N.getValueType(),
+                                             N.getOperand(0).getValueType(), VT,
                                              N.getOpcode());
 
   assert(N.getOpcode() == ISD::BUILD_VECTOR && "expected BUILD_VECTOR");
-  EVT VT = N.getValueType();
   SDLoc dl(N);
-  unsigned EltSize = VT.getScalarSizeInBits() / 2;
-  unsigned NumElts = VT.getVectorNumElements();
-  MVT TruncVT = MVT::getIntegerVT(EltSize);
   SmallVector<SDValue, 8> Ops;
   for (unsigned i = 0; i != NumElts; ++i) {
     ConstantSDNode *C = cast<ConstantSDNode>(N.getOperand(i));
     const APInt &CInt = C->getAPIntValue();
     // Element types smaller than 32 bits are not legal, so use i32 elements.
     // The values are implicitly truncated so sext vs. zext doesn't matter.
     Ops.push_back(DAG.getConstant(CInt.zextOrTrunc(32), dl, MVT::i32));
   }
-  return DAG.getBuildVector(MVT::getVectorVT(TruncVT, NumElts), dl, Ops);
+  return DAG.getBuildVector(TruncVT, dl, Ops);
 }
 
 static bool isSignExtended(SDValue N, SelectionDAG &DAG) {
   return N.getOpcode() == ISD::SIGN_EXTEND ||
          N.getOpcode() == ISD::ANY_EXTEND ||
          isExtendedBUILD_VECTOR(N, DAG, true);
 }
 
if (((IsN0SExt && IsN1ZExt) || (IsN0ZExt && IsN1SExt)) &&
     }
   }
 
   // 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);
-    if (DAG.MaskedValueIsZero(IsN0ZExt ? N1 : N0, Mask)) {
-      EVT HalfVT;
-      switch (VT.getSimpleVT().SimpleTy) {
-      case MVT::v2i64:
-        HalfVT = MVT::v2i32;
-        break;
-      case MVT::v4i32:
-        HalfVT = MVT::v4i16;
-        break;
-      case MVT::v8i16:
-        HalfVT = MVT::v8i8;
-        break;
-      default:
-        return 0;
-      }
-      // Truncate and then extend the result.
-      SDValue NewExt =
-          DAG.getNode(ISD::TRUNCATE, DL, HalfVT, IsN0ZExt ? N1 : N0);
-      NewExt = DAG.getZExtOrTrunc(NewExt, DL, VT);
-      if (IsN0ZExt)
-        N1 = NewExt;
-      else
-        N0 = NewExt;
+    if (DAG.MaskedValueIsZero(IsN0ZExt ? N1 : N0, Mask))
       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;