[AArch64][SVE] Only fold frame indexes referencing SVE objects into SVE loads/stores

[bsmith]

Summary

Currently we always fold frame indexes into SVE load/store instructions,
however these instructions can only encode VL scaled offests. This means
that when we are accessing a fixed length stack object with these
instructions, the folded in frame index gets pulled back out during frame
lowering. This can cause issues when we have no spare registers and no
emergency spill slot.

Rather than causing issues like this, don't fold in frame indexes that
reference fixed length objects.

Fixes: #55041

[efriedma]

Is FrameIndex 0 actually reliably located in a position that doesn't have an offset? That seems unlikely to me, especially if variable-size allocation becomes involved.

Alternatively, we could fix the code that decides whether we need an emergency spill slot. But I guess there's not really much point to attempting to fold the frame indexes in the first place.

[bsmith]

In D124457#3474976, @efriedma wrote:

Is FrameIndex 0 actually reliably located in a position that doesn't have an offset? That seems unlikely to me, especially if variable-size allocation becomes involved.

Hm, you might be right. It seems a shame to regress cases where the index would be zero, but I'm not sure there is way to tell at this point if the offset will resolve to zero?

Alternatively, we could fix the code that decides whether we need an emergency spill slot. But I guess there's not really much point to attempting to fold the frame indexes in the first place.

I considered that, the issue though is finding the correct condition to do this in without being far too conservative, I think you'd have to do it by traversing the function finding all SVE load/stores and checking what they access, which doesn't seem very nice for an issue we are creating for ourselves.

[bsmith]

In D124457#3476866, @bsmith wrote:

In D124457#3474976, @efriedma wrote:

Is FrameIndex 0 actually reliably located in a position that doesn't have an offset? That seems unlikely to me, especially if variable-size allocation becomes involved.

Hm, you might be right. It seems a shame to regress cases where the index would be zero, but I'm not sure there is way to tell at this point if the offset will resolve to zero?

Alternatively, we could fix the code that decides whether we need an emergency spill slot. But I guess there's not really much point to attempting to fold the frame indexes in the first place.

I considered that, the issue though is finding the correct condition to do this in without being far too conservative, I think you'd have to do it by traversing the function finding all SVE load/stores and checking what they access, which doesn't seem very nice for an issue we are creating for ourselves.

Thinking about this some more, in the interest of not regressing and given this isn't a correctness issue, perhaps we can leave it like this and 'assume' frame index zero has a zero offset, if it is not it will just get pulled back out like it does now (hence no correctness issue), it just may still be possible (although probably incredibly hard) to produce a crash due to lack of emergency spill slot. If it does turn out to be an issue still we can always come back and remove this assumption if necessary.

[bsmith]

• Update test to use update_llc_test_checks.py script

[paulwalker-arm]

I'm inclined to agree given this patch is more mitigation rather than an actual fix. Is it not possible to include a minimised version of the original crashing test case? just as a no crash test. I ask because otherwise somebody could easily remove this mitigation without realising the mistake.

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

5105

I know we have the if (N.getOpcode() != ISD::ADD) test below but I think it's better to have an explicit return false for invalid values for FI.

[bsmith]

• Add testcase protecting original issue
• Add explict false return for invalid frame index

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

 /// Base + OffImm * sizeof(MemVT) for Min >= OffImm <= Max
 /// where Root is the memory access using N for its address.
 template <int64_t Min, int64_t Max>
 bool AArch64DAGToDAGISel::SelectAddrModeIndexedSVE(SDNode *Root, SDValue N,
                                                    SDValue &Base,
                                                    SDValue &OffImm) {
   const EVT MemVT = getMemVTFromNode(*(CurDAG->getContext()), Root);
   const DataLayout &DL = CurDAG->getDataLayout();
+  const MachineFrameInfo &MFI = MF->getFrameInfo();
 
   if (N.getOpcode() == ISD::FrameIndex) {
     int FI = cast<FrameIndexSDNode>(N)->getIndex();
+    // We can only encode VL scaled offsets, so only fold in frame indexes
+    // referencing SVE objects.
+    if (FI == 0 || MFI.getStackID(FI) == TargetStackID::ScalableVector) {
       Base = CurDAG->getTargetFrameIndex(FI, TLI->getPointerTy(DL));
       OffImm = CurDAG->getTargetConstant(0, SDLoc(N), MVT::i64);
       return true;
     }
 
+    return false;
+  }
+
   if (MemVT == EVT())
     return false;
 
   if (N.getOpcode() != ISD::ADD)
     return false;
 
   SDValue VScale = N.getOperand(1);
   if (VScale.getOpcode() != ISD::VSCALE)
     return false;
 
   TypeSize TS = MemVT.getSizeInBits();
   int64_t MemWidthBytes = static_cast<int64_t>(TS.getKnownMinSize()) / 8;
   int64_t MulImm = cast<ConstantSDNode>(VScale.getOperand(0))->getSExtValue();
 
   if ((MulImm % MemWidthBytes) != 0)
     return false;
 
   int64_t Offset = MulImm / MemWidthBytes;
   if (Offset < Min || Offset > Max)
     return false;
 
   Base = N.getOperand(0);
   if (Base.getOpcode() == ISD::FrameIndex) {
     int FI = cast<FrameIndexSDNode>(Base)->getIndex();
+    // We can only encode VL scaled offsets, so only fold in frame indexes
+    // referencing SVE objects.
+    if (FI == 0 || MFI.getStackID(FI) == TargetStackID::ScalableVector)
       Base = CurDAG->getTargetFrameIndex(FI, TLI->getPointerTy(DL));
   }
 
   OffImm = CurDAG->getTargetConstant(Offset, SDLoc(N), MVT::i64);
   return true;
 }
 
 /// Select register plus register addressing mode for SVE, with scaled
 /// offset.

