Add JIT memory counters

Author: pchintalapudi

src/cgmemmgr.cpp

@@ -866,6 +866,8 @@ uint8_t *RTDyldMemoryManagerJL::allocateCodeSection(uintptr_t Size,
     code_allocated = true;
 #endif
     total_allocated += Size;
+    jl_timing_counter_inc(JL_TIMING_COUNTER_JITSize, Size);
+    jl_timing_counter_inc(JL_TIMING_COUNTER_JITCodeSize, Size);
     if (exe_alloc)
         return (uint8_t*)exe_alloc->alloc(Size, Alignment);
     return SectionMemoryManager::allocateCodeSection(Size, Alignment, SectionID,
@@ -879,6 +881,8 @@ uint8_t *RTDyldMemoryManagerJL::allocateDataSection(uintptr_t Size,
                                                     bool isReadOnly)
 {
     total_allocated += Size;
+    jl_timing_counter_inc(JL_TIMING_COUNTER_JITSize, Size);
+    jl_timing_counter_inc(JL_TIMING_COUNTER_JITDataSize, Size);
     if (!isReadOnly)
         return (uint8_t*)rw_alloc.alloc(Size, Alignment);
     if (ro_alloc)

src/jitlayers.cpp

@@ -851,10 +851,27 @@ class JLMemoryUsagePlugin : public ObjectLinkingLayer::Plugin {
                           jitlink::PassConfiguration &Config) override {
         Config.PostAllocationPasses.push_back([this](jitlink::LinkGraph &G) {
             size_t graph_size = 0;
+            size_t code_size = 0;
+            size_t data_size = 0;
             for (auto block : G.blocks()) {
                 graph_size += block->getSize();
             }
+            for (auto &section : G.sections()) {
+                size_t secsize = 0;
+                for (auto block : section.blocks()) {
+                    secsize += block->getSize();
+                }
+                if ((section.getMemProt() & orc::MemProt::Exec) == orc::MemProt::None) {
+                    data_size += secsize;
+                } else {
+                    code_size += secsize;
+                }
+                graph_size += secsize;
+            }
             this->total_size.fetch_add(graph_size, std::memory_order_relaxed);
+            jl_timing_counter_inc(JL_TIMING_COUNTER_JITSize, graph_size);
+            jl_timing_counter_inc(JL_TIMING_COUNTER_JITCodeSize, code_size);
+            jl_timing_counter_inc(JL_TIMING_COUNTER_JITDataSize, data_size);
             return Error::success();
         });
     }

src/timing.c

@@ -108,6 +108,9 @@ void jl_init_timing(void)
     // We reference these by enum indexing and then asking for the name, since that allows the compiler
     // to catch name mismatches.
     TracyCPlotConfig(jl_timing_counters[JL_TIMING_COUNTER_HeapSize].tracy_counter.name, TracyPlotFormatMemory, /* rectilinear */ 0, /* fill */ 1, /* color */ 0);
+    TracyCPlotConfig(jl_timing_counters[JL_TIMING_COUNTER_JITSize].tracy_counter.name, TracyPlotFormatMemory, /* rectilinear */ 0, /* fill */ 1, /* color */ 0);
+    TracyCPlotConfig(jl_timing_counters[JL_TIMING_COUNTER_JITCodeSize].tracy_counter.name, TracyPlotFormatMemory, /* rectilinear */ 0, /* fill */ 1, /* color */ 0);
+    TracyCPlotConfig(jl_timing_counters[JL_TIMING_COUNTER_JITDataSize].tracy_counter.name, TracyPlotFormatMemory, /* rectilinear */ 0, /* fill */ 1, /* color */ 0);
 #endif
 }
 

src/timing.h

@@ -194,6 +194,9 @@ void jl_timing_puts(jl_timing_block_t *cur_block, const char *str);
 #define JL_TIMING_COUNTERS \
         X(Invalidations) \
         X(HeapSize) \
+        X(JITSize) \
+        X(JITCodeSize) \
+        X(JITDataSize) \
 
 
 enum jl_timing_owners {