codegen: slightly optimize gc-frame allocation (#58794)

Try to avoid allocating frames for some very simple function that only
have the safepoint on entry and don't define any values themselves.

Author: vtjnash

src/cgutils.cpp

@@ -341,11 +341,11 @@ static void find_perm_offsets(jl_datatype_t *typ, SmallVectorImpl<unsigned> &res
 }
 
 // load a pointer to N inlined_roots into registers (as a SmallVector)
-static llvm::SmallVector<Value*,0> load_gc_roots(jl_codectx_t &ctx, Value *inline_roots_ptr, size_t npointers, bool isVolatile=false)
+static llvm::SmallVector<Value*,0> load_gc_roots(jl_codectx_t &ctx, Value *inline_roots_ptr, size_t npointers, MDNode *tbaa, bool isVolatile=false)
 {
     SmallVector<Value*,0> gcroots(npointers);
     Type *T_prjlvalue = ctx.types().T_prjlvalue;
-    auto roots_ai = jl_aliasinfo_t::fromTBAA(ctx, ctx.tbaa().tbaa_gcframe);
+    auto roots_ai = jl_aliasinfo_t::fromTBAA(ctx, tbaa);
     for (size_t i = 0; i < npointers; i++) {
         auto *ptr = ctx.builder.CreateAlignedLoad(T_prjlvalue, emit_ptrgep(ctx, inline_roots_ptr, i * sizeof(jl_value_t*)), Align(sizeof(void*)), isVolatile);
         roots_ai.decorateInst(ptr);

src/codegen.cpp

@@ -5054,7 +5054,7 @@ static jl_cgval_t emit_call_specfun_other(jl_codectx_t &ctx, bool is_opaque_clos
             break;
         case jl_returninfo_t::SRet:
             assert(result);
-            retval = mark_julia_slot(result, jlretty, NULL, ctx.tbaa().tbaa_gcframe, load_gc_roots(ctx, return_roots, returninfo.return_roots));
+            retval = mark_julia_slot(result, jlretty, NULL, ctx.tbaa().tbaa_gcframe, load_gc_roots(ctx, return_roots, returninfo.return_roots, ctx.tbaa().tbaa_gcframe));
             break;
         case jl_returninfo_t::Union: {
             Value *box = ctx.builder.CreateExtractValue(call, 0);
@@ -5603,7 +5603,7 @@ static jl_cgval_t emit_varinfo(jl_codectx_t &ctx, jl_varinfo_t &vi, jl_sym_t *va
                 T_prjlvalue = AT->getElementType();
             }
             assert(T_prjlvalue == ctx.types().T_prjlvalue);
-            v.inline_roots = load_gc_roots(ctx, varslot, nroots, vi.isVolatile);
+            v.inline_roots = load_gc_roots(ctx, varslot, nroots, ctx.tbaa().tbaa_gcframe, vi.isVolatile);
         }
         if (vi.usedUndef) {
             assert(vi.defFlag);
@@ -6927,7 +6927,7 @@ static void emit_specsig_to_specsig(
                 auto tracked = CountTrackedPointers(et);
                 SmallVector<Value*,0> roots;
                 if (tracked.count && !tracked.all) {
-                    roots = load_gc_roots(ctx, &*AI, tracked.count);
+                    roots = load_gc_roots(ctx, &*AI, tracked.count, ctx.tbaa().tbaa_const);
                     ++AI;
                 }
                 myargs[i] = mark_julia_slot(arg_v, jt, NULL, ctx.tbaa().tbaa_const, roots);
@@ -8511,7 +8511,7 @@ static jl_llvm_functions_t
             ctx.spvals_ptr = &*AI++;
         }
     }
-    // step 6. set up GC frame and special arguments
+    // step 6a. set up special arguments and attributes
     Function::arg_iterator AI = f->arg_begin();
     SmallVector<AttributeSet, 0> attrs(f->arg_size()); // function declaration attributes
 
@@ -8558,7 +8558,11 @@ static jl_llvm_functions_t
         attrs[Arg->getArgNo()] = AttributeSet::get(Arg->getContext(), param);
     }
 
+    // step 6b. Setup the GC frame and entry safepoint before any loads
     allocate_gc_frame(ctx, b0);
+    if (params.safepoint_on_entry && JL_FEAT_TEST(ctx, safepoint_on_entry))
+        emit_gc_safepoint(ctx.builder, ctx.types().T_size, get_current_ptls(ctx), ctx.tbaa().tbaa_const);
+
     Value *last_age = NULL;
     Value *world_age_field = NULL;
     if (ctx.is_opaque_closure) {
@@ -8716,7 +8720,14 @@ static jl_llvm_functions_t
             SmallVector<Value*,0> roots;
             auto tracked = CountTrackedPointers(llvmArgType);
             if (tracked.count && !tracked.all) {
-                roots = load_gc_roots(ctx, &*AI, tracked.count);
+                Argument *RootArg = &*AI;
+                roots = load_gc_roots(ctx, RootArg, tracked.count, ctx.tbaa().tbaa_const);
+                AttrBuilder param(ctx.builder.getContext(), f->getAttributes().getParamAttrs(Arg->getArgNo()));
+                param.addAttribute(Attribute::NonNull);
+                param.addAttribute(Attribute::NoUndef);
+                param.addDereferenceableAttr(tracked.count * sizeof(void*));
+                param.addAlignmentAttr(alignof(void*));
+                attrs[RootArg->getArgNo()] = AttributeSet::get(Arg->getContext(), param);
                 ++AI;
             }
             theArg = mark_julia_slot(Arg, argType, NULL, ctx.tbaa().tbaa_const, roots); // this argument is by-pointer
@@ -9026,11 +9037,7 @@ static jl_llvm_functions_t
 
     Instruction &prologue_end = ctx.builder.GetInsertBlock()->back();
 
-    // step 11a. Emit the entry safepoint
-    if (params.safepoint_on_entry && JL_FEAT_TEST(ctx, safepoint_on_entry))
-        emit_gc_safepoint(ctx.builder, ctx.types().T_size, get_current_ptls(ctx), ctx.tbaa().tbaa_const);
-
-    // step 11b. Do codegen in control flow order
+    // step 11. Do codegen in control flow order
     SmallVector<int, 0> workstack;
     DenseMap<size_t, BasicBlock*> BB;
     DenseMap<size_t, BasicBlock*> come_from_bb;

src/llvm-gc-interface-passes.h

@@ -251,11 +251,13 @@ struct BBState {
     // These get updated during dataflow
     LargeSparseBitVector LiveIn;
     LargeSparseBitVector LiveOut;
-    SmallVector<int, 0> Safepoints;
-    int TopmostSafepoint = -1;
+    // auto Safepoints = std::range(LastSafepoint, FirstSafepoint);
     bool HasSafepoint = false;
-    // Have we gone through this basic block in our local scan yet?
-    bool Done = false;
+    // This lets us refine alloca tracking to avoid creating GC frames in
+    // some simple functions that only have the initial safepoint.
+    int FirstSafepoint = -1;
+    int LastSafepoint = -1;
+    int FirstSafepointAfterFirstDef = -1;
 };
 
 struct State {
@@ -292,21 +294,18 @@ struct State {
     // of its uses need to preserve the values listed in the map value.
     std::map<Instruction *, SmallVector<int, 0>> GCPreserves;
 
-    // The assignment of numbers to safepoints. The indices in the map
-    // are indices into the next three maps which store safepoint properties
-    std::map<Instruction *, int> SafepointNumbering;
+    // The assignment of numbers to safepoints. These have the same ordering as
+    // LiveSets, LiveIfLiveOut, and CalleeRoots.
+    SmallVector<Instruction*, 0> SafepointNumbering;
 
-    // Reverse mapping index -> safepoint
-    SmallVector<Instruction *, 0> ReverseSafepointNumbering;
-
-    // Instructions that can return twice. For now, all values live at these
-    // instructions will get their own, dedicated GC frame slots, because they
-    // have unobservable control flow, so we can't be sure where they're
-    // actually live. All of these are also considered safepoints.
-    SmallVector<Instruction *, 0> ReturnsTwice;
+    // Safepoint number of instructions that can return twice. For now, all
+    // values live at these instructions will get their own, dedicated GC frame
+    // slots, because they have unobservable control flow, so we can't be sure
+    // where they're actually live.
+    SmallVector<int, 0> ReturnsTwice;
 
     // The set of values live at a particular safepoint
-    SmallVector< LargeSparseBitVector , 0> LiveSets;
+    SmallVector<LargeSparseBitVector, 0> LiveSets;
     // Those values that - if live out from our parent basic block - are live
     // at this safepoint.
     SmallVector<SmallVector<int, 0>> LiveIfLiveOut;
@@ -332,7 +331,7 @@ struct LateLowerGCFrame:  private JuliaPassContext {
     Value *pgcstack;
     Function *smallAllocFunc;
 
-    void MaybeNoteDef(State &S, BBState &BBS, Value *Def, const ArrayRef<int> &SafepointsSoFar,
+    bool MaybeNoteDef(State &S, BBState &BBS, Value *Def,
                       SmallVector<int, 1> &&RefinedPtr = SmallVector<int, 1>());
     void NoteUse(State &S, BBState &BBS, Value *V, LargeSparseBitVector &Uses, Function &F);
     void NoteUse(State &S, BBState &BBS, Value *V, Function &F) {