llvm/test/CodeGen/AArch64/sve-fixed-length-frame-offests-crash.ll

+; RUN: llc < %s | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+; Ensure we don't crash by trying to fold fixed length frame indexes into
+; loads/stores that don't support an appropriate addressing mode, hence creating
+; too many extra vregs during frame lowering, when we don't have an emergency
+; spill slot.
+
+define dso_local void @func1(i64* %v1, i64* %v2, i64* %v3, i64* %v4, i64* %v5, i64* %v6, i64* %v7, i64* %v8,
+                             i64* %v9, i64* %v10, i64* %v11, i64* %v12, i64* %v13, i64* %v14,  i64* %v15, i64* %v16,
+                             i64* %v17, i64* %v18, i64* %v19, i64* %v20, i64* %v21, i64* %v22, i64* %v23, i64* %v24,
+                             i64* %v25, i64* %v26, i64* %v27, i64* %v28, i64* %v29, i64* %v30, i64* %v31, i64* %v32,
+                             i64* %v33, i64* %v34, i64* %v35, i64* %v36, i64* %v37, i64* %v38, i64* %v39, i64* %v40,
+                             i64* %v41, i64* %v42, i64* %v43, i64* %v44, i64* %v45, i64* %v46, i64* %v47, i64* %v48,
+                             i64 %v49) #0 {
+; CHECK-LABEL: func1
+  tail call void @func2(i64* %v1, i64* %v2, i64* %v3, i64* %v4, i64* %v5, i64* %v6, i64* %v7, i64* %v8,
+                        i64* %v9, i64* %v10, i64* %v11, i64* %v12, i64* undef, i64* %v14, i64* %v15, i64* %v16,
+                        i64* %v17, i64* %v18, i64* %v19, i64* %v20, i64* %v21, i64* %v22, i64* %v23, i64* %v24,
+                        i64* %v25, i64* %v26, i64* %v27, i64* %v28, i64* %v29, i64* %v30, i64* undef, i64* undef,
+                        i64* undef, i64* undef, i64* undef, i64* undef, i64* %v37, i64* %v38, i64* %v39, i64* %v40,
+                        i64* %v41, i64* %v42, i64* %v43, i64* %v44, i64* %v45, i64* undef, i64* %v47, i64* %v48,
+                        i64 undef)
+  ret void
+}
+
+declare dso_local void @func2(i64*, i64*, i64*, i64*, i64*, i64*, i64*, i64*,
+                              i64*, i64*, i64*, i64*, i64*, i64*, i64*, i64*,
+                              i64*, i64*, i64*, i64*, i64*, i64*, i64*, i64*,
+                              i64*, i64*, i64*, i64*, i64*, i64*, i64*, i64*,
+                              i64*, i64*, i64*, i64*, i64*, i64*, i64*, i64*,
+                              i64*, i64*, i64*, i64*, i64*, i64*, i64*, i64*,
+                              i64)
+
+attributes #0 = { "target-features"="+sve" vscale_range(2,2) }

llvm/test/CodeGen/AArch64/sve-fixed-length-frame-offests.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -debug-only=isel < %s 2>&1 | FileCheck %s
+
+; REQUIRES: asserts
+
+target triple = "aarch64-unknown-linux-gnu"
+
+; Ensure that only no offset frame indexes are folded into SVE load/stores when
+; accessing fixed width objects.
+define void @foo(<8 x i64>* %a) #0 {
+; CHECK-LABEL: foo:
+; CHECK:       SelectionDAG has 14 nodes:
+; CHECK-NEXT:    t0: ch = EntryToken
+; CHECK-NEXT:    t12: nxv2i1 = PTRUE_D TargetConstant:i32<31>
+; CHECK-NEXT:    t2: i64,ch = CopyFromReg t0, Register:i64 %0
+; CHECK-NEXT:    t18: nxv2i64,ch = LD1D_IMM<Mem:(volatile load (s512) from %ir.a)> t12, t2, TargetConstant:i64<0>, t0
+; CHECK-NEXT:    t8: i64 = ADDXri TargetFrameIndex:i64<1>, TargetConstant:i32<0>, TargetConstant:i32<0>
+; CHECK-NEXT:    t17: ch = ST1D_IMM<Mem:(volatile store (s512) into %ir.r0)> t18, t12, TargetFrameIndex:i64<0>, TargetConstant:i64<0>, t0
+; CHECK-NEXT:    t16: ch = ST1D_IMM<Mem:(volatile store (s512) into %ir.r1)> t18, t12, t8, TargetConstant:i64<0>, t17
+; CHECK-NEXT:    t10: ch = RET_ReallyLR t16
+; CHECK-EMPTY:
+entry:
+  %r0 = alloca <8 x i64>
+  %r1 = alloca <8 x i64>
+  %r = load volatile <8 x i64>, <8 x i64>* %a
+  store volatile <8 x i64> %r, <8 x i64>* %r0
+  store volatile <8 x i64> %r, <8 x i64>* %r1
+  ret void
+}
+
+attributes #0 = { nounwind "target-features"="+sve" vscale_range(4,4) }