Put more values into global variables

CHANGES.md

@@ -10,6 +10,7 @@
 * Compiler: improved shape computation (#2198)
 * Add the --build-config and --apply-build-config flags (#2177)
 * Runtime/wasm: optimized some bigstring primitives (#2144)
+* Put more values into global variables (#2211)
 
 ## Bug fixes
 * Compiler: fix reference unboxing (#2210)

compiler/lib-wasm/code_generation.ml

@@ -433,6 +433,72 @@ let is_small_constant e =
   | W.GlobalGet name -> global_is_constant name
   | _ -> return false
 
+let rec check_is_constant st (e : W.expression) =
+  match e with
+  | Const _ | RefFunc _ | RefNull _ -> true
+  | RefI31 e -> check_is_constant st e
+  | GlobalGet x -> (
+      match Var.Map.find_opt x st.context.constant_globals with
+      | Some { init = Some _; _ } -> true
+      | Some { init = None; _ } | None -> false)
+  | ArrayNewFixed (_, args) | StructNew (_, args) ->
+      List.for_all ~f:(fun x -> check_is_constant st x) args
+  | RefCast (_, e) | ExternConvertAny e | AnyConvertExtern e -> check_is_constant st e
+  | UnOp _
+  | BinOp _
+  | I32WrapI64 _
+  | I64ExtendI32 _
+  | F32DemoteF64 _
+  | F64PromoteF32 _
+  | LocalGet _
+  | LocalTee _
+  | BlockExpr _
+  | Call _
+  | Seq _
+  | Pop _
+  | Call_ref _
+  | I31Get _
+  | ArrayNew _
+  | ArrayNewData _
+  | ArrayGet _
+  | ArrayLen _
+  | StructGet _
+  | RefTest _
+  | RefEq _
+  | Br_on_cast _
+  | Br_on_cast_fail _
+  | Br_on_null _
+  | IfExpr _
+  | Try _ -> false
+
+let is_constant_expression e st = check_is_constant st e, st
+
+let partial_const_init x e st =
+  match e with
+  | W.ArrayNewFixed (ty, args) ->
+      let placeholder =
+        match (Var.Hashtbl.find st.context.types ty).typ with
+        | Array { typ = Value (Ref { typ = Eq; _ }); _ } -> W.RefI31 (W.Const (I32 0l))
+        | Array { typ = Value F64; _ } -> W.Const (F64 0.)
+        | _ -> assert false
+      in
+      let init_args, patches, _ =
+        List.fold_left
+          ~f:(fun (init_args, patches, i) a ->
+            if check_is_constant st a
+            then a :: init_args, patches, i + 1
+            else
+              ( placeholder :: init_args
+              , W.ArraySet (ty, W.GlobalGet x, W.Const (I32 (Int32.of_int i)), a)
+                :: patches
+              , i + 1 ))
+          ~init:([], [], 0)
+          args
+      in
+      assert (not (List.is_empty patches));
+      Some (W.ArrayNewFixed (ty, List.rev init_args), List.rev patches), st
+  | _ -> None, st
+
 let load x =
   let* x = var x in
   match x with
@@ -576,34 +642,47 @@ let rec store ?(always = false) ?typ x e =
         let* b = should_make_global x in
         if b
         then
-          let* () =
-            let* b = global_is_registered x in
-            if b
-            then return ()
+          let* b = global_is_registered x in
+          if b
+          then instr (GlobalSet (x, e))
+          else
+            let* typ =
+              match typ with
+              | Some typ -> return typ
+              | None -> (
+                  if always
+                  then value_type
+                  else
+                    let* typ = expression_type e in
+                    match typ with
+                    | None -> value_type
+                    | Some typ -> return typ)
+            in
+            let* is_const = if always then return false else is_constant_expression e in
+            if is_const
+            then
+              let* () = register_constant x (W.GlobalGet x) in
+              register_global x { mut = false; typ } e
             else
-              let* typ =
-                match typ with
-                | Some typ -> return typ
-                | None -> (
-                    if always
-                    then value_type
-                    else
-                      let* typ = expression_type e in
-                      match typ with
-                      | None -> value_type
-                      | Some typ -> return typ)
-              in
-              let* default, typ', cast = default_value typ in
-              let* () =
-                register_constant
-                  x
-                  (match cast with
-                  | Some typ -> W.RefCast (typ, W.GlobalGet x)
-                  | None -> W.GlobalGet x)
-              in
-              register_global ~constant:true x { mut = true; typ = typ' } default
-          in
-          instr (GlobalSet (x, e))
+              let* partial = if always then return None else partial_const_init x e in
+              match partial with
+              | Some (init_expr, patches) ->
+                  let* () = register_constant x (W.GlobalGet x) in
+                  let* () = register_global x { mut = false; typ } init_expr in
+                  instrs patches
+              | None ->
+                  let* default, typ', cast = default_value typ in
+                  let* () =
+                    register_constant
+                      x
+                      (match cast with
+                      | Some typ -> W.RefCast (typ, W.GlobalGet x)
+                      | None -> W.GlobalGet x)
+                  in
+                  let* () =
+                    register_global ~constant:true x { mut = true; typ = typ' } default
+                  in
+                  instr (GlobalSet (x, e))
         else
           let* typ =
             match typ with

compiler/lib-wasm/globalize.ml

@@ -37,6 +37,7 @@ type st =
   ; visited_variables : int Code.Var.Map.t
   ; globals : Code.Var.Set.t
   ; closures : Closure_conversion.closure Code.Var.Map.t
+  ; constants : Code.Var.Set.t
   }
 
 let threshold = 1000
@@ -80,7 +81,8 @@ let traverse_expression x e st =
         ~f:(fun st x -> use x st)
         ~init:st
         (Code.Var.Map.find x st.closures).Closure_conversion.free_variables
-  | Constant _ | Special _ -> st
+  | Constant _ -> { st with constants = Code.Var.Set.add x st.constants }
+  | Special _ -> st
   | Prim (_, args) ->
       List.fold_left
         ~f:(fun st a ->
@@ -104,6 +106,35 @@ let traverse_block p st pc =
   let st = List.fold_left ~f:(fun st x -> declare x st) ~init:st b.Code.params in
   List.fold_left ~f:(fun st i -> traverse_instruction st i) ~init:st b.Code.body
 
+let available x st = Code.Var.Set.mem x st.globals || Code.Var.Set.mem x st.constants
+
+let propagate_instruction st i =
+  match i with
+  | Code.Let (x, Block (_, a, _, _)) when not (Code.Var.Set.mem x st.globals) ->
+      (* Globalize a block when most of its fields are available
+         (global or constant). Available fields go into the global's
+         initializer; the rest are patched via [array.set] in the
+         function body. The [+1] keeps 2-field cons cells eligible
+         (one non-available field is allowed), which matters for
+         cascading: globalizing an inner block makes its variable
+         available for outer blocks. *)
+      let non_available =
+        Array.fold_right ~f:(fun v n -> if available v st then n else n + 1) a ~init:0
+      in
+      if 3 * non_available <= Array.length a + 1 then globalize st x else st
+  | Code.Let (x, Closure _) when not (Code.Var.Set.mem x st.globals) -> (
+      match Code.Var.Map.find x st.closures with
+      | { free_variables; _ } ->
+          if List.for_all ~f:(fun v -> available v st) free_variables
+          then globalize st x
+          else st
+      | exception Not_found -> st)
+  | _ -> st
+
+let propagate_block p st pc =
+  let b = Code.Addr.Map.find pc p.Code.blocks in
+  List.fold_left ~f:(fun st i -> propagate_instruction st i) ~init:st b.Code.body
+
 let f p g closures =
   let l = Structure.blocks_in_reverse_post_order g in
   let in_loop = Freevars.find_loops_in_closure p p.Code.start in
@@ -116,7 +147,15 @@ let f p g closures =
         ; visited_variables = Code.Var.Map.empty
         ; globals = Code.Var.Set.empty
         ; closures
+        ; constants = Code.Var.Set.empty
         }
       l
   in
+  let st =
+    List.fold_left
+      ~f:(fun st pc ->
+        if Code.Addr.Map.mem pc in_loop then st else propagate_block p st pc)
+      ~init:st
+      l
+  in
   st.globals