Jib: Optimize set_slice for ARM v8.5

3 files changed, 70 insertions, 43 deletions

diff --git a/src/jib/c_backend.ml b/src/jib/c_backend.ml
index 94425384..cb02d17d 100644
--- a/src/jib/c_backend.ml
+++ b/src/jib/c_backend.ml
@@ -528,22 +528,6 @@ let analyze_primop' ctx id args typ =
      | _ -> no_change
      end
 
-  | "sail_unsigned", [AV_C_fragment (frag, vtyp, _)] ->
-     begin match destruct_vector ctx.tc_env vtyp with
-     | Some (Nexp_aux (Nexp_constant n, _), _, _)
-          when Big_int.less_equal n (Big_int.of_int 63) && is_stack_typ ctx typ ->
-        AE_val (AV_C_fragment (F_call ("fast_unsigned", [frag]), typ, ctyp_of_typ ctx typ))
-     | _ -> no_change
-     end
-
-  | "sail_signed", [AV_C_fragment (frag, vtyp, _)] ->
-     begin match destruct_vector ctx.tc_env vtyp with
-     | Some (Nexp_aux (Nexp_constant n, _), _, _)
-          when Big_int.less_equal n (Big_int.of_int 64) && is_stack_typ ctx typ ->
-        AE_val (AV_C_fragment (F_call ("fast_signed", [frag; v_int (Big_int.to_int n)]), typ, ctyp_of_typ ctx typ))
-     | _ -> no_change
-     end
-
   | "neg_int", [AV_C_fragment (frag, _, _)] ->
      AE_val (AV_C_fragment (F_op (v_int 0, "-", frag), typ, CT_fint 64))
 
@@ -1100,6 +1084,22 @@ let rec sgen_ctyp_name = function
   | CT_ref ctyp -> "ref_" ^ sgen_ctyp_name ctyp
   | CT_poly -> "POLY" (* c_error "Tried to generate code for non-monomorphic type" *)
 
+let sgen_mask n =
+  if n = 0 then
+    "UINT64_C(0)"
+  else if n <= 64 then
+    let chars_F = String.make (n / 4) 'F' in
+    let first = match (n mod 4) with
+      | 0 -> ""
+      | 1 -> "1"
+      | 2 -> "3"
+      | 3 -> "7"
+      | _ -> assert false
+    in
+    "UINT64_C(0x" ^ first ^ chars_F ^ ")"
+  else
+    failwith "Tried to create a mask literal for a vector greater than 64 bits."
+
 let rec sgen_cval = function
   | V_id (id, ctyp) -> string_of_name id
   | V_ref (id, _) -> "&" ^ string_of_name id
@@ -1129,6 +1129,8 @@ let rec sgen_cval = function
 and sgen_call op cvals =
   let open Printf in
   match op, cvals with
+  | Bit_to_bool, [v] ->
+     sprintf "((bool) %s)" (sgen_cval v)
   | Bnot, [v] -> "!(" ^ sgen_cval v ^ ")"
   | List_hd, [v] ->
      sprintf "(%s).hd" ("*" ^ sgen_cval v)
@@ -1160,6 +1162,14 @@ and sgen_call op cvals =
      sprintf "(%s + %s)" (sgen_cval v1) (sgen_cval v2)
   | Isub, [v1; v2] ->
      sprintf "(%s - %s)" (sgen_cval v1) (sgen_cval v2)
+  | Unsigned 64, [vec] ->
+     sprintf "((mach_int) %s)" (sgen_cval vec)
+  | Signed 64, [vec] ->
+     begin match cval_ctyp vec with
+     | CT_fbits (n, _) ->
+        sprintf "fast_signed(%s, %d)" (sgen_cval vec) n
+     | _ -> assert false
+     end
   | Bvand, [v1; v2] ->
      begin match cval_ctyp v1 with
      | CT_fbits _ ->
@@ -1232,6 +1242,12 @@ and sgen_call op cvals =
         sprintf "sslice(%s.bits, %s, %s)" (sgen_cval vec) (sgen_cval start) (sgen_cval len)
      | _ -> assert false
      end
+  | Set_slice, [vec; start; slice] ->
+     begin match cval_ctyp vec, cval_ctyp slice with
+     | CT_fbits (n, _), CT_fbits (m, _) ->
+        sprintf "((%s & ~(%s << %s)) | (%s << %s))" (sgen_cval vec) (sgen_mask m) (sgen_cval start) (sgen_cval slice) (sgen_cval start)
+     | _ -> assert false
+     end
   | Zero_extend n, [v] ->
      begin match cval_ctyp v with
      | CT_fbits _ -> sgen_cval v
@@ -1269,8 +1285,8 @@ and sgen_call op cvals =
         sprintf "append_ss(%s, %s)" (sgen_cval v1) (sgen_cval v2)
      | _ -> assert false
      end
-  | _, vs ->
-     sprintf "%s(%s)" (string_of_op op) (Util.string_of_list ", " sgen_cval cvals)
+  | _, _ ->
+     failwith "Could not generate cval primop"
 
 let sgen_cval_param cval =
   match cval_ctyp cval with
