In progress: getting kernel logging to work again on the cuda backend

lukstafi · lukstafi · commit 2c2690bf488a · 2025-05-20T20:47:35.000+02:00
Signed-off-by: lukstafi &lt;lukstafi@users.noreply.github.com&gt;
diff --git a/arrayjit/lib/cuda_backend.ml b/arrayjit/lib/cuda_backend.ml
@@ -264,12 +264,12 @@ end) : Ir.Backend_impl.Lowered_backend = struct
 
       let use_host_memory = None
       let procs = Input.procs
+
       let full_printf_support =
         not @@ Bool.of_string
         @@ Utils.get_global_arg ~default:"false" ~arg_name:"prefer_backend_uniformity"
     end)
 
-    let logs_to_stdout = true
     let main_kernel_prefix = "extern \"C\" __global__"
 
     let kernel_prep_line =
@@ -450,7 +450,7 @@ end) : Ir.Backend_impl.Lowered_backend = struct
       let open PPrint in
       let format_string_literal =
         !Utils.captured_log_prefix ^ "%d: "
-        ^ String.substr_replace_all base_message_literal ~pattern:"\n" ~with_:"$"
+        ^ String.substr_replace_all base_message_literal ~pattern:"\n" ~with_:"$" ^ "\\n"
       in
       let all_args =
         match log_param_c_expr_doc with
@@ -568,7 +568,6 @@ end) : Ir.Backend_impl.Lowered_backend = struct
       [%log "launching the kernel"];
       (* Stdio.printf "launching %s\n" name; *)
       (if Utils.debug_log_from_routines () then
-         (* FIXME: this needs to use Utils.log_debug_routine_logs. *)
          Utils.add_log_processor ~prefix:log_id_prefix @@ fun log_contents ->
          Utils.log_debug_routine_logs ~log_contents ~stream_name);
       S.launch_kernel func ~grid_dim_x:1 ~block_dim_x:1 ~shared_mem_bytes:0 stream.runner args;
diff --git a/test/dune b/test/dune
@@ -40,6 +40,7 @@
    "--ocannl_debug_log_to_stream_files=true")))
 
 (rule
+ (deps "log_files/%{read:config/ocannl_backend.txt}-0-0.log")
  (target "%{read:config/ocannl_backend.txt}-0-0.log.actual")
  (action
   (progn
diff --git a/test/micrograd_demo_logging.ml b/test/micrograd_demo_logging.ml
@@ -32,7 +32,8 @@ let () =
   List.iter ~f:(Option.iter ~f:(fun diff -> Train.set_hosted diff.Tensor.grad)) [ a.diff; b.diff ];
   let update = Train.grad_update g in
   let step = Train.to_routine (module Backend) ctx IDX.empty update.fwd_bprop in
+  Utils.capture_stdout_logs @@ fun () ->
   Train.run step;
   Tensor.print ~with_code:false ~with_grad:false `Default g;
   Tensor.print ~with_code:false ~with_grad:true `Default a;
-  Tensor.print ~with_code:false ~with_grad:true `Default b
+  Tensor.print ~with_code:false ~with_grad:true `Default b
