Inline single getters i.e. "views"

lukstafi · lukstafi · commit 3b98f9bcbb2e · 2025-08-12T16:51:08.000+02:00
diff --git a/arrayjit/lib/low_level.ml b/arrayjit/lib/low_level.ml
@@ -119,6 +119,7 @@ type traced_array = {
   mutable read_before_write : bool;
   mutable read_only : bool;
   mutable is_scalar_constexpr : bool;
+  mutable is_accessing : bool;
   mutable is_complex : bool;
 }
 [@@deriving sexp_of]
@@ -149,6 +150,7 @@ let get_node store tn =
         read_before_write = false;
         read_only = false;
         is_scalar_constexpr = false;
+        is_accessing = false;
         is_complex = false;
       })
 
@@ -180,18 +182,19 @@ let is_constexpr_comp traced_store llsc =
   in
   loop llsc
 
-let is_complex_comp traced_store llsc =
+let is_accessing_comp traced_store llsc =
   let rec loop llsc =
     match llsc with
     | Get_local { tn; _ } | Local_scope { id = { tn; _ }; _ } ->
         let traced = get_node traced_store tn in
-        traced.is_complex
+        traced.is_accessing
     | Get (tn, _) ->
         let traced = get_node traced_store tn in
         not traced.is_scalar_constexpr
     | Get_merge_buffer (tn, _) ->
         let traced = get_node traced_store tn in
-        not traced.is_scalar_constexpr
+        traced.is_accessing <- true;
+        true
     | Ternop (_, v1, v2, v3) -> loop v1 || loop v2 || loop v3
     | Binop (_, v1, v2) -> loop v1 || loop v2
     | Unop (_, v) -> loop v
@@ -200,6 +203,20 @@ let is_complex_comp traced_store llsc =
   in
   loop llsc
 
+let is_complex_comp traced_store llsc =
+  let accessing = is_accessing_comp traced_store in
+  match llsc with
+  | Get_local { tn; _ } | Local_scope { id = { tn; _ }; _ } ->
+      let traced = get_node traced_store tn in
+      traced.is_complex
+  | Get _ -> false
+  | Get_merge_buffer _ -> false
+  | Ternop (_, v1, v2, v3) -> accessing v1 || accessing v2 || accessing v3
+  | Binop (_, v1, v2) -> accessing v1 || accessing v2
+  | Unop (_, v) -> accessing v
+  | Constant _ -> false
+  | Embed_index _ -> false
+
 let is_scalar_dims tn = Array.for_all ~f:(( = ) 1) @@ Lazy.force tn.Tn.dims
 
 let visit_llc traced_store ~merge_node_id reverse_node_map ~max_visits llc =
@@ -234,6 +251,7 @@ let visit_llc traced_store ~merge_node_id reverse_node_map ~max_visits llc =
         let traced : traced_array = get_node traced_store tn in
         if Hash_set.is_empty traced.assignments && Hashtbl.is_empty traced.accesses then (
           traced.zero_initialized <- true;
+          traced.is_accessing <- false;
           traced.is_complex <- false;
           if is_scalar_dims tn then traced.is_scalar_constexpr <- true);
         traced.zeroed_out <- true
@@ -246,8 +264,9 @@ let visit_llc traced_store ~merge_node_id reverse_node_map ~max_visits llc =
         then traced.is_scalar_constexpr <- is_constexpr_comp traced_store llsc
           (* Note: this prevents detection if the same constant is assigned inside a loop. *)
         else if not @@ Hash_set.is_empty traced.assignments then traced.is_scalar_constexpr <- false;
