track-allocation: charge entry cost to caller (#34391)

Don't include one-time costs (JIT compilation) so that warm-up isn't generally required.
And adjust codegen emission to charge call entry costs to the caller.

fixes #11753
fixes #19981
fixes #21871
fixes #34054
close #18595

Author: vtjnash

src/codegen.cpp

@@ -320,6 +320,7 @@ Function *juliapersonality_func;
 #endif
 #endif
 static Function *diff_gc_total_bytes_func;
+static Function *sync_gc_total_bytes_func;
 static Function *jlarray_data_owner_func;
 static GlobalVariable *jlgetworld_global;
 
@@ -1799,17 +1800,18 @@ static void coverageAllocLine(StringRef filename, int line)
 
 static logdata_t mallocData;
 
-static void mallocVisitLine(jl_codectx_t &ctx, StringRef filename, int line)
+static void mallocVisitLine(jl_codectx_t &ctx, StringRef filename, int line, Value *sync)
 {
     assert(!imaging_mode);
     if (filename == "" || filename == "none" || filename == "no file" || filename == "<missing>" || line < 0)
         return;
-    Value *addend = ctx.builder.CreateCall(prepare_call(diff_gc_total_bytes_func), {});
+    Value *addend = sync
+        ? ctx.builder.CreateCall(prepare_call(sync_gc_total_bytes_func), {sync})
+        : ctx.builder.CreateCall(prepare_call(diff_gc_total_bytes_func), {});
     visitLine(ctx, mallocData[filename], line, addend, "bytecnt");
 }
 
-// Resets the malloc counts. Needed to avoid including memory usage
-// from JITting.
+// Resets the malloc counts.
 extern "C" JL_DLLEXPORT void jl_clear_malloc_data(void)
 {
     logdata_t::iterator it = mallocData.begin();
@@ -1826,7 +1828,7 @@ extern "C" JL_DLLEXPORT void jl_clear_malloc_data(void)
             }
         }
     }
-    jl_gc_sync_total_bytes();
+    jl_gc_sync_total_bytes(0);
 }
 
 extern "C" int isabspath(const char *in);
@@ -6374,12 +6376,15 @@ static std::unique_ptr<Module> emit_function(
         }
         new_lineinfo.clear();
     };
-    auto mallocVisitStmt = [&] (unsigned dbg) {
-        if (!do_malloc_log(mod_is_user_mod) || dbg == 0)
+    auto mallocVisitStmt = [&] (unsigned dbg, Value *sync) {
+        if (!do_malloc_log(mod_is_user_mod) || dbg == 0) {
+            if (do_malloc_log(true) && sync)
+                ctx.builder.CreateCall(prepare_call(sync_gc_total_bytes_func), {sync});
             return;
+        }
         while (linetable.at(dbg).inlined_at)
             dbg = linetable.at(dbg).inlined_at;
-        mallocVisitLine(ctx, ctx.file, linetable.at(dbg).line);
+        mallocVisitLine(ctx, ctx.file, linetable.at(dbg).line, sync);
     };
     if (coverage_mode != JL_LOG_NONE) {
         // record all lines that could be covered
@@ -6450,8 +6455,10 @@ static std::unique_ptr<Module> emit_function(
                 current_lineinfo.push_back(1);
         }
     }
-    if (do_malloc_log(mod_is_user_mod))
-        mallocVisitLine(ctx, ctx.file, toplineno);
+    Value *sync_bytes = nullptr;
+    if (do_malloc_log(true))
+        sync_bytes = ctx.builder.CreateCall(prepare_call(diff_gc_total_bytes_func), {});
+
     find_next_stmt(0);
     while (cursor != -1) {
         int32_t debuginfoloc = ((int32_t*)jl_array_data(src->codelocs))[cursor];
@@ -6563,7 +6570,7 @@ static std::unique_ptr<Module> emit_function(
                 }
             }
 
-            mallocVisitStmt(debuginfoloc);
+            mallocVisitStmt(debuginfoloc, sync_bytes);
             if (toplevel)
                 ctx.builder.CreateStore(last_age, ctx.world_age_field);
             assert(type_is_ghost(retty) || returninfo.cc == jl_returninfo_t::SRet ||
@@ -6605,7 +6612,7 @@ static std::unique_ptr<Module> emit_function(
             jl_value_t *cond = args[0];
             int lname = jl_unbox_long(args[1]);
             Value *isfalse = emit_condition(ctx, cond, "if");
-            mallocVisitStmt(debuginfoloc);
+            mallocVisitStmt(debuginfoloc, nullptr);
             come_from_bb[cursor+1] = ctx.builder.GetInsertBlock();
             workstack.push_back(lname - 1);
             BasicBlock *ifnot = BB[lname];
@@ -6643,7 +6650,7 @@ static std::unique_ptr<Module> emit_function(
         }
         else {
             emit_stmtpos(ctx, stmt, cursor);
-            mallocVisitStmt(debuginfoloc);
+            mallocVisitStmt(debuginfoloc, nullptr);
         }
         find_next_stmt(cursor + 1);
     }
@@ -7683,6 +7690,13 @@ static void init_julia_llvm_env(Module *m)
                          "jl_gc_diff_total_bytes", m);
     add_named_global(diff_gc_total_bytes_func, &jl_gc_diff_total_bytes);
 
+    sync_gc_total_bytes_func =
+        Function::Create(FunctionType::get(T_int64, {T_int64}, false),
+                         Function::ExternalLinkage,
+                         "jl_gc_sync_total_bytes", m);
+    add_named_global(sync_gc_total_bytes_func, &jl_gc_sync_total_bytes);
+
+
     std::vector<Type*> array_owner_args(0);
     array_owner_args.push_back(T_prjlvalue);
     jlarray_data_owner_func =

src/gc.c

@@ -2789,9 +2789,13 @@ JL_DLLEXPORT int64_t jl_gc_diff_total_bytes(void)
     return newtb - oldtb;
 }
 
