[AArch64] Ensure we do not access illegal operands in tryCombineMULLWithUZP1

[dmgreen]

Summary

https://github.com/llvm/llvm-project/issues/65015 shows a case where tryCombineMULLWithUZP1 could attempt to look at the wrong operand of another user instruction. This adds an extra else as if we don't find the right opcode, we don't need to check the operands.

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

22984

Minor nit: add braces around "if" body so it more obviously matches the "else".

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static SDValue tryCombineMULLWithUZP1(SDNode *N,
   }
 
   if (!ExtractLow || !ExtractLow->hasOneUse())
     HasFoundMULLow = false;
 
   // Check ExtractLow's user.
   if (HasFoundMULLow) {
     SDNode *ExtractLowUser = *ExtractLow.getNode()->use_begin();
     if (ExtractLowUser->getOpcode() != N->getOpcode())
       HasFoundMULLow = false;
-
+    else {
       if (ExtractLowUser->getOperand(0) == ExtractLow) {
         if (ExtractLowUser->getOperand(1).getOpcode() == ISD::TRUNCATE)
           TruncLow = ExtractLowUser->getOperand(1);
         else
           HasFoundMULLow = false;
       } else {
         if (ExtractLowUser->getOperand(0).getOpcode() == ISD::TRUNCATE)
           TruncLow = ExtractLowUser->getOperand(0);
         else
           HasFoundMULLow = false;
       }
     }
+  }
 
   // If the truncate's operand is BUILD_VECTOR with DUP, do not combine the op
   // with uzp1.
   // You can see the regressions on test/CodeGen/AArch64/aarch64-smull.ll
   EVT TruncHighVT = TruncHigh.getValueType();
   EVT UZP1VT = TruncHighVT.getDoubleNumVectorElementsVT(*DAG.getContext());
   SDValue TruncLowOp =
       HasFoundMULLow ? TruncLow.getOperand(0) : DAG.getUNDEF(UZP1VT);

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

entry:
   %12 = shufflevector <8 x i16> %0, <8 x i16> poison, <4 x i32> <i32 4, i32 5, i32 6, i32 7>
   %13 = tail call <4 x i32> @llvm.aarch64.neon.umull.v4i32(<4 x i16> %12, <4 x i16> %9)
   %14 = add <4 x i32> %11, %13
   %15 = sub <4 x i32> %1, %14
   store <4 x i32> %15, ptr %2, align 16
   ret void
 }
 
+define <2 x i32> @do_stuff(<2 x i64> %0, <2 x i64> %1) {
+; CHECK-LABEL: do_stuff:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uzp1 v0.4s, v0.4s, v0.4s
+; CHECK-NEXT:    smull2 v0.2d, v1.4s, v0.4s
+; CHECK-NEXT:    xtn v0.2s, v0.2d
+; CHECK-NEXT:    add v0.2s, v0.2s, v1.2s
+; CHECK-NEXT:    ret
+  %bc.1 = bitcast <2 x i64> %1 to <4 x i32>
+  %trunc.0 = trunc <2 x i64> %0 to <2 x i32>
+  %shuff.hi = shufflevector <4 x i32> %bc.1, <4 x i32> zeroinitializer, <2 x i32> <i32 2, i32 3>
+  %shuff.lo = shufflevector <4 x i32> %bc.1, <4 x i32> zeroinitializer, <2 x i32> <i32 0, i32 1>
+  %smull = tail call <2 x i64> @llvm.aarch64.neon.smull.v2i64(<2 x i32> %shuff.hi, <2 x i32> %trunc.0)
+  %trunc.smull = trunc <2 x i64> %smull to <2 x i32>
+  %final = add <2 x i32> %trunc.smull, %shuff.lo
+  ret <2 x i32> %final
+}
+
 declare <8 x i16> @llvm.aarch64.neon.pmull.v8i16(<8 x i8>, <8 x i8>)
 declare <8 x i16> @llvm.aarch64.neon.smull.v8i16(<8 x i8>, <8 x i8>)
 declare <8 x i16> @llvm.aarch64.neon.umull.v8i16(<8 x i8>, <8 x i8>)
 declare <4 x i32> @llvm.aarch64.neon.smull.v4i32(<4 x i16>, <4 x i16>)
 declare <4 x i32> @llvm.aarch64.neon.umull.v4i32(<4 x i16>, <4 x i16>)
+declare <2 x i64> @llvm.aarch64.neon.smull.v2i64(<2 x i32>, <2 x i32>)