[SVE] Fix VLS FMA generation at CodeGenOpt::Aggressive

[cameron.mcinally]

Summary

@paulwalker-arm Here is a hastily prepared patch of the problem we discussed earlier today. I'm a little out of the VLS loop, so perhaps there is a better way to handle this...

For VLS lowering, the DAGCombiner is not matching fixed width vector FMAs at CodeGenOpt::Aggressive, since the fixed width FMA matching at CodeGenOpt::Aggressive is done in the MachineCombiner. However, when the MachineCombiner runs, we have already lowered the fixed width vectors to scalable vectors, so FMAs are not generated at all. This patch corrects that by allowing the DAGCombiner to match FMAs if we are using VLS.

[madhur13490]

llvm/test/CodeGen/AArch64/sve-fixed-length-fp-fma.ll

81

Do we want to have a test where %a and %b are globals?

[cameron.mcinally]

llvm/test/CodeGen/AArch64/sve-fixed-length-fp-fma.ll

81

I haven't looked at VLS tests in a while, so that may have changed, but IINM this is the standard format for those tests. See sve-fixed-length*.ll for reference.

[cameron.mcinally]

Fix clang-format warning and add the missing '+' to "+sve".

[paulwalker-arm]

Thanks @cameron.mcinally, I see what you mean now.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12572

Does !useSVEForFixedLengthVectorVT(VT) also work? That way we maintain existing behaviour for NEON sized vectors.

llvm/test/CodeGen/AArch64/sve-fixed-length-fp-fma.ll

17

Given this is a DAG combine I don't see much value in testing all combinations of -aarch64-sve-vector-bits-min so perhaps just have a single RUN line using the maximum value. That way you don't need VBYTES and can use update_llc_test_checks.py to generate the CHECK lines.

[cameron.mcinally]

Updated Diff for @paulwalker-arm's reviews...

[paulwalker-arm]

Perhaps worth adding matching half/fp16 tests to sve-fixed-length-fp-fma.ll but otherwise looks good.

[cameron.mcinally]

Perhaps worth adding matching half/fp16 tests to sve-fixed-length-fp-fma.ll but otherwise looks good.

Added and submitted. Will address any post-commit reviews as needed.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

case Type::DoubleTyID:
     return true;
   default:
     return false;
   }
 }
 
 bool AArch64TargetLowering::generateFMAsInMachineCombiner(
     EVT VT, CodeGenOpt::Level OptLevel) const {
-  return (OptLevel >= CodeGenOpt::Aggressive) && !VT.isScalableVector();
+  return (OptLevel >= CodeGenOpt::Aggressive) && !VT.isScalableVector() &&
+          !Subtarget->useSVEForFixedLengthVectors();
 }
 
 const MCPhysReg *
 AArch64TargetLowering::getScratchRegisters(CallingConv::ID) const {
   // LR is a callee-save register, but we must treat it as clobbered by any call
   // site. Hence we include LR in the scratch registers, which are in turn added
   // as implicit-defs for stackmaps and patchpoints.
   static const MCPhysReg ScratchRegs[] = {

llvm/test/CodeGen/AArch64/sve-fixed-length-fp-fma.ll

+; RUN: llc -O3 -aarch64-sve-vector-bits-min=128  < %s | FileCheck %s -D#VBYTES=16  -check-prefix=NO_SVE
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=256  < %s | FileCheck %s -D#VBYTES=32  -check-prefixes=CHECK
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=384  < %s | FileCheck %s -D#VBYTES=32  -check-prefixes=CHECK
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=512  < %s | FileCheck %s -D#VBYTES=64  -check-prefixes=CHECK
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=640  < %s | FileCheck %s -D#VBYTES=64  -check-prefixes=CHECK
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=768  < %s | FileCheck %s -D#VBYTES=64  -check-prefixes=CHECK
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=896  < %s | FileCheck %s -D#VBYTES=64  -check-prefixes=CHECK
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=1024 < %s | FileCheck %s -D#VBYTES=128
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=1152 < %s | FileCheck %s -D#VBYTES=128
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=1280 < %s | FileCheck %s -D#VBYTES=128
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=1408 < %s | FileCheck %s -D#VBYTES=128
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=1536 < %s | FileCheck %s -D#VBYTES=128
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=1664 < %s | FileCheck %s -D#VBYTES=128
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=1792 < %s | FileCheck %s -D#VBYTES=128
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=1920 < %s | FileCheck %s -D#VBYTES=128
+; RUN: llc -O3 -aarch64-sve-vector-bits-min=2048 < %s | FileCheck %s -D#VBYTES=256
+
+; VBYTES represents the useful byte size of a vector register from the code
+; generator's point of view. It is clamped to power-of-2 values because
+; only power-of-2 vector lengths are considered legal, regardless of the
+; user specified vector length.
+
+target triple = "aarch64-unknown-linux-gnu"
+
+; Don't use SVE when its registers are no bigger than NEON.
+; NO_SVE-NOT: ptrue
+
+;
+; FMA
+;
+
+; Don't use SVE for 64-bit vectors.
+define <2 x float> @fma_v2f32(<2 x float> %op1, <2 x float> %op2, <2 x float> %op3) #0 {
+; CHECK-LABEL: fma_v2f32:
+; CHECK: fmla v2.2s, v1.2s, v0.2s
+; CHECK: ret
+  %mul = fmul contract <2 x float> %op1, %op2
+  %res = fadd contract <2 x float> %mul, %op3
+  ret <2 x float> %res
+}
+
+; Don't use SVE for 128-bit vectors.
+define <4 x float> @fma_v4f32(<4 x float> %op1, <4 x float> %op2, <4 x float> %op3) #0 {
+; CHECK-LABEL: fma_v4f32:
+; CHECK: fmla v2.4s, v1.4s, v0.4s
+; CHECK: ret
+  %mul = fmul contract <4 x float> %op1, %op2
+  %res = fadd contract <4 x float> %mul, %op3
+  ret <4 x float> %res
+}
+
+define void @fma_v8f32(<8 x float>* %a, <8 x float>* %b, <8 x float>* %c) #0 {
+; CHECK-LABEL: fma_v8f32:
+; CHECK: ptrue [[PG:p[0-9]]].s, vl[[#min(div(VBYTES,4),8)]]
+; CHECK-DAG: ld1w { [[OP1:z[0-9]]].s }, [[PG]]/z, [x0]
+; CHECK-DAG: ld1w { [[OP2:z[0-9]]].s }, [[PG]]/z, [x1]
+; CHECK-DAG: ld1w { [[OP3:z[0-9]]].s }, [[PG]]/z, [x2]
+; CHECK: fmad [[OP1]].s, [[PG]]/m, [[OP2]].s, [[OP3]].s
+; CHECK: st1w { [[OP1]].s }, [[PG]], [x0]
+; CHECK: ret
+  %op1 = load <8 x float>, <8 x float>* %a
+  %op2 = load <8 x float>, <8 x float>* %b
+  %op3 = load <8 x float>, <8 x float>* %c
+  %mul = fmul contract <8 x float> %op1, %op2
+  %res = fadd contract <8 x float> %mul, %op3
+  store <8 x float> %res, <8 x float>* %a
+  ret void
+}
+
+; NOTE: Check lines only cover the first VBYTES because the fma_v#f16 tests
+; already cover the general legalisation cases.
+define void @fma_v16f32(<16 x float>* %a, <16 x float>* %b, <16 x float>* %c) #0 {
+; CHECK-LABEL: fma_v16f32:
+; CHECK: ptrue [[PG:p[0-9]]].s, vl[[#min(div(VBYTES,4),16)]]
+; CHECK-DAG: ld1w { [[OP1:z[0-9]]].s }, [[PG]]/z, [x0]
+; CHECK-DAG: ld1w { [[OP2:z[0-9]]].s }, [[PG]]/z, [x1]
+; CHECK-DAG: ld1w { [[OP3:z[0-9]]].s }, [[PG]]/z, [x2]
+; CHECK: fmad [[OP1]].s, [[PG]]/m, [[OP2]].s, [[OP3]].s
+; CHECK: st1w { [[OP1]].s }, [[PG]], [x0]
+; CHECK: ret
+  %op1 = load <16 x float>, <16 x float>* %a
+  %op2 = load <16 x float>, <16 x float>* %b
+  %op3 = load <16 x float>, <16 x float>* %c
+  %mul = fmul contract <16 x float> %op1, %op2
+  %res = fadd contract <16 x float> %mul, %op3
+  store <16 x float> %res, <16 x float>* %a
+  ret void
+}
+
+; NOTE: Check lines only cover the first VBYTES because the fma_v#f16 tests
+; already cover the general legalisation cases.
+define void @fma_v32f32(<32 x float>* %a, <32 x float>* %b, <32 x float>* %c) #0 {
+; CHECK-LABEL: fma_v32f32:
+; CHECK: ptrue [[PG:p[0-9]]].s, vl[[#min(div(VBYTES,4),32)]]
+; CHECK-DAG: ld1w { [[OP1:z[0-9]]].s }, [[PG]]/z, [x0]
+; CHECK-DAG: ld1w { [[OP2:z[0-9]]].s }, [[PG]]/z, [x1]
+; CHECK-DAG: ld1w { [[OP3:z[0-9]]].s }, [[PG]]/z, [x2]
+; CHECK: fmad [[OP1]].s, [[PG]]/m, [[OP2]].s, [[OP3]].s
+; CHECK: st1w { [[OP1]].s }, [[PG]], [x0]
+; CHECK: ret
+  %op1 = load <32 x float>, <32 x float>* %a
+  %op2 = load <32 x float>, <32 x float>* %b
+  %op3 = load <32 x float>, <32 x float>* %c
+  %mul = fmul contract <32 x float> %op1, %op2
+  %res = fadd contract <32 x float> %mul, %op3
+  store <32 x float> %res, <32 x float>* %a
+  ret void
+}
+
+; NOTE: Check lines only cover the first VBYTES because the fma_v#f16 tests
+; already cover the general legalisation cases.
+define void @fma_v64f32(<64 x float>* %a, <64 x float>* %b, <64 x float>* %c) #0 {
+; CHECK-LABEL: fma_v64f32:
+; CHECK: ptrue [[PG:p[0-9]]].s, vl[[#min(div(VBYTES,4),64)]]
+; CHECK-DAG: ld1w { [[OP1:z[0-9]]].s }, [[PG]]/z, [x0]
+; CHECK-DAG: ld1w { [[OP2:z[0-9]]].s }, [[PG]]/z, [x1]
+; CHECK-DAG: ld1w { [[OP3:z[0-9]]].s }, [[PG]]/z, [x2]
+; CHECK: fmad [[OP1]].s, [[PG]]/m, [[OP2]].s, [[OP3]].s
+; CHECK: st1w { [[OP1]].s }, [[PG]], [x0]
+; CHECK: ret
+  %op1 = load <64 x float>, <64 x float>* %a
+  %op2 = load <64 x float>, <64 x float>* %b
+  %op3 = load <64 x float>, <64 x float>* %c
+  %mul = fmul contract <64 x float> %op1, %op2
+  %res = fadd contract <64 x float> %mul, %op3
+  store <64 x float> %res, <64 x float>* %a
+  ret void
+}
+
+; Don't use SVE for 64-bit vectors.
+define <1 x double> @fma_v1f64(<1 x double> %op1, <1 x double> %op2, <1 x double> %op3) #0 {
+; CHECK-LABEL: fma_v1f64:
+; CHECK: fmadd d0, d0, d1, d2
+; CHECK: ret
+  %mul = fmul contract <1 x double> %op1, %op2
+  %res = fadd contract <1 x double> %mul, %op3
+  ret <1 x double> %res
+}
+
+; Don't use SVE for 128-bit vectors.
+define <2 x double> @fma_v2f64(<2 x double> %op1, <2 x double> %op2, <2 x double> %op3) #0 {
+; CHECK-LABEL: fma_v2f64:
+; CHECK: fmla v2.2d, v1.2d, v0.2d
+; CHECK: ret
+  %mul = fmul contract <2 x double> %op1, %op2
+  %res = fadd contract <2 x double> %mul, %op3
+  ret <2 x double> %res
+}
+
+define void @fma_v4f64(<4 x double>* %a, <4 x double>* %b, <4 x double>* %c) #0 {
+; CHECK-LABEL: fma_v4f64:
+; CHECK: ptrue [[PG:p[0-9]]].d, vl[[#min(div(VBYTES,8),4)]]
+; CHECK-DAG: ld1d { [[OP1:z[0-9]]].d }, [[PG]]/z, [x0]
+; CHECK-DAG: ld1d { [[OP2:z[0-9]]].d }, [[PG]]/z, [x1]
+; CHECK-DAG: ld1d { [[OP3:z[0-9]]].d }, [[PG]]/z, [x2]
+; CHECK: fmad [[OP1]].d, [[PG]]/m, [[OP2]].d, [[OP3]].d
+; CHECK: st1d { [[OP1]].d }, [[PG]], [x0]
+; CHECK: ret
+  %op1 = load <4 x double>, <4 x double>* %a
+  %op2 = load <4 x double>, <4 x double>* %b
+  %op3 = load <4 x double>, <4 x double>* %c
+  %mul = fmul contract <4 x double> %op1, %op2
+  %res = fadd contract <4 x double> %mul, %op3
+  store <4 x double> %res, <4 x double>* %a
+  ret void
+}
+
+; NOTE: Check lines only cover the first VBYTES because the fma_v#f16 tests
+; already cover the general legalisation cases.
+define void @fma_v8f64(<8 x double>* %a, <8 x double>* %b, <8 x double>* %c) #0 {
+; CHECK-LABEL: fma_v8f64:
+; CHECK: ptrue [[PG:p[0-9]]].d, vl[[#min(div(VBYTES,8),8)]]
+; CHECK-DAG: ld1d { [[OP1:z[0-9]]].d }, [[PG]]/z, [x0]
+; CHECK-DAG: ld1d { [[OP2:z[0-9]]].d }, [[PG]]/z, [x1]
+; CHECK-DAG: ld1d { [[OP3:z[0-9]]].d }, [[PG]]/z, [x2]
+; CHECK: fmad [[OP1]].d, [[PG]]/m, [[OP2]].d, [[OP3]].d
+; CHECK: st1d { [[OP1]].d }, [[PG]], [x0]
+; CHECK: ret
+  %op1 = load <8 x double>, <8 x double>* %a
+  %op2 = load <8 x double>, <8 x double>* %b
+  %op3 = load <8 x double>, <8 x double>* %c
+  %mul = fmul contract <8 x double> %op1, %op2
+  %res = fadd contract <8 x double> %mul, %op3
+  store <8 x double> %res, <8 x double>* %a
+  ret void
+}
+
+; NOTE: Check lines only cover the first VBYTES because the fma_v#f16 tests
+; already cover the general legalisation cases.
+define void @fma_v16f64(<16 x double>* %a, <16 x double>* %b, <16 x double>* %c) #0 {
+; CHECK-LABEL: fma_v16f64:
+; CHECK: ptrue [[PG:p[0-9]]].d, vl[[#min(div(VBYTES,8),16)]]
+; CHECK-DAG: ld1d { [[OP1:z[0-9]]].d }, [[PG]]/z, [x0]
+; CHECK-DAG: ld1d { [[OP2:z[0-9]]].d }, [[PG]]/z, [x1]
+; CHECK-DAG: ld1d { [[OP3:z[0-9]]].d }, [[PG]]/z, [x2]
+; CHECK: fmad [[OP1:z[0-9]]].d, [[PG]]/m, [[OP2]].d, [[OP3]].d
+; CHECK: st1d { [[OP1]].d }, [[PG]], [x0]
+; CHECK: ret
+  %op1 = load <16 x double>, <16 x double>* %a
+  %op2 = load <16 x double>, <16 x double>* %b
+  %op3 = load <16 x double>, <16 x double>* %c
+  %mul = fmul contract <16 x double> %op1, %op2
+  %res = fadd contract <16 x double> %mul, %op3
+  store <16 x double> %res, <16 x double>* %a
+  ret void
+}
+
+; NOTE: Check lines only cover the first VBYTES because the fma_v#f16 tests
+; already cover the general legalisation cases.
+define void @fma_v32f64(<32 x double>* %a, <32 x double>* %b, <32 x double>* %c) #0 {
+; CHECK-LABEL: fma_v32f64:
+; CHECK: ptrue [[PG:p[0-9]]].d, vl[[#min(div(VBYTES,8),32)]]
+; CHECK-DAG: ld1d { [[OP1:z[0-9]]].d }, [[PG]]/z, [x0]
+; CHECK-DAG: ld1d { [[OP2:z[0-9]]].d }, [[PG]]/z, [x1]
+; CHECK-DAG: ld1d { [[OP3:z[0-9]]].d }, [[PG]]/z, [x2]
+; CHECK: fmad [[OP1:z[0-9]]].d, [[PG]]/m, [[OP2]].d, [[OP3]].d
+; CHECK: st1d { [[OP1]].d }, [[PG]], [x0]
+; CHECK: ret
+  %op1 = load <32 x double>, <32 x double>* %a
+  %op2 = load <32 x double>, <32 x double>* %b
+  %op3 = load <32 x double>, <32 x double>* %c
+  %mul = fmul contract <32 x double> %op1, %op2
+  %res = fadd contract <32 x double> %mul, %op3
+  store <32 x double> %res, <32 x double>* %a
+  ret void
+}
+
+attributes #0 = { "target-features"="sve" }