CodeGen: support plugin IR passes immediately before ISel

.github/workflows/goallc-ci.yml

@@ -110,6 +110,7 @@ jobs:
             "$build/test/CodeGen/Generic/goobj-entry-stackmap-cfg.ll"
             "$build/test/CodeGen/Generic/goobj-entry-stackmap-sentinel.ll"
             "$build/test/CodeGen/Generic/goobj-stack-check-policy.ll"
+            "$build/test/CodeGen/X86/branchfolding-inlineasmbr-successor.mir"
             "$build/test/CodeGen/X86/go-callconv.ll"
             "$build/test/CodeGen/X86/go-stack-byval.ll"
             "$build/test/CodeGen/X86/go-gc-write-barrier.ll"

llvm/include/llvm/BinaryFormat/GoObj.h

@@ -46,21 +46,24 @@ inline constexpr uint64_t EntryArgsStackMapID = 0x476f4e6f53706c74ULL;
 //
 //   ordinary-deopt*, BEGIN, protocol-length, record-count,
 //     (RECORD-TAG, record-length, direct-base, byte-offset, byte-size,
-//     alignment,
-//      pointer-size, bit-count, word-bits, word-count, bitmap-word*)*,
+//      alignment, pointer-size, contents-live, bit-count, word-bits,
+//      word-count, bitmap-word*)*,
 //   END, protocol-length
 //
 // Protocol length counts BEGIN through END and excludes its trailing duplicate;
-// record length counts RECORD-TAG through the final bitmap word. A matching
-// direct alloca in the statepoint gc-live operands says that the record's
+// record length counts RECORD-TAG through the final bitmap word.
+// contents-live is zero or one and independently says whether the record's
 // contents contribute to that callsite's ArgsPointerMaps or
-// LocalsPointerMaps, according to the alloca's frame region. An unmatched
-// record identifies a function-level native Go StackObject; argument/result
-// objects use a non-negative argp-relative offset and local objects use a
-// negative varp-relative offset. The producer must repeat that same layout at
-// every ordinary statepoint. The direct address itself remains a rematerialized
-// frame index, not a bitmap slot. The first contract has no version, requires a
-// whole alloca at byte offset zero, and uses 64-bit bitmap words. Bit N, stored
+// LocalsPointerMaps, according to the alloca's frame region. The direct alloca
+// may still occur in the statepoint gc-live operands when only its frame
+// address needs gc.relocate rematerialization; that occurrence never makes the
+// contents live by itself. A record with contents-live zero identifies a
+// function-level native Go StackObject; argument/result objects use a
+// non-negative argp-relative offset and local objects use a negative
+// varp-relative offset. The producer must repeat that same layout at every
+// ordinary statepoint. The direct address itself remains a rematerialized frame
+// index, not a bitmap slot. The first contract has no version, requires a whole
+// alloca at byte offset zero, and uses 64-bit bitmap words. Bit N, stored
 // low-bit first, describes the pointer-sized slot at direct-base + byte-offset
 // + N * pointer-size. Padding bits must be zero.
 // These tags are intentionally small enough to remain inline StackMaps

llvm/include/llvm/CodeGen/TargetPassConfig.h

@@ -113,6 +113,10 @@ class LLVM_ABI TargetPassConfig : public ImmutablePass {
   /// but may still change instruction sizes before target branch relaxation.
   std::vector<std::function<Pass *()>> PreBranchRelaxationPassFactories;
 
+  /// Plugin-provided LLVM IR passes that run after all standard codegen IR
+  /// preparation and immediately before instruction selection.
+  std::vector<std::function<Pass *()>> PreISelPassFactories;
+
   /// Set the StartAfter, StartBefore and StopAfter passes to allow running only
   /// a portion of the normal code-gen pass sequence.
   ///
@@ -228,6 +232,12 @@ class LLVM_ABI TargetPassConfig : public ImmutablePass {
     PreBranchRelaxationPassFactories.push_back(std::move(Factory));
   }
 
+  /// Register a late LLVM IR pass after CodeGenPrepare and the remaining
+  /// standard IR preparation passes, but before instruction selection.
+  void addPreISelPass(std::function<Pass *()> Factory) {
+    PreISelPassFactories.push_back(std::move(Factory));
+  }
+
   /// Allow the target to enable a specific standard pass by default.
   void enablePass(AnalysisID PassID) { substitutePass(PassID, PassID); }
 

llvm/lib/CodeGen/BranchFolding.cpp

@@ -1441,11 +1441,13 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
         Pred->ReplaceUsesOfBlockWith(MBB, &*FallThrough);
       }
       // Add rest successors of MBB to successors of FallThrough. Those
