Rename float_t -> scalar_t

lukstafi · lukstafi · commit 20809e9f95c6 · 2025-07-26T08:17:14.000+02:00
diff --git a/arrayjit/lib/c_syntax.ml b/arrayjit/lib/c_syntax.ml
@@ -576,7 +576,7 @@ module C_syntax (B : C_syntax_config) = struct
         in
         if PPrint.is_empty local_defs then assignment else local_defs ^^ hardline ^^ assignment
 
-  and pp_float (prec : Ops.prec) (vcomp : Low_level.float_t) : PPrint.document * PPrint.document =
+  and pp_float (prec : Ops.prec) (vcomp : Low_level.scalar_t) : PPrint.document * PPrint.document =
     (* Returns (local definitions, value expression) *)
     let open PPrint in
     match vcomp with
@@ -658,7 +658,7 @@ module C_syntax (B : C_syntax_config) = struct
         let expr = group (B.unop_syntax prec op expr_v) in
         (defs, expr)
 
-  and debug_float (prec : Ops.prec) (value : Low_level.float_t) :
+  and debug_float (prec : Ops.prec) (value : Low_level.scalar_t) :
       PPrint.document
       * [ `Accessor of Indexing.axis_index array * int array | `Value of PPrint.document ] list =
     (* Returns (value expression doc, list of arguments for printf) *)
diff --git a/arrayjit/lib/low_level.ml b/arrayjit/lib/low_level.ml
@@ -35,18 +35,18 @@ type t =
   | Seq of t * t
   | For_loop of { index : Indexing.symbol; from_ : int; to_ : int; body : t; trace_it : bool }
   | Zero_out of Tn.t
-  | Set of { tn : Tn.t; idcs : Indexing.axis_index array; llv : float_t; mutable debug : string }
-  | Set_local of scope_id * float_t
+  | Set of { tn : Tn.t; idcs : Indexing.axis_index array; llv : scalar_t; mutable debug : string }
+  | Set_local of scope_id * scalar_t
 [@@deriving sexp_of, equal]
 
-and float_t =
+and scalar_t =
   | Local_scope of { id : scope_id; body : t; orig_indices : Indexing.axis_index array }
   | Get_local of scope_id
   | Get of Tn.t * Indexing.axis_index array
   | Get_merge_buffer of Tn.t * Indexing.axis_index array
-  | Ternop of Ops.ternop * float_t * float_t * float_t
-  | Binop of Ops.binop * float_t * float_t
-  | Unop of Ops.unop * float_t
+  | Ternop of Ops.ternop * scalar_t * scalar_t * scalar_t
+  | Binop of Ops.binop * scalar_t * scalar_t
+  | Unop of Ops.unop * scalar_t
   | Constant of float
   | Embed_index of Indexing.axis_index
 [@@deriving sexp_of, equal, compare]
@@ -530,7 +530,7 @@ let inline_computation ~id computations_table traced static_indices call_args =
       | Set_local (id, llv) -> Some (Set_local (id, loop_float env llv))
       | Comment _ -> Some llc
       | Staged_compilation _ -> Some llc
-    and loop_float env llv : float_t =
+    and loop_float env llv : scalar_t =
       match llv with
       | Constant _ -> llv
       | Get (tn, indices) when Tn.equal tn traced.tn ->
@@ -564,7 +564,7 @@ let inline_computation ~id computations_table traced static_indices call_args =
 
 let optimize_integer_pow = ref true
 
-let rec unroll_pow ~(base : float_t) ~(exp : int) : float_t =
+let rec unroll_pow ~(base : scalar_t) ~(exp : int) : scalar_t =
   if exp < 0 then unroll_pow ~base:(Binop (Div, Constant 1., base)) ~exp:(Int.neg exp)
   else if exp = 0 then Constant 1.
   else Fn.apply_n_times ~n:(exp - 1) (fun accu -> Binop (Mul, base, accu)) base
@@ -603,7 +603,7 @@ let virtual_llc computations_table traced_store reverse_node_map static_indices
     | Set_local (id, llv) -> Set_local (id, loop_float ~process_for llv)
     | Comment _ -> llc
     | Staged_compilation _ -> llc
-  and loop_float ~process_for (llv : float_t) : float_t =
+  and loop_float ~process_for (llv : scalar_t) : scalar_t =
     match llv with
     | Constant _ -> llv
     | Get (tn, _) when Set.mem process_for tn ->
@@ -680,7 +680,7 @@ let cleanup_virtual_llc reverse_node_map ~static_indices (llc : t) : t =
         Some (Set_local (id, loop_float ~balanced ~env_dom llv))
     | Comment _ -> Some llc
     | Staged_compilation _ -> Some llc
