Wasm: specialization of number comparisons

Author: vouillon

compiler/lib-wasm/generate.ml

@@ -110,13 +110,6 @@ module Generate (Target : Target_sig.S) = struct
       ; "caml_erf_float", (`Pure, [ Float ], Float)
       ; "caml_erfc_float", (`Pure, [ Float ], Float)
       ; "caml_float_compare", (`Pure, [ Float; Float ], Int)
-      ; "caml_greaterthan", (`Mutable, [ Value; Value ], Int)
-      ; "caml_greaterequal", (`Mutable, [ Value; Value ], Int)
-      ; "caml_lessthan", (`Mutable, [ Value; Value ], Int)
-      ; "caml_lessequal", (`Mutable, [ Value; Value ], Int)
-      ; "caml_equal", (`Mutable, [ Value; Value ], Int)
-      ; "caml_notequal", (`Mutable, [ Value; Value ], Int)
-      ; "caml_compare", (`Mutable, [ Value; Value ], Int)
       ];
     h
 
@@ -285,6 +278,39 @@ module Generate (Target : Target_sig.S) = struct
               (transl_prim_arg ctx ?typ:tz z)
         | _ -> invalid_arity name l ~expected:3)
 
+  let register_comparison name cmp_int cmp_boxed_int cmp_float =
+    register_prim name `Mutable (fun ctx _ l ->
+        match l with
+        | [ x; y ] -> (
+            match get_type ctx x, get_type ctx y with
+            | Int _, Int _ -> cmp_int ctx x y
+            | Number (Int32, _), Number (Int32, _) ->
+                let x = transl_prim_arg ctx ~typ:(Number (Int32, Unboxed)) x in
+                let y = transl_prim_arg ctx ~typ:(Number (Int32, Unboxed)) y in
+                int32_bin_op cmp_boxed_int x y
+            | Number (Nativeint, _), Number (Nativeint, _) ->
+                let x = transl_prim_arg ctx ~typ:(Number (Nativeint, Unboxed)) x in
+                let y = transl_prim_arg ctx ~typ:(Number (Nativeint, Unboxed)) y in
+                nativeint_bin_op cmp_boxed_int x y
+            | Number (Int64, _), Number (Int64, _) ->
+                let x = transl_prim_arg ctx ~typ:(Number (Int64, Unboxed)) x in
+                let y = transl_prim_arg ctx ~typ:(Number (Int64, Unboxed)) y in
+                int64_bin_op cmp_boxed_int x y
+            | Number (Float, _), Number (Float, _) ->
+                let x = transl_prim_arg ctx ~typ:(Number (Float, Unboxed)) x in
+                let y = transl_prim_arg ctx ~typ:(Number (Float, Unboxed)) y in
+                float_bin_op cmp_float x y
+            | _ ->
+                let* f =
+                  register_import
+                    ~name
+                    (Fun { W.params = [ Type.value; Type.value ]; result = [ I32 ] })
+                in
+                let* x = transl_prim_arg ctx x in
+                let* y = transl_prim_arg ctx y in
+                return (W.Call (f, [ x; y ])))
+        | _ -> invalid_arity name l ~expected:2)
+
   let () =
     register_bin_prim
       "caml_floatarray_unsafe_get"
@@ -1087,7 +1113,83 @@ module Generate (Target : Target_sig.S) = struct
       ~ty:(Int Normalized)
       (fun i j -> Arith.((j < i) - (i < j)));
     register_prim "%js_array" `Pure (fun ctx _ l ->
-        Memory.allocate ~tag:0 (expression_list (fun x -> transl_prim_arg ctx x) l))
+        Memory.allocate ~tag:0 (expression_list (fun x -> transl_prim_arg ctx x) l));
+    register_comparison
+      "caml_greaterthan"
+      (fun ctx x y -> translate_int_comparison ctx (fun y x -> Arith.(x < y)) x y)
+      (Gt S)
+      Gt;
+    register_comparison
+      "caml_greaterequal"
+      (fun ctx x y -> translate_int_comparison ctx (fun y x -> Arith.(x <= y)) x y)
+      (Ge S)
+      Ge;
+    register_comparison
+      "caml_lessthan"
+      (fun ctx x y -> translate_int_comparison ctx Arith.( < ) x y)
+      (Lt S)
+      Lt;
+    register_comparison
+      "caml_lessequal"
+      (fun ctx x y -> translate_int_comparison ctx Arith.( <= ) x y)
+      (Le S)
+      Le;
+    register_comparison
+      "caml_equal"
+      (fun ctx x y -> translate_int_equality ctx ~negate:false x y)
+      Eq
+      Eq;
+    register_comparison
+      "caml_notequal"
+      (fun ctx x y -> translate_int_equality ctx ~negate:true x y)
+      Ne
+      Ne;
+    register_prim "caml_compare" `Mutable (fun ctx _ l ->
+        match l with
+        | [ x; y ] -> (
+            match get_type ctx x, get_type ctx y with
+            | Int _, Int _ ->
+                let x' = transl_prim_arg ctx ~typ:(Int Normalized) x in
+                let y' = transl_prim_arg ctx ~typ:(Int Normalized) y in
+                Arith.((y' < x') - (x' < y'))
+            | Number (Int32, _), Number (Int32, _)
+            | Number (Nativeint, _), Number (Nativeint, _) ->
+                let* f =
+                  register_import
+                    ~name:"caml_int32_compare"
+                    (Fun { W.params = [ I32; I32 ]; result = [ I32 ] })
+                in
+                let* x' = transl_prim_arg ctx ~typ:(Number (Int32, Unboxed)) x in
+                let* y' = transl_prim_arg ctx ~typ:(Number (Int32, Unboxed)) y in
+                return (W.Call (f, [ x'; y' ]))
+            | Number (Int64, _), Number (Int64, _) ->
+                let* f =
+                  register_import
+                    ~name:"caml_int64_compare"
+                    (Fun { W.params = [ I64; I64 ]; result = [ I32 ] })
+                in
+                let* x' = transl_prim_arg ctx ~typ:(Number (Int64, Unboxed)) x in
+                let* y' = transl_prim_arg ctx ~typ:(Number (Int64, Unboxed)) y in
+                return (W.Call (f, [ x'; y' ]))
+            | Number (Float, _), Number (Float, _) ->
+                let* f =
+                  register_import
+                    ~name:"caml_float_compare"
+                    (Fun { W.params = [ F64; F64 ]; result = [ I32 ] })
+                in
+                let* x' = transl_prim_arg ctx ~typ:(Number (Float, Unboxed)) x in
+                let* y' = transl_prim_arg ctx ~typ:(Number (Float, Unboxed)) y in
+                return (W.Call (f, [ x'; y' ]))
+            | _ ->
+                let* f =
+                  register_import
+                    ~name:"caml_compare"
+                    (Fun { W.params = [ Type.value; Type.value ]; result = [ I32 ] })
+                in
+                let* x' = transl_prim_arg ctx x in
+                let* y' = transl_prim_arg ctx y in
+                return (W.Call (f, [ x'; y' ])))
+        | _ -> invalid_arity "caml_compare" l ~expected:2)
 
   let unboxed_type ty : W.value_type option =
     match ty with

compiler/lib-wasm/typing.ml

@@ -588,6 +588,24 @@ let primitives_with_unboxed_parameters =
     ];
   h
 
+let type_specialized_primitive types name args =
+  match name with
+  | "caml_greaterthan"
+  | "caml_greaterequal"
+  | "caml_lessthan"
+  | "caml_lessequal"
+  | "caml_equal"
+  | "caml_notequal"
+  | "caml_compare" -> (
+      match List.map ~f:(arg_type ~approx:types) args with
+      | [ Int _; Int _ ]
+      | [ Number (Int32, _); Number (Int32, _) ]
+      | [ Number (Int64, _); Number (Int64, _) ]
+      | [ Number (Nativeint, _); Number (Nativeint, _) ]
+      | [ Number (Float, _); Number (Float, _) ] -> true
+      | _ -> false)
+  | _ -> false
+
 let box_numbers p st types =
   (* We box numbers eagerly if the boxed value is ever used. *)
   let should_box = Var.ISet.empty () in
@@ -634,7 +652,9 @@ let box_numbers p st types =
                       then List.iter ~f:box args
                   | Block (tag, lst, _, _) -> if tag <> 254 then Array.iter ~f:box lst
                   | Prim (Extern s, args) ->
-                      if not (String.Hashtbl.mem primitives_with_unboxed_parameters s)
+                      if
+                        (not (String.Hashtbl.mem primitives_with_unboxed_parameters s))
+                        || type_specialized_primitive types s args
                       then
                         List.iter
                           ~f:(fun a ->
@@ -660,6 +680,12 @@ let box_numbers p st types =
         ())
     ()
 
+let print_opt typ f e =
+  match e with
+  | Prim (Extern name, args) when type_specialized_primitive typ name args ->
+      Format.fprintf f " OPT"
+  | _ -> ()
+
 type t =
   { types : typ Var.Tbl.t
   ; return_types : typ Var.Hashtbl.t
@@ -689,7 +715,13 @@ let f ~global_flow_state ~global_flow_info ~fun_info ~deadcode_sentinal p =
       Format.err_formatter
       (fun _ i ->
         match i with
-        | Instr (Let (x, _)) -> Format.asprintf "{%a}" Domain.print (Var.Tbl.get types x)
+        | Instr (Let (x, e)) ->
+            Format.asprintf
+              "{%a}%a"
+              Domain.print
+              (Var.Tbl.get types x)
+              (print_opt types)
+              e
         | _ -> "")
       p);
   let return_types = Var.Hashtbl.create 128 in