flambda-backend: Improved reinterpret casts for integers and floats (#…

…2686)

bytecomp/bytegen.ml

@@ -161,7 +161,8 @@ let preserve_tailcall_for_prim = function
   | Pprobe_is_enabled _ | Pobj_dup
   | Pctconst _ | Pbswap16 | Pbbswap _ | Pint_as_pointer _
   | Patomic_exchange | Patomic_cas | Patomic_fetch_add | Patomic_load _
-  | Pdls_get ->
+  | Pdls_get | Preinterpret_tagged_int63_as_unboxed_int64
+  | Preinterpret_unboxed_int64_as_tagged_int63 ->
       false
 
 (* Add a Kpop N instruction in front of a continuation *)
@@ -583,6 +584,20 @@ let comp_primitive stack_info p sz args =
   | Punboxed_int32_array_set_128 _ | Punboxed_int64_array_set_128 _
   | Punboxed_nativeint_array_set_128 _ ->
     fatal_error "128-bit load/store is not supported in bytecode mode."
+  | Preinterpret_tagged_int63_as_unboxed_int64 ->
+    if not (Target_system.is_64_bit ())
+    then
+      Misc.fatal_error
+        "Preinterpret_tagged_int63_as_unboxed_int64 can only be used on 64-bit \
+         targets";
+    Kccall("caml_reinterpret_tagged_int63_as_unboxed_int64", 1)
+  | Preinterpret_unboxed_int64_as_tagged_int63 ->
+    if not (Target_system.is_64_bit ())
+    then
+      Misc.fatal_error
+        "Preinterpret_unboxed_int64_as_tagged_int63 can only be used on 64-bit \
+         targets";
+    Kccall("caml_reinterpret_unboxed_int64_as_tagged_int63", 1)
   (* The cases below are handled in [comp_expr] before the [comp_primitive] call
      (in the order in which they appear below),
      so they should never be reached in this function. *)

lambda/lambda.ml

@@ -305,6 +305,8 @@ type primitive =
   | Pbox_float of boxed_float * alloc_mode
   | Punbox_int of boxed_integer
   | Pbox_int of boxed_integer * alloc_mode
+  | Preinterpret_unboxed_int64_as_tagged_int63
+  | Preinterpret_tagged_int63_as_unboxed_int64
   (* Jane Street extensions *)
   | Parray_to_iarray
   | Parray_of_iarray
@@ -1805,7 +1807,16 @@ let primitive_may_allocate : primitive -> alloc_mode option = function
   | Patomic_exchange
   | Patomic_cas
   | Patomic_fetch_add
-  | Pdls_get -> None
+  | Pdls_get
+  | Preinterpret_unboxed_int64_as_tagged_int63 -> None
+  | Preinterpret_tagged_int63_as_unboxed_int64 ->
+    if !Clflags.native_code then None
+    else
+      (* We don't provide a locally-allocating version of this primitive
+         since it would only apply to bytecode, and code requiring performance
+         at a level where these primitives are necessary is very likely going
+         to be native. *)
+      Some alloc_heap
 
 let constant_layout: constant -> layout = function
   | Const_int _ | Const_char _ -> Pvalue Pintval
@@ -1991,6 +2002,8 @@ let primitive_result_layout (p : primitive) =
   | Patomic_cas
   | Patomic_fetch_add
   | Pdls_get -> layout_any_value
+  | Preinterpret_tagged_int63_as_unboxed_int64 -> layout_unboxed_int64
+  | Preinterpret_unboxed_int64_as_tagged_int63 -> layout_int
 
 let compute_expr_layout free_vars_kind lam =
   let rec compute_expr_layout kinds = function

lambda/lambda.mli

@@ -290,6 +290,17 @@ type primitive =
   | Pbox_float of boxed_float * alloc_mode
   | Punbox_int of boxed_integer
   | Pbox_int of boxed_integer * alloc_mode
+  | Preinterpret_unboxed_int64_as_tagged_int63
+  | Preinterpret_tagged_int63_as_unboxed_int64
+    (** At present [Preinterpret_unboxed_int64_as_tagged_int63] and
+        [Preinterpret_tagged_int63_as_unboxed_int64] will cause a fatal error
+        if the target system is 32-bit bytecode.
+
+        The [Preinterpret_tagged_int63_as_unboxed_int64] primitive is the
+        identity on machine words.  The
+        [Preinterpret_unboxed_int64_as_tagged_int63] compiles to logical OR
+        with 1 on machine words, to ensure that the tag bit is always set in
+        the result, just in case it was not in the incoming unboxed int64. *)
   (* Jane Street extensions *)
   | Parray_to_iarray (* Unsafely reinterpret a mutable array as an immutable
                         one; O(1) *)

lambda/printlambda.ml

@@ -799,6 +799,10 @@ let primitive ppf = function
   | Parray_to_iarray -> fprintf ppf "array_to_iarray"
   | Parray_of_iarray -> fprintf ppf "array_of_iarray"
   | Pget_header m -> fprintf ppf "get_header%s" (alloc_kind m)
+  | Preinterpret_tagged_int63_as_unboxed_int64 ->
+      fprintf ppf "reinterpret_tagged_int63_as_unboxed_int64"
+  | Preinterpret_unboxed_int64_as_tagged_int63 ->
+      fprintf ppf "reinterpret_unboxed_int64_as_tagged_int63"
 
 let name_of_primitive = function
   | Pbytes_of_string -> "Pbytes_of_string"
@@ -964,6 +968,10 @@ let name_of_primitive = function
   | Parray_of_iarray -> "Parray_of_iarray"
   | Parray_to_iarray -> "Parray_to_iarray"
   | Pget_header _ -> "Pget_header"
+  | Preinterpret_tagged_int63_as_unboxed_int64 ->
+      "Preinterpret_tagged_int63_as_unboxed_int64"
+  | Preinterpret_unboxed_int64_as_tagged_int63 ->
+      "Preinterpret_unboxed_int64_as_tagged_int63"
 
 let zero_alloc_attribute ppf check =
   match check with

lambda/tmc.ml

@@ -900,6 +900,8 @@ let rec choice ctx t =
     | Pisint _ | Pisout
     | Pignore
     | Pcompare_ints | Pcompare_floats _ | Pcompare_bints _
+    | Preinterpret_tagged_int63_as_unboxed_int64
+    | Preinterpret_unboxed_int64_as_tagged_int63
 
     (* we don't handle effect or DLS primitives *)
     | Prunstack | Pperform | Presume | Preperform | Pdls_get

lambda/translprim.ml

@@ -809,6 +809,10 @@ let lookup_primitive loc ~poly_mode ~poly_sort pos p =
     | "%box_int32" -> Primitive(Pbox_int (Pint32, mode), 1)
     | "%unbox_int64" -> Primitive(Punbox_int Pint64, 1)
     | "%box_int64" -> Primitive(Pbox_int (Pint64, mode), 1)
+    | "%reinterpret_tagged_int63_as_unboxed_int64" ->
+      Primitive(Preinterpret_tagged_int63_as_unboxed_int64, 1)
+    | "%reinterpret_unboxed_int64_as_tagged_int63" ->
+      Primitive(Preinterpret_unboxed_int64_as_tagged_int63, 1)
     | s when String.length s > 0 && s.[0] = '%' ->
        raise(Error(loc, Unknown_builtin_primitive s))
     | _ -> External lambda_prim
@@ -1559,6 +1563,9 @@ let lambda_primitive_needs_event_after = function
   | Punboxed_nativeint_array_set_128 _
   | Prunstack | Pperform | Preperform | Presume
   | Pbbswap _ | Pobj_dup | Pget_header _ -> true
+  (* [Preinterpret_tagged_int63_as_unboxed_int64] has to allocate in
+     bytecode, because int64# is actually represented as a boxed value. *)
+  | Preinterpret_tagged_int63_as_unboxed_int64 -> true
 
   | Pbytes_to_string | Pbytes_of_string
   | Parray_to_iarray | Parray_of_iarray
@@ -1584,6 +1591,7 @@ let lambda_primitive_needs_event_after = function
   | Pintofbint _ | Pctconst _ | Pbswap16 | Pint_as_pointer _ | Popaque _
   | Pdls_get
   | Pobj_magic _ | Punbox_float _ | Punbox_int _
+  | Preinterpret_unboxed_int64_as_tagged_int63
   (* These don't allocate in bytecode; they're just identity functions: *)
   | Pbox_float (_, _) | Pbox_int _
     -> false

lambda/value_rec_compiler.ml

@@ -395,7 +395,9 @@ let compute_static_size lam =
     | Pbox_int (_, _)
     | Pfloatoffloat32 _
     | Pfloat32offloat _
-    | Pget_header _ ->
+    | Pget_header _
+    | Preinterpret_tagged_int63_as_unboxed_int64
+    | Preinterpret_unboxed_int64_as_tagged_int63 ->
         dynamic_size lam
   in
   compute_expression_size Ident.Map.empty lam

runtime/ints.c

@@ -864,3 +864,24 @@ CAMLprim value caml_nativeint_of_string(value s)
 {
   return caml_copy_nativeint(caml_nativeint_of_string_unboxed(s));
 }
+
+CAMLprim value caml_reinterpret_tagged_int63_as_unboxed_int64(value i)
+{
+  // This should only be called on 64-bit targets.
+  // This stub is only used for bytecode, so in fact the "unboxed_int64"
+  // is to be boxed.
+  CAMLassert(sizeof(value) == 8);
+  CAMLassert(Is_long(i));
+  return caml_copy_int64((int64_t) i);
+}
+
+CAMLprim value caml_reinterpret_unboxed_int64_as_tagged_int63(value i)
+{
+  // This should only be called on 64-bit targets.
+  // This stub is only used for bytecode, so in fact the "unboxed_int64"
+  // is boxed.
+  CAMLassert(sizeof(value) == 8);
+  CAMLassert(Is_block(i));
+  CAMLassert(Tag_val(i) == Custom_tag);
+  return (value) (Int64_val(i) | 1L);
+}

runtime4/ints.c

@@ -864,3 +864,24 @@ CAMLprim value caml_nativeint_of_string(value s)
 {
   return caml_copy_nativeint(caml_nativeint_of_string_unboxed(s));
 }
+
+CAMLprim value caml_reinterpret_tagged_int63_as_unboxed_int64(value i)
+{
+  // This should only be called on 64-bit targets.
+  // This stub is only used for bytecode, so in fact the "unboxed_int64"
+  // is to be boxed.
+  CAMLassert(sizeof(value) == 8);
+  CAMLassert(Is_long(i));
+  return caml_copy_int64((int64_t) i);
+}
+
+CAMLprim value caml_reinterpret_unboxed_int64_as_tagged_int63(value i)
+{
+  // This should only be called on 64-bit targets.
+  // This stub is only used for bytecode, so in fact the "unboxed_int64"
+  // is boxed.
+  CAMLassert(sizeof(value) == 8);
+  CAMLassert(Is_block(i));
+  CAMLassert(Tag_val(i) == Custom_tag);
+  return (value) (Int64_val(i) | 1L);
+}

testsuite/tests/reinterpret-casts/i64_i63_reinterpret.ml

@@ -0,0 +1,66 @@
+(* TEST *)
+
+(* Tests for reinterpret-cast primitives tagged_int63 <-> unboxed_int64 *)
+
+[@@@ocaml.flambda_o3]
+
+external i63_to_i64 : int -> int64# =
+  "%reinterpret_tagged_int63_as_unboxed_int64"
+external i64_to_i63 : int64# -> int =
+  "%reinterpret_unboxed_int64_as_tagged_int63"
+
+external box_int64 : int64# -> (int64[@local_opt]) = "%box_int64"
+
+(* List functions that can be unrolled *)
+
+let rec map f = function
+  | [] -> []
+  | x::xs -> (f x) :: map f xs
+
+let[@loop never] rec iter2 f l1 l2 =
+  match (l1, l2) with
+    ([], []) -> ()
+  | (a1::l1, a2::l2) -> f a1 a2; iter2 f l1 l2
+  | (_, _) -> invalid_arg "iter2"
+
+let examples = [min_int; -1; 0; 1; max_int]
+
+let examples_int64 =
+  (map [@unrolled 6])
+    (fun i -> Int64.logor (Int64.mul (Int64.of_int i) 2L) 1L) examples
+
+(* This simplifies to a function that just returns unit (and does nothing
+   else)! *)
+let[@inline never] test () =
+  (iter2 [@unrolled 6]) (fun i i64 ->
+      let i64_unboxed = i63_to_i64 i in
+      let i64' = box_int64 i64_unboxed in
+      if not (Int64.equal i64 i64') then
+        failwith (Printf.sprintf
+          "i63_to_i64 failure on 0x%x: got 0x%Lx, expected 0x%Lx" i i64' i64);
+      let i' = i64_to_i63 i64_unboxed in
+      if not (Int.equal i i') then
+        failwith (Printf.sprintf
+          "i64_to_i63 failure on 0x%Lx -> 0x%x: expected 0x%x" i64 i' i))
+    examples examples_int64
+
+(* This version checks the Cmm compilation of the primitives. *)
+let[@inline never] test_opaque () =
+  List.iter2 (fun i i64 ->
+      let i64_unboxed =
+        Sys.opaque_identity (i63_to_i64 (Sys.opaque_identity i))
+      in
+      let i64' = box_int64 i64_unboxed in
+      if not (Int64.equal i64 i64') then
+        failwith (Printf.sprintf
+          "i63_to_i64 failure on 0x%x: got 0x%Lx, expected 0x%Lx" i i64' i64);
+      let i' = i64_to_i63 i64_unboxed in
+      if not (Int.equal i i') then
+        failwith (Printf.sprintf
+          "i64_to_i63 failure on 0x%Lx -> 0x%x: expected 0x%x" i64 i' i))
+    examples examples_int64
+
+let () =
+  test ();
+  test_opaque ();
+  print_endline "ok"

testsuite/tests/reinterpret-casts/i64_i63_reinterpret.reference

@@ -0,0 +1 @@
+ok

typing/primitive.ml

@@ -615,6 +615,11 @@ let prim_has_valid_reprs ~loc prim =
     | "%unbox_int64" ->
       exactly [Same_as_ocaml_repr Value; Same_as_ocaml_repr Bits64]
 
+    | "%reinterpret_tagged_int63_as_unboxed_int64" ->
+      exactly [Same_as_ocaml_repr Value; Same_as_ocaml_repr Bits64]
+    | "%reinterpret_unboxed_int64_as_tagged_int63" ->
+      exactly [Same_as_ocaml_repr Bits64; Same_as_ocaml_repr Value]
+
     (* Bigstring primitives *)
     | "%caml_bigstring_get32#" ->
       exactly [