[ISel] Commute FMUL and inserting zero into vector lane

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -26268,11 +26268,51 @@ static SDValue removeRedundantInsertVectorElt(SDNode *N) {
   return ExtractVec;
 }
 
+static SDValue commuteInsertVectorEltFMul(SDNode *N, SelectionDAG &DAG) {
+  assert(N->getOpcode() == ISD::INSERT_VECTOR_ELT && "Unexpected node!");
+  SDValue InsertVec = N->getOperand(0);
+  SDValue InsertVal = N->getOperand(1);
+  SDValue InsertIdx = N->getOperand(2);
+
+  // Only handle constant 0 insertion...
+  if (!(isNullConstant(InsertVal) || isNullFPConstant(InsertVal)))
+    return SDValue();
+  // ... into the result of an FMUL.
+  if (InsertVec.getOpcode() != ISD::FMUL)
+    return SDValue();
+
+  // Insert into the operand of FMUL instead.
+  SDValue FMulOp = InsertVec.getOperand(0);
+  SDValue FMulOp2 = InsertVec.getOperand(1);
+
+  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);
+  if (FMulOp == FMulOp2)
+    FMulOp2 = InsertOp;
+  SDValue FMul = DAG.getNode(ISD::FMUL, SDLoc(InsertVec),
+                             InsertVec.getValueType(), InsertOp, FMulOp2);
+  DAG.ReplaceAllUsesWith(N, &FMul);
+  return FMul;
+}
+
 static SDValue
 performInsertVectorEltCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI) {
   if (SDValue Res = removeRedundantInsertVectorElt(N))
     return Res;
 
+  if (SDValue Res = commuteInsertVectorEltFMul(N, DCI.DAG))
+    return Res;
+
   return performPostLD1Combine(N, DCI, true);
 }
 

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

@@ -1186,7 +1186,53 @@ define double @fmul_lane_d(double %A, <2 x double> %vec) nounwind {
   ret double %res
 }
 
+define <4 x float> @fmul_insert_zero(<4 x float> %A, <4 x float> %B) {
+; CHECK-LABEL: fmul_insert_zero:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov.s v0[3], wzr
+; CHECK-NEXT:    fmul.4s v0, v0, v1
+; CHECK-NEXT:    ret
+  %mul = fmul <4 x float> %A, %B
+  %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_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: