[AArch64] Use simd mov to materialize big fp constants

[Allen]

Summary

mov w8, #1325400064 + fmov s0, w8 ==> movi v0.2s, 0x4f, lsl 24

Fix https://github.com/llvm/llvm-project/issues/53651

[Allen]

This comment was removed by Allen.

[Allen]

fix precommit format

[dmgreen]

Using the MOVI to materialize fp constants sounds like a good idea.

But I don't think this should be starting from a CopyToReg. CopyToReg usually means we are copying the result into the return value of a function (or into a call argument). Copying into s0 for example. CopyFromReg is used to get the value from input physical registers from arguments.

An example like this, with fadd, would preferably use the movi lowering too: https://godbolt.org/z/bc99E85qz

If the FPConstant is still in the DAG at the point that we run instruction selection (I think it should be but it may be dependant on what AArch64TargetLowering::isFPImmLegal returns), then it may be best to add a tablegen pattern for cases that we can materialize with the integer movi instruction. From here: https://github.com/llvm/llvm-project/blob/ecff9b65b54c7a4bd79ca2af157c81595678f0ee/llvm/lib/Target/AArch64/AArch64InstrInfo.td#L1545, with a different fpimm to limit it to cases where the MOVi can be generated, and an EXTRACT_SUBREG to get the s reg. If that doesn't work (the custom ImmLeaf's might be a bit much) then perhaps something in AArch64ISelDAGToDAG.cpp?

[Allen]

In D120452#3343727, @dmgreen wrote:

Using the MOVI to materialize fp constants sounds like a good idea.

Thanks very much for detail, refactor with above new method.

[Allen]

update test cases

[fhahn]

llvm/test/Transforms/LoopVectorize/AArch64/remat-const-float-simd.ll

3

It looks like the test is completely unrelated to the loop-vectorize pass. It should probably be either moved to CodeGen/AArch64/ or added to one of the existing tests there.

[dmgreen]

llvm/lib/Target/AArch64/AArch64InstrFormats.td

1210 ↗

(On Diff #411714)

This likely doesn't need to be an Operand, it can just be a FPImmLeaf

llvm/lib/Target/AArch64/AArch64InstrInfo.td

1551 ↗

(On Diff #411714)

This doesn't need to define a new instruction, it can use the existing MOVI (which I believe is probably called MOVIv2i32?)

llvm/test/Transforms/LoopVectorize/AArch64/remat-const-float-simd.ll

2

This test file should be in CodeGen/AArch64

[Allen]

a) move case remat-const-float-simd.ll to CodeGen/AArch64
b) use MOVIv2i32 , and delete unneeded instruction MOVI2s_ns
c) delete Operand<f32>, and just be a FPImmLeaf

[Allen]

In D120452#3348692, @Allen wrote:

a) move case remat-const-float-simd.ll to CodeGen/AArch64
b) use MOVIv2i32 , and delete unneeded instruction MOVI2s_ns
c) delete Operand<f32>, and just be a FPImmLeaf

Address all comment

llvm/lib/Target/AArch64/AArch64InstrFormats.td

1210 ↗

