[AArch64] splat before (f)mul to allow mul-by-element isel

[spatel]

Summary

A splat of a vector multiply (either integer or FP) can be turned into a multiply-by-element:

splat (mul X, Y), Lane --> mul (splat X, Lane), (splat Y, Lane) --> mul-by-element (splat X, Lane), Y.[Lane]

These patterns showed up as an ARM regression in D60214, but we have this transform in IR, so it's an existing problem IIUC.

The constant cases look better, but I'm not sure if this is a win if both operands are variables.

[efriedma]

The constant cases look better, but I'm not sure if this is a win if both operands are variables.

The key really isn't whether one of the operands is a constant; it's whether the operand is free or cheap to splat. Constants are usually free to splat. A splat is free to splat... although I guess that's unlikely to come up in practice given other optimizations. A load is cheap to splat (the addressing mode for ld1r is very limited, so you're likely adding an extra instruction for address computation). A loop-invariant operand is likely cheap to splat, but I don't think there's any way to handle that in SelectionDAG at the moment. And of course, a multiply with an operand that's free to splat is itself free to splat, recursively. Probably worth adding a testcase with more than one multiply.

For two arbitrary variables, it's basically neutral, like you've noted; splatting an operand has the same cost as splatting a result.

It's worth noting that at least on some chips, 128-bit multiplies have half the throughput of scalar and 64-bit vector multiplies. But I don't think that directly affects this patch; you probably wouldn't want to add extra instructions just to make a multiply smaller.

[spatel]

Abandoning - implemented a limited version of splat re-ordering for DAGCombiner with D79886.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Lane == 0 && V1.getOpcode() == ISD::SCALAR_TO_VECTOR)
       return DAG.getNode(AArch64ISD::DUP, dl, V1.getValueType(),
                          V1.getOperand(0));
     // Test if V1 is a BUILD_VECTOR and the lane being referenced is a non-
     // constant. If so, we can just reference the lane's definition directly.
     if (V1.getOpcode() == ISD::BUILD_VECTOR &&
         !isa<ConstantSDNode>(V1.getOperand(Lane)))
       return DAG.getNode(AArch64ISD::DUP, dl, VT, V1.getOperand(Lane));
 
