[CodeGen] Make CreateStackTemporary work for scalable vectors

[david-arm]

Summary

In CreateStackTemporary we were promoting alignments beyond the
stack alignment, which I have fixed in this patch. In addition, we
need to set the stack id explicitly for scalable vectors so that
we allocate the correct amount of stack using vscale.

I added a test to

CodeGen/AArch64/sve-insert-element.ll

that tries to insert an element into an illegal scalable vector
type that involves creating temporary stack objects.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

2020

Can you split this patch up into two patches:

• one that prevents the stack temporary to exceed the stack alignment (which also needs tests for fixed-width vectors). Please also add @efriedma as reviewer to these patches, since he recently fixed a very similar problem in D79532.
• another to make createStackTemporary support scalable types.

[david-arm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

2020

Sure. It will take me a long time to get a patch ready for the stack alignment fix anyway due to it breaking 64 tests!

[david-arm]

This patch will be split into two patches in this submission order:

1. A patch to fix the alignment and fix up all of the tests, which will take time since at least some of the tests were using the alignment as a convenient way to defend bug fixes for instruction selection, etc.
2. A patch to add scalable vector support for CreateStackTemporary.

I've decided to abandon this revision for now and will create new ones later.

[efriedma]

I would split the patch into the following steps:

1. Introduce the new CreateStackTemporary(TypeSize, Align) API, and refactor the code.
2. Set the StackID in CreateStackTemporary
3. Fix the APIs that take a VT to choose more appropriate alignment. This is the hard part; probably requires some work to make the interaction between the alignment assumed by the users, vs. the actual alignment, work correctly. Probably makes sense to do this last; you might end up splitting this into multiple patches and/or change users of the API.

llvm/include/llvm/CodeGen/SelectionDAG.h

public:
   }
 
   /// To be invoked on an SDNode that is slated to be erased. This
   /// function mirrors \c llvm::salvageDebugInfo.
   void salvageDebugInfo(SDNode &N);
 
   void dump() const;
 
+  /// Create a stack temporary based on the size in bytes and the alignment
+  SDValue CreateStackTemporary(TypeSize Bytes, Align Alignment);
+
   /// Create a stack temporary, suitable for holding the specified value type.
   /// If minAlign is specified, the slot size will have at least that alignment.
   SDValue CreateStackTemporary(EVT VT, unsigned minAlign = 1);
 
   /// Create a stack temporary suitable for holding either of the specified
   /// value types.
   SDValue CreateStackTemporary(EVT VT1, EVT VT2);
 

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

 #include "llvm/CodeGen/MachineConstantPool.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/RuntimeLibcalls.h"
 #include "llvm/CodeGen/SelectionDAGAddressAnalysis.h"
 #include "llvm/CodeGen/SelectionDAGNodes.h"
 #include "llvm/CodeGen/SelectionDAGTargetInfo.h"
+#include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/CodeGen/ValueTypes.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugInfoMetadata.h"
SDValue SelectionDAG::expandVACopy(SDNode *Node) {
   const Value *VS = cast<SrcValueSDNode>(Node->getOperand(4))->getValue();
   SDValue Tmp1 =
       getLoad(TLI.getPointerTy(getDataLayout()), dl, Node->getOperand(0),
               Node->getOperand(2), MachinePointerInfo(VS));
   return getStore(Tmp1.getValue(1), dl, Tmp1, Node->getOperand(1),
                   MachinePointerInfo(VD));
 }
 