(On Diff #411714)

Yes, verified ok!

llvm/lib/Target/AArch64/AArch64InstrInfo.td

1551 ↗

(On Diff #411714)

Thanks, delete unneeded instruction MOVI2s_ns

1551 ↗

(On Diff #411608)

It's strange that here error with the following info, as it is ok on my local linux ?

AArch64InstrInfo.td:1551:1: error: Type set is empty for each HW mode in 'MOVIv2s_ns' def MOVIv2s_ns : BaseSIMDModifiedImmVectorShift<1, 1, 0b10, V128, "movi", ".2s",

[Allen]

ping ?

[dmgreen]

Thanks for the updates. There are other forms of MOVI that might also be useful for some fp immediates. The one we have here (shift an i8 by 24) I imagine is the most useful, but the others could be good in places too. The MOVN's too. I'm not sure what the best way to generalise this would be. As far as I can tell it needs lots of checks to the various isAdvSIMDModImmTypeXYZ methods wherever it is? This seems like a good beginning at least.

llvm/lib/Target/AArch64/AArch64InstrFormats.td

1210 ↗

(On Diff #411786)

It's likely worth calling this out as a "AdvSIMDModImmType4" constant somehow. Maybe call it fpimm32SIMDModImmType4? Same for the XForm.

llvm/lib/Target/AArch64/AArch64InstrInfo.td

6146 ↗

(On Diff #411786)

It is quite uncommon to not have NEON, but can you add a predicate for it:

let Predicates = [HasNEON] in {

It might be worth adding a run line without neon (-mattr=-neon) for the new remat test case too, to show the difference.

llvm/test/CodeGen/AArch64/remat-const-float-simd.ll

7 ↗

(On Diff #411786)

0x7fffffff -> 2147483648, as the 0x7fffffff gets rounded.

[Allen]

update testcase 0x7fffffff with 2147483648 as gets rounded

[dmgreen]

Thanks. LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.h

bool functionArgumentNeedsConsecutiveRegisters(
       Type *Ty, CallingConv::ID CallConv, bool isVarArg,
       const DataLayout &DL) const override;
 
   /// Used for exception handling on Win64.
   bool needsFixedCatchObjects() const override;
 
   bool fallBackToDAGISel(const Instruction &Inst) const override;
 
+  SDValue LowerCopyToReg(SDValue Op, SelectionDAG &DAG) const;
+
   /// SVE code generation for fixed length vectors does not custom lower
   /// BUILD_VECTOR. This makes BUILD_VECTOR legalisation a source of stores to
   /// merge. However, merging them creates a BUILD_VECTOR that is just as
   /// illegal as the original, thus leading to an infinite legalisation loop.
   /// NOTE: Once BUILD_VECTOR is legal or can be custom lowered for all legal
   /// vector types this override can be removed.
   bool mergeStoresAfterLegalization(EVT VT) const override;
 

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
   // When comparing vectors the result sets the different elements in the
   // vector to all-one or all-zero.
   setBooleanVectorContents(ZeroOrNegativeOneBooleanContent);
 
   // Set up the register classes.
   addRegisterClass(MVT::i32, &AArch64::GPR32allRegClass);
   addRegisterClass(MVT::i64, &AArch64::GPR64allRegClass);
 
+  setOperationAction(ISD::CopyToReg, MVT::Other, Custom);
+
   if (Subtarget->hasLS64()) {
     addRegisterClass(MVT::i64x8, &AArch64::GPR64x8ClassRegClass);
     setOperationAction(ISD::LOAD, MVT::i64x8, Custom);
     setOperationAction(ISD::STORE, MVT::i64x8, Custom);
   }
 
   if (Subtarget->hasFPARMv8()) {
     addRegisterClass(MVT::f16, &AArch64::FPR16RegClass);
SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
                                               SelectionDAG &DAG) const {
   LLVM_DEBUG(dbgs() << "Custom lowering: ");
   LLVM_DEBUG(Op.dump());
 
   switch (Op.getOpcode()) {
   default:
     llvm_unreachable("unimplemented operand");
     return SDValue();
+  case ISD::CopyToReg:
+    return LowerCopyToReg(Op, DAG);
   case ISD::BITCAST:
     return LowerBITCAST(Op, DAG);
   case ISD::GlobalAddress:
     return LowerGlobalAddress(Op, DAG);
   case ISD::GlobalTLSAddress:
     return LowerGlobalTLSAddress(Op, DAG);
   case ISD::SETCC:
   case ISD::STRICT_FSETCC:
case ISD::CTLZ:
     return LowerToPredicatedOp(Op, DAG, AArch64ISD::CTLZ_MERGE_PASSTHRU);
   case ISD::CTTZ:
     return LowerCTTZ(Op, DAG);
   case ISD::VECTOR_SPLICE:
     return LowerVECTOR_SPLICE(Op, DAG);
   }
 }
 
+static SDValue
+tryAdvSIMDModImm32(unsigned NewOp, SDValue Op, SelectionDAG &DAG,
+                   const APInt &Bits, const SDValue *LHS = nullptr);
+static bool resolveBuildVector(BuildVectorSDNode *BVN, APInt &CnstBits,
+                               APInt &UndefBits);
+
+SDValue AArch64TargetLowering::LowerCopyToReg(SDValue Op,
+                                              SelectionDAG &DAG) const {
+  SDValue Chain = Op->getOperand(0);
+  SDValue LHS = Op->getOperand(1);
+  SDValue RHS = Op->getOperand(2);
+  ConstantFPSDNode *CFP = dyn_cast<ConstantFPSDNode>(RHS);
+  EVT VT = LHS->getValueType(0);
+  if (!CFP || !LHS->hasNUsesOfValue(2,0) ||
+      VT.getSimpleVT().SimpleTy != MVT::f32)
+    return SDValue();
+
+  const APFloat &FPVal = CFP->getValueAPF();
+  const APInt ImmInt = FPVal.bitcastToAPInt();
+  uint64_t Imm = ImmInt.getZExtValue();
+  // Skip getFP32Imm as related value already deal with fmov.
+  if (AArch64_AM::getFP32Imm(ImmInt) != -1 || FPVal.isPosZero() ||
+      !AArch64_AM::isAdvSIMDModImmType4(Imm << 32 | Imm))
+    return SDValue();
+
+  SDLoc dl(Op);
+  MVT VecTy = MVT::v2f32;
+  APInt DefBits(VecTy.getSizeInBits(), 0);
+  APInt UndefBits(VecTy.getSizeInBits(), 0);
+  SDValue Parts[2];
+  for (int Elt=0; Elt< 2; Elt++)
+    Parts[Elt] = RHS;
+  SDValue DupVal = DAG.getNode(ISD::BUILD_VECTOR, dl, MVT::v2f32, Parts);
+
+  BuildVectorSDNode *BVN = cast<BuildVectorSDNode>(DupVal.getNode());
+  if (!resolveBuildVector(BVN, DefBits, UndefBits))
+    return SDValue();
+
+  SDValue NewOp = tryAdvSIMDModImm32(AArch64ISD::MOVIshift,
+                                     DupVal, DAG, DefBits);
+  SDValue Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, dl, VT, NewOp,
+                            DAG.getConstant(0, dl, MVT::i32));
+
+  DAG.ReplaceAllUsesWith(RHS, Elt);
+  return Op;
+}
+
 bool AArch64TargetLowering::mergeStoresAfterLegalization(EVT VT) const {
   return !Subtarget->useSVEForFixedLengthVectors();
 }
 
 bool AArch64TargetLowering::useSVEForFixedLengthVectorVT(
     EVT VT, bool OverrideNEON) const {
   if (!VT.isFixedLengthVector())
     return false;
if (Bits.getHiBits(64) == Bits.getLoBits(64)) {
     }
   }
 
   return SDValue();
 }
 
 // Try 32-bit splatted SIMD immediate.
 static SDValue tryAdvSIMDModImm32(unsigned NewOp, SDValue Op, SelectionDAG &DAG,
                                   const APInt &Bits,
-                                  const SDValue *LHS = nullptr) {
+                                  const SDValue *LHS) {
   if (Bits.getHiBits(64) == Bits.getLoBits(64)) {
     uint64_t Value = Bits.zextOrTrunc(64).getZExtValue();
     EVT VT = Op.getValueType();
     MVT MovTy = (VT.getSizeInBits() == 128) ? MVT::v4i32 : MVT::v2i32;
     bool isAdvSIMDModImm = false;
     uint64_t Shift;
 
     if ((isAdvSIMDModImm = AArch64_AM::isAdvSIMDModImmType1(Value))) {

llvm/test/Transforms/LoopVectorize/AArch64/remat-const-float-simd.ll

+; RUN: llc < %s -mtriple=aarch64-none-linux-gnu -verify-machineinstrs | FileCheck %s
+
+; Check that float(0x7fffffff) can be rematerialized with simd instruction
+target triple = "aarch64-unknown-linux-gnu"
+
+; float foo(void) { return float(0x7fffffff); }
+define float @foo() {
+; CHECK: movi   v0.2s, #79, lsl #24
+entry:
+  ret float 0x41E0000000000000
+}
+
+