From 88a0627003c45ddac304b7be933c93caae8ae6b3 Mon Sep 17 00:00:00 2001
From: pchintalapudi <34727397+pchintalapudi@users.noreply.github.com>
Date: Sat, 26 Nov 2022 04:58:10 -0500
Subject: [PATCH] Fix and simplify inference timing logic (#47711)

* Fix and simplify inference timing logic

* Reduce task struct size
---
 src/gf.c    | 14 ++++++++------
 src/julia.h |  4 ++--
 src/task.c  |  2 ++
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/src/gf.c b/src/gf.c
index 0bce672ca729c..0e98f2a140d4a 100644
--- a/src/gf.c
+++ b/src/gf.c
@@ -3419,18 +3419,20 @@ int jl_has_concrete_subtype(jl_value_t *typ)
 
 JL_DLLEXPORT void jl_typeinf_timing_begin(void)
 {
-    if (jl_atomic_load_relaxed(&jl_measure_compile_time_enabled)) {
-        jl_task_t *ct = jl_current_task;
-        if (ct->inference_start_time == 0 && ct->reentrant_inference == 1)
-            ct->inference_start_time = jl_hrtime();
+    jl_task_t *ct = jl_current_task;
+    if (ct->reentrant_inference == 1) {
+        ct->inference_start_time = jl_hrtime();
     }
 }
 
 JL_DLLEXPORT void jl_typeinf_timing_end(void)
 {
     jl_task_t *ct = jl_current_task;
-    if (ct->inference_start_time != 0 && ct->reentrant_inference == 1) {
-        jl_atomic_fetch_add_relaxed(&jl_cumulative_compile_time, (jl_hrtime() - ct->inference_start_time));
+    if (ct->reentrant_inference == 1) {
+        if (jl_atomic_load_relaxed(&jl_measure_compile_time_enabled)) {
+            uint64_t inftime = jl_hrtime() - ct->inference_start_time;
+            jl_atomic_fetch_add_relaxed(&jl_cumulative_compile_time, inftime);
+        }
         ct->inference_start_time = 0;
     }
 }
diff --git a/src/julia.h b/src/julia.h
index 1ec6fe2bd39bf..981e6a0ee8232 100644
--- a/src/julia.h
+++ b/src/julia.h
@@ -1940,8 +1940,8 @@ typedef struct _jl_task_t {
     void *stkbuf; // malloc'd memory (either copybuf or stack)
     size_t bufsz; // actual sizeof stkbuf
     uint64_t inference_start_time; // time when inference started
-    unsigned int reentrant_inference; // How many times we've reentered inference
-    unsigned int reentrant_codegen; // How many times we've reentered codegen
+    uint16_t reentrant_inference; // How many times we've reentered inference
+    uint16_t reentrant_codegen; // How many times we've reentered codegen
     unsigned int copy_stack:31; // sizeof stack for copybuf
     unsigned int started:1;
 } jl_task_t;
diff --git a/src/task.c b/src/task.c
index 81b90a832e2dd..a5ebc1ce26005 100644
--- a/src/task.c
+++ b/src/task.c
@@ -940,6 +940,7 @@ JL_DLLEXPORT jl_task_t *jl_new_task(jl_function_t *start, jl_value_t *completion
     t->world_age = ct->world_age;
     t->reentrant_codegen = 0;
     t->reentrant_inference = 0;
+    t->inference_start_time = 0;
 
 #ifdef COPY_STACKS
     if (!t->copy_stack) {
@@ -1527,6 +1528,7 @@ jl_task_t *jl_init_root_task(jl_ptls_t ptls, void *stack_lo, void *stack_hi)
     ct->world_age = 1; // OK to run Julia code on this task
     ct->reentrant_codegen = 0;
     ct->reentrant_inference = 0;
+    ct->inference_start_time = 0;
     ptls->root_task = ct;
     jl_atomic_store_relaxed(&ptls->current_task, ct);
     JL_GC_PROMISE_ROOTED(ct);