-SDValue SelectionDAG::CreateStackTemporary(EVT VT, unsigned minAlign) {
-  MachineFrameInfo &MFI = getMachineFunction().getFrameInfo();
-  unsigned ByteSize = VT.getStoreSize();
-  Type *Ty = VT.getTypeForEVT(*getContext());
-  unsigned StackAlign =
-      std::max((unsigned)getDataLayout().getPrefTypeAlignment(Ty), minAlign);
+SDValue SelectionDAG::CreateStackTemporary(TypeSize Bytes, Align Alignment) {
+  // Scalable vectors need a special StackID to distinguish them from other
+  // (fixed size) stack objects.
+  unsigned StackId = 0;
+  if (Bytes.isScalable()) {
+    const TargetFrameLowering *TFI = MF->getSubtarget().getFrameLowering();
+    StackId = TFI->getStackIDForScalableVectors();
+  }
 
-  int FrameIdx = MFI.CreateStackObject(ByteSize, StackAlign, false);
+  MachineFrameInfo &MFI = MF->getFrameInfo();
+  int FrameIdx =
+      MFI.CreateStackObject(Bytes, Alignment, false, nullptr, StackId);
   return getFrameIndex(FrameIdx, TLI->getFrameIndexTy(getDataLayout()));
 }
 
+SDValue SelectionDAG::CreateStackTemporary(EVT VT, unsigned minAlign) {
+  const TargetFrameLowering *TFI = MF->getSubtarget().getFrameLowering();
+  Type *Ty = VT.getTypeForEVT(*getContext());
+  const Align PrefAlign = getDataLayout().getPrefTypeAlign(Ty);
+  const Align StackAlign = TFI->getStackAlign();
+  const Align Alignment =
+      std::max(std::min(PrefAlign, StackAlign), Align(minAlign));
+  return CreateStackTemporary(VT.getStoreSize(), Alignment);
+}
+
 SDValue SelectionDAG::CreateStackTemporary(EVT VT1, EVT VT2) {
-  unsigned Bytes = std::max(VT1.getStoreSize(), VT2.getStoreSize());
+  const TypeSize Bytes = std::max(VT1.getStoreSize(), VT2.getStoreSize());
   Type *Ty1 = VT1.getTypeForEVT(*getContext());
   Type *Ty2 = VT2.getTypeForEVT(*getContext());
-  const DataLayout &DL = getDataLayout();
-  unsigned Align =
-      std::max(DL.getPrefTypeAlignment(Ty1), DL.getPrefTypeAlignment(Ty2));
-
-  MachineFrameInfo &MFI = getMachineFunction().getFrameInfo();
-  int FrameIdx = MFI.CreateStackObject(Bytes, Align, false);
-  return getFrameIndex(FrameIdx, TLI->getFrameIndexTy(getDataLayout()));
+  const DataLayout DL = getDataLayout();
+  Align Alignment =
+      std::max(DL.getPrefTypeAlign(Ty1), DL.getPrefTypeAlign(Ty2));
+  return CreateStackTemporary(Bytes, Alignment);
 }
 
 SDValue SelectionDAG::FoldSetCC(EVT VT, SDValue N1, SDValue N2,
                                 ISD::CondCode Cond, const SDLoc &dl) {
   EVT OpVT = N1.getValueType();
 
   // These setcc operations always fold.
   switch (Cond) {

llvm/test/CodeGen/AArch64/sve-insert-element.ll

 }
 
 define <vscale x 16 x i8> @test_lane0_undef_16xi8(i8 %a) {
 ; CHECK-LABEL: test_lane0_undef_16xi8
 ; CHECK:       fmov s0, w0
   %b = insertelement <vscale x 16 x i8> undef, i8 %a, i32 0
   ret <vscale x 16 x i8> %b
 }
+
+; TODO: The code generated for this function is still incorrect, hence there
+; is only a single CHECK line to ensure we use 'addvl' for scalable stack
+; objects.
+define <vscale x 32 x i8> @test_lanex_32xi8(<vscale x 32 x i8> %a, i32 %x) {
+; CHECK-LABEL: test_lanex_32xi8
+; CHECK:       addvl sp, sp, #-2
+  %b = insertelement <vscale x 32 x i8> %a, i8 30, i32 %x
+  ret <vscale x 32 x i8> %b
+}