[X86][VARARG] Assign MMO earlier to avoid prolog insert point been sunk across VASTART_SAVE_XMM_REGS

[pengfei]

Summary

The changes in D80163 defered the assignment of MachineMemOperand (MMO)
until the X86ExpandPseudo pass. This will result in crash due to prolog
insert point been sunk across the pseudo instruction VASTART_SAVE_XMM_REGS.

Moving the assignment to the creation of the node can avoid the problem.

[avl]

It looks good to me. Though, wait for one more reviewer, please.

[rnk]

lgtm

[pengfei]

Thanks @avl and @rnk

llvm/lib/Target/X86/X86ExpandPseudo.cpp

void X86ExpandPseudo::ExpandVastartSaveXmmRegs(
   Func->insert(EntryBlkIter, TailBlk);
 
   // Transfer the remainder of EntryBlk and its successor edges to TailBlk.
   TailBlk->splice(TailBlk->begin(), EntryBlk,
                   std::next(MachineBasicBlock::iterator(VAStartPseudoInstr)),
                   EntryBlk->end());
   TailBlk->transferSuccessorsAndUpdatePHIs(EntryBlk);
 
-  int64_t FrameIndex = VAStartPseudoInstr->getOperand(1).getImm();
-  Register BaseReg;
-  uint64_t FrameOffset =
-      X86FL->getFrameIndexReference(*Func, FrameIndex, BaseReg).getFixed();
-  uint64_t VarArgsRegsOffset = VAStartPseudoInstr->getOperand(2).getImm();
+  uint64_t FrameOffset = VAStartPseudoInstr->getOperand(4).getImm();
+  uint64_t VarArgsRegsOffset = VAStartPseudoInstr->getOperand(6).getImm();
 
   // TODO: add support for YMM and ZMM here.
   unsigned MOVOpc = STI->hasAVX() ? X86::VMOVAPSmr : X86::MOVAPSmr;
 
   // In the XMM save block, save all the XMM argument registers.
-  for (int64_t OpndIdx = 3, RegIdx = 0;
+  for (int64_t OpndIdx = 7, RegIdx = 0;
        OpndIdx < VAStartPseudoInstr->getNumOperands() - 1;
        OpndIdx++, RegIdx++) {
-
-    int64_t Offset = FrameOffset + VarArgsRegsOffset + RegIdx * 16;
-
-    MachineMemOperand *MMO = Func->getMachineMemOperand(
-        MachinePointerInfo::getFixedStack(*Func, FrameIndex, Offset),
-        MachineMemOperand::MOStore,
-        /*Size=*/16, Align(16));
-
-    BuildMI(GuardedRegsBlk, DL, TII->get(MOVOpc))
-        .addReg(BaseReg)
-        .addImm(/*Scale=*/1)
-        .addReg(/*IndexReg=*/0)
-        .addImm(/*Disp=*/Offset)
-        .addReg(/*Segment=*/0)
-        .addReg(VAStartPseudoInstr->getOperand(OpndIdx).getReg())
-        .addMemOperand(MMO);
+    auto NewMI = BuildMI(GuardedRegsBlk, DL, TII->get(MOVOpc));
+    for (int i = 0; i < X86::AddrNumOperands; ++i) {
+      if (i == X86::AddrDisp)
+        NewMI.addImm(FrameOffset + VarArgsRegsOffset + RegIdx * 16);
+      else
+        NewMI.add(VAStartPseudoInstr->getOperand(i + 1));
+    }
+    NewMI.addReg(VAStartPseudoInstr->getOperand(OpndIdx).getReg());
     assert(Register::isPhysicalRegister(
         VAStartPseudoInstr->getOperand(OpndIdx).getReg()));
   }
 
   // The original block will now fall through to the GuardedRegsBlk.
   EntryBlk->addSuccessor(GuardedRegsBlk);
   // The GuardedRegsBlk will fall through to the TailBlk.
   GuardedRegsBlk->addSuccessor(TailBlk);

llvm/lib/Target/X86/X86ISelLowering.h

enum NodeType : unsigned {
     // Masked version of above.
     // SRC, PASSTHRU, MASK
     MCVTNEPS2BF16,
 
     // Dot product of BF16 pairs to accumulated into
     // packed single precision.
     DPBF16PS,
 
-    // Save xmm argument registers to the stack, according to %al. An operator
-    // is needed so that this can be expanded with control flow.
-    VASTART_SAVE_XMM_REGS,
-
     // A stack checking function call. On Windows it's _chkstk call.
     DYN_ALLOCA,
 
     // For allocating variable amounts of stack space when using
     // segmented stacks. Check if the current stacklet has enough space, and
     // falls back to heap allocation if not.
     SEG_ALLOCA,
 
enum NodeType : unsigned {
     AESDEC128KL,
     AESENC256KL,
     AESDEC256KL,
     AESENCWIDE128KL,
     AESDECWIDE128KL,
     AESENCWIDE256KL,
     AESDECWIDE256KL,
 
+    // Save xmm argument registers to the stack, according to %al. An operator
+    // is needed so that this can be expanded with control flow.
+    VASTART_SAVE_XMM_REGS,
+
     // WARNING: Do not add anything in the end unless you want the node to
     // have memop! In fact, starting from FIRST_TARGET_MEMORY_OPCODE all
     // opcodes will be thought as target memory ops!
   };
   } // end namespace X86ISD
 
   namespace X86 {
     /// Current rounding mode is represented in bits 11:10 of FPSR. These

llvm/lib/Target/X86/X86ISelLowering.cpp

for (SDValue Val : LiveGPRs) {
       Offset += 8;
     }
 
     // Now store the XMM (fp + vector) parameter registers.
     if (!LiveXMMRegs.empty()) {
       SmallVector<SDValue, 12> SaveXMMOps;
       SaveXMMOps.push_back(Chain);
       SaveXMMOps.push_back(ALVal);
-      SaveXMMOps.push_back(
-          DAG.getTargetConstant(FuncInfo->getRegSaveFrameIndex(), DL, MVT::i32));
+      SaveXMMOps.push_back(RSFIN);
       SaveXMMOps.push_back(
           DAG.getTargetConstant(FuncInfo->getVarArgsFPOffset(), DL, MVT::i32));
       llvm::append_range(SaveXMMOps, LiveXMMRegs);
-      MemOps.push_back(DAG.getNode(X86ISD::VASTART_SAVE_XMM_REGS, DL,
-                                   MVT::Other, SaveXMMOps));
+      MachineMemOperand *StoreMMO =
+          DAG.getMachineFunction().getMachineMemOperand(
+              MachinePointerInfo::getFixedStack(
+                  DAG.getMachineFunction(), FuncInfo->getRegSaveFrameIndex(),
+                  Offset),
+              MachineMemOperand::MOStore, 128, Align(16));
+      MemOps.push_back(DAG.getMemIntrinsicNode(X86ISD::VASTART_SAVE_XMM_REGS,
+                                               DL, DAG.getVTList(MVT::Other),
+                                               SaveXMMOps, MVT::i8, StoreMMO));
     }
 
     if (!MemOps.empty())
       Chain = DAG.getNode(ISD::TokenFactor, DL, MVT::Other, MemOps);
   }
 }
 
 void VarArgsLoweringHelper::forwardMustTailParameters(SDValue &Chain) {

llvm/lib/Target/X86/X86InstrCompiler.td

def ADJCALLSTACKUP64   : I<0, Pseudo, (outs), (ins i32imm:$amt1, i32imm:$amt2),
                            Requires<[IsLP64]>;
 }
 def : Pat<(X86callseq_start timm:$amt1, timm:$amt2),
         (ADJCALLSTACKDOWN64 i32imm:$amt1, i32imm:$amt2, 0)>, Requires<[IsLP64]>;
 
 let SchedRW = [WriteSystem] in {
 
 // x86-64 va_start lowering magic.
-let hasSideEffects = 1, Defs = [EFLAGS] in {
+let hasSideEffects = 1, mayStore = 1, Defs = [EFLAGS] in {
 def VASTART_SAVE_XMM_REGS : I<0, Pseudo,
                               (outs),
-                              (ins GR8:$al,
-                                   i32imm:$regsavefi, i32imm:$offset,
-                                   variable_ops),
-                              "#VASTART_SAVE_XMM_REGS $al, $regsavefi, $offset",
-                              [(X86vastart_save_xmm_regs GR8:$al,
-                                                         timm:$regsavefi,
-                                                         timm:$offset),
+                              (ins GR8:$al, i8mem:$regsavefi, variable_ops),
+                              "#VASTART_SAVE_XMM_REGS $al, $regsavefi",
+                              [(X86vastart_save_xmm_regs GR8:$al, addr:$regsavefi),
                                (implicit EFLAGS)]>;
 }
 
 let usesCustomInserter = 1, Defs = [EFLAGS] in {
 // The VAARG_64 and VAARG_X32 pseudo-instructions take the address of the
 // va_list, and place the address of the next argument into a register.
 let Defs = [EFLAGS] in {
 def VAARG_64 : I<0, Pseudo,

llvm/lib/Target/X86/X86InstrInfo.td

 def SDT_X86CallSeqEnd   : SDCallSeqEnd<[SDTCisVT<0, i32>,
                                         SDTCisVT<1, i32>]>;
 
 def SDT_X86Call   : SDTypeProfile<0, -1, [SDTCisVT<0, iPTR>]>;
 
 def SDT_X86NtBrind : SDTypeProfile<0, -1, [SDTCisVT<0, iPTR>]>;
 
 def SDT_X86VASTART_SAVE_XMM_REGS : SDTypeProfile<0, -1, [SDTCisVT<0, i8>,
-                                                         SDTCisVT<1, iPTR>,
-                                                         SDTCisVT<2, iPTR>]>;
+                                                         SDTCisPtrTy<1>]>;
 
 def SDT_X86VAARG : SDTypeProfile<1, -1, [SDTCisPtrTy<0>,
                                          SDTCisPtrTy<1>,
                                          SDTCisVT<2, i32>,
                                          SDTCisVT<3, i8>,
                                          SDTCisVT<4, i32>]>;
 
 def SDTX86RepStr  : SDTypeProfile<0, 1, [SDTCisVT<0, OtherVT>]>;

 def X86retflag : SDNode<"X86ISD::RET_FLAG", SDTX86Ret,
                         [SDNPHasChain, SDNPOptInGlue, SDNPVariadic]>;
 def X86iret : SDNode<"X86ISD::IRET", SDTX86Ret,
                         [SDNPHasChain, SDNPOptInGlue]>;
 
 def X86vastart_save_xmm_regs :
                  SDNode<"X86ISD::VASTART_SAVE_XMM_REGS",
                         SDT_X86VASTART_SAVE_XMM_REGS,
-                        [SDNPHasChain, SDNPVariadic]>;
+                        [SDNPHasChain, SDNPMayStore, SDNPMemOperand, SDNPVariadic]>;
 def X86vaarg64 :
                  SDNode<"X86ISD::VAARG_64", SDT_X86VAARG,
                         [SDNPHasChain, SDNPMayLoad, SDNPMayStore,
                          SDNPMemOperand]>;
 def X86vaargx32 :
                  SDNode<"X86ISD::VAARG_X32", SDT_X86VAARG,
                         [SDNPHasChain, SDNPMayLoad, SDNPMayStore,
                          SDNPMemOperand]>;

llvm/test/CodeGen/X86/vaargs-prolog-insert.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=x86_64 < %s | FileCheck %s
 
 ; Check the prolog won't be sunk across the save of CSRs.
 define void @reduce(i32, i32, i32, i32, i32, i32, ...) nounwind {
 ; CHECK-LABEL: reduce:
 ; CHECK:       # %bb.0:
+; CHECK-NEXT:    subq $56, %rsp
 ; CHECK-NEXT:    testb %al, %al
 ; CHECK-NEXT:    je .LBB0_4
 ; CHECK-NEXT:  # %bb.3:
 ; CHECK-NEXT:    movaps %xmm0, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movaps %xmm1, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movaps %xmm2, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movaps %xmm3, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movaps %xmm4, -{{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movaps %xmm5, (%rsp)
 ; CHECK-NEXT:    movaps %xmm6, {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:    movaps %xmm7, {{[0-9]+}}(%rsp)
 ; CHECK-NEXT:  .LBB0_4:
 ; CHECK-NEXT:    xorl %eax, %eax
 ; CHECK-NEXT:    testb %al, %al
 ; CHECK-NEXT:    jne .LBB0_2
 ; CHECK-NEXT:  # %bb.1:
-; CHECK-NEXT:    subq $56, %rsp
 ; CHECK-NEXT:    leaq -{{[0-9]+}}(%rsp), %rax
 ; CHECK-NEXT:    movq %rax, 16
 ; CHECK-NEXT:    leaq {{[0-9]+}}(%rsp), %rax
 ; CHECK-NEXT:    movq %rax, 8
 ; CHECK-NEXT:    movl $48, 4
 ; CHECK-NEXT:    movl $48, 0
-; CHECK-NEXT:    addq $56, %rsp
 ; CHECK-NEXT:  .LBB0_2:
+; CHECK-NEXT:    addq $56, %rsp
 ; CHECK-NEXT:    retq
   br i1 undef, label %8, label %7
 
 7:                                                ; preds = %6
   call void @llvm.va_start(i8* null)
   br label %8
 
 8:                                                ; preds = %7, %6
   ret void
 }
 
 declare void @llvm.va_start(i8*)
 declare void @llvm.va_end(i8*)