From bbefd5713ff55de3a23e4bdc7c15edd74cbc5e7a Mon Sep 17 00:00:00 2001
From: Bjorn Pettersson <bjorn.a.pettersson@ericsson.com>
Date: Fri, 9 Aug 2024 17:47:56 +0200
Subject: [PATCH] [TargetLowering] Handle vector types in expandFixedPointMul
 (#102635)

In TargetLowering::expandFixedPointMul when expanding fixed point
multiplication, and when using a widened MUL as strategy for the
lowering, there was a bug resulting in assertion failures like this:
   Assertion `VT.isVector() == N1.getValueType().isVector() &&
   "SIGN_EXTEND result type type should be vector iff the operand "
   "type is vector!"' failed.

Problem was that we did not consider that VT could be a vector type
when setting up the WideVT. This patch should fix that bug.
---
 llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp |  3 +++
 llvm/test/CodeGen/AArch64/smul_fix.ll            | 14 ++++++++++++++
 llvm/test/CodeGen/AArch64/umul_fix.ll            | 14 ++++++++++++++
 3 files changed, 31 insertions(+)

diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index 866fac092792f..db8a64c52ecf6 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -10782,6 +10782,9 @@ TargetLowering::expandFixedPointMul(SDNode *Node, SelectionDAG &DAG) const {
   unsigned LoHiOp = Signed ? ISD::SMUL_LOHI : ISD::UMUL_LOHI;
   unsigned HiOp = Signed ? ISD::MULHS : ISD::MULHU;
   EVT WideVT = EVT::getIntegerVT(*DAG.getContext(), VTSize * 2);
+  if (VT.isVector())
+    WideVT =
+        EVT::getVectorVT(*DAG.getContext(), WideVT, VT.getVectorElementCount());
   if (isOperationLegalOrCustom(LoHiOp, VT)) {
     SDValue Result = DAG.getNode(LoHiOp, dl, DAG.getVTList(VT, VT), LHS, RHS);
     Lo = Result.getValue(0);
diff --git a/llvm/test/CodeGen/AArch64/smul_fix.ll b/llvm/test/CodeGen/AArch64/smul_fix.ll
index 7bd80a00ec81b..50a189df2be53 100644
--- a/llvm/test/CodeGen/AArch64/smul_fix.ll
+++ b/llvm/test/CodeGen/AArch64/smul_fix.ll
@@ -137,3 +137,17 @@ define <4 x i64> @vec3(<4 x i64> %x, <4 x i64> %y) nounwind {
   %tmp = call <4 x i64> @llvm.smul.fix.v4i64(<4 x i64> %x, <4 x i64> %y, i32 32)
   ret <4 x i64> %tmp
 }
+
+define <4 x i16> @widemul(<4 x i16> %x, <4 x i16> %y) nounwind {
+; CHECK-LABEL: widemul:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    smull v0.4s, v0.4h, v1.4h
+; CHECK-NEXT:    shrn v1.4h, v0.4s, #16
+; CHECK-NEXT:    xtn v2.4h, v0.4s
+; CHECK-NEXT:    add v1.4h, v1.4h, v1.4h
+; CHECK-NEXT:    shl v0.4h, v1.4h, #13
+; CHECK-NEXT:    usra v0.4h, v2.4h, #2
+; CHECK-NEXT:    ret
+  %tmp = call <4 x i16> @llvm.smul.fix.v4i16(<4 x i16> %x, <4 x i16> %y, i32 2)
+  ret <4 x i16> %tmp
+}
diff --git a/llvm/test/CodeGen/AArch64/umul_fix.ll b/llvm/test/CodeGen/AArch64/umul_fix.ll
index 6ec6832993841..9f4da82dd74b6 100644
--- a/llvm/test/CodeGen/AArch64/umul_fix.ll
+++ b/llvm/test/CodeGen/AArch64/umul_fix.ll
@@ -145,3 +145,17 @@ define <4 x i64> @vec3(<4 x i64> %x, <4 x i64> %y) nounwind {
   %tmp = call <4 x i64> @llvm.umul.fix.v4i64(<4 x i64> %x, <4 x i64> %y, i32 32)
   ret <4 x i64> %tmp
 }
+
+define <4 x i16> @widemul(<4 x i16> %x, <4 x i16> %y) nounwind {
+; CHECK-LABEL: widemul:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    umull v0.4s, v0.4h, v1.4h
+; CHECK-NEXT:    shrn v1.4h, v0.4s, #16
+; CHECK-NEXT:    xtn v2.4h, v0.4s
+; CHECK-NEXT:    add v1.4h, v1.4h, v1.4h
+; CHECK-NEXT:    shl v0.4h, v1.4h, #11
+; CHECK-NEXT:    usra v0.4h, v2.4h, #4
+; CHECK-NEXT:    ret
+  %tmp = call <4 x i16> @llvm.umul.fix.v4i16(<4 x i16> %x, <4 x i16> %y, i32 4)
+  ret <4 x i16> %tmp
+}