-void jl_gc_sync_total_bytes(void)
+JL_DLLEXPORT int64_t jl_gc_sync_total_bytes(int64_t offset)
 {
-    jl_gc_get_total_bytes(&last_gc_total_bytes);
+    int64_t oldtb = last_gc_total_bytes;
+    int64_t newtb;
+    jl_gc_get_total_bytes(&newtb);
+    last_gc_total_bytes = newtb - offset;
+    return newtb - oldtb;
 }
 
 JL_DLLEXPORT int64_t jl_gc_live_bytes(void)

src/gf.c

@@ -2146,6 +2146,7 @@ STATIC_INLINE jl_value_t *_jl_invoke(jl_value_t *F, jl_value_t **args, uint32_t
         }
         codeinst = codeinst->next;
     }
+    int64_t last_alloc = jl_options.malloc_log ? jl_gc_diff_total_bytes() : 0;
     int last_errno = errno;
 #ifdef _OS_WINDOWS_
     DWORD last_error = GetLastError();
@@ -2155,6 +2156,8 @@ STATIC_INLINE jl_value_t *_jl_invoke(jl_value_t *F, jl_value_t **args, uint32_t
     SetLastError(last_error);
 #endif
     errno = last_errno;
+    if (jl_options.malloc_log)
+        jl_gc_sync_total_bytes(last_alloc); // discard allocation count from compilation
     jl_value_t *res = codeinst->invoke(F, args, nargs, codeinst);
     return verify_type(res);
 }
