In progress: preparations for threefry, get rid of File_mapped

lukstafi · lukstafi · commit 04aebe83f689 · 2025-07-05T19:55:37.000+02:00
diff --git a/arrayjit/lib/assignments.ml b/arrayjit/lib/assignments.ml
@@ -362,9 +362,7 @@ let to_doc ?name ?static_indices () c =
     | Access (Low_level.Merge_buffer { source }) -> string (ident source ^ ".merge")
     | Access (Low_level.External_unsafe { ptr; prec; dims = _ }) ->
         string (Ops.ptr_to_string_hum ptr prec)
-    | Access (Low_level.File_mapped (file, file_prec)) ->
-        string ("file_mapped(\"" ^ file ^ "\", " ^ Ops.prec_string file_prec ^ ")")
-    | Access (Low_level.Uint4x32_to_prec_uniform { source; prec = target_prec }) ->
+    | Access (Low_level.Uint4x32_to_prec_uniform { source; target_prec; target_dims = _ }) ->
         string ("uint4x32_to_" ^ Ops.prec_string target_prec ^ "_uniform(" ^ ident source ^ ")")
     | Slice { batch_idx; sliced } ->
         string (ident sliced ^ " @| " ^ Indexing.symbol_ident batch_idx.static_symbol)
diff --git a/arrayjit/lib/c_syntax.ml b/arrayjit/lib/c_syntax.ml
@@ -617,23 +617,14 @@ module C_syntax (B : C_syntax_config) = struct
           ^^ offset_doc ^^ string "))" ^^ string postfix
         in
         (empty, expr)
-    | Access (Low_level.File_mapped (file, source_prec), Some idcs) ->
-        let prefix, postfix = B.convert_precision ~from:source_prec ~to_:prec in
-        let expr =
-          string prefix
-          ^^ string ("file_mapped_data_" ^ file ^ "[")
-          ^^ pp_array_offset (idcs, [||])
-          ^^ string "]" ^^ string postfix
-        in
-        (empty, expr)
-    | Access (Low_level.Uint4x32_to_prec_uniform { source; prec = source_prec }, Some idcs) ->
+    | Access (Low_level.Uint4x32_to_prec_uniform { source; target_prec; target_dims }, Some idcs) ->
         let tn = source in
-        let prefix, postfix = B.convert_precision ~from:source_prec ~to_:prec in
-        let offset_doc = pp_array_offset (idcs, Lazy.force tn.dims) in
+        let prefix, postfix = B.convert_precision ~from:target_prec ~to_:prec in
+        let offset_doc = pp_array_offset (idcs, Lazy.force target_dims) in
         let source_ident = string (get_ident tn) in
         let expr =
           string prefix
-          ^^ string ("uint4x32_to_" ^ Ops.prec_string source_prec ^ "_uniform(")
+          ^^ string ("uint4x32_to_" ^ Ops.prec_string target_prec ^ "_uniform(")
           ^^ source_ident ^^ brackets offset_doc ^^ string ")" ^^ string postfix
         in
         (empty, expr)
@@ -736,34 +727,20 @@ module C_syntax (B : C_syntax_config) = struct
           string prefix ^^ string ("external[%u]{=" ^ B.float_log_style ^ "}") ^^ string postfix
         in
         (expr_doc, [ `Accessor (idcs, dims_val); `Value access_doc ])
-    | Access (Low_level.File_mapped (file, source_prec), Some idcs) ->
-        let prefix, postfix = B.convert_precision ~from:source_prec ~to_:prec in
-        let access_doc =
-          string prefix
-          ^^ string ("file_mapped_data_" ^ file ^ "[")
-          ^^ pp_array_offset (idcs, [||])
-          ^^ string "]" ^^ string postfix
-        in
-        let expr_doc =
-          string prefix
-          ^^ string ("file_mapped_" ^ file ^ "[%u]{=" ^ B.float_log_style ^ "}")
-          ^^ string postfix
-        in
-        (expr_doc, [ `Accessor (idcs, [||]); `Value access_doc ])
-    | Access (Low_level.Uint4x32_to_prec_uniform { source; prec = source_prec }, Some idcs) ->
+    | Access (Low_level.Uint4x32_to_prec_uniform { source; target_prec; target_dims }, Some idcs) ->
         let tn = source in