-  and loop_float ~balanced ~env_dom (llv : float_t) : float_t =
+  and loop_float ~balanced ~env_dom (llv : scalar_t) : scalar_t =
     let loop = loop_float ~balanced ~env_dom in
     match llv with
     | Constant _ -> llv
@@ -770,7 +770,7 @@ let simplify_llc llc =
     | Set_local (id, llv) -> Set_local (id, loop_float llv)
     | Comment _ -> llc
     | Staged_compilation _ -> llc
-  and loop_float (llv : float_t) : float_t =
+  and loop_float (llv : scalar_t) : scalar_t =
     let local_scope_body, llv' =
       match llv with
       | Local_scope opts ->
@@ -838,9 +838,9 @@ let simplify_llc llc =
     | Binop (Div, Binop (Div, llv1, llv2), llv3) ->
         loop_float @@ Binop (Div, llv1, Binop (Mul, llv2, llv3))
     | Binop (ToPowOf, llv1, llv2) -> (
-        let v1 : float_t = loop_float llv1 in
-        let v2 : float_t = loop_float llv2 in
-        let result : float_t = Binop (ToPowOf, v1, v2) in
+        let v1 : scalar_t = loop_float llv1 in
+        let v2 : scalar_t = loop_float llv2 in
+        let result : scalar_t = Binop (ToPowOf, v1, v2) in
         if not !optimize_integer_pow then result
         else
           match v2 with
diff --git a/arrayjit/lib/low_level.mli b/arrayjit/lib/low_level.mli
@@ -16,31 +16,32 @@ type scope_id = Scope_id.t = { tn : Tnode.t; scope_id : int }
 
 (** {2 Low-level representation} *)
 
-(** Cases: [t] -- code, [float_t] -- single number at some precision. *)
+(** Cases: [t] -- code, [scalar_t] -- single number at some precision. *)
 type t =
   | Noop
   | Comment of string
   | Staged_compilation of (unit -> PPrint.document)
   | 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
+  | Set of { tn : Tnode.t; idcs : Indexing.axis_index array; llv : scalar_t; mutable debug : string }
+  | Set_local of scope_id * scalar_t
 [@@deriving sexp_of, equal]
 
-and float_t =
+and scalar_t =
   | Local_scope of { id : scope_id; body : t; orig_indices : Indexing.axis_index array }
   | Get_local of scope_id
   | Get of Tnode.t * Indexing.axis_index array
   | Get_merge_buffer of Tnode.t * Indexing.axis_index array
-  | Ternop of Ops.ternop * float_t * float_t * float_t
-  | Binop of Ops.binop * float_t * float_t
-  | Unop of Ops.unop * float_t
+  | Ternop of Ops.ternop * scalar_t * scalar_t * scalar_t
+  | Binop of Ops.binop * scalar_t * scalar_t
+  | Unop of Ops.unop * scalar_t
   | Constant of float
   | Embed_index of Indexing.axis_index
+  | Vec of int * scalar_t
 [@@deriving sexp_of, equal, compare]
 
-val apply_op : Ops.op -> float_t array -> float_t
+val apply_op : Ops.op -> scalar_t array -> scalar_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
diff --git a/arrayjit/lib/lowering_and_inlining.md b/arrayjit/lib/lowering_and_inlining.md
@@ -14,22 +14,22 @@ type t =
   | 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 { array : Tnode.t; idcs : Indexing.axis_index array; llv : float_t; mutable debug : string }
-  | Set_local of scope_id * float_t
+  | Set of { array : Tnode.t; idcs : Indexing.axis_index array; llv : scalar_t; mutable debug : string }
+  | Set_local of scope_id * scalar_t
 
-and float_t =
+and scalar_t =
   | Local_scope of { id : scope_id; body : t; orig_indices : Indexing.axis_index array }
   | Get_local of scope_id
   | Access of dedicated_access * Indexing.axis_index array option
   | Get of Tnode.t * Indexing.axis_index array
-  | Ternop of Ops.ternop * float_t * float_t * float_t
-  | Binop of Ops.binop * float_t * float_t
-  | Unop of Ops.unop * float_t
+  | Ternop of Ops.ternop * scalar_t * scalar_t * scalar_t
+  | Binop of Ops.binop * scalar_t * scalar_t
+  | Unop of Ops.unop * scalar_t
   | Constant of float
   | Embed_index of Indexing.axis_index
 ```
 
-`t` represents code/statements while `float_t` represents scalar expressions. The `trace_it` flag in `For_loop` indicates whether the loop should be traced for optimization (its initial segment will be unrolled for analysis).
+`t` represents code/statements while `scalar_t` represents scalar expressions. The `trace_it` flag in `For_loop` indicates whether the loop should be traced for optimization (its initial segment will be unrolled for analysis).
 
 ## Translation from Assignments
 
