Track caml_alloc_custom_mem allocations with Statmemprof

Changes

@@ -5,6 +5,9 @@ Working version
 
 ### Runtime system:
 
+- #??: Track custom blocks (e.g. Bigarray) with Statmemprof
+  (Stephen Dolan, review by ??)
+
 ### Code generation and optimizations:
 
 ### Standard library:

runtime/caml/memprof.h

@@ -27,6 +27,7 @@ extern void caml_memprof_set_suspended(int);
 extern value caml_memprof_handle_postponed_exn(void);
 
 extern void caml_memprof_track_alloc_shr(value block);
+extern void caml_memprof_track_custom(value block, mlsize_t bytes);
 extern void caml_memprof_track_young(uintnat wosize, int from_caml,
                                      int nallocs, unsigned char* alloc_lens);
 extern void caml_memprof_track_interned(header_t* block, header_t* blockend);

runtime/custom.c

@@ -24,6 +24,7 @@
 #include "caml/memory.h"
 #include "caml/mlvalues.h"
 #include "caml/signals.h"
+#include "caml/memprof.h"
 
 uintnat caml_custom_major_ratio = Custom_major_ratio_def;
 uintnat caml_custom_minor_ratio = Custom_minor_ratio_def;
@@ -102,7 +103,9 @@ CAMLexport value caml_alloc_custom_mem(struct custom_operations * ops,
     Bsize_wsize (Caml_state->stat_heap_wsz) / 150 * caml_custom_major_ratio;
   mlsize_t max_minor =
     Bsize_wsize (Caml_state->minor_heap_wsz) / 100 * caml_custom_minor_ratio;
-  return alloc_custom_gen (ops, bsz, mem, max_major, mem_minor, max_minor);
+  value v = alloc_custom_gen (ops, bsz, mem, max_major, mem_minor, max_minor);
+  caml_memprof_track_custom(v, mem);
+  return v;
 }
 
 struct custom_operations_list {

runtime/memprof.c

@@ -765,6 +765,25 @@ void caml_memprof_track_alloc_shr(value block)
   check_action_pending();
 }
 
+void caml_memprof_track_custom(value block, mlsize_t bytes)
+{
+  uintnat n_samples;
+  value callstack = 0;
+  CAMLassert(Is_young(block) || Is_in_heap(block));
+
+  if (lambda == 0 || local->suspended) return;
+
+  n_samples = rand_binom(Wsize_bsize(bytes));
+  if (n_samples == 0) return;
+
+  callstack = capture_callstack_postponed();
+  if (callstack == 0) return;
+
+  new_tracked(n_samples, Wsize_bsize(bytes), 0, Is_young(block),
+              block, callstack);
+  check_action_pending();
+}
+
 /* Shifts the next sample in the minor heap by [n] words. Essentially,
    this tells the sampler to ignore the next [n] words of the minor
    heap. */

testsuite/tests/statmemprof/custom.ml

@@ -0,0 +1,44 @@
+(* TEST *)
+
+open Gc.Memprof
+
+let bigstring_create sz =
+  Bigarray.Array1.create Bigarray.char Bigarray.c_layout sz
+
+let keep = ref []
+
+let test sampling_rate =
+  let size = 256 in
+  let iters = 100_000 in
+  let size_words = size / (Sys.word_size / 8) in
+  let alloc = ref 0 and collect = ref 0 and promote = ref 0 in
+  let tracker =
+    { null_tracker with
+      alloc_minor = (fun info ->
+        if info.size < size_words then None
+        else begin
+          alloc := !alloc + info.n_samples;
+          Some info.n_samples
+        end);
+      promote = (fun ns ->
+        promote := !promote + ns; None);
+      dealloc_minor = (fun ns ->
+        collect := !collect + ns) } in
+  start ~sampling_rate tracker;
+  for i = 1 to iters do
+    let str = Sys.opaque_identity bigstring_create size in
+    if i mod 10 = 0 then keep := str :: !keep
+  done;
+  keep := [];
+  Gc.full_major ();
+  stop ();
+  assert (!alloc = !promote + !collect);
+  let iters = float_of_int iters and size_words = float_of_int size_words in
+  (* see comballoc.ml for notes on precision *)
+  Printf.printf "%.2f %.1f\n"
+    ((float_of_int !alloc /. iters) /. size_words)
+    ((float_of_int !promote /. iters) /. size_words *. 10.)
+
+
+let () =
+  [0.01; 0.5; 0.17] |> List.iter test

testsuite/tests/statmemprof/custom.reference

@@ -0,0 +1,3 @@
+0.01 0.0
+0.50 0.5
+0.17 0.2