-        let prefix, postfix = B.convert_precision ~from:source_prec ~to_:prec in
-        let dims = Lazy.force tn.dims in
+        let prefix, postfix = B.convert_precision ~from:target_prec ~to_:prec in
+        let dims = Lazy.force target_dims in
         let offset_doc = pp_array_offset (idcs, dims) in
         let source_ident = string (get_ident tn) in
         let access_doc =
           string prefix
-          ^^ string ("uint4x32_to_" ^ Ops.prec_string source_prec ^ "_uniform(")
+          ^^ string ("uint4x32_to_" ^ Ops.prec_string target_prec ^ "_uniform(")
           ^^ source_ident ^^ brackets offset_doc ^^ string ")" ^^ string postfix
         in
         let expr_doc =
           string prefix
-          ^^ string ("uint4x32_to_" ^ Ops.prec_string source_prec ^ "_uniform(")
+          ^^ string ("uint4x32_to_" ^ Ops.prec_string target_prec ^ "_uniform(")
           ^^ source_ident
           ^^ brackets (string "%u")
           ^^ string "){=" ^^ string B.float_log_style ^^ string "}" ^^ string postfix
diff --git a/arrayjit/lib/low_level.ml b/arrayjit/lib/low_level.ml
@@ -10,16 +10,12 @@ let _get_local_debug_runtime = Utils.get_local_debug_runtime
 
 type dedicated_access =
   | C_function of string
-  | External_unsafe of {
-      ptr : Ops.voidptr;
-      prec : (Ops.prec[@equal.ignore] [@compare.ignore]);
-      dims : int array Lazy.t;
-    }
+  | External_unsafe of { ptr : Ops.voidptr; prec : Ops.prec; dims : int array Lazy.t }
   | Merge_buffer of { source : Tnode.t }
-  | File_mapped of string * Ops.prec
   | Uint4x32_to_prec_uniform of {
       source : Tnode.t;
-      prec : (Ops.prec[@equal.ignore] [@compare.ignore]);
+      target_prec : Ops.prec;
+      target_dims : int array Lazy.t;
     }
 [@@deriving sexp_of, equal, compare]
 
@@ -100,9 +96,7 @@ let virtualize_settings =
   let max_tracing_dim =
     Int.of_string @@ Utils.get_global_arg ~arg_name:"virtualize_max_tracing_dim" ~default:"5"
   in
-  let enable_device_only =
-    Utils.get_global_flag ~default:true ~arg_name:"enable_device_only"
-  in
+  let enable_device_only = Utils.get_global_flag ~default:true ~arg_name:"enable_device_only" in
   let inline_scalar_constexprs =
     Utils.get_global_flag ~default:true ~arg_name:"inline_scalar_constexprs"
   in
@@ -1083,16 +1077,11 @@ let to_doc_cstyle ?name ?static_indices () llc =
     | Access (Merge_buffer { source }, None) -> doc_ident source ^^ string ".merge"
     | Access (Merge_buffer { source }, Some idcs) ->
         group (doc_ident source ^^ string ".merge" ^^ brackets (pp_indices idcs))
-    | Access (File_mapped (file, prec), None) ->
-        string ("file_mapped(\"" ^ file ^ "\", " ^ Ops.prec_string prec ^ ")")
-    | Access (File_mapped (file, prec), Some idcs) ->
-        string ("file_mapped(\"" ^ file ^ "\", " ^ Ops.prec_string prec ^ ")")
-        ^^ brackets (pp_indices idcs)
-    | Access (Uint4x32_to_prec_uniform { source; prec }, None) ->
-        string ("uint4x32_to_" ^ Ops.prec_string prec ^ "_uniform(")
+    | Access (Uint4x32_to_prec_uniform { source; target_prec; target_dims = _ }, None) ->
+        string ("uint4x32_to_" ^ Ops.prec_string target_prec ^ "_uniform(")
         ^^ doc_ident source ^^ string ")"
