Use liveness info to speed up allocation

CHANGELOG.md

@@ -20,6 +20,10 @@
   `-ec`, `-oec`, `-oecarray` and `-CT` command-line options are deprecated
   ([PR #914](https://github.com/jasmin-lang/jasmin/pull/914)).
 
+- The “allocation” pass now uses the liveness information to reduce the sizes
+  of the tables it uses internally; it should be faster on large functions
+  ([PR #965](https://github.com/jasmin-lang/jasmin/pull/965)).
+
 # Jasmin 2024.07.2 — Nancy, 2024-11-21
 
 ## New features

compiler/src/compile.ml

@@ -243,6 +243,45 @@ let compile (type reg regx xreg rflag cond asm_op extra_op)
     (get_annot fn).stack_zero_strategy
   in
 
+  (* This implements an analysis returning the set of variables becoming dead
+     after each instruction. It is based on the liveness analysis available
+     in Liveness. *)
+  let dead_vars_fd (f : _ func) =
+    let hvars = Hashtbl.create 97 in
+    let live = Liveness.live_fd false f in
+    let rec analyze (i : _ ginstr) =
+      begin match i.i_desc with
+      | Cif (_, c1, c2) -> List.iter analyze c1; List.iter analyze c2
+      | Cfor (_, _, c) -> List.iter analyze c
+      | Cwhile (_, c, _, c') -> List.iter analyze c; List.iter analyze c'
+      | _ -> ()
+      end;
+      let (in_set, out_set) = i.i_info in
+      let s = Conv.csv_of_sv (Sv.diff in_set out_set) in
+      Hashtbl.add hvars i.i_loc s
+    in
+    List.iter analyze live.f_body;
+
+    fun ii ->
+      let loc, _ = ii in
+      try Hashtbl.find hvars loc with
+      | Not_found ->
+          hierror ~loc:(Lmore loc) ~kind:"compilation error" ~internal:true
+            "dead_vars_fd: location not found"
+  in
+
+  (* We expose a version of dead_vars_fd for _ufun_decl. *)
+  let dead_vars_ufd (f : _ Expr._ufun_decl) =
+    let f = Conv.fdef_of_cufdef f in
+    dead_vars_fd f
+  in
+
+  (* We expose a version of dead_vars_fd for _sfun_decl. *)
+  let dead_vars_sfd (f : _ Expr._sfun_decl) =
+    let _, f = Conv.fdef_of_csfdef f in
+    dead_vars_fd f
+  in
+
   let cparams =
     {
       Compiler.rename_fd;
@@ -279,6 +318,8 @@ let compile (type reg regx xreg rflag cond asm_op extra_op)
       Compiler.fresh_var_ident = Conv.fresh_var_ident;
       Compiler.slh_info;
       Compiler.stack_zero_info = szs_of_fn;
+      Compiler.dead_vars_ufd;
+      Compiler.dead_vars_sfd;
     }
   in
 

compiler/src/conv.ml

@@ -55,6 +55,14 @@ let var_of_cvar cv =
   assert (cty_of_ty v.v_ty = cv.Var.vtype);
   v
 
+let csv_of_sv s =
+  let open Var0.SvExtra.Sv in
+  Sv.fold (fun x cs -> add (Obj.magic (cvar_of_var x)) cs) s empty
+
+let sv_of_csv cs =
+  let open Var0.SvExtra.Sv in
+  fold (fun x s -> Sv.add (var_of_cvar (Obj.magic x)) s) cs Sv.empty
+
 (* ------------------------------------------------------------------------ *)
 
 let cvari_of_vari v =

compiler/src/conv.mli

@@ -37,6 +37,9 @@ val cvar_of_var :  var -> Var0.Var.var
 val var_of_cvar :  Var0.Var.var -> var
 val vari_of_cvari :  Expr.var_i -> var L.located
 
+val csv_of_sv : Sv.t -> Var0.SvExtra.Sv.t
+val sv_of_csv : Var0.SvExtra.Sv.t -> Sv.t
+
 val lval_of_clval :  Expr.lval -> Prog.lval
 
 val cexpr_of_expr :  expr -> Expr.pexpr

compiler/src/liveness.ml

@@ -50,7 +50,14 @@ and live_d weak d (s_o: Sv.t) =
     let s2, c2 = live_c weak c2 s_o in
     Sv.union (vars_e e) (Sv.union s1 s2), s_o, Cif(e, c1, c2)
 
-  | Cfor _ -> assert false (* Should have been removed before *)
+  | Cfor (x, (_dir, e1, e2 as r), c) ->
+    let rec loop s_o =
+      let s_i, c = live_c weak c s_o in
+      let s_i = Sv.remove (L.unloc x) s_i in
+      if Sv.subset s_i s_o then s_o, c
+      else loop (Sv.union s_i s_o) in
+    let s_i, c = loop s_o in
+    Sv.union (vars_es [ e1; e2 ]) s_i, s_o, Cfor (x, r, c)
 
   | Cwhile(a,c,e,c') ->
     let ve = (vars_e e) in

proofs/compiler/allocation.v

@@ -391,6 +391,14 @@ Module M.
     by have := @mergeP_mset r1 r2;SvD.fsetdec.
   Qed.
 
+  Lemma remove_incl r x : incl (remove r x) r.
+  Proof.
+    apply /inclP; split=> //= y id _.
+    rewrite removeP.
+    case: get => // id'.
+    by case: eq_op.
+  Qed.
+
   End WSW.
   Arguments add {wsw} m x id h.
   Arguments set {wsw} m x id h.
@@ -543,73 +551,94 @@ Context
   {asm_op syscall_state : Type}
   {asmop:asmOp asm_op}.
 
+Section FUNCTION.
+
+(* This is liveness info in the target program. It returns the set of variables
+   that become dead at the instruction identified by the instr_info. We
+   remove them from the map, since they will not appear in the remainder
+   of the program. This helps keeping the map small. *)
+Context (dead_vars : instr_info -> Sv.t).
+
+(* TODO: further optimizations
+   - It seems that merge_aux/merge are costly. Can they be optimized? Maybe
+     we could make a safe over-approximation (e.g. take the union for the
+     reverse map [mid]).
+   - At some point, we are in quasi-SSA form. Variables are mostly defined once
+     (except in conditions and loops?). This can probably be used to reduce
+     the amount of information we need to remember. *)
+
 Fixpoint check_i (i1 i2:instr_r) r :=
-  match i1, i2 with
-  | Cassgn x1 _ ty1 e1, Cassgn x2 _ ty2 e2 =>
-    Let _ := assert (ty1 == ty2) (alloc_error "bad type in assignment") in
-    check_e e1 e2 r >>= check_lval (Some (ty2,e2)) x1 x2
-
-  | Copn xs1 _ o1 es1, Copn xs2 _ o2 es2 =>
-    Let _ := assert (o1 == o2) (alloc_error "operators not equals") in
-    check_es es1 es2 r >>= check_lvals xs1 xs2
-
-  | Csyscall xs1 o1 es1, Csyscall xs2 o2 es2 =>
-    Let _ := assert (o1 == o2) (alloc_error "syscall not equals") in
-    check_es es1 es2 r >>= check_lvals xs1 xs2
-
-  | Ccall x1 f1 arg1, Ccall x2 f2 arg2 =>
-    Let _ := assert (f1 == f2) (alloc_error "functions not equals") in
-    check_es arg1 arg2 r >>= check_lvals x1 x2
-
-  | Cif e1 c11 c12, Cif e2 c21 c22 =>
-    Let re := check_e e1 e2 r in
-    Let r1 := fold2 E.fold2 check_I c11 c21 re in
-    Let r2 := fold2 E.fold2 check_I c12 c22 re in
-    ok (M.merge r1 r2)
-
-  | Cfor x1 (d1,lo1,hi1) c1, Cfor x2 (d2,lo2,hi2) c2 =>
-    Let _ := assert (d1 == d2) (alloc_error "loop directions not equals") in
-    Let rhi := check_e lo1 lo2 r >>=check_e hi1 hi2 in
-    let check_c r :=
-      check_var x1 x2 r >>=
-      fold2 E.fold2 check_I c1 c2 in
-    loop check_c Loop.nb rhi
-
-  | Cwhile a1 c1 e1 c1', Cwhile a2 c2 e2 c2' =>
-    let check_c r :=
-      Let r := fold2 E.fold2 check_I c1 c2 r in
-      Let re := check_e e1 e2 r in
-      Let r' := fold2 E.fold2 check_I c1' c2' re in
-      ok (re, r') in
-    Let r := loop2 check_c Loop.nb r in
-    ok r
+    match i1, i2 with
+    | Cassgn x1 _ ty1 e1, Cassgn x2 _ ty2 e2 =>
+      Let _ := assert (ty1 == ty2) (alloc_error "bad type in assignment") in
+      check_e e1 e2 r >>= check_lval (Some (ty2,e2)) x1 x2
+
+    | Copn xs1 _ o1 es1, Copn xs2 _ o2 es2 =>
+      Let _ := assert (o1 == o2) (alloc_error "operators not equals") in
+      check_es es1 es2 r >>= check_lvals xs1 xs2
 
-  | _, _ => Error (alloc_error "instructions not equals")
-  end
+    | Csyscall xs1 o1 es1, Csyscall xs2 o2 es2 =>
+      Let _ := assert (o1 == o2) (alloc_error "syscall not equals") in
+      check_es es1 es2 r >>= check_lvals xs1 xs2
+
+    | Ccall x1 f1 arg1, Ccall x2 f2 arg2 =>
+      Let _ := assert (f1 == f2) (alloc_error "functions not equals") in
+      check_es arg1 arg2 r >>= check_lvals x1 x2
+
+    | Cif e1 c11 c12, Cif e2 c21 c22 =>
+      Let re := check_e e1 e2 r in
+      Let r1 := fold2 E.fold2 check_I c11 c21 re in
+      Let r2 := fold2 E.fold2 check_I c12 c22 re in
+      ok (M.merge r1 r2)
+
+    | Cfor x1 (d1,lo1,hi1) c1, Cfor x2 (d2,lo2,hi2) c2 =>
+      Let _ := assert (d1 == d2) (alloc_error "loop directions not equals") in
+      Let rhi := check_e lo1 lo2 r >>=check_e hi1 hi2 in
+      let check_c r :=
+        check_var x1 x2 r >>=
+        fold2 E.fold2 check_I c1 c2 in
+      loop check_c Loop.nb rhi
+
+    | Cwhile a1 c1 e1 c1', Cwhile a2 c2 e2 c2' =>
+      let check_c r :=
+        Let r := fold2 E.fold2 check_I c1 c2 r in
+        Let re := check_e e1 e2 r in
+        Let r' := fold2 E.fold2 check_I c1' c2' re in
+        ok (re, r') in
+      Let r := loop2 check_c Loop.nb r in
+      ok r
+
+    | _, _ => Error (alloc_error "instructions not equals")
+    end
 
 with check_I i1 i2 r :=
   match i1, i2 with
-  | MkI _ i1, MkI ii i2 => add_iinfo ii (check_i i1 i2 r)
+  | MkI _ i1, MkI ii i2 =>
+    Let m := add_iinfo ii (check_i i1 i2 r) in
+    ok (Sv.fold (fun x acc => M.remove acc x) (dead_vars ii) m)
   end.
 
 Definition check_cmd := fold2 E.fold2 check_I.
 
+End FUNCTION.
+
 Section PROG.
 
 Context {pT: progT}.
 
 Variable (init_alloc : extra_fun_t -> extra_prog_t -> extra_prog_t -> cexec M.t).
 Variable (check_f_extra: M.t → extra_fun_t → extra_fun_t → seq var_i → seq var_i → cexec M.t).
+Variable (dead_vars_fd : fun_decl -> instr_info -> Sv.t).
 
 Definition check_fundef (ep1 ep2 : extra_prog_t) (f1 f2: funname * fundef) (_:Datatypes.unit) :=
-  let (f1,fd1) := f1 in
-  let (f2,fd2) := f2 in
-  add_funname f1 (add_finfo fd1.(f_info) (
-    Let _ := assert [&& f1 == f2, fd1.(f_tyin) == fd2.(f_tyin) & fd1.(f_tyout) == fd2.(f_tyout) ]
+  let (fn1,fd1) := f1 in
+  let (fn2,fd2) := f2 in
+  add_funname fn1 (add_finfo fd1.(f_info) (
+    Let _ := assert [&& fn1 == fn2, fd1.(f_tyin) == fd2.(f_tyin) & fd1.(f_tyout) == fd2.(f_tyout) ]
                         (E.error "functions not equal") in
     Let r := init_alloc fd1.(f_extra) ep1 ep2 in
     Let r := check_f_extra r fd1.(f_extra) fd2.(f_extra) fd1.(f_params) fd2.(f_params) in
-    Let r := check_cmd fd1.(f_body) fd2.(f_body) r in
+    Let r := check_cmd (dead_vars_fd f2) fd1.(f_body) fd2.(f_body) r in
     let es1 := map Plvar fd1.(f_res) in
     let es2 := map Plvar fd2.(f_res) in
     Let _r := check_es es1 es2 r in