diff --git a/src/jib/jib_compile.ml b/src/jib/jib_compile.ml
index b76e65b7..a5a411f5 100644
--- a/src/jib/jib_compile.ml
+++ b/src/jib/jib_compile.ml
@@ -439,6 +439,12 @@ let optimize_call l ctx clexp id args arg_ctyps ret_ctyp =
        let start = ngensym () in
        [iinit (CT_fint 64) start (List.nth args 1);
         icopy l clexp (V_call (Slice n, [List.nth args 0; V_id (start, CT_fint 64)]))]
+    | "sail_unsigned", [CT_fbits _], CT_fint 64 ->
+       [icopy l clexp (V_call (Unsigned 64, [List.nth args 0]))]
+    | "sail_signed", [CT_fbits _], CT_fint 64 ->
+       [icopy l clexp (V_call (Signed 64, [List.nth args 0]))]
+    | "set_slice", [_; _; CT_fbits (n, _); CT_fint 64; CT_fbits (m, _)], CT_fbits (n', _) when n = n' ->
+       [icopy l clexp (V_call (Set_slice, [List.nth args 2; List.nth args 3; List.nth args 4]))]
     | _, _, _ ->
        call ()
     end
diff --git a/src/jib/jib_util.ml b/src/jib/jib_util.ml
index 441f8bd1..1c2e9530 100644
--- a/src/jib/jib_util.ml
+++ b/src/jib/jib_util.ml
@@ -276,31 +276,34 @@ let string_of_name ?deref_current_exception:(dce=true) ?zencode:(zencode=true) =
      "current_exception" ^ ssa_num n
 
 let string_of_op = function
-  | Bnot -> "not"
-  | List_hd -> "hd"
-  | List_tl -> "tl"
-  | Bit_to_bool -> "bit_to_bool"
-  | Eq -> "eq"
-  | Neq -> "neq"
-  | Bvnot -> "bvnot"
-  | Bvor -> "bvor"
-  | Bvand -> "bvand"
-  | Bvxor -> "bvxor"
-  | Bvadd -> "bvadd"
-  | Bvsub -> "bvsub"
-  | Bvaccess -> "bvaccess"
-  | Ilt -> "lt"
-  | Igt -> "gt"
-  | Ilteq -> "lteq"
-  | Igteq -> "gteq"
-  | Iadd -> "iadd"
-  | Isub -> "isub"
-  | Zero_extend n -> "zero_extend::<" ^ string_of_int n ^ ">"
-  | Sign_extend n -> "sign_extend::<" ^ string_of_int n ^ ">"
-  | Slice n -> "slice::<" ^ string_of_int n ^ ">"
-  | Sslice n -> "sslice::<" ^ string_of_int n ^ ">"
-  | Replicate n -> "replicate::<" ^ string_of_int n ^ ">"
-  | Concat -> "concat"
+  | Bnot -> "@not"
+  | List_hd -> "@hd"
+  | List_tl -> "@tl"
+  | Bit_to_bool -> "@bit_to_bool"
+  | Eq -> "@eq"
+  | Neq -> "@neq"
+  | Bvnot -> "@bvnot"
+  | Bvor -> "@bvor"
+  | Bvand -> "@bvand"
+  | Bvxor -> "@bvxor"
+  | Bvadd -> "@bvadd"
+  | Bvsub -> "@bvsub"
+  | Bvaccess -> "@bvaccess"
+  | Ilt -> "@lt"
+  | Igt -> "@gt"
+  | Ilteq -> "@lteq"
+  | Igteq -> "@gteq"
+  | Iadd -> "@iadd"
+  | Isub -> "@isub"
+  | Unsigned n -> "@unsigned::<" ^ string_of_int n ^ "@>"
+  | Signed n -> "@signed::<" ^ string_of_int n ^ "@>"
+  | Zero_extend n -> "@zero_extend::<" ^ string_of_int n ^ "@>"
+  | Sign_extend n -> "@sign_extend::<" ^ string_of_int n ^ "@>"
+  | Slice n -> "@slice::<" ^ string_of_int n ^ "@>"
+  | Sslice n -> "@sslice::<" ^ string_of_int n ^ "@>"
+  | Replicate n -> "@replicate::<" ^ string_of_int n ^ "@>"
+  | Set_slice -> "@set_slice"
+  | Concat -> "@concat"
 
 let rec string_of_cval = function
   | V_id (id, ctyp) -> string_of_name ~zencode:false id
@@ -908,6 +911,7 @@ let rec infer_call op vs =
   | (Bvor | Bvand | Bvxor | Bvadd | Bvsub), [v; _] -> cval_ctyp v
   | (Ilt | Igt | Ilteq | Igteq), _ -> CT_bool
   | (Iadd | Isub), _ -> CT_fint 64
+  | (Unsigned n | Signed n), _ -> CT_fint n
   | (Zero_extend n | Sign_extend n), [v] ->
      begin match cval_ctyp v with
      | CT_fbits (_, ord) | CT_sbits (_, ord) | CT_lbits ord ->
@@ -926,6 +930,7 @@ let rec infer_call op vs =
         CT_sbits (n, ord)
      | _ -> Reporting.unreachable Parse_ast.Unknown __POS__ "Invalid type for extract argument"
      end
+  | Set_slice, [vec; _; _] -> cval_ctyp vec
   | Replicate n, [vec] ->
      begin match cval_ctyp vec with
      | CT_fbits (m, ord) -> CT_fbits (n * m, ord)