-      // successors are not directly reachable via MBB, so it should be
-      // landing-pad.
+      // successors are not represented by the analyzed branch, so they must
+      // be EH pads or inlineasm_br indirect targets.
       for (auto SI = MBB->succ_begin(), SE = MBB->succ_end(); SI != SE; ++SI)
         if (*SI != &*FallThrough && !FallThrough->isSuccessor(*SI)) {
-          assert((*SI)->isEHPad() && "Bad CFG");
+          assert(((*SI)->isEHPad() ||
+                  (*SI)->isInlineAsmBrIndirectTarget()) &&
+                 "Bad CFG");
           FallThrough->copySuccessor(MBB, SI);
         }
       // If MBB was the target of a jump table, update jump tables to go to the
@@ -1710,12 +1712,14 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
               DidChange = true;
               PMBB->ReplaceUsesOfBlockWith(MBB, CurTBB);
               // Add rest successors of MBB to successors of CurTBB. Those
-              // successors are not directly reachable via MBB, so it should be
-              // landing-pad.
+              // successors are not represented by the analyzed branch, so
+              // they must be EH pads or inlineasm_br indirect targets.
               for (auto SI = MBB->succ_begin(), SE = MBB->succ_end(); SI != SE;
                    ++SI)
                 if (*SI != CurTBB && !CurTBB->isSuccessor(*SI)) {
-                  assert((*SI)->isEHPad() && "Bad CFG");
+                  assert(((*SI)->isEHPad() ||
+                          (*SI)->isInlineAsmBrIndirectTarget()) &&
+                         "Bad CFG");
                   CurTBB->copySuccessor(MBB, SI);
                 }
               // If this change resulted in PMBB ending in a conditional

llvm/lib/CodeGen/SelectionDAG/StatepointLowering.cpp

