GoObj: mark live stack-object pointer fields

llvm/include/llvm/IR/Intrinsics.td

@@ -1129,6 +1129,13 @@ def int_call_preallocated_teardown : DefaultAttrsIntrinsic<[], [llvm_token_ty]>;
 def int_callbr_landingpad : Intrinsic<[llvm_any_ty], [LLVMMatchType<0>],
                                       [IntrNoMerge]>;
 
+// Model the out-of-band control-flow edge from a Go defer registration site to
+// the function's deferreturn path. This intrinsic is only valid as the callee
+// of a callbr with one indirect destination and emits no machine instruction.
+def int_go_defer_edge : Intrinsic<[], [],
+    [IntrNoMem, IntrHasSideEffects, IntrNoCallback, IntrNoSync, IntrNoFree,
+     IntrNoDuplicate, IntrNoMerge]>;
+
 def int_structured_gep
     : DefaultAttrsIntrinsic<[llvm_anyptr_ty],
                             [LLVMMatchType<0>, llvm_vararg_ty],

llvm/include/llvm/MC/MCContext.h

@@ -234,6 +234,9 @@ class MCContext {
   /// Go object symbol flags keyed by MC symbol.
   DenseMap<const MCSymbol *, std::pair<uint8_t, uint8_t>> GoObjSymbolFlags;
 
+  /// Go object function ID and function flags keyed by MC symbol.
+  DenseMap<const MCSymbol *, std::pair<uint8_t, uint8_t>> GoObjFunctionInfos;
+
   /// Native Go content-addressable identity hashes keyed by MC symbol.
   DenseMap<const MCSymbol *, std::string> GoObjSymbolContentHashes;
 
@@ -709,6 +712,19 @@ class MCContext {
     return It->second;
   }
 
+  void setGoObjFunctionInfo(const MCSymbol *Sym, uint8_t FuncID,
+                            uint8_t FuncFlag) {
+    GoObjFunctionInfos[Sym] = {FuncID, FuncFlag};
+  }
+
+  std::optional<std::pair<uint8_t, uint8_t>>
+  getGoObjFunctionInfo(const MCSymbol *Sym) const {
+    auto It = GoObjFunctionInfos.find(Sym);
+    if (It == GoObjFunctionInfos.end())
+      return std::nullopt;
+    return It->second;
+  }
+
   void setGoObjSymbolContentHash(const MCSymbol *Sym, std::string Hash) {
     GoObjSymbolContentHashes[Sym] = std::move(Hash);
   }

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -898,6 +898,23 @@ getGoObjSymbolFlags(const GlobalObject *GO) {
   return std::make_pair(Flag, Flag2);
 }
 
+static std::optional<std::pair<uint8_t, uint8_t>>
+getGoObjFunctionInfo(const Function &F) {
+  const MDNode *MD = F.getMetadata("goobj.func.info");
+  if (!MD)
+    return std::nullopt;
+  if (MD->getNumOperands() != 2)
+    report_fatal_error("expected !goobj.func.info to have two operands");
+
+  auto ReadByte = [&](unsigned I) -> uint8_t {
+    const auto *CI = mdconst::dyn_extract<ConstantInt>(MD->getOperand(I));
+    if (!CI || CI->getType()->getIntegerBitWidth() != 8)
+      report_fatal_error("expected !goobj.func.info operands to be i8");
+    return static_cast<uint8_t>(CI->getZExtValue());
+  };
+  return std::make_pair(ReadByte(0), ReadByte(1));
+}
+
 static std::optional<std::string>
 getGoObjSymbolContentHash(const GlobalVariable *GV) {
   const MDNode *MD = GV->getMetadata("goobj.content_hash");
@@ -3686,6 +3703,9 @@ void AsmPrinter::SetupMachineFunction(MachineFunction &MF) {
             getGoObjSymbolFlags(&F))
       OutContext.setGoObjSymbolFlags(CurrentFnSym, Flags->first,
                                      Flags->second);
+    if (std::optional<std::pair<uint8_t, uint8_t>> Info =
+            getGoObjFunctionInfo(F))
+      OutContext.setGoObjFunctionInfo(CurrentFnSym, Info->first, Info->second);
     if (goabi::isGoABIInternalCallingConv(F.getCallingConv()))
       OutContext.setGoObjSymbolABI(CurrentFnSym, GoObj::SymABIInternal);
     else if (goabi::isGoABI0CallingConv(F.getCallingConv()))

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

@@ -3133,13 +3133,14 @@ bool IRTranslator::translateCallBr(const User &U,
   MachineBasicBlock *CallBrMBB = &MIRBuilder.getMBB();
 
   Intrinsic::ID IID = I.getIntrinsicID();
+  const bool IsGoDeferEdge = IID == Intrinsic::go_defer_edge;
   if (I.isInlineAsm()) {
     // FIXME: inline asm is not yet supported for callbr in GlobalISel. As soon
     // as we add support, we need to handle the indirect asm targets, see
     // SelectionDAGBuilder::visitCallBr().
     return false;
   }
-  if (!translateIntrinsic(I, IID, MIRBuilder))
+  if (!IsGoDeferEdge && !translateIntrinsic(I, IID, MIRBuilder))
     return false;
 
   // Retrieve successors.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -3627,14 +3627,15 @@ void SelectionDAGBuilder::visitCallBrIntrinsic(const CallBrInst &I) {
 
 void SelectionDAGBuilder::visitCallBr(const CallBrInst &I) {
   MachineBasicBlock *CallBrMBB = FuncInfo.MBB;
+  const bool IsGoDeferEdge = I.getIntrinsicID() == Intrinsic::go_defer_edge;
 
   if (I.isInlineAsm()) {
     // Deopt bundles are lowered in LowerCallSiteWithDeoptBundle, and we don't
     // have to do anything here to lower funclet bundles.
     failForInvalidBundles(I, "callbrs",
                           {LLVMContext::OB_deopt, LLVMContext::OB_funclet});
     visitInlineAsm(I);
-  } else {
+  } else if (!IsGoDeferEdge) {
     assert(!I.hasOperandBundles() &&
            "Can't have operand bundles for intrinsics");
     visitCallBrIntrinsic(I);
@@ -3653,7 +3654,7 @@ void SelectionDAGBuilder::visitCallBr(const CallBrInst &I) {
   // this changes, we might need to enhance
   // Target->setIsInlineAsmBrIndirectTarget or add something similar for
   // intrinsic indirect branches.
-  if (I.isInlineAsm()) {
+  if (I.isInlineAsm() || IsGoDeferEdge) {
     for (BasicBlock *Dest : I.getIndirectDests()) {
       MachineBasicBlock *Target = FuncInfo.getMBB(Dest);
       Target->setIsInlineAsmBrIndirectTarget();

llvm/lib/IR/Verifier.cpp

@@ -3484,9 +3484,7 @@ void Verifier::visitIndirectBrInst(IndirectBrInst &BI) {
 }
 
 static bool isSupportedCallBrIntrinsic(Intrinsic::ID ID) {
-  // Currently we only support callbr for amdgcn.kill. Add more checks here as
-  // needed.
-  return isAMDGPUCallBrIntrinsic(ID);
+  return ID == Intrinsic::go_defer_edge || isAMDGPUCallBrIntrinsic(ID);
 }
 
 void Verifier::visitCallBrInst(CallBrInst &CBI) {
@@ -6923,6 +6921,20 @@ void Verifier::visitIntrinsicCall(Intrinsic::ID ID, CallBase &Call) {
           &Call);
     break;
   }
+  case Intrinsic::go_defer_edge: {
+    const auto *CBR = dyn_cast<CallBrInst>(&Call);
+    Check(CBR, "llvm.go.defer.edge must be used with callbr", &Call);
+    if (!CBR)
+      break;
+    Check(CBR->getNumIndirectDests() == 1,
+          "llvm.go.defer.edge callbr must have exactly one indirect "
+          "destination",
+          &Call);
+    if (CBR->getNumIndirectDests() == 1)
+      Check(CBR->getDefaultDest() != CBR->getIndirectDest(0),
+            "llvm.go.defer.edge callbr destinations must be distinct", &Call);
+    break;
+  }
   case Intrinsic::structured_gep: {
     // Parser should refuse those 2 cases.
     assert(Call.arg_size() >= 1);

llvm/lib/MC/GoObjObjectWriter.cpp

@@ -253,14 +253,15 @@ SmallString<0> makeConstantPCTab(int32_t Value, uint64_t CodeSize,
 }
 
 SmallString<0> makeFuncInfoData(uint32_t ArgSize, uint32_t StackSize,
+                                uint8_t FuncID, uint8_t FuncFlag,
                                 ArrayRef<uint32_t> Files, int32_t StartLine) {
   SmallString<0> Data;
   raw_svector_ostream OS(Data);
   support::endian::Writer W(OS, llvm::endianness::little);
   W.write<uint32_t>(ArgSize);   // Args.
   W.write<uint32_t>(StackSize); // Locals.
-  W.write<uint8_t>(0);          // FuncIDNormal.
-  W.write<uint8_t>(0);          // No FuncFlag bits.
+  W.write<uint8_t>(FuncID);
+  W.write<uint8_t>(FuncFlag);
   W.write<uint8_t>(0);
   W.write<uint8_t>(0);
   W.write<uint32_t>(static_cast<uint32_t>(StartLine));
@@ -1566,10 +1567,14 @@ uint64_t GoObjObjectWriter::writeObject() {
         return LHS.PC < RHS.PC;
       });
 
+      auto [FuncID, FuncFlag] =
+          Asm->getContext()
+              .getGoObjFunctionInfo(Symbols[I].Symbol)
+              .value_or(std::make_pair(uint8_t(0), uint8_t(0)));
       uint32_t FuncInfoSym = addAuxCarrierSymbol(
           Symbols, GoObj::DefinedSymbolBlock::Symdef,
-          makeFuncInfoData(ArgSize, FrameLayout.FuncInfoLocalsSize,
-                           LineInfo.Files, LineInfo.StartLine));
+          makeFuncInfoData(ArgSize, FrameLayout.FuncInfoLocalsSize, FuncID,
+                           FuncFlag, LineInfo.Files, LineInfo.StartLine));
       uint32_t PcspSym = getOrAddHashedAuxCarrierSymbol(
           Symbols, AuxCarrierIndexes, 'P',
           PCSPEntries.empty()

llvm/lib/MC/MCContext.cpp

@@ -171,6 +171,7 @@ void MCContext::reset() {
   GoObjSymbolArgSizes.clear();
   GoObjSymbolHasFramePointers.clear();
   GoObjSymbolFlags.clear();
+  GoObjFunctionInfos.clear();
   GoObjSymbolContentHashes.clear();
   GoObjRelocOverrides.clear();
   GoObjWeakRelocs.clear();

llvm/test/CodeGen/Generic/go-defer-edge.ll

@@ -0,0 +1,49 @@
+; RUN: opt -passes='default<O2>' -S %s | FileCheck %s --check-prefix=OPT
+; RUN: llc -O0 -fast-isel=0 -verify-machineinstrs -mtriple=aarch64-unknown-linux-gnu %s -o - | FileCheck %s --check-prefix=AARCH64
+; RUN: llc -O0 -fast-isel=0 -verify-machineinstrs -mtriple=x86_64-unknown-linux-gnu %s -o - | FileCheck %s --check-prefix=X86
+; RUN: llc -O0 -fast-isel=0 -mtriple=aarch64-unknown-linux-gnu -stop-after=finalize-isel %s -o - | FileCheck %s --check-prefix=SDAG
+; RUN: llc -O0 -global-isel=1 -mtriple=aarch64-unknown-linux-gnu -stop-after=irtranslator %s -o - | FileCheck %s --check-prefix=GISEL
+
+declare void @llvm.go.defer.edge()
+declare void @runtime.deferreturn()
+
+define void @go_defer_edge() {
+; OPT-LABEL: define void @go_defer_edge()
+; OPT: callbr void @llvm.go.defer.edge()
+; OPT-NEXT: to label %{{.*}} [label %recover]
+; OPT: recover:
+; OPT: call void @runtime.deferreturn()
+;
+; AARCH64-LABEL: go_defer_edge:
+; AARCH64-NOT: llvm.go.defer.edge
+; AARCH64: ret
+; AARCH64: runtime.deferreturn
+; AARCH64: ret
+;
+; X86-LABEL: go_defer_edge:
+; X86-NOT: llvm.go.defer.edge
+; X86: retq
+; X86: runtime.deferreturn
+; X86: retq
+;
+; SDAG-LABEL: name: go_defer_edge
+; SDAG-NOT: llvm.go.defer.edge
+; SDAG: bb.{{[0-9]+}}.entry:
+; SDAG: successors: %bb.{{[0-9]+}}{{.*}}%bb.{{[0-9]+}}
+; SDAG: bb.{{[0-9]+}}.recover (inlineasm-br-indirect-target)
+;
+; GISEL-LABEL: name: go_defer_edge
+; GISEL-NOT: G_INTRINSIC intrinsic(@llvm.go.defer.edge)
+; GISEL: bb.{{[0-9]+}}.entry:
+; GISEL: successors: %bb.{{[0-9]+}}{{.*}}%bb.{{[0-9]+}}
+; GISEL: bb.{{[0-9]+}}.recover (inlineasm-br-indirect-target)
+entry:
+  callbr void @llvm.go.defer.edge() to label %normal [label %recover]
+
+normal:
+  ret void
+
+recover:
+  call void @runtime.deferreturn()
+  ret void
+}

llvm/test/CodeGen/X86/goobj-function-info.ll

@@ -0,0 +1,19 @@
+; RUN: llc -mtriple=x86_64-unknown-linux-goobj -filetype=obj < %s -o %t.o
+; RUN: %python %S/../../MC/GoObj/Inputs/dump-goobj.py %t.o | FileCheck %s
+; RUN: llc -mtriple=aarch64-unknown-linux-goobj -filetype=obj < %s -o %t.arm64.o
+; RUN: %python %S/../../MC/GoObj/Inputs/dump-goobj.py %t.arm64.o | FileCheck %s
+
+define goabiinternal void @normal() {
+  ret void
+}
+
+define goabiinternal void @wrapper() !goobj.func.info !0 {
+  ret void
+}
+
+!0 = !{i8 23, i8 1}
+
+; CHECK: symdef 0: normal
+; CHECK: symdef 1: wrapper
+; CHECK: aux 0.0: type=funcinfo target= args=0 locals={{[0-9]+}} funcid=0 funcflag=0
+; CHECK: aux 1.8: type=funcinfo target= args=0 locals={{[0-9]+}} funcid=23 funcflag=1

llvm/test/MC/GoObj/Inputs/dump-goobj.py

@@ -257,9 +257,13 @@ def symbol_data(pkg_index, sym_index):
             aux_name = AUX_TYPES.get(aux_type, str(aux_type))
             extra = ""
             payload = symbol_data(pkg_index, sym_index)
-            if aux_type == 1 and payload is not None and len(payload) >= 8:
+            if aux_type == 1 and payload is not None and len(payload) >= 10:
                 args, locals_ = struct.unpack_from("<II", payload, 0)
-                extra = f" args={args} locals={locals_}"
+                func_id, func_flag = struct.unpack_from("<BB", payload, 8)
+                extra = (
+                    f" args={args} locals={locals_} funcid={func_id}"
+                    f" funcflag={func_flag}"
+                )
             if aux_type == 2 and payload is not None:
                 extra = f" data={payload.hex()}"
             if aux_type in (7, 8, 9, 10, 11) and payload is not None:

llvm/test/Verifier/go-defer-edge.ll

@@ -0,0 +1,72 @@
+; RUN: split-file %s %t
+; RUN: opt -passes=verify %t/valid.ll -disable-output
+; RUN: not opt -passes=verify %t/ordinary-call.ll -disable-output 2>&1 | FileCheck %s --check-prefix=ORDINARY
+; RUN: not opt -passes=verify %t/no-indirect-dest.ll -disable-output 2>&1 | FileCheck %s --check-prefix=NO-DEST
+; RUN: not opt -passes=verify %t/multiple-indirect-dests.ll -disable-output 2>&1 | FileCheck %s --check-prefix=MULTIPLE-DESTS
+; RUN: not opt -passes=verify %t/same-dest.ll -disable-output 2>&1 | FileCheck %s --check-prefix=SAME-DEST
+
+;--- valid.ll
+declare void @llvm.go.defer.edge()
+
+define void @valid() {
+entry:
+  callbr void @llvm.go.defer.edge() to label %normal [label %recover]
+
+normal:
+  ret void
+
+recover:
+  ret void
+}
+
+;--- ordinary-call.ll
+; ORDINARY: llvm.go.defer.edge must be used with callbr
+declare void @llvm.go.defer.edge()
+
+define void @ordinary_call() {
+entry:
+  call void @llvm.go.defer.edge()
+  ret void
+}
+
+;--- no-indirect-dest.ll
+; NO-DEST: llvm.go.defer.edge callbr must have exactly one indirect destination
+declare void @llvm.go.defer.edge()
+
+define void @no_indirect_dest() {
+entry:
+  callbr void @llvm.go.defer.edge() to label %normal []
+
+normal:
+  ret void
+}
+
+;--- multiple-indirect-dests.ll
+; MULTIPLE-DESTS: llvm.go.defer.edge callbr must have exactly one indirect destination
+declare void @llvm.go.defer.edge()
+
+define void @multiple_indirect_dests() {
+entry:
+  callbr void @llvm.go.defer.edge() to label %normal [label %recover1, label %recover2]
+
+normal:
+  ret void
+
+recover1:
+  ret void
+
+recover2:
+  ret void
+}
+
+;--- same-dest.ll
+; SAME-DEST: llvm.go.defer.edge callbr destinations must be distinct
+declare void @llvm.go.defer.edge()
+
+define void @same_dest() {
+entry:
+  callbr void @llvm.go.defer.edge() to label %dest [label %dest]
+
+dest:
+  ret void
+}