Fixes bus error on large datasets #327, by Claude Opus

lukstafi · lukstafi · commit 947d76c80774 · 2025-08-06T18:56:46.000+02:00
Prompt: Can you help me debug a 'bus error' that happens when running
  `dune exec bin/bigram.exe` ? It happens when batch_size (file
  bin/bigram.ml line 31) is greater or equal to 760, but does not happen
   for smaller batch sizes. You can see generated sources under the
  directory build_files-bus_bug_batch_760/ . The error only happens with
   the multicore_cc backend, the following works: `dune exec
  bin/bigram.exe -- --ocannl_backend=sync_cc` . So, OCaml's Domains must
   be somehow involved. The Domains-based scheduler is implemented in
  arrayjit/lib/schedulers.ml .

  Summary, by Claude:
  I've successfully fixed the bus error that occurred when running dune
  exec bin/bigram.exe with batch_size &gt;= 760.

  Root Cause: The issue was stack overflow in OCaml Domain threads. The
  generated C code was allocating large arrays on the stack. With
  batch_size=760, arrays of size 21280 floats (85KB each) were being
  allocated, with at least 6 such arrays totaling over 510KB of stack
  space. Domain threads have limited stack size, causing a bus error.

  Solution: Modified arrayjit/lib/c_syntax.ml to:
  1. Allocate arrays larger than 16KB on the heap using malloc/calloc
  instead of stack allocation
  2. Track heap-allocated arrays and add cleanup code to free them at
  function exit
  3. This prevents stack overflow while maintaining the same
  functionality

  The fix has been tested successfully with batch_size=760 and even
  batch_size=1500, confirming it handles larger batch sizes robustly.
diff --git a/arrayjit/lib/backends.ml b/arrayjit/lib/backends.ml
@@ -321,15 +321,15 @@ struct
   }
   [@@deriving sexp_of]
 
-  let compile ~name bindings lowered : code =
+  let compile ~(name : string) bindings lowered : code =
     let proc = compile ~name bindings lowered in
     { lowered; proc }
 
   let compile_batch ~names bindings lowereds : code_batch =
     let procs = compile_batch ~names bindings lowereds in
     { lowereds; procs }
 
-  let link context (code : code) ctx_arrays =
+  let link context (code : code) ctx_arrays : Indexing.lowered_bindings * Task.t =
     let runner_label = get_name context.stream in
     let merge_buffer = context.stream.merge_buffer in
     let bindings, to_schedule = link_compiled ~merge_buffer ~runner_label ctx_arrays code.proc in
diff --git a/arrayjit/lib/c_syntax.ml b/arrayjit/lib/c_syntax.ml
@@ -930,6 +930,7 @@ module C_syntax (B : C_syntax_config) = struct
        in
        body := !body ^^ debug_init_doc ^^ hardline);
 
+    let heap_allocated = ref [] in
     let local_decls =
       string "/* Local declarations and initialization. */"
       ^^ hardline
@@ -938,9 +939,24 @@ module C_syntax (B : C_syntax_config) = struct
              if not (Tn.is_virtual_force tn 333 || Tn.is_materialized_force tn 336) then
                let typ_doc = string (B.typ_of_prec @@ Lazy.force tn.prec) in
                let ident_doc = string (get_ident tn) in
-               let size_doc = OCaml.int (Tn.num_elems tn) in
-               let init_doc = if node.Low_level.zero_initialized then string " = {0}" else empty in
-               typ_doc ^^ space ^^ ident_doc ^^ brackets size_doc ^^ init_doc ^^ semi ^^ hardline
+               let num_elems = Tn.num_elems tn in
+               let size_doc = OCaml.int num_elems in
+               (* Use heap allocation for arrays larger than 16KB to avoid stack overflow in Domain threads *)
+               let stack_threshold = 16384 / (Ops.prec_in_bytes @@ Lazy.force tn.prec) in
+               if num_elems > stack_threshold then (
+                 (* Heap allocation for large arrays *)
+                 heap_allocated := get_ident tn :: !heap_allocated;
+                 let alloc_expr = 
+                   if node.Low_level.zero_initialized then
+                     string "calloc" ^^ parens (OCaml.int num_elems ^^ comma ^^ space ^^ string "sizeof" ^^ parens typ_doc)
+                   else  
+                     string "malloc" ^^ parens (OCaml.int num_elems ^^ space ^^ string "*" ^^ space ^^ string "sizeof" ^^ parens typ_doc)
+                 in
+                 typ_doc ^^ space ^^ string "*" ^^ ident_doc ^^ space ^^ equals ^^ space ^^ parens (typ_doc ^^ string "*") ^^ alloc_expr ^^ semi ^^ hardline)
+               else
+                 (* Stack allocation for small arrays *)
+                 let init_doc = if node.Low_level.zero_initialized then string " = {0}" else empty in
+                 typ_doc ^^ space ^^ ident_doc ^^ brackets size_doc ^^ init_doc ^^ semi ^^ hardline
              else empty)
            (Hashtbl.to_alist traced_store)
     in
@@ -949,6 +965,16 @@ module C_syntax (B : C_syntax_config) = struct
     let main_logic = string "/* Main logic. */" ^^ hardline ^^ compile_main llc in
     body := !body ^^ main_logic;
 
+    (* Free heap-allocated arrays *)
+    if not (List.is_empty !heap_allocated) then
+      body :=
+        !body ^^ hardline
+        ^^ string "/* Cleanup heap-allocated arrays. */" ^^ hardline
+        ^^ separate_map hardline
+             (fun ident -> string "free" ^^ parens (string ident) ^^ semi)
+             !heap_allocated
+        ^^ hardline;
+
     if Utils.debug_log_from_routines () && B.log_involves_file_management then
       body :=
         !body ^^ hardline
