Use epoll instead of select to avoid the FD_SETSIZE limit

lukstafi · lukstafi · commit ba1c2319fb6b · 2024-09-09T22:57:08.000+02:00
But `pipes_cc` is super slow compared to `cc`, I might just get rid of it.
diff --git a/arrayjit.opam b/arrayjit.opam
@@ -25,6 +25,7 @@ depends: [
   "stdio"
   "num"
   "saturn_lockfree"
+  "poll"
   "ppxlib"
   "ppx_jane"
   "ppx_expect"
diff --git a/arrayjit/lib/dune b/arrayjit/lib/dune
@@ -12,6 +12,7 @@
   ctypes
   ctypes.foreign
   saturn_lockfree
+  poll
   (select
    gcc_backend.ml
    from
diff --git a/arrayjit/lib/utils.ml b/arrayjit/lib/utils.ml
@@ -556,6 +556,18 @@ type waiter = {
 }
 (** Note: this waiter is meant for sequential waiting. *)
 
+let poll pipe_inp =
+  let poll = Poll.create () in
+  Poll.set poll pipe_inp Poll.Event.read;
+  let rec wait () =
+    Poll.clear poll;
+    match Poll.wait poll (Poll.Timeout.after 10_000_000L) with
+    | `Timeout -> ()
+    | `Ok -> ()
+    | exception Unix.Unix_error (Unix.EINTR, _, _) -> wait ()
+  in
+  wait
+
 let waiter ~name () =
   let is_open = Atomic.make true in
   (* TODO: since OCaml 5.2, consider [make_contended]? *)
@@ -564,6 +576,7 @@ let waiter ~name () =
   let pipe_inp, pipe_out =
     try Unix.pipe ~cloexec:true () with e -> Exn.reraise e @@ "waiter " ^ name ^ ": Unix.pipe"
   in
+  let poll = poll pipe_inp in
   let await ~keep_waiting =
     let%track_l_sexp rec wait () : bool =
       if Atomic.compare_and_set is_released true false then (
@@ -574,10 +587,7 @@ let waiter ~name () =
         assert (n = 1);
         true)
       else if keep_waiting () then
-        let _, _, _ =
-          try Unix.select [ pipe_inp ] [] [] 5.0
-          with e -> Exn.reraise e @@ "waiter " ^ name ^ ": Unix.select"
-        in
+        let () = try poll () with e -> Exn.reraise e @@ "waiter " ^ name ^ ": poll error" in
         wait ()
       else false
     in
diff --git a/bin/moons_benchmark.ml b/bin/moons_benchmark.ml
@@ -171,7 +171,7 @@ let cuda_benchmarks =
       List.concat_map [ 1; 3; 6; 12 (* ; 16; 32; 64 *) ] ~f:(fun num_devices ->
           List.concat_map [ 64; 128 (* ; 256 *) ] ~f:(fun batch_size ->
               List.concat_map [ 0; 1 (* ; 2; 3; 4 *) ] ~f:(fun seed ->
-                  List.concat_map [ (* "gccjit" ; *) "cc"; "cuda" ] ~f:(fun backend_name ->
+                  List.concat_map [ (* "gccjit" ; *) "pipes_cc" (* ; "cc"; "cuda" *) ] ~f:(fun backend_name ->
                       [
                         classify_moons ~seed ~on_device:true ~inlining_cutoff ~num_devices
                           ~batch_size ~backend_name ~value_prec:CDSL.single ~grad_prec:CDSL.single;
diff --git a/dune-project b/dune-project
@@ -66,6 +66,7 @@
   stdio
   num
   saturn_lockfree
+  poll
   ppxlib
   ppx_jane
   ppx_expect

Original file line number	Diff line number	Diff line change
`@@ -171,7 +171,7 @@ let cuda_benchmarks =`
`171`	`171`	`List.concat_map [ 1; 3; 6; 12 (* ; 16; 32; 64 *) ] ~f:(fun num_devices ->`
`172`	`172`	`List.concat_map [ 64; 128 (* ; 256 *) ] ~f:(fun batch_size ->`
`173`	`173`	`List.concat_map [ 0; 1 (* ; 2; 3; 4 *) ] ~f:(fun seed ->`
`174`		`- List.concat_map [ (* "gccjit" ; *) "cc"; "cuda" ] ~f:(fun backend_name ->`
	`174`	`+ List.concat_map [ (* "gccjit" ; ) "pipes_cc" ( ; "cc"; "cuda" *) ] ~f:(fun backend_name ->`
`175`	`175`	`[`
`176`	`176`	`classify_moons ~seed ~on_device:true ~inlining_cutoff ~num_devices`
`177`	`177`	`~batch_size ~backend_name ~value_prec:CDSL.single ~grad_prec:CDSL.single;`