+    // Splat multiply operands to allow selecting this as an (f)mul by element:
+    // splat (mul X, Y), Lane --> mul (splat X, Lane), (splat Y, Lane)
+    if (V1.hasOneUse() &&
+        (V1.getOpcode() == ISD::FMUL || V1.getOpcode() == ISD::MUL)) {
+      SDValue SplatX = DAG.getVectorShuffle(VT, dl, V1.getOperand(0),
+                                            DAG.getUNDEF(VT), ShuffleMask);
+      SDValue SplatY = DAG.getVectorShuffle(VT, dl, V1.getOperand(1),
+                                            DAG.getUNDEF(VT), ShuffleMask);
+      return DAG.getNode(V1.getOpcode(), dl, VT, SplatX, SplatY);
+    }
+
     // Otherwise, duplicate from the lane of the input vector.
     unsigned Opcode = getDUPLANEOp(V1.getValueType().getVectorElementType());
 
     // SelectionDAGBuilder may have "helpfully" already extracted or conatenated
     // to make a vector of the same size as this SHUFFLE. We can ignore the
     // extract entirely, and canonicalise the concat using WidenVector.
     if (V1.getOpcode() == ISD::EXTRACT_SUBVECTOR) {
       Lane += cast<ConstantSDNode>(V1.getOperand(1))->getZExtValue();

llvm/test/CodeGen/AArch64/mul_by_elt.ll

; CHECK-NEXT:    ret
   %mul = fmul <4 x float> %splat, <float 3.0, float 3.0, float 3.0, float 3.0>
   ret <4 x float> %mul
 }
 
 define <4 x float> @splat0_after_fmul_constant(<4 x float> %a) {
 ; CHECK-LABEL: splat0_after_fmul_constant:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    fmov v1.4s, #3.00000000
-; CHECK-NEXT:    fmul v0.4s, v0.4s, v1.4s
-; CHECK-NEXT:    dup v0.4s, v0.s[0]
+; CHECK-NEXT:    fmul v0.4s, v1.4s, v0.s[0]
 ; CHECK-NEXT:    ret
   %mul = fmul <4 x float> %a, <float 3.0, float 42.0, float 3.0, float 3.0>
   %splat = shufflevector <4 x float> %mul, <4 x float> undef, <4 x i32> zeroinitializer
   ret <4 x float> %splat
 }
 
 ; Try different type and splat lane.
 
 define <2 x double> @splat1_before_fmul_constant(<2 x double> %a) {
 ; CHECK-LABEL: splat1_before_fmul_constant:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    fmov v1.2d, #5.00000000
 ; CHECK-NEXT:    fmul v0.2d, v1.2d, v0.d[1]
 ; CHECK-NEXT:    ret
   %splat = shufflevector <2 x double> %a, <2 x double> undef, <2 x i32> <i32 1, i32 1>
   %mul = fmul <2 x double> %splat, <double 5.0, double 5.0>
   ret <2 x double> %mul
 }
 
 define <2 x double> @splat1_after_fmul_constant(<2 x double> %a) {
 ; CHECK-LABEL: splat1_after_fmul_constant:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    fmov v1.2d, #5.00000000
-; CHECK-NEXT:    fmul v0.2d, v0.2d, v1.2d
-; CHECK-NEXT:    dup v0.2d, v0.d[1]
+; CHECK-NEXT:    fmul v0.2d, v1.2d, v0.d[1]
 ; CHECK-NEXT:    ret
   %mul = fmul <2 x double> %a, <double -1.0, double 5.0>
   %splat = shufflevector <2 x double> %mul, <2 x double> undef, <2 x i32> <i32 1, i32 1>
   ret <2 x double> %splat
 }
 
 ; 2 variable operands
 
 define <2 x double> @splat1_before_fmul(<2 x double> %a, <2 x double> %b) {
 ; CHECK-LABEL: splat1_before_fmul:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    fmul v0.2d, v0.2d, v1.2d
 ; CHECK-NEXT:    dup v0.2d, v0.d[1]
+; CHECK-NEXT:    fmul v0.2d, v0.2d, v1.d[1]
 ; CHECK-NEXT:    ret
   %splata = shufflevector <2 x double> %a, <2 x double> undef, <2 x i32> <i32 1, i32 1>
   %splatb = shufflevector <2 x double> %b, <2 x double> undef, <2 x i32> <i32 1, i32 1>
   %mul = fmul <2 x double> %splata, %splatb
   ret <2 x double> %mul
 }
 
 define <2 x double> @splat1_after_fmul(<2 x double> %a, <2 x double> %b) {
 ; CHECK-LABEL: splat1_after_fmul:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    fmul v0.2d, v0.2d, v1.2d
 ; CHECK-NEXT:    dup v0.2d, v0.d[1]
+; CHECK-NEXT:    fmul v0.2d, v0.2d, v1.d[1]
 ; CHECK-NEXT:    ret
   %mul = fmul <2 x double> %a, %b
   %splat = shufflevector <2 x double> %mul, <2 x double> undef, <2 x i32> <i32 1, i32 1>
   ret <2 x double> %splat
 }
 
 ; Integer multiply
 
 define <4 x i32> @splat2_before_mul_constant(<4 x i32> %a) {
 ; CHECK-LABEL: splat2_before_mul_constant:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v1.4s, #3
 ; CHECK-NEXT:    mul v0.4s, v1.4s, v0.s[2]
 ; CHECK-NEXT:    ret
   %splat = shufflevector <4 x i32> %a, <4 x i32> undef, <4 x i32> <i32 2, i32 undef, i32 2, i32 undef>
   %mul = mul <4 x i32> %splat, <i32 3, i32 3, i32 3, i32 3>
   ret <4 x i32> %mul
 }
 
 define <4 x i32> @splat2_after_mul_constant(<4 x i32> %a) {
 ; CHECK-LABEL: splat2_after_mul_constant:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    movi v1.4s, #3
-; CHECK-NEXT:    mul v0.4s, v0.4s, v1.4s
-; CHECK-NEXT:    dup v0.4s, v0.s[2]
+; CHECK-NEXT:    mul v0.4s, v1.4s, v0.s[2]
 ; CHECK-NEXT:    ret
   %mul = mul <4 x i32> %a, <i32 3, i32 3, i32 3, i32 3>
   %splat = shufflevector <4 x i32> %mul, <4 x i32> undef, <4 x i32> <i32 2, i32 undef, i32 2, i32 undef>
   ret <4 x i32> %splat
 }
 
 define <8 x i16> @splat1_before_mul(<8 x i16> %a, <8 x i16> %b) {
 ; CHECK-LABEL: splat1_before_mul:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mul v0.8h, v0.8h, v1.8h
 ; CHECK-NEXT:    dup v0.8h, v0.h[1]
+; CHECK-NEXT:    mul v0.8h, v0.8h, v1.h[1]
 ; CHECK-NEXT:    ret
   %splata = shufflevector <8 x i16> %a, <8 x i16> undef, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
   %splatb = shufflevector <8 x i16> %b, <8 x i16> undef, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
   %mul = mul <8 x i16> %splata, %splatb
   ret <8 x i16> %mul
 }
 
 define <8 x i16> @splat1_after_mul(<8 x i16> %a, <8 x i16> %b) {
 ; CHECK-LABEL: splat1_after_mul:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mul v0.8h, v0.8h, v1.8h
 ; CHECK-NEXT:    dup v0.8h, v0.h[1]
+; CHECK-NEXT:    mul v0.8h, v0.8h, v1.h[1]
 ; CHECK-NEXT:    ret
   %mul = mul <8 x i16> %a, %b
   %splat = shufflevector <8 x i16> %mul, <8 x i16> undef, <8 x i32> <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
   ret <8 x i16> %splat
 }
-