GoObj: mark asynchronous-unsafe functions

llvm/include/llvm/BinaryFormat/GoObj.h

@@ -43,6 +43,10 @@ inline constexpr int64_t AllocaPtrMapEndMagic = 0x414c4c43;   // "ALLC"
 inline constexpr int64_t AllocaPtrMapRecordTag = 0x5054524d;  // "PTRM"
 inline constexpr uint32_t AllocaPtrMapBitmapWordBits = 64;
 
+// Special values in PCDATA_UnsafePoint (PCDATA stream 0).
+inline constexpr int32_t UnsafePointSafe = -1;
+inline constexpr int32_t UnsafePointUnsafe = -2;
+
 inline constexpr char Magic[] = {'\0', 'g', 'o', '1', '2', '0', 'l', 'd'};
 inline constexpr uint32_t MagicSize = sizeof(Magic);
 inline constexpr uint32_t FingerprintSize = 8;

llvm/include/llvm/MC/MCContext.h

@@ -244,6 +244,10 @@ class MCContext {
   DenseMap<const MCSymbol *, std::vector<GoObjPCSPEntry>>
       GoObjSymbolPCSPEntries;
 
+  /// Whether a Go object function is unsafe for asynchronous preemption over
+  /// its complete PC range.
+  DenseMap<const MCSymbol *, bool> GoObjSymbolAsyncUnsafe;
+
   /// Go object statepoint stack maps keyed by function MC symbol.
   DenseMap<const MCSymbol *, std::vector<GoObjStackMapEntry>>
       GoObjSymbolStackMapEntries;
@@ -777,6 +781,14 @@ class MCContext {
     return &It->second;
   }
 
+  void setGoObjSymbolAsyncUnsafe(const MCSymbol *Sym, bool AsyncUnsafe) {
+    GoObjSymbolAsyncUnsafe[Sym] = AsyncUnsafe;
+  }
+
+  bool isGoObjSymbolAsyncUnsafe(const MCSymbol *Sym) const {
+    return GoObjSymbolAsyncUnsafe.lookup(Sym);
+  }
+
   void addGoObjSymbolStackMapEntry(const MCSymbol *Sym,
                                    GoObjStackMapEntry Entry) {
     GoObjSymbolStackMapEntries[Sym].push_back(std::move(Entry));

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -3520,6 +3520,8 @@ void AsmPrinter::SetupMachineFunction(MachineFunction &MF) {
         getGoObjArgSize(F, MF.getDataLayout(), TM.getTargetTriple()));
     OutContext.setGoObjSymbolHasFramePointer(
         CurrentFnSym, MF.getSubtarget().getFrameLowering()->hasFP(MF));
+    OutContext.setGoObjSymbolAsyncUnsafe(CurrentFnSym,
+                                         F.hasFnAttribute("go-async-unsafe"));
   }
 
   CurrentFnSymForSize = CurrentFnSym;

llvm/lib/MC/GoObjObjectWriter.cpp

@@ -1291,6 +1291,11 @@ uint64_t GoObjObjectWriter::writeObject() {
         });
       }
 
+      int32_t InitialUnsafePointValue =
+          Asm->getContext().isGoObjSymbolAsyncUnsafe(Symbols[I].Symbol)
+              ? GoObj::UnsafePointUnsafe
+              : GoObj::UnsafePointSafe;
+
       GoObjFuncDebugLines &LineInfo = FuncDebugLines[I];
       if (!LineInfo.hasLines())
         LineInfo.Files.push_back(GetFallbackFile());
@@ -1345,7 +1350,7 @@ uint64_t GoObjObjectWriter::writeObject() {
           Symbols, AuxCarrierIndexes, 'P', StackMapIndex);
       uint32_t UnsafePointSym = getOrAddHashedAuxCarrierSymbol(
           Symbols, AuxCarrierIndexes, 'P',
-          makeConstantPCTab(-1, CodeSize, PCQuantum));
+          makeConstantPCTab(InitialUnsafePointValue, CodeSize, PCQuantum));
 
       Symbols[I].Auxiliaries.push_back({GoObj::AuxFuncInfo, FuncInfoSym});
       Symbols[I].Auxiliaries.push_back({GoObj::AuxFuncdata, ArgsMapSym});

llvm/lib/MC/MCContext.cpp

@@ -179,6 +179,7 @@ void MCContext::reset() {
   GoObjSymbolSizes.clear();
   GoObjGotypeTargets.clear();
   GoObjSymbolPCSPEntries.clear();
+  GoObjSymbolAsyncUnsafe.clear();
   GoObjSymbolStackMapEntries.clear();
   Symbols.clear();
   Allocator.Reset();

llvm/test/CodeGen/AArch64/goobj-function-unsafe.ll

@@ -0,0 +1,29 @@
+; REQUIRES: aarch64-registered-target
+; RUN: opt -passes='default<O2>' -S < %s | FileCheck %s --check-prefix=OPT
+; RUN: llc -mtriple=aarch64-apple-darwin-goobj -verify-machineinstrs \
+; RUN:   -enable-shrink-wrap=true -filetype=obj < %s -o %t.o
+; RUN: %python %S/../../MC/GoObj/Inputs/dump-goobj.py %t.o | \
+; RUN:   FileCheck %s --check-prefix=OBJ
+
+define goabiinternal i64 @main.safe(ptr %p) {
+entry:
+  %v = load i64, ptr %p, align 8
+  ret i64 %v
+}
+
+define goabiinternal i64 @main.async_unsafe(ptr %p) #0 {
+; OPT-LABEL: define goabiinternal i64 @main.async_unsafe
+; OPT-SAME: #[[UNSAFE:[0-9]+]]
+; OPT: %v = load i64, ptr %p, align 8
+entry:
+  %v = load i64, ptr %p, align 8
+  ret i64 %v
+}
+
+attributes #0 = { "go-async-unsafe" }
+; OPT: attributes #[[UNSAFE]] = { {{.*}}"go-async-unsafe"{{.*}} }
+
+; OBJ: symdef 0: main.safe
+; OBJ: symdef 1: main.async_unsafe
+; OBJ: aux 0.6: type=pcdata target= pc=[0-2:-1]
+; OBJ: aux 1.14: type=pcdata target= pc=[0-2:-2]