@@ -2281,6 +2284,7 @@ STATIC_INLINE jl_method_instance_t *jl_lookup_generic_(jl_value_t *F, jl_value_t
         mfunc = entry->func.linfo;
     }
     else {
+        int64_t last_alloc = jl_options.malloc_log ? jl_gc_diff_total_bytes() : 0;
         JL_LOCK(&mt->writelock);
         // cache miss case
         JL_TIMING(METHOD_LOOKUP_SLOW);
@@ -2289,6 +2293,8 @@ STATIC_INLINE jl_method_instance_t *jl_lookup_generic_(jl_value_t *F, jl_value_t
         mfunc = jl_mt_assoc_by_type(mt, tt, /*cache*/1, world);
         JL_GC_POP();
         JL_UNLOCK(&mt->writelock);
+        if (jl_options.malloc_log)
+            jl_gc_sync_total_bytes(last_alloc); // discard allocation count from compilation
         if (mfunc == NULL) {
 #ifdef JL_TRACE
             if (error_en)
@@ -2382,6 +2388,7 @@ static jl_value_t *jl_gf_invoke_by_method(jl_method_t *method, jl_value_t *gf, j
         mfunc = tm->func.linfo;
     }
     else {
+        int64_t last_alloc = jl_options.malloc_log ? jl_gc_diff_total_bytes() : 0;
         jl_svec_t *tpenv = jl_emptysvec;
         jl_tupletype_t *tt = NULL;
         JL_GC_PUSH2(&tpenv, &tt);
@@ -2398,6 +2405,8 @@ static jl_value_t *jl_gf_invoke_by_method(jl_method_t *method, jl_value_t *gf, j
         mfunc = cache_method(NULL, &method->invokes, (jl_value_t*)method, tt, method, 1, tpenv);
         JL_UNLOCK(&method->writelock);
         JL_GC_POP();
+        if (jl_options.malloc_log)
+            jl_gc_sync_total_bytes(last_alloc); // discard allocation count from compilation
     }
     JL_GC_PROMISE_ROOTED(mfunc);
     size_t world = jl_get_ptls_states()->world_age;

src/julia_internal.h

@@ -315,7 +315,7 @@ JL_DLLEXPORT void *jl_gc_counted_malloc(size_t sz);
 JL_DLLEXPORT void JL_NORETURN jl_throw_out_of_memory_error(void);
 
 JL_DLLEXPORT int64_t jl_gc_diff_total_bytes(void);
-void jl_gc_sync_total_bytes(void);
+JL_DLLEXPORT int64_t jl_gc_sync_total_bytes(int64_t offset);
 void jl_gc_track_malloced_array(jl_ptls_t ptls, jl_array_t *a) JL_NOTSAFEPOINT;
 void jl_gc_count_allocd(size_t sz) JL_NOTSAFEPOINT;
 void jl_gc_run_all_finalizers(jl_ptls_t ptls);

test/cmdlineargs.jl

@@ -264,10 +264,31 @@ let exename = `$(Base.julia_cmd()) --startup-file=no`
     @test readchomp(`$exename -E "Base.JLOptions().malloc_log != 0"`) == "false"
     @test readchomp(`$exename -E "Base.JLOptions().malloc_log != 0" --track-allocation=none`) == "false"
 
-    @test readchomp(`$exename -E "Base.JLOptions().malloc_log != 0"
-        --track-allocation`) == "true"
-    @test readchomp(`$exename -E "Base.JLOptions().malloc_log != 0"
-        --track-allocation=user`) == "true"
+    @test readchomp(`$exename -E "Base.JLOptions().malloc_log != 0" --track-allocation`) == "true"
+    @test readchomp(`$exename -E "Base.JLOptions().malloc_log != 0" --track-allocation=user`) == "true"
+    mktempdir() do dir
+        helperdir = joinpath(@__DIR__, "testhelpers")
+        inputfile = joinpath(helperdir, "allocation_file.jl")
+        pid = readchomp(`$exename -E "getpid()" -L $inputfile --track-allocation=user`)
+        memfile = "$inputfile.$pid.mem"
+        got = readlines(memfile)
+        rm(memfile)
+        @test popfirst!(got) == "        0 g(x) = x + 123456"
+        @test popfirst!(got) == "        - function f(x)"
+        @test popfirst!(got) == "       80     []"
+        if Sys.WORD_SIZE == 64
+            @test popfirst!(got) == "       32     Base.invokelatest(g, 0)"
+            @test popfirst!(got) == "       48     Base.invokelatest(g, x)"
+        else
+            @test popfirst!(got) == "       24     Base.invokelatest(g, 0)"
+            @test popfirst!(got) == "       36     Base.invokelatest(g, x)"
+        end
+        @test popfirst!(got) == "       80     []"
+        @test popfirst!(got) == "        - end"
+        @test popfirst!(got) == "        - f(1.23)"
+        @test isempty(got) || got
+    end
+
 
     # --optimize
     @test readchomp(`$exename -E "Base.JLOptions().opt_level"`) == "2"

test/testhelpers/allocation_file.jl

@@ -0,0 +1,8 @@
+g(x) = x + 123456
+function f(x)
+    []
+    Base.invokelatest(g, 0)
+    Base.invokelatest(g, x)
+    []
+end
+f(1.23)