-        if first_visit then
-          traced.is_complex <- traced.is_complex || is_complex_comp traced_store llsc;
+        if first_visit then (
+          traced.is_accessing <- traced.is_accessing || is_accessing_comp traced_store llsc;
+          traced.is_complex <- traced.is_complex || is_complex_comp traced_store llsc);
         Hash_set.add traced.assignments (lookup env idcs);
         Array.iter idcs ~f:(function
           | Fixed_idx _ -> ()
@@ -265,7 +284,9 @@ let visit_llc traced_store ~merge_node_id reverse_node_map ~max_visits llc =
         let traced : traced_array = get_node traced_store tn in
         (* Vector operations cannot be scalar constexpr *)
         traced.is_scalar_constexpr <- false;
-        if first_visit then traced.is_complex <- false;
+        if first_visit then (
+          traced.is_accessing <- traced.is_accessing || is_accessing_comp traced_store arg;
+          traced.is_complex <- traced.is_complex || not (is_constexpr_comp traced_store arg));
         (* Mark all positions that will be written to *)
         for i = 0 to length - 1 do
           let pos_idcs = Array.copy idcs in
diff --git a/arrayjit/lib/low_level.mli b/arrayjit/lib/low_level.mli
@@ -90,9 +90,12 @@ type traced_array = {
       (** 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. *)
-  mutable is_complex : bool;
+  mutable is_accessing : bool;
       (** False only if the tensor node is built from index embeddings and scalar constant
           expressions. *)
+  mutable is_complex : bool;
+      (** True only if the tensor node is built acciessing computations that are not a single
+          getter. *)
 }
 [@@deriving sexp_of]
 
diff --git a/arrayjit/lib/lowering_and_inlining.md b/arrayjit/lib/lowering_and_inlining.md
@@ -117,7 +117,7 @@ The optimization behavior is controlled by `virtualize_settings`:
 - `max_tracing_dim`: Maximum dimension size for loop unrolling during analysis
 - `enable_device_only`: Whether to prefer device-only storage when possible
 - `inline_scalar_constexprs`: Whether to inline scalar constant expressions regardless of accesses
-- `inline_simple_computations`: Currently, whether to inline computations built from index embeddings and scalar constant expressions, regardless of accesses
+- `inline_simple_computations`: Currently, whether to inline computations built from either single getters, or index embeddings and scalar constant expressions, regardless of accesses
 
 ## Memory Mode Management
 
diff --git a/ocannl_config.example b/ocannl_config.example
@@ -101,7 +101,8 @@ enable_device_only=true
 # If true, scalar constant expressions will always be inlined.
 inline_scalar_constexprs=true
 
-# If true, if the tensor node is built from index embeddings and scalar constant expressions, regardless of accesses, it will be inlined.
+# If true, if the tensor node is built from either single getters, or index embeddings and
+# scalar constant expressions, regardless of accesses, it will be inlined.
 inline_simple_computations=true
 
 # The random number library. Options: `stdlib` -- `Base.Random`;
diff --git a/test/einsum/moons_demo_variant.expected b/test/einsum/moons_demo_variant.expected
@@ -22,8 +22,8 @@ n17 threefry4x32 as threefry4x32: Virt/15; uint4x32 prec 4; mem in bytes: <not-i
 n18 threefry4x32 as threefry4x32: Virt/15; uint4x32 prec 4; mem in bytes: <not-in-yet>
 n19 w2 as w2: Host&stream/412410; single prec 1x16; mem in bytes: <not-in-yet>
 n20 grad_w2 as w2.grad: Local/26046; single prec 1x16; mem in bytes: <not-in-yet>
-n21 @|_moons_input as moons_input: Local/1046; single prec 10x2; mem in bytes: <not-in-yet>
-n24 @|_moons_class as moons_class: Local/1046; single prec 10x1; mem in bytes: <not-in-yet>
+n21 @|_moons_input as moons_input: Virt/15; single prec 10x2; mem in bytes: <not-in-yet>
+n24 @|_moons_class as moons_class: Virt/15; single prec 10x1; mem in bytes: <not-in-yet>
 n27 * as n27: Local/1046; single prec 10x16; mem in bytes: <not-in-yet>
 n28 grad_* as n27.grad: Local/1046; single prec 10x16; mem in bytes: <not-in-yet>
 n29 + as n29: Local/1046; single prec 10x16; mem in bytes: <not-in-yet>
diff --git a/test/training/bigram.expected b/test/training/bigram.expected
diff --git a/test/training/bigram.ml b/test/training/bigram.ml
