Factor out Indexing.Pp_helpers, more interface files

lukstafi · lukstafi · commit 1f2a22b2f4fe · 2024-10-19T21:28:19.000+02:00
diff --git a/CHANGES.md b/CHANGES.md
@@ -2,6 +2,7 @@
 
 ### Added
 
+- Interface files for `Backends` and `Low_level`.
 - TODO: stream-to-stream synchronization functionality, with lazy per-tensor-node synchronization.
 - Fixed #245: tracking of used memory.
 
@@ -16,7 +17,7 @@
 - Moved the multicore backend from a `device = stream` model to a single device model.
 - Got rid of `unsafe_cleanup`.
 - Got rid of `subordinal`.
-- Removed dependency on `dore`, broke up dependency on `ppx_jane`.
+- Removed dependency on `core`, broke up dependency on `ppx_jane`.
 - TODO: Built per-tensor-node stream-to-stream synchronization into device-to-device copying functions, removed obsolete blocking synchronizations.
 
 ### Fixed
diff --git a/arrayjit/lib/backends.ml b/arrayjit/lib/backends.ml
@@ -805,15 +805,12 @@ end
 module Cuda_backend : Backend_types.Backend = Lowered_backend ((
   Cuda_backend : Backend_types.Lowered_backend))
 
-(** Initializes the backend, and if it was already initialized, performs garbage collection. *)
 let reinitialize (module Backend : Backend_types.Backend) config =
   if not @@ Backend.is_initialized () then Backend.initialize config
   else (
     Stdlib.Gc.full_major ();
     Backend.initialize config)
 
-(** Reinitializes and returns a backend corresponding to [backend_name], or if omitted, selected via
-    the global [backend] setting. See {!reinitialize}. *)
 let fresh_backend ?backend_name ?(config = Only_devices_parallel) () =
   let backend =
     match
diff --git a/arrayjit/lib/backends.mli b/arrayjit/lib/backends.mli
@@ -0,0 +1,22 @@
+(** {1 The collection of the execution backends} *)
+
+open Base
+
+val sync_suggested_num_streams : int ref
+
+module Cc_backend : Backend_types.Backend
+module Sync_cc_backend : Backend_types.Backend
+module Gccjit_backend : Backend_types.Backend
+module Sync_gccjit_backend : Backend_types.Backend
+module Cuda_backend : Backend_types.Backend
+
+val reinitialize : (module Backend_types.Backend) -> Backend_types.Types.config -> unit
+(** Initializes the backend, and if it was already initialized, performs garbage collection. *)
+
+val fresh_backend :
+  ?backend_name:string ->
+  ?config:Backend_types.Types.config ->
+  unit ->
+  (module Backend_types.Backend)
+(** Reinitializes and returns a backend corresponding to [backend_name], or if omitted, selected via
+    the global [backend] setting. See {!reinitialize}. *)
diff --git a/arrayjit/lib/c_syntax.ml b/arrayjit/lib/c_syntax.ml
@@ -38,16 +38,7 @@ struct
   let pp_zero_out ppf tn =
     Stdlib.Format.fprintf ppf "@[<2>memset(%s, 0, %d);@]@ " (get_ident tn) @@ Tn.size_in_bytes tn
 
-  (* let pp_semi ppf () = Stdlib.Format.fprintf ppf ";@ " *)
-  let pp_comma ppf () = Stdlib.Format.fprintf ppf ",@ "
-
-  (* let pp_symbol ppf sym = Stdlib.Format.fprintf ppf "%s" @@ Indexing.symbol_ident sym *)
-  let pp_index ppf sym = Stdlib.Format.fprintf ppf "%s" @@ Indexing.symbol_ident sym
-
-  let pp_index_axis ppf = function
-    | Indexing.Iterator it -> pp_index ppf it
-    | Fixed_idx i when i < 0 -> Stdlib.Format.fprintf ppf "(%d)" i
-    | Fixed_idx i -> Stdlib.Format.fprintf ppf "%d" i
+  open Indexing.Pp_helpers
 
   let pp_array_offset ppf (idcs, dims) =
     let open Stdlib.Format in
@@ -57,9 +48,9 @@ struct
     done;
     for i = 0 to Array.length idcs - 1 do
       let dim = dims.(i) in
-      if i = 0 then fprintf ppf "%a" pp_index_axis idcs.(i)
-      else if i = Array.length idcs - 1 then fprintf ppf " * %d + %a" dim pp_index_axis idcs.(i)
-      else fprintf ppf " * %d +@ %a@;<0 -1>)@]" dim pp_index_axis idcs.(i)
+      if i = 0 then fprintf ppf "%a" pp_axis_index idcs.(i)
+      else if i = Array.length idcs - 1 then fprintf ppf " * %d + %a" dim pp_axis_index idcs.(i)
+      else fprintf ppf " * %d +@ %a@;<0 -1>)@]" dim pp_axis_index idcs.(i)
     done
 
   let array_offset_to_string (idcs, dims) =
@@ -113,15 +104,15 @@ struct
           fprintf ppf "@[<v 0>%a@]" (pp_print_list pp_ll)
             (List.filter [ c1; c2 ] ~f:(function Noop -> false | _ -> true))
       | For_loop { index = i; from_; to_; body; trace_it = _ } ->
-          fprintf ppf "@[<2>for (int@ %a = %d;@ %a <= %d;@ ++%a) {@ " pp_index i from_ pp_index i
-            to_ pp_index i;
+          fprintf ppf "@[<2>for (int@ %a = %d;@ %a <= %d;@ ++%a) {@ " pp_symbol i from_ pp_symbol i
+            to_ pp_symbol i;
           if Utils.debug_log_from_routines () then
             if B.logs_to_stdout then
               fprintf ppf {|printf(@[<h>"%s%%d: index %a = %%d\n",@] log_id, %a);@ |}
-                !Utils.captured_log_prefix pp_index i pp_index i
+                !Utils.captured_log_prefix pp_symbol i pp_symbol i
             else
-              fprintf ppf {|fprintf(log_file,@ @[<h>"index %a = %%d\n",@] %a);@ |} pp_index i
-                pp_index i;
+              fprintf ppf {|fprintf(log_file,@ @[<h>"index %a = %%d\n",@] %a);@ |} pp_symbol i
+                pp_symbol i;
           fprintf ppf "%a@;<1 -2>}@]@," pp_ll body
       | Zero_out tn ->
           let traced = Low_level.(get_node traced_store tn) in
@@ -235,7 +226,7 @@ struct
           fprintf ppf "%s%.16g%s" prefix c postfix
       | Embed_index idx ->
           let prefix, postfix = B.convert_precision ~from:Ops.double ~to_:prec in
-          fprintf ppf "%s%a%s" prefix pp_index_axis idx postfix
+          fprintf ppf "%s%a%s" prefix pp_axis_index idx postfix
       | Binop (Arg1, v1, _v2) -> loop ppf v1
       | Binop (Arg2, _v1, v2) -> loop ppf v2
       | Binop (op, v1, v2) ->
diff --git a/arrayjit/lib/indexing.ml b/arrayjit/lib/indexing.ml
@@ -175,3 +175,21 @@ let derive_index ~product_syms ~(projection : axis_index array) =
       | it -> Second it)
   in
   fun ~product -> Array.map positions ~f:(function First p -> product.(p) | Second it -> it)
+
+module Pp_helpers = struct
+  let pp_comma ppf () = Stdlib.Format.fprintf ppf ",@ "
+  let pp_symbol ppf sym = Stdlib.Format.fprintf ppf "%s" @@ symbol_ident sym
+
+  let pp_static_symbol ppf { static_symbol; static_range } =
+    match static_range with
+    | None -> pp_symbol ppf static_symbol
+    | Some range -> Stdlib.Format.fprintf ppf "%a : [0..%d]" pp_symbol static_symbol (range - 1)
+
+  let pp_axis_index ppf idx =
+    match idx with
+    | Iterator sym -> pp_symbol ppf sym
+    | Fixed_idx i -> Stdlib.Format.fprintf ppf "%d" i
+
+  let pp_indices ppf idcs =
+    Stdlib.Format.pp_print_list ~pp_sep:pp_comma pp_axis_index ppf @@ Array.to_list idcs
+end
diff --git a/arrayjit/lib/low_level.ml b/arrayjit/lib/low_level.ml
@@ -1,7 +1,6 @@
 open Base
 
 module Lazy = Utils.Lazy
-(** The code for operating on n-dimensional arrays. *)
 
 module Nd = Ndarray
 module Tn = Tnode
@@ -26,15 +25,12 @@ end
 type scope_id = Scope_id.t = { tn : Tn.t; scope_id : int }
 [@@deriving sexp_of, equal, hash, compare]
 
-(** *** Low-level representation. *)
-
 let get_scope =
   let uid = ref 0 in
   fun tn ->
     Int.incr uid;
     { tn; scope_id = !uid }
 
-(** Cases: [t] -- code, [float_t] -- single number at some precision. *)
 type t =
   | Noop
   | Comment of string
@@ -71,12 +67,6 @@ let rec unflat_lines = function
   | Noop :: tl -> unflat_lines tl
   | llc :: tl -> Seq (llc, unflat_lines tl)
 
-let comment_to_name =
-  let nonliteral = Str.regexp {|[^a-zA-Z0-9_]|} in
-  Str.global_replace nonliteral "_"
-
-(** *** Optimization *** *)
-
 type virtualize_settings = {
   mutable enable_device_only : bool;
   mutable max_visits : int;
@@ -102,31 +92,18 @@ let virtualize_settings =
 type visits =
   | Visits of int
   | Recurrent
-      (** A [Recurrent] visit is when there is an access prior to any assignment in an update. *)
 [@@deriving sexp, equal, variants]
 
 type traced_array = {
   tn : Tn.t;
   mutable computations : (Indexing.axis_index array option * t) list;
-      (** The computations (of the tensor node) are retrieved for optimization just as they are
-          populated, so that the inlined code corresponds precisely to the changes to the arrays
-          that would happen up till that point. Within the code blocks paired with an index tuple,
-          all assignments and accesses must happen via the index tuple; if this is not the case for
-          some assignment, the node cannot be virtual. Currently, we only allow for-loop symbols in
-          assignment indices of virtual nodes. *)
   assignments : int array Hash_set.t;
   accesses : (int array, visits) Hashtbl.t;
-      (** For dynamic indexes, we take a value of 0. This leads to an overestimate of visits, which
-          is safe. *)
   mutable zero_initialized : bool;
   mutable zeroed_out : bool;
   mutable read_before_write : bool;
-      (** The node is read before it is written (i.e. it is recurrent). *)
   mutable read_only : bool;
   mutable is_scalar_constexpr : bool;
-      (** True only if the tensor node has all axes of dimension 1, is either zeroed-out or assigned
-          before accessed, is assigned at most once, and from an expression involving only constants
-          or tensor nodes that were at the time is_scalar_constexpr. *)
 }
 [@@deriving sexp_of]
 
@@ -144,22 +121,6 @@ let get_node store tn =
         is_scalar_constexpr = false;
       })
 
-let partition_tf_with_comment cs ~f =
-  let both = Array.map cs ~f:(fun c -> if f c then Either.First c else Either.Second c) in
-  let trues =
-    Array.filter_map both ~f:(function
-      | First x -> Some x
-      | Second (Comment _ as x) -> Some x
-      | Second _ -> None)
-  in
-  let falses =
-    Array.filter_map both ~f:(function
-      | First (Comment _ as x) -> Some x
-      | First _ -> None
-      | Second x -> Some x)
-  in
-  (trues, falses)
-
 let visit ~is_assigned old =
   if not is_assigned then Recurrent
   else
@@ -801,21 +762,8 @@ let%diagn2_sexp optimize_proc static_indices llc =
   { traced_store; llc; merge_node }
 
 let code_hum_margin = ref 100
-let pp_comma ppf () = Stdlib.Format.fprintf ppf ",@ "
-let pp_symbol ppf sym = Stdlib.Format.fprintf ppf "%s" @@ Indexing.symbol_ident sym
-
-let pp_static_symbol ppf { Indexing.static_symbol; static_range } =
-  match static_range with
-  | None -> pp_symbol ppf static_symbol
-  | Some range -> Stdlib.Format.fprintf ppf "%a : [0..%d]" pp_symbol static_symbol (range - 1)
-
-let pp_index ppf idx =
-  match idx with
-  | Indexing.Iterator sym -> pp_symbol ppf sym
-  | Fixed_idx i -> Stdlib.Format.fprintf ppf "%d" i
 
-let pp_indices ppf idcs =
-  Stdlib.Format.pp_print_list ~pp_sep:pp_comma pp_index ppf @@ Array.to_list idcs
+open Indexing.Pp_helpers
 
 let fprint_function_header ?name ?static_indices () ppf =
   let open Stdlib.Format in
@@ -924,7 +872,7 @@ let fprint_hum ?name ?static_indices () ppf llc =
         fprintf ppf "@[<2>%a.merge[@,%a]@]" pp_ident tn pp_indices idcs
     | Get (tn, idcs) -> fprintf ppf "@[<2>%a[@,%a]@]" pp_ident tn pp_indices idcs
     | Constant c -> fprintf ppf "%.16g" c
-    | Embed_index idx -> pp_index ppf idx
+    | Embed_index idx -> pp_axis_index ppf idx
     | Binop (Arg1, v1, _v2) -> pp_float prec ppf v1
     | Binop (Arg2, _v1, v2) -> pp_float prec ppf v2
     | Binop (op, v1, v2) ->
diff --git a/arrayjit/lib/low_level.mli b/arrayjit/lib/low_level.mli
@@ -0,0 +1,121 @@
+(** {1 A for-loop-based array language and backend-agnostic optimization} *)
+
+open Base
+
+module Scope_id : sig
+  type t = { tn : Tnode.t; scope_id : int } [@@deriving sexp_of, equal, hash, compare]
+  type comparator_witness
+
+  val comparator : (t, comparator_witness) Base.Comparator.t
+end
+
+type scope_id = Scope_id.t = { tn : Tnode.t; scope_id : int }
+[@@deriving sexp_of, equal, hash, compare]
+
+(** {2 Low-level representation} *)
+
+(** Cases: [t] -- code, [float_t] -- single number at some precision. *)
+type t =
+  | Noop
+  | Comment of string
+  | Staged_compilation of (unit -> unit)
+  | Seq of t * t
+  | For_loop of { index : Indexing.symbol; from_ : int; to_ : int; body : t; trace_it : bool }
+  | Zero_out of Tnode.t
+  | Set of { tn : Tnode.t; idcs : Indexing.axis_index array; llv : float_t; mutable debug : string }
+  | Set_local of scope_id * float_t
+[@@deriving sexp_of, equal]
+
+and float_t =
+  | Local_scope of { id : scope_id; body : t; orig_indices : Indexing.axis_index array }
+  | Get_local of scope_id
+  | Get_global of Ops.global_identifier * Indexing.axis_index array option
+  | Get of Tnode.t * Indexing.axis_index array
+  | Binop of Ops.binop * float_t * float_t
+  | Unop of Ops.unop * float_t
+  | Constant of float
+  | Embed_index of Indexing.axis_index
+[@@deriving sexp_of, equal, compare]
+
+val binop : op:Ops.binop -> rhs1:float_t -> rhs2:float_t -> float_t
+val unop : op:Ops.unop -> rhs:float_t -> float_t
+val flat_lines : t list -> t list
+val unflat_lines : t list -> t
+val loop_over_dims : int array -> body:(Indexing.axis_index array -> t) -> t
+
+(** {2 Optimization} *)
+
+type virtualize_settings = {
+  mutable enable_device_only : bool;
+  mutable max_visits : int;
+  mutable max_tracing_dim : int;
+  mutable inline_scalar_constexprs : bool;
+}
+
+val virtualize_settings : virtualize_settings
+
+type visits =
+  | Visits of int
+  | Recurrent
+      (** A [Recurrent] visit is when there is an access prior to any assignment in an update. *)
+[@@deriving sexp, equal, variants]
+
+type traced_array = {
+  tn : Tnode.t;
+  mutable computations : (Indexing.axis_index array option * t) list;
+      (** The computations (of the tensor node) are retrieved for optimization just as they are
+          populated, so that the inlined code corresponds precisely to the changes to the arrays
+          that would happen up till that point. Within the code blocks paired with an index tuple,
+          all assignments and accesses must happen via the index tuple; if this is not the case for
+          some assignment, the node cannot be virtual. Currently, we only allow for-loop symbols in
+          assignment indices of virtual nodes. *)
+  assignments : int array Base.Hash_set.t;
+  accesses : (int array, visits) Base.Hashtbl.t;
+  mutable zero_initialized : bool;
+  mutable zeroed_out : bool;
+  mutable read_before_write : bool;
+      (** The node is read before it is written (i.e. it is recurrent). *)
+  mutable read_only : bool;
+  mutable is_scalar_constexpr : bool;
+      (** True only if the tensor node has all axes of dimension 1, is either zeroed-out or assigned
+          before accessed, is assigned at most once, and from an expression involving only constants
+          or tensor nodes that were at the time is_scalar_constexpr. *)
+}
+[@@deriving sexp_of]
+
+val get_node : (Tnode.t, traced_array) Base.Hashtbl.t -> Tnode.t -> traced_array
+val optimize_integer_pow : bool ref
+
+type traced_store = (Tnode.t, traced_array) Base.Hashtbl.t [@@deriving sexp_of]
+
+type optimized = { traced_store : traced_store; llc : t; merge_node : Tnode.t option }
+[@@deriving sexp_of]
+
+val optimize_proc :
+  unoptim_ll_source:Stdlib.Format.formatter option ->
+  ll_source:Stdlib.Format.formatter option ->
+  name:string ->
+  Indexing.static_symbol list ->
+  t ->
+  optimized
+
+(** {2 Printing} *)
+
+val code_hum_margin : int ref
+
+val fprint_function_header :
+  ?name:string ->
+  ?static_indices:Indexing.static_symbol list ->
+  unit ->
+  Stdlib.Format.formatter ->
+  unit
+
+val get_ident_within_code : ?no_dots:bool -> t array -> Tnode.t -> string
+
+val fprint_hum :
+  ?name:string ->
+  ?static_indices:Indexing.static_symbol list ->
+  unit ->
+  Stdlib.Format.formatter ->
+  t ->
+  unit
