fixup! [ISel] Commute FMUL and inserting zero into vector lane

Author: juliannagele

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -26283,16 +26283,24 @@ static SDValue commuteInsertVectorEltFMul(SDNode *N, SelectionDAG &DAG) {
 
   // Insert into the operand of FMUL instead.
   SDValue FMulOp = InsertVec.getOperand(0);
+  SDValue FMulOp2 = InsertVec.getOperand(1);
 
-  if (!InsertVec.hasOneUse() || !FMulOp.hasOneUse())
+  if (!InsertVec.hasOneUse())
     return SDValue();
 
+  if (!InsertVec->isOnlyUserOf(FMulOp.getNode())) {
+    if (!InsertVec->isOnlyUserOf(FMulOp2.getNode()))
+      return SDValue();
+    std::swap(FMulOp, FMulOp2);
+  }
+
   SDValue InsertOp =
       DAG.getNode(ISD::INSERT_VECTOR_ELT, SDLoc(N), FMulOp.getValueType(),
                   FMulOp, InsertVal, InsertIdx);
-  SDValue FMul =
-      DAG.getNode(ISD::FMUL, SDLoc(InsertVec), InsertVec.getValueType(),
-                  InsertOp, InsertVec.getOperand(1));
+  if (FMulOp == FMulOp2)
+    FMulOp2 = InsertOp;
+  SDValue FMul = DAG.getNode(ISD::FMUL, SDLoc(InsertVec),
+                             InsertVec.getValueType(), InsertOp, FMulOp2);
   DAG.ReplaceAllUsesWith(N, &FMul);
   return FMul;
 }

llvm/test/CodeGen/AArch64/arm64-vmul.ll

@@ -1197,6 +1197,42 @@ define <4 x float> @fmul_insert_zero(<4 x float> %A, <4 x float> %B) {
   ret <4 x float> %mul_set_lane
 }
 
+define <4 x float> @fmul_insert_zero_same(<4 x float> %A) {
+; CHECK-LABEL: fmul_insert_zero_same:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov.s v0[3], wzr
+; CHECK-NEXT:    fmul.4s v0, v0, v0
+; CHECK-NEXT:    ret
+  %mul = fmul <4 x float> %A, %A
+  %mul_set_lane = insertelement <4 x float> %mul, float 0.000000e+00, i64 3
+  ret <4 x float> %mul_set_lane
+}
+
+define <4 x float> @fmul_insert_zero1(<4 x float> %A, <4 x float> %B, <4 x float> %C) {
+; CHECK-LABEL: fmul_insert_zero1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    fsub.4s v0, v2, v0
+; CHECK-NEXT:    mov.s v1[3], wzr
+; CHECK-NEXT:    fmul.4s v0, v1, v0
+; CHECK-NEXT:    ret
+  %sub = fsub <4 x float> %C, %A
+  %mul = fmul <4 x float> %B, %sub
+  %mul_set_lane = insertelement <4 x float> %mul, float 0.000000e+00, i64 3
+  ret <4 x float> %mul_set_lane
+}
+
+define <4 x float> @fmul_insert_zero2(<4 x float> %A, <4 x float> %B) {
+; CHECK-LABEL: fmul_insert_zero2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov.s v0[3], wzr
+; CHECK-NEXT:    fmul.4s v0, v0, v1
+; CHECK-NEXT:    fsub.4s v0, v1, v0
+; CHECK-NEXT:    ret
+  %mul = fmul <4 x float> %B, %A
+  %mul_set_lane = insertelement <4 x float> %mul, float 0.000000e+00, i64 3
+  %sub = fsub <4 x float> %B, %mul_set_lane
+  ret <4 x float> %sub
+}
 
 define <2 x float> @fmulx_lane_2s(<2 x float> %A, <2 x float> %B) nounwind {
 ; CHECK-LABEL: fmulx_lane_2s: