[CodeGen] Mark setjmp/catchret MBBs address-taken

[JosephTremoulet]

Summary

This ensures that BranchFolding (and similar) won't remove these blocks.

Also allow AsmPrinter::EmitBasicBlockStart to process MBBs which are
address-taken but do not have BBs that are address-taken, since otherwise
its call to getAddrLabelSymbolTableToEmit would fail an assertion on such
blocks. I audited the other callers of getAddrLabelSymbolTableToEmit
(and getAddrLabelSymbol); they all have BBs known to be address-taken
except for the call through getAddrLabelSymbol from
WinException::create32bitRef; that call is actually now unreachable, so
I've removed it and updated the signature of create32bitRef.

This fixes PR25168.

[rnk]

I'd actually like this transform to work, and I don't want WinEH to get entangled with BlockAddress support, which is what hasAddressTaken is for.

My idea was to create a pseudo instr during frame lowering that defs EAX/RAX and takes no operands. During X86MCInstrLowering, we would scan forwards to the CATCHRET terminator, get the destination MBB operand, and print its symbol.

[JosephTremoulet]

I'd actually like this transform to work,

In the case that we hit in the bug, we actually get the same codegen despite the MBB being marked address-taken:

• Tail-merge can still reduce the address-taken block to a jmp
• We can still position the address-taken block before its succ, leaving it empty
• We just can't remove the empty block anymore, so we use a different label for the same codegen

Of course, I agree we wouldn't always be able to get the same codegen. I believe the setjmp testcase added here left around a block with a jmp, owing I guess to the successor having multiple preds...

I don't want WinEH to get entangled with BlockAddress support, which is what hasAddressTaken is for.

I was actually trying to be careful *not* to use BlockAddress, so that nothing is address-taken at the IR level. I take it you're saying that distinction isn't helpful?

Happy to consider alternatives:

1. BranchFolding knows how to rewrite jump tables; perhaps something patterned after that to allow it to rewrite the lea/mov for a catchret would be better? I considered that but since the lea/mov is in code and the jump tables are data it seemed like it would be a harder problem.
2. I'm not sure I understand why we need to split out the lea/mov at PEI. Could we persist a CATCHRET with explicit target block through to X86MCInsrtLowering, and there find the appropriate place to stick the lea/mov?

My idea was to create a pseudo instr during frame lowering that defs EAX/RAX and takes no operands. During X86MCInstrLowering, we would scan forwards to the CATCHRET terminator, get the destination MBB operand, and print its symbol.

Since the same catch could have multiple catchrets targeting different successors, I'd be worried about the pseudos for different targets getting merged together somehow and then we wouldn't be able to expand the merged pseudo to something correct. Does that sound right, or am I misunderstanding?

Thanks.

[rnk]

lgtm
I looked at the setjmp test case, and that convinced me that your approach here is better.

lib/CodeGen/AsmPrinter/AsmPrinter.cpp

2483

There should be a comment here explaining that MBBs can have their address taken as part of CodeGen without having their address taken in IR.

[JosephTremoulet]

rebase, incorporate feedback

[JosephTremoulet]

update newly-added tests

lib/CodeGen/AsmPrinter/AsmPrinter.cpp

void AsmPrinter::EmitBasicBlockStart(const MachineBasicBlock &MBB) const {
   // reference the block.  It is possible that there is more than one label
   // here, because multiple LLVM BB's may have been RAUW'd to this block after
   // the references were generated.
   if (MBB.hasAddressTaken()) {
     const BasicBlock *BB = MBB.getBasicBlock();
     if (isVerbose())
       OutStreamer->AddComment("Block address taken");
 
+    if (BB->hasAddressTaken())
       for (MCSymbol *Sym : MMI->getAddrLabelSymbolToEmit(BB))
         OutStreamer->EmitLabel(Sym);
   }
 
   // Print some verbose block comments.
   if (isVerbose()) {
     if (const BasicBlock *BB = MBB.getBasicBlock())
       if (BB->hasName())
         OutStreamer->AddComment("%" + BB->getName());
     emitBasicBlockLoopComments(MBB, LI, *this);

lib/CodeGen/AsmPrinter/WinException.h

 
 #ifndef LLVM_LIB_CODEGEN_ASMPRINTER_WIN64EXCEPTION_H
 #define LLVM_LIB_CODEGEN_ASMPRINTER_WIN64EXCEPTION_H
 
 #include "EHStreamer.h"
 
 namespace llvm {
 class Function;
+class GlobalValue;
 class MachineFunction;
 class MCExpr;
 class Value;
 struct WinEHFuncInfo;
 
 class LLVM_LIBRARY_VISIBILITY WinException : public EHStreamer {
   /// Per-function flag to indicate if personality info should be emitted.
   bool shouldEmitPersonality = false;
void computeIP2StateTable(
       SmallVectorImpl<std::pair<const MCExpr *, int>> &IPToStateTable);
 
   /// Emits the label used with llvm.x86.seh.recoverfp, which is used by
   /// outlined funclets.
   void emitEHRegistrationOffsetLabel(const WinEHFuncInfo &FuncInfo,
                                      StringRef FLinkageName);
 
   const MCExpr *create32bitRef(const MCSymbol *Value);
-  const MCExpr *create32bitRef(const Value *V);
+  const MCExpr *create32bitRef(const GlobalValue *GV);
   const MCExpr *getLabelPlusOne(const MCSymbol *Label);
   const MCExpr *getOffset(const MCSymbol *OffsetOf, const MCSymbol *OffsetFrom);
   const MCExpr *getOffsetPlusOne(const MCSymbol *OffsetOf,
                                  const MCSymbol *OffsetFrom);
 
   /// Gets the offset that we should use in a table for a stack object with the
   /// given index. For targets using CFI (Win64, etc), this is relative to the
   /// established SP at the end of the prologue. For targets without CFI (Win32

lib/CodeGen/AsmPrinter/WinException.cpp

const MCExpr *WinException::create32bitRef(const MCSymbol *Value) {
   if (!Value)
     return MCConstantExpr::create(0, Asm->OutContext);
   return MCSymbolRefExpr::create(Value, useImageRel32
                                             ? MCSymbolRefExpr::VK_COFF_IMGREL32
                                             : MCSymbolRefExpr::VK_None,
                                  Asm->OutContext);
 }
 
-const MCExpr *WinException::create32bitRef(const Value *V) {
-  if (!V)
+const MCExpr *WinException::create32bitRef(const GlobalValue *GV) {
+  if (!GV)
     return MCConstantExpr::create(0, Asm->OutContext);
-  if (const auto *GV = dyn_cast<GlobalValue>(V))
-    return create32bitRef(Asm->getSymbol(GV));
-  return create32bitRef(MMI->getAddrLabelSymbol(cast<BasicBlock>(V)));
 }
 
 const MCExpr *WinException::getLabelPlusOne(const MCSymbol *Label) {
   return MCBinaryExpr::createAdd(create32bitRef(Label),
                                  MCConstantExpr::create(1, Asm->OutContext),
                                  Asm->OutContext);
 }
 

lib/Target/X86/X86FrameLowering.cpp

if (STI.is64Bit()) {
           .addMBB(RestoreMBB)
           .addReg(0);
     } else {
       // MOV32ri $RestoreMBB, %eax
       BuildMI(MBB, FirstCSPop, DL, TII.get(X86::MOV32ri))
           .addReg(ReturnReg)
           .addMBB(RestoreMBB);
     }
+    // Record that we've taken the address of RestoreMBB and no longer just
+    // reference it in a terminator.
+    RestoreMBB->setHasAddressTaken();
   }
 
   if (MBBI != MBB.end())
     DL = MBBI->getDebugLoc();
 
   // If there is an ADD32ri or SUB32ri of ESP immediately before this
   // instruction, merge the two instructions.
   if (NumBytes || MFI->hasVarSizedObjects())

lib/Target/X86/X86ISelLowering.cpp

X86TargetLowering::emitEHSjLjSetJmp(MachineInstr *MI,
 
   MVT PVT = getPointerTy(MF->getDataLayout());
   assert((PVT == MVT::i64 || PVT == MVT::i32) &&
          "Invalid Pointer Size!");
 
   // For v = setjmp(buf), we generate
   //
   // thisMBB:
-  //  buf[LabelOffset] = restoreMBB
+  //  buf[LabelOffset] = restoreMBB <-- takes address of restoreMBB
   //  SjLjSetup restoreMBB
   //
   // mainMBB:
   //  v_main = 0
   //
   // sinkMBB:
   //  v = phi(main, restore)
   //
   // restoreMBB:
   //  if base pointer being used, load it from frame
   //  v_restore = 1
 
   MachineBasicBlock *thisMBB = MBB;
   MachineBasicBlock *mainMBB = MF->CreateMachineBasicBlock(BB);
   MachineBasicBlock *sinkMBB = MF->CreateMachineBasicBlock(BB);
   MachineBasicBlock *restoreMBB = MF->CreateMachineBasicBlock(BB);
   MF->insert(I, mainMBB);
   MF->insert(I, sinkMBB);
   MF->push_back(restoreMBB);
+  restoreMBB->setHasAddressTaken();
 
   MachineInstrBuilder MIB;
 
   // Transfer the remainder of BB and its successor edges to sinkMBB.
   sinkMBB->splice(sinkMBB->begin(), MBB,
                   std::next(MachineBasicBlock::iterator(MI)), MBB->end());
   sinkMBB->transferSuccessorsAndUpdatePHIs(MBB);
 

test/CodeGen/X86/late-address-taken.ll

+; RUN: llc -mtriple=x86_64-pc-windows-msvc < %s | FileCheck %s
+
+; Repro cases from PR25168
+
+; test @catchret - catchret target is not address-taken until PEI
+; splits it into lea/mov followed by ret.  Make sure the MBB is
+; handled, both by tempting BranchFolding to merge it with %early_out
+; and delete it, and by checking that we emit a proper reference
+; to it in the LEA
+
+declare void @ProcessCLRException()
+declare void @f()
+
+define void @catchret(i1 %b) personality void ()* @ProcessCLRException {
+entry:
+  br i1 %b, label %body, label %early_out
+early_out:
+  ret void
+body:
+  invoke void @f()
+          to label %exit unwind label %catch.pad
+catch.pad:
+  %catch = catchpad [i32 33554467]
+          to label %catch.body unwind label %catch.end
+catch.body:
+  catchret %catch to label %exit
+catch.end:
+  catchendpad unwind to caller
+exit:
+  ret void
+}
+; CHECK-LABEL: catchret:  # @catchret
+; CHECK: [[Exit:^[^ :]+]]: # Block address taken
+; CHECK-NEXT:              # %exit
+; CHECK: # %catch.pad
+; CHECK: .seh_endprolog
+; CHECK: leaq [[Exit]](%rip), %rax
+; CHECK: retq # CATCHRET
+
+
+; test @setjmp - similar to @catchret, but the MBB in question
+; is the one generated when the setjmp's block is split
+
+@buf = internal global [5 x i8*] zeroinitializer
+declare i8* @llvm.frameaddress(i32) nounwind readnone
+declare i8* @llvm.stacksave() nounwind
+declare i32 @llvm.eh.sjlj.setjmp(i8*) nounwind
+declare void @llvm.eh.sjlj.longjmp(i8*) nounwind
+
+define void @setjmp(i1 %b) nounwind {
+entry:
+  br i1 %b, label %early_out, label %sj
+early_out:
+  ret void
+sj:
+  %fp = call i8* @llvm.frameaddress(i32 0)
+  store i8* %fp, i8** getelementptr inbounds ([5 x i8*], [5 x i8*]* @buf, i64 0, i64 0), align 16
+  %sp = call i8* @llvm.stacksave()
+  store i8* %sp, i8** getelementptr inbounds ([5 x i8*], [5 x i8*]* @buf, i64 0, i64 2), align 16
+  call i32 @llvm.eh.sjlj.setjmp(i8* bitcast ([5 x i8*]* @buf to i8*))
+  ret void
+}
+; CHECK-LABEL: setjmp: # @setjmp
+; CHECK: # %sj
+; CHECK: leaq [[Label:\..+]](%rip), %[[Reg:.+]]{{$}}
+; CHECK-NEXT: movq %[[Reg]], buf
+; CHECK: {{^}}[[Label]]:  # Block address taken
+; CHECK-NEXT:              # %sj

test/CodeGen/X86/win-catchpad-csrs.ll

 ; X64: .seh_endprologue
 ; X64: callq getint
 ; X64: callq getint
 ; X64: callq getint
 ; X64: callq getint
 ; X64: callq useints
 ; X64: movl $1, %ecx
 ; X64: callq f
-; X64: [[contbb:\.LBB0_[0-9]+]]: # %try.cont
+; X64: [[contbb:\.LBB0_[0-9]+]]: # Block address taken
+; X64-NEXT:                      # %try.cont
 ; X64: addq $40, %rsp
 ; X64: popq %rbp
 ; X64: retq
 
 ; X64: "?catch$[[catch1bb:[0-9]+]]@?0?try_catch_catch@4HA":
 ; X64: LBB0_[[catch1bb]]: # %catch.dispatch{{$}}
 ; X64: movq %rdx, 16(%rsp)
 ; X64: pushq %rbp

test/CodeGen/X86/win-catchpad.ll

 ; X86: movl $0, -{{[0-9]+}}(%ebp)
 ; X86: leal -[[local_offs:[0-9]+]](%ebp), %[[addr_reg:[a-z]+]]
 ; X86-DAG: movl %[[addr_reg]], 4(%esp)
 ; X86-DAG: movl $1, (%esp)
 ; X86: calll _f
 ; X86: [[contbb:LBB0_[0-9]+]]: # %try.cont
 ; X86: retl
 
-; X86: [[restorebb1:LBB0_[0-9]+]]: # %invoke.cont.2
+; X86: [[restorebb1:LBB0_[0-9]+]]: # Block address taken
+; X86-NEXT:                        # %invoke.cont.2
 ; X86: movl -16(%ebp), %esp
 ; X86: addl $12, %ebp
 ; X86: jmp [[contbb]]
 
 ; FIXME: These should be de-duplicated.
-; X86: [[restorebb2:LBB0_[0-9]+]]: # %invoke.cont.3
+; X86: [[restorebb2:LBB0_[0-9]+]]: # Block address taken
+; X86-NEXT:                        # %invoke.cont.3
 ; X86: movl -16(%ebp), %esp
 ; X86: addl $12, %ebp
 ; X86: jmp [[contbb]]
 
 ; X86: "?catch$[[catch1bb:[0-9]+]]@?0?try_catch_catch@4HA":
 ; X86: LBB0_[[catch1bb]]: # %catch.dispatch{{$}}
 ; X86: pushl %ebp
 ; X86: subl $8, %esp

 ; X64: .seh_stackalloc 48
 ; X64: leaq 48(%rsp), %rbp
 ; X64: .seh_setframe 5, 48
 ; X64: .seh_endprologue
 ; X64: .Ltmp0
 ; X64-DAG: leaq -[[local_offs:[0-9]+]](%rbp), %rdx
 ; X64-DAG: movl $1, %ecx
 ; X64: callq f
-; X64: [[contbb:\.LBB0_[0-9]+]]: # %try.cont
+; X64: [[contbb:\.LBB0_[0-9]+]]: # Block address taken
+; X64-NEXT:                      # %try.cont
 ; X64: addq $48, %rsp
 ; X64: popq %rbp
 ; X64: retq
 
 ; X64: "?catch$[[catch1bb:[0-9]+]]@?0?try_catch_catch@4HA":
 ; X64: LBB0_[[catch1bb]]: # %catch.dispatch{{$}}
 ; X64: movq %rdx, 16(%rsp)
 ; X64: pushq %rbp

 }
 
 ; X86-LABEL: _branch_to_normal_dest:
 ; X86: calll _f
 
 ; X86: [[contbb:LBB1_[0-9]+]]: # %try.cont
 ; X86: retl
 
-; X86: [[restorebb:LBB1_[0-9]+]]: # %catch.done
+; X86: [[restorebb:LBB1_[0-9]+]]: # Block address taken
+; X86-NEXT:                       # %catch.done
 ; X86: movl -16(%ebp), %esp
 ; X86: addl $12, %ebp
 ; X86: jmp [[contbb]]
 
 ; X86: "?catch$[[catchdispbb:[0-9]+]]@?0?branch_to_normal_dest@4HA":
 ; X86: LBB1_[[catchdispbb]]: # %catch.dispatch{{$}}
 ; X86: pushl %ebp
 ; X86: subl $8, %esp

 ; X64: subq $48, %rsp
 ; X64: .seh_stackalloc 48
 ; X64: leaq 48(%rsp), %rbp
 ; X64: .seh_setframe 5, 48
 ; X64: .seh_endprologue
 ; X64: .Ltmp[[before_call:[0-9]+]]:
 ; X64: callq f
 ; X64: .Ltmp[[after_call:[0-9]+]]:
-; X64: [[contbb:\.LBB1_[0-9]+]]: # %try.cont
+; X64: [[contbb:\.LBB1_[0-9]+]]: # Block address taken
+; X64-NEXT:                      # %try.cont
 ; X64: addq $48, %rsp
 ; X64: popq %rbp
 ; X64: retq
 
 ; X64: "?catch$[[catchbb:[0-9]+]]@?0?branch_to_normal_dest@4HA":
 ; X64: LBB1_[[catchbb]]: # %catch.dispatch{{$}}
 ; X64: movq %rdx, 16(%rsp)
 ; X64: pushq %rbp