@@ -107,14 +107,14 @@ getArgumentValueOffset(const Value *V, const DataLayout &DL) {
 
 static SDValue getStatepointGCValue(const Value *V,
                                     SelectionDAGBuilder &Builder) {
-  // A typed byval argument denotes its incoming Go stack home, not a heap
-  // pointer stored in that home. Always use the canonical fixed frame index
-  // for this address. In particular, do not let an earlier gc.relocate or a
-  // larger live set turn the address into an ordinary pointer spill: stack
-  // growth rematerializes frame-index addresses, while the separate object
-  // layout describes which words in the home are GC roots.
+  // A typed byval or goret argument denotes a fixed Go stack home, not a heap
+  // pointer stored in that home. Always use the canonical frame index for this
+  // address. In particular, do not let an earlier gc.relocate or a larger live
+  // set turn the address into an ordinary pointer spill: stack growth
+  // rematerializes frame-index addresses, while separate object metadata
+  // describes which words in the home are GC roots.
   if (const auto *Arg = dyn_cast<Argument>(V);
-      Arg && Arg->hasByValAttr() &&
+      Arg && (Arg->hasByValAttr() || Arg->hasGoRetAttr()) &&
       goabi::isGoCallingConv(
           Builder.DAG.getMachineFunction().getFunction().getCallingConv())) {
     int FI = Builder.FuncInfo.getArgumentFrameIndex(Arg);
@@ -698,7 +698,9 @@ lowerStatepointMetaArgs(SmallVectorImpl<SDValue> &Ops,
 
     if (auto ArgValue = getArgumentValueOffset(V, Builder.DAG.getDataLayout());
         ArgValue &&
-        !(V == ArgValue->first && ArgValue->first->hasByValAttr())) {
+        !(V == ArgValue->first &&
+          (ArgValue->first->hasByValAttr() ||
+           ArgValue->first->hasGoRetAttr()))) {
       uint64_t Size = PtrSD.getValueType().getStoreSize().getKnownMinValue();
       int FI = Builder.FuncInfo.getArgumentValueHome(ArgValue->first,
                                                      ArgValue->second, Size);
@@ -1064,7 +1066,9 @@ SDValue SelectionDAGBuilder::LowerAsSTATEPOINT(
                    DAG.getMachineFunction().getFunction().getCallingConv()) &&
                isa<FrameIndexSDNode>(SDV) &&
                (isa<AllocaInst>(V) ||
-                (isa<Argument>(V) && cast<Argument>(V)->hasByValAttr() &&
+                (isa<Argument>(V) &&
+                 (cast<Argument>(V)->hasByValAttr() ||
+                  cast<Argument>(V)->hasGoRetAttr()) &&
                  FuncInfo.getArgumentFrameIndex(cast<Argument>(V)) ==
                      cast<FrameIndexSDNode>(SDV)->getIndex()))) {
       Record.type = RecordType::FrameIndexRemat;
@@ -1224,14 +1228,15 @@ SelectionDAGBuilder::LowerStatepoint(const GCStatepointInst &I,
   // pointers passed to deopt are base pointers; relaxing that assumption
   // would require relatively large changes to how we represent relocations.
   for (Value *V : I.deopt_operands()) {
-    // GoALLC uses direct static allocas and typed byval parameters as
-    // frame-layout carriers for its per-object pointer maps. The object is a
-    // GC root only when its base is also present in the explicit gc-live
-    // bundle; treating the deopt carrier as a root would make an inactive
-    // lifetime scan uninitialized or dead storage.
+    // GoALLC uses direct static allocas and typed byval/goret parameters as
+    // fixed-frame metadata carriers. The object is a GC root only when its
+    // base is also present in the explicit gc-live bundle; treating the deopt
+    // carrier as a root would make inactive storage or the frame address
+    // itself enter the pointer map.
     const auto *Arg = dyn_cast<Argument>(V);
     if (GFI->getStrategy().getName() == "goallc" &&
-        (isa<AllocaInst>(V) || (Arg && Arg->hasByValAttr())))
+        (isa<AllocaInst>(V) ||
+         (Arg && (Arg->hasByValAttr() || Arg->hasGoRetAttr()))))
       continue;
     if (!isGCValue(V, *this))
       continue;

llvm/lib/CodeGen/TargetPassConfig.cpp

@@ -982,6 +982,9 @@ void TargetPassConfig::addISelPrepare() {
   addPass(createSafeStackPass());
   addPass(createStackProtectorPass());
 
+  for (const auto &Factory : PreISelPassFactories)
+    addPass(Factory());
+
   if (PrintISelInput)
     addPass(createPrintFunctionPass(
         dbgs(), "\n\n*** Final LLVM Code input to ISel ***\n"));

llvm/lib/MC/GoObjObjectWriter.cpp

@@ -415,6 +415,7 @@ struct GoObjAllocaPtrMapRecord {
   uint64_t ByteSize;
   uint64_t Alignment;
   uint64_t PointerSize;
+  bool ContentsLive;
   uint64_t BitCount;
   SmallVector<uint64_t, 4> BitmapWords;
 };
@@ -562,28 +563,32 @@ parseAllocaPtrMapRecords(const MCContext::GoObjStackMapEntry &Entry) {
       Deopts[ProtocolStart + 2], "record count");
   size_t Cursor = ProtocolStart + 3;
   size_t RecordsEnd = ProtocolStart + ProtocolLength - 1;
-  if (RecordCount > (RecordsEnd - Cursor) / 10)
+  if (RecordCount > (RecordsEnd - Cursor) / 11)
     report_fatal_error("GoObj alloca ptrmap record count is invalid");
   SmallVector<GoObjAllocaPtrMapRecord, 4> Records;
   Records.reserve(static_cast<size_t>(RecordCount));
   for (uint64_t RecordIndex = 0; RecordIndex != RecordCount; ++RecordIndex) {
-    if (Cursor > RecordsEnd || RecordsEnd - Cursor < 10)
+    if (Cursor > RecordsEnd || RecordsEnd - Cursor < 11)
       report_fatal_error("GoObj alloca ptrmap record header is malformed");
     if (!IsConstant(Deopts[Cursor], GoObj::AllocaPtrMapRecordTag))
       report_fatal_error("GoObj alloca ptrmap record tag is invalid");
     uint64_t RecordLength =
         getNonnegativeAllocaPtrMapConstant(Deopts[Cursor + 1], "record length");
     uint64_t WordCount = getNonnegativeAllocaPtrMapConstant(
-        Deopts[Cursor + 9], "bitmap word count");
-    if (WordCount > RecordsEnd - Cursor - 10 ||
-        RecordLength != 10 + WordCount || RecordLength > RecordsEnd - Cursor)
+        Deopts[Cursor + 10], "bitmap word count");
+    if (WordCount > RecordsEnd - Cursor - 11 ||
+        RecordLength != 11 + WordCount || RecordLength > RecordsEnd - Cursor)
       report_fatal_error("GoObj alloca ptrmap record length is invalid");
 
     const auto &Base = Deopts[Cursor + 2];
     if (Base.Type != MCContext::GoObjStackMapLocation::Direct)
       report_fatal_error(
           "GoObj alloca ptrmap base is not a direct frame location");
-    uint64_t WordBits = getNonnegativeAllocaPtrMapConstant(Deopts[Cursor + 8],
+    uint64_t ContentsLive = getNonnegativeAllocaPtrMapConstant(
+        Deopts[Cursor + 7], "contents-live flag");
+    if (ContentsLive > 1)
+      report_fatal_error("GoObj alloca ptrmap contents-live flag is invalid");
+    uint64_t WordBits = getNonnegativeAllocaPtrMapConstant(Deopts[Cursor + 9],
                                                            "bitmap word width");
     if (WordBits != GoObj::AllocaPtrMapBitmapWordBits)
       report_fatal_error("GoObj alloca ptrmap bitmap word width is invalid");
@@ -594,12 +599,13 @@ parseAllocaPtrMapRecords(const MCContext::GoObjStackMapEntry &Entry) {
         getNonnegativeAllocaPtrMapConstant(Deopts[Cursor + 4], "byte size"),
         getNonnegativeAllocaPtrMapConstant(Deopts[Cursor + 5], "alignment"),
         getNonnegativeAllocaPtrMapConstant(Deopts[Cursor + 6], "pointer size"),
-        getNonnegativeAllocaPtrMapConstant(Deopts[Cursor + 7], "bit count"),
+        ContentsLive != 0,
+        getNonnegativeAllocaPtrMapConstant(Deopts[Cursor + 8], "bit count"),
         {}};
     Record.BitmapWords.reserve(WordCount);
     for (uint64_t Word = 0; Word != WordCount; ++Word)
       Record.BitmapWords.push_back(static_cast<uint64_t>(
-          getAllocaPtrMapConstant(Deopts[Cursor + 10 + Word], "bitmap word")));
+          getAllocaPtrMapConstant(Deopts[Cursor + 11 + Word], "bitmap word")));
     Records.push_back(std::move(Record));
     Cursor += RecordLength;
   }
@@ -687,7 +693,7 @@ GoObjStatepointStackMaps makeStatepointStackMaps(
     GoObjAllocaPtrMapRecord Layout;
     goobj::StackMapSlotKind Kind;
     uint64_t SeenOrdinaryEntries = 0;
-    bool SawUnmatchedGCLive = false;
+    bool SawInactiveContents = false;
   };
   SmallVector<FunctionAllocaRecord, 4> FunctionAllocaRecords;
   uint64_t OrdinaryEntryCount =
@@ -805,7 +811,7 @@ GoObjStatepointStackMaps makeStatepointStackMaps(
       }
       AllocaRanges.push_back({RangeStart, RangeEnd});
 
-      bool IsActive = llvm::any_of(
+      bool HasDirectGCLive = llvm::any_of(
           GCLiveLocations,
           [&](const MCContext::GoObjStackMapLocation &Location) {
             MCContext::GoObjStackMapLocation Normalized =
@@ -816,6 +822,9 @@ GoObjStatepointStackMaps makeStatepointStackMaps(
                    Normalized.DwarfRegNum == RecordBase.DwarfRegNum &&
                    Normalized.Offset == RecordBase.Offset;
           });
+      if (Record.ContentsLive && !HasDirectGCLive)
+        report_fatal_error(
+            "GoObj live alloca ptrmap has no direct gc-live base");
 
       uint64_t PaddingBits = Record.BitmapWords.size() * 64 - Record.BitCount;
       if (PaddingBits && (Record.BitmapWords.back() >> (64 - PaddingBits)) != 0)
@@ -842,9 +851,10 @@ GoObjStatepointStackMaps makeStatepointStackMaps(
         if ((Record.BitmapWords[Bit / 64] & (uint64_t(1) << (Bit % 64))) == 0)
           continue;
         HasPointer = true;
-        // The layout record is interpretation-free. A matching direct
-        // gc-live base means this pointer word is live at this callsite.
-        if (IsActive) {
+        // gc-live may carry the direct frame base solely for gc.relocate.
+        // Only the producer's independent contents-live bit makes this
+        // pointer word live at the callsite.
+        if (Record.ContentsLive) {
           DenseSet<uint32_t> &AllocaPointerBits =
               Slot.Kind == goobj::StackMapSlotKind::Args
                   ? AllocaArgsPointerBits
@@ -869,13 +879,13 @@ GoObjStatepointStackMaps makeStatepointStackMaps(
           });
       if (FunctionRecord == FunctionAllocaRecords.end()) {
         FunctionAllocaRecords.push_back(
-            {NormalizedRecord, *RecordKind, 1, !IsActive});
+            {NormalizedRecord, *RecordKind, 1, !Record.ContentsLive});
       } else {
         if (FunctionRecord->Kind != *RecordKind)
           report_fatal_error(
               "GoObj alloca ptrmap frame region changes between statepoints");
         ++FunctionRecord->SeenOrdinaryEntries;
-        FunctionRecord->SawUnmatchedGCLive |= !IsActive;
+        FunctionRecord->SawInactiveContents |= !Record.ContentsLive;
       }
     }
 
@@ -1022,7 +1032,7 @@ GoObjStatepointStackMaps makeStatepointStackMaps(
 
   SmallVector<GoObjStatepointStackMaps::StackObject, 4> FunctionStackObjects;
   for (const FunctionAllocaRecord &FunctionRecord : FunctionAllocaRecords) {
-    if (!FunctionRecord.SawUnmatchedGCLive)
+    if (!FunctionRecord.SawInactiveContents)
       continue;
     if (FunctionRecord.SeenOrdinaryEntries != OrdinaryEntryCount)
       report_fatal_error(

llvm/test/CodeGen/AArch64/go-statepoint-stack-args.ll

@@ -3,6 +3,7 @@
 ; RUN:   -stop-after=finalize-isel < %s | FileCheck %s
 
 %aggregate = type { ptr addrspace(1), i64, ptr addrspace(1) }
+%goret.home = type { ptr, i64, ptr }
 
 declare goabiinternal void @safepoint()
 
@@ -148,10 +149,30 @@ entry:
   ret ptr addrspace(1) %home.relocated
 }
 
+define goabiinternal { i64, i64, i64, i64, i64, i64, i64, i64,
+                       i64, i64, i64, i64, i64, i64, i64 }
+    @goret_home_address(
+        ptr addrspace(1) goret(%goret.home) align 8 "goretindex"="15" %result)
+    #0 gc "statepoint-example" {
+entry:
+  %token = call goabiinternal token (i64, i32, ptr, i32, i32, ...)
+      @llvm.experimental.gc.statepoint.p0(
+          i64 9, i32 0, ptr elementtype(void ()) @safepoint,
+          i32 0, i32 0, i32 0, i32 0)
+      [ "gc-live"(ptr addrspace(1) %result) ]
+  %result.relocated = call ptr addrspace(1)
+      @llvm.experimental.gc.relocate.p1(
+      token %token, i32 0, i32 0)
+  store ptr null, ptr addrspace(1) %result.relocated, align 8
+  ret { i64, i64, i64, i64, i64, i64, i64, i64,
+        i64, i64, i64, i64, i64, i64, i64 } zeroinitializer
+}
+
 declare token @llvm.experimental.gc.statepoint.p0(
     i64 immarg, i32 immarg, ptr, i32 immarg, i32 immarg, ...)
 declare ptr addrspace(1) @llvm.experimental.gc.relocate.p1(
     token, i32 immarg, i32 immarg)
+attributes #0 = { "go_results_tuple" }
 
 ; CHECK-LABEL: name: scalar_stack_arg
 ; CHECK: fixedStack:
@@ -213,3 +234,15 @@ declare ptr addrspace(1) @llvm.experimental.gc.relocate.p1(
 ; CHECK-SAME: 2, 2, 0, %fixed-stack.[[BYVAL_HOME]], 0, 1, 8, %stack.[[HEAP_SLOT]], 0,
 ; CHECK-NEXT: ADJCALLSTACKUP
 ; CHECK-NEXT: {{%[0-9]+}}:gpr64sp = ADDXri %fixed-stack.[[BYVAL_HOME]], 0, 0
+
+; CHECK-LABEL: name: goret_home_address
+; CHECK: fixedStack:
+; CHECK: - { id: [[GORET_HOME:[0-9]+]], type: default, offset: 8, size: 24,
+; CHECK: stack: []
+; CHECK: STATEPOINT 9,
+; CHECK-SAME: 2, 1, 0, %fixed-stack.[[GORET_HOME]], 0,
+; CHECK-SAME: 2, 1, 0, %fixed-stack.[[GORET_HOME]], 0,
+; CHECK-NEXT: ADJCALLSTACKUP
+; CHECK-NEXT: [[GORET_FRAME:%[0-9]+]]:gpr64sp = ADDXri %fixed-stack.[[GORET_HOME]], 0, 0
+; CHECK-NEXT: [[GORET_ADDR:%[0-9]+]]:gpr64sp = COPY [[GORET_FRAME]]
+; CHECK: STRXui {{.*}}, [[GORET_ADDR]], 0