-    | Access (Uint4x32_to_prec_uniform { source; prec }, Some idcs) ->
-        string ("uint4x32_to_" ^ Ops.prec_string prec ^ "_uniform(")
+    | Access (Uint4x32_to_prec_uniform { source; target_prec; target_dims = _ }, Some idcs) ->
+        string ("uint4x32_to_" ^ Ops.prec_string target_prec ^ "_uniform(")
         ^^ doc_ident source ^^ string ")"
         ^^ brackets (pp_indices idcs)
     | Get (tn, idcs) -> group (doc_ident tn ^^ brackets (pp_indices idcs))
@@ -1170,16 +1159,11 @@ let to_doc ?name ?static_indices () llc =
     | Access (Merge_buffer { source }, None) -> doc_ident source ^^ string ".merge"
     | Access (Merge_buffer { source }, Some idcs) ->
         group (doc_ident source ^^ string ".merge" ^^ brackets (pp_indices idcs))
-    | Access (File_mapped (file, prec), None) ->
-        string ("file_mapped(\"" ^ file ^ "\", " ^ Ops.prec_string prec ^ ")")
-    | Access (File_mapped (file, prec), Some idcs) ->
-        string ("file_mapped(\"" ^ file ^ "\", " ^ Ops.prec_string prec ^ ")")
-        ^^ brackets (pp_indices idcs)
-    | Access (Uint4x32_to_prec_uniform { source; prec }, None) ->
-        string ("uint4x32_to_" ^ Ops.prec_string prec ^ "_uniform(")
+    | Access (Uint4x32_to_prec_uniform { source; target_prec; target_dims = _ }, None) ->
+        string ("uint4x32_to_" ^ Ops.prec_string target_prec ^ "_uniform(")
         ^^ doc_ident source ^^ string ")"
-    | Access (Uint4x32_to_prec_uniform { source; prec }, Some idcs) ->
-        string ("uint4x32_to_" ^ Ops.prec_string prec ^ "_uniform(")
+    | Access (Uint4x32_to_prec_uniform { source; target_prec; target_dims = _ }, Some idcs) ->
+        string ("uint4x32_to_" ^ Ops.prec_string target_prec ^ "_uniform(")
         ^^ doc_ident source ^^ string ")"
         ^^ brackets (pp_indices idcs)
     | Get (tn, idcs) -> group (doc_ident tn ^^ brackets (pp_indices idcs))
diff --git a/arrayjit/lib/low_level.mli b/arrayjit/lib/low_level.mli
@@ -7,24 +7,20 @@ open Base
 (** A dedicated access that might need to be implemented differently for each backend. *)
 type dedicated_access =
   | C_function of string  (** Calls a no-argument or indices-arguments C function. *)
-  | External_unsafe of {
-      ptr : Ops.voidptr;
-      prec : (Ops.prec[@equal.ignore] [@compare.ignore]);
-      dims : int array Lazy.t;
-    }
+  | External_unsafe of { ptr : Ops.voidptr; prec : Ops.prec; dims : int array Lazy.t }
   | Merge_buffer of { source : Tnode.t }
       (** Each device has at most one merge buffer, which is re-used, and re-allocated as needed, by
           merge operations. The merge buffer is associated with the source node of the device's most
           recent [device_to_device ~into_merge_buffer:true] operation. *)
-  | File_mapped of string * Ops.prec
-      (** Reads the data using [Unix.openfile] and [Unix.map_file]. *)
   | Uint4x32_to_prec_uniform of {
       source : Tnode.t;
-      prec : (Ops.prec[@equal.ignore] [@compare.ignore]);
+      target_prec : Ops.prec;
+      target_dims : int array Lazy.t;
     }
       (** Converts the given Uint4x32 to the given precision in a bit-efficient manner. For random
           bits, the result is uniform over the range of the precision for integer precisions, and
-          over the range [[0.0, 1.0)] for floating point precisions. *)
+          over the range \[0.0, 1.0) for floating point precisions. When used in an access pattern,
+          the indices are converted to a byte offset depending on the given precision. *)
 [@@deriving sexp_of, equal, compare]
 
 module Scope_id : sig
diff --git a/arrayjit/lib/ops.ml b/arrayjit/lib/ops.ml
@@ -343,7 +343,15 @@ let interpret_binop op v1 v2 =
   (* | Shr -> v1 / (int_pow 2. @@ to_int v2) *)
   | Or -> if v1 <> 0. || v2 <> 0. then 1. else 0.
   | And -> if v1 <> 0. && v2 <> 0. then 1. else 0.
-  | Threefry4x32 -> invalid_arg "interpret_binop: Threefry4x32 requires hardware implementation"
+  | Threefry4x32 ->
+      (* NOTE: the purpose of this code is to reflect the reference implementation that all backends
+         should implement. Due to precision constraints and the fact that threefry is inherently a
+         non-numerical, bit-level operation, this code reinterprets the arguments and the result as
+         the lower 64 bits of the 128-bit values; with the upper 64 bits being 0 for the arguments
+         and ignored for the result. This agrees with the Bigarray reinterpretation of the
+         [Uint4x32] precision as the [Complex.t] type with the real part exposed. *)
+      (* FIXME: NOT IMPLEMENTED YET *)
+      failwith "NOT IMPLEMENTED YET: Threefry4x32"
 
 let interpret_unop op v =
   let open Float in
@@ -461,7 +469,9 @@ let binop_c_syntax prec v =
   (* | Shr, _ -> ("((", ") / exp2(", "))") *)
   | Or, _ -> ("(", " ||", ")")
   | And, _ -> ("(", " &&", ")")
-  | Threefry4x32, _ -> ("threefry4x32(", ",", ")")
+  | Threefry4x32, _ ->
+      (* This corresponds to the pure C implementation in arrayjit_stubs.c. *)
+      ("arrayjit_threefry4x32(", ",", ")")
 
 let is_assign_op = function
   | Arg1 | Mod | Threefry4x32 (* | Shl | Shr *) | Cmplt | Cmpeq | Cmpne -> false
diff --git a/lib/shape.ml b/lib/shape.ml
@@ -442,18 +442,11 @@ let%debug4_sexp get_inequalities ({ shape = cur_sh; logic; id = _ } as _upd : up
   | Terminal (`Fetch (Access (C_function _))) -> (Row.dim_map_empty, mark_terminal ())
   | Terminal (`Fetch (Access (External_unsafe _))) -> (Row.dim_map_empty, mark_terminal ())
   | Terminal (`Fetch (Access (Merge_buffer _))) -> (Row.dim_map_empty, mark_terminal ())
-  | Terminal (`Fetch (Access (Uint4x32_to_prec_uniform _))) -> (Row.dim_map_empty, mark_terminal ())
-  | Terminal (`Fetch (Access (File_mapped (filename, prec)))) ->
-      let fd = Unix.openfile filename [ Unix.O_RDONLY ] 0o640 in
-      let len = Unix.lseek fd 0 Unix.SEEK_END / Ir.Ops.prec_in_bytes prec in
-      Unix.close fd;
-      ( Row.dim_map_empty,
-        Rows_constr
-          {
-            r = [ cur_sh.batch; cur_sh.output; cur_sh.input ];
-            constr = Total_elems { nominator = len; divided_by = dim_var_set_empty };
-          }
-        :: mark_terminal () )
+  | Terminal (`Fetch (Access (Uint4x32_to_prec_uniform _))) ->
+      (* FIXME: NOT IMPLEMENTED YET -- we need to propagate the precision-adjusted dimensions
+         between the source tensor and the target tensor. This is tricky because the dimensions
+         are not known at the time of the shape inference. *)
+      (Row.dim_map_empty, mark_terminal ())
   | Terminal (`Fetch (Slice { sliced = tn; batch_idx = _ })) ->
       if Lazy.is_val tn.dims then
         ( dim_map_empty,
diff --git a/lib/shape.mli b/lib/shape.mli
@@ -149,8 +149,7 @@ type logic =
           [s1], hence the name. *)
   | Broadcast_tern of ternary_type * t * t * t  (** Matches the shapes for a ternary operation. *)
   | Terminal of [ `Data of Ir.Assignments.init_data | `Fetch of Ir.Assignments.fetch_op ]
-      (** Extracts any available shape information from the initialization. E.g. for
-          [`Fetch (File_mapped fn)], opens the file [fn] to check its length. *)
+      (** Extracts any available shape information from the initialization. *)
 [@@deriving equal, sexp_of]
 
 type update_id [@@deriving equal, compare, hash, sexp]
diff --git a/lib/shape_inference.md b/lib/shape_inference.md
@@ -126,8 +126,7 @@ type logic =
       (** Permutes the axes of a shape. One case of [Transpose] is to swap inputs with outputs of [s1],
       hence the name. *)
   | Terminal of [ `Data of Ir.Assignments.init_data | `Fetch of Ir.Assignments.fetch_op ]
-      (** Extracts any available shape information from the initialization. E.g.
-      for [File_mapped fn], opens the file [fn] to check its length. *)
+      (** Extracts any available shape information from the initialization, e.g. the number of elements. *)
 ```
 
 ### Non-tensor-like constraints
@@ -208,7 +207,7 @@ There is an important and intentional difference between `dims` in the `arrayjit
 Other important functions in the `Shape` module.
 
 * `einsum_slot_spec_to_dims_bio ~generative` parses an einsum spec for a single shape, returns the three rows and a mapping from axis (`dim`) variables to indices where the einsum specifies fixed indexing. When `generative` is true for the kind of a row, when an axis has a fixed projection to dimension 0, the axis is not a variable added to the fixed indexing mapping, but is instead dimension-1 (solved). The "generative" rows are the ones with no initial user-provided shape information. This is just a heuristic to avoid surprises where a tensor axis with only dimension 0 populated gets inferred a bigger dimension size -- it might be revisited in the future.
-* `get_inequalities` builds row inequalities by pairing the rows of the current shape (as `cur`) with the rows of sub-shapes (as `subr`). It also derives a batch row constraint for terminals initialized with `Constant_fill values` and `File_mapped (filename, prec)` (where the file is scanned to get its length). For `Batch_slice` (the `@|` operation) it waits till the batch row variables (if any) are solved, and derives row equations (not inequalities) between the current shape and the sub-shape, with `cur_sh.batch.dims` expanded to account for the slicing / indexing. For einsum specs, it derives inequalities, roughly: _current shape ≥ lhs spec shape_, and _rhs spec shape ≥ sub-shape_.
+* `get_inequalities` builds row inequalities by pairing the rows of the current shape (as `cur`) with the rows of sub-shapes (as `subr`). It also derives a batch row constraint for terminals initialized with `Constant_fill values`. For `Batch_slice` (the `@|` operation) it waits till the batch row variables (if any) are solved, and derives row equations (not inequalities) between the current shape and the sub-shape, with `cur_sh.batch.dims` expanded to account for the slicing / indexing. For einsum specs, it derives inequalities, roughly: _current shape ≥ lhs spec shape_, and _rhs spec shape ≥ sub-shape_.
 * `propagate_shapes` gets and then solves the inequalities, using a global state for the environment. It udpates the shapes in-place with the partial solution. It is invoked twice for each `update_step`: first during the bottom-up process of building tensors, and then in reverse order from `finish_inference`.
 * `finish_inference` is called right before some projections or array dimensions are required (typically, because of jitting). It performs a second round of `propagate_shapes`, and then once again attempts to solve any remaining constraints that `propagate_shapes` didn't solve. Then it "closes the shapes": substitutes out remaining shape variables by their LUBs if any, or dimension-1 / `Broadcastable` (no-more-axes). Then it resets the environment state, since the shapes are now guaranteed to not have variables.
 * `derive_projections` starts by freshening the `proj_id`s in the `update_step`. Then it generates and solves shape inequalities, and then generates and solves projection equations, and constructs the `projections` record.
