Merge branch 'master' of https://github.com/ahrefs/ocannl

lukstafi · lukstafi · commit f28a0990d33e · 2025-11-28T19:51:00.000+01:00
diff --git a/arrayjit/lib/metal_backend.ml b/arrayjit/lib/metal_backend.ml
@@ -40,7 +40,7 @@ module Device_config = struct
   type runner = {
     queue : Me.CommandQueue.t;
     event : Me.SharedEvent.t; (* Use SharedEvent for signalling *)
-    counter : ullong; (* Next value to signal *)
+    mutable counter : ullong; (* Next value to signal *)
   }
   [@@deriving sexp_of]
 
@@ -212,6 +212,7 @@ end) : Ir.Backend_impl.Lowered_backend = struct
     let shared_event = stream.runner.event in
     let counter = stream.runner.counter in
     let next_value = Unsigned.ULLong.add counter Unsigned.ULLong.one in
+    stream.runner.counter <- next_value;
     let command_buffer = Me.CommandBuffer.on_queue queue in
     Me.CommandBuffer.encode_signal_event command_buffer
       (Me.SharedEvent.super shared_event)
@@ -220,10 +221,10 @@ end) : Ir.Backend_impl.Lowered_backend = struct
     { shared = shared_event; value = next_value }
 
   let await stream =
-    let queue = stream.runner.queue in
-    let command_buffer = Me.CommandBuffer.on_queue queue in
-    Me.CommandBuffer.commit command_buffer;
-    Me.CommandBuffer.wait_until_completed command_buffer;
+    (* Signal an event after all current work and wait for it.
+       This ensures all previously submitted command buffers complete. *)
+    let event = all_work stream in
+    sync event;
     (* Process captured logs if any *)
     if Utils.debug_log_from_routines () then
       match Hashtbl.find stream_logs stream.stream_id with
diff --git a/test/operations/dune b/test/operations/dune
@@ -377,6 +377,17 @@
  (preprocess
   (pps ppx_here ppx_ocannl ppx_expect)))
 
+(test
+ (name small_literal_tensor)
+ (package neural_nets_lib)
+ (deps
+  ocannl_config
+  (env_var OCANNL_BACKEND))
+ (modules small_literal_tensor)
+ (libraries base ocannl stdio)
+ (preprocess
+  (pps ppx_here ppx_ocannl ppx_expect)))
+
 (test
  (name zero2hero_1of7_exec)
  (package neural_nets_lib)
diff --git a/test/operations/small_literal_tensor.expected b/test/operations/small_literal_tensor.expected
@@ -0,0 +1,12 @@
+Retrieving commandline, environment, or config file variable ocannl_log_level
+Found 0, in the config file
+HERE: test/operations/small_literal_tensor.ml:11:21
+[0]: 1,2,3,4,5..._hey shape 1:3->0:2  [
+   1.00 , 2.00 , 3.00
+  ;  4.00 , 5.00 , 6.00
+]
+HERE: test/operations/small_literal_tensor.ml:14:21
+[1]: 1,2,3,4,5..._hoo shape 0:2|1:3  [|
+  [ 1.00 ; 2.00 ; 3.00 ]
+  ; [ 4.00 ; 5.00 ; 6.00 ]
+|]
diff --git a/test/operations/small_literal_tensor.ml b/test/operations/small_literal_tensor.ml
@@ -0,0 +1,14 @@
+open! Base
+open! Ocannl
+open! Nn_blocks.DSL_modules
+
+let () =
+  Tensor.unsafe_reinitialize ();
+  let ctx = Context.auto () in
+  let%op hey = [ (1, 2, 3); (4, 5, 6) ] in
+  let ctx = Train.forward_once ctx hey in
+  (* ignore (failwith @@ Tn.debug_memory_mode hey.value.memory_mode); *)
+  Train.printf ~here:[%here] ~with_code:false ~with_grad:false ~style:`Inline hey;
+  let%op hoo = [| [ 1; 2; 3 ]; [ 4; 5; 6 ] |] in
+  let _ctx = Train.forward_once ctx hoo in
+  Train.printf ~here:[%here] ~with_code:false ~with_grad:false ~style:`Inline hoo
diff --git a/test/operations/test_param_shape_error.expected b/test/operations/test_param_shape_error.expected
@@ -1,6 +1,18 @@
 Retrieving commandline, environment, or config file variable ocannl_log_level
 Found 0, in the config file
-Testing default parameter 1D -- corner case
+Testing default lone parameter -- corner case
+┌───────────────────────────────────────┐
+│[5]: w_o shape 0:4                     │
+│┌┬────────────────────────────────────┐│
+│││axis 0                              ││
+│├┼────────────────────────────────────┤│
+│││ 8.38e-1  1.38e-1  2.61e-1  7.74e-1 ││
+│└┴────────────────────────────────────┘│
+└───────────────────────────────────────┘
+grad_w_o <not-hosted>
+
+ERROR: Should have raised an exception
+Testing lone parameter 1D -- corner case
 ┌───────────────────┐
 │[5]: w_o shape 0:1 │
 │┌┬─────────┐       │
@@ -25,3 +37,7 @@ Testing default affine operation with propagated dimensions
 grad_w <not-hosted>
 Testing default affine operation with unknown input dimensions
 Got expected error: You forgot to specify the hidden dimension(s) 1
+Testing default bias parameter
+Got expected error: You forgot to specify the hidden dimension(s) 4
+Testing bias parameter 1D
+Got expected error: You forgot to specify the hidden dimension(s) 1
diff --git a/test/operations/test_param_shape_error.ml b/test/operations/test_param_shape_error.ml
@@ -13,8 +13,8 @@ let default_lone_param () =
     Stdio.print_endline "\nERROR: Should have raised an exception"
   with Row.Shape_error (msg, _) -> Stdio.printf "Got acceptable error: %s\n" msg
 
-let default_param_1d () =
-  Stdio.printf "Testing default parameter 1D -- corner case\n";
+let lone_param_1d () =
+  Stdio.printf "Testing lone parameter 1D -- corner case\n";
   Tensor.unsafe_reinitialize ();
   (* This should raise an error because we have a parameter with unspecified dimensions *)
   try
@@ -57,9 +57,33 @@ let default_affine_op_unknown_input () =
     ()
   with Row.Shape_error (msg, _) -> Stdio.printf "Got expected error: %s\n" msg
 
+let default_bias_param () =
+  Stdio.printf "Testing default bias parameter\n";
+  Tensor.unsafe_reinitialize ();
+  (* This should raise an error because we have a parameter with unspecified dimensions *)
+  try
+    let%op w_o = { x } + { y } in
+    let _ctx : Context.t = Train.init_params (Context.auto ()) Train.IDX.empty w_o in
+    Train.printf w_o;
+    Stdio.print_endline "\nERROR: Should have raised an exception"
+  with Row.Shape_error (msg, _) -> Stdio.printf "Got expected error: %s\n" msg
+
+let default_bias_param_1d () =
+  Stdio.printf "Testing bias parameter 1D\n";
+  Tensor.unsafe_reinitialize ();
+  (* This should raise an error because we have a parameter with unspecified dimensions *)
+  try
+    let%op w_o = { x = uniform1 () } + { y = uniform1 () } in
+    let _ctx : Context.t = Train.init_params (Context.auto ()) Train.IDX.empty w_o in
+    Train.printf w_o;
+    ()
+  with Row.Shape_error (msg, _) -> Stdio.printf "Got expected error: %s\n" msg
+
 let () =
-  ignore default_lone_param;
-  default_param_1d ();
+  default_lone_param ();
+  lone_param_1d ();
   default_linear_op ();
   default_affine_op_propagated ();
-  default_affine_op_unknown_input ()
+  default_affine_op_unknown_input ();
+  default_bias_param ();
+  default_bias_param_1d ()
