lambda: Add support for new atomic primitives.

Uses of existing atomic primitives %atomic_foo, which act on
single-field references, are now translated into %atomic_foo_field,
which act on a pointer and an offset -- passed as separate arguments.

In particular, note that the arity of the internal Lambda primitive
    Patomic_load
increases by one with this patchset. (Initially we renamed it into
    Patomic_load_field
but this creates a lot of churn for no clear benefits.)

We also support primitives of the form %atomic_foo_loc, which
expects a pair of a pointer and an offset (as a single argument),
as we proposed in the RFC on atomic fields
  ocaml/RFCs#39
(but there is no language-level support for atomic record fields yet)

Co-authored-by: Clément Allain <[email protected]>

Author: 2 people

bytecomp/bytegen.ml

@@ -656,7 +656,7 @@ let comp_primitive stack_info p sz args =
   | Parray_of_iarray -> Kccall("caml_array_of_iarray", 1)
   | Pget_header _ -> Kccall("caml_get_header", 1)
   | Pobj_dup -> Kccall("caml_obj_dup", 1)
-  | Patomic_load _ -> Kccall("caml_atomic_load", 1)
+  | Patomic_load _ -> Kccall("caml_atomic_load_field", 2)
   | Patomic_set _
   | Patomic_exchange _ -> Kccall("caml_atomic_exchange", 2)
   | Patomic_compare_exchange _ -> Kccall("caml_atomic_compare_exchange", 3)

lambda/translprim.ml

@@ -114,6 +114,24 @@ type loc_kind =
   | Loc_POS
   | Loc_FUNCTION
 
+type atomic_kind =
+  | Ref   (* operation on an atomic reference (takes only a pointer) *)
+  | Field (* operation on an atomic field (takes a pointer and an offset) *)
+  | Loc   (* operation on a first-class field (takes a (pointer, offset) pair *)
+
+type atomic_op =
+  | Load
+  | Set
+  | Exchange
+  | Compare_exchange
+  | Cas
+  | Faa
+  | Add
+  | Sub
+  | Lxor
+  | Land
+  | Lor
+
 type prim =
   | Primitive of Lambda.primitive * int
   | External of Lambda.external_call_description
@@ -134,6 +152,7 @@ type prim =
   | Poke of Lambda.peek_or_poke option
     (* For [Peek] and [Poke] the [option] is [None] until the primitive
        specialization code (below) has been run. *)
+  | Atomic of atomic_op * atomic_kind
   | Unsupported of Lambda.primitive
 
 let units_with_used_primitives = Hashtbl.create 7
@@ -897,22 +916,26 @@ let lookup_primitive loc ~poly_mode ~poly_sort pos p =
     | "%unbox_vec128" -> Primitive(Punbox_vector Boxed_vec128, 1)
     | "%box_vec128" -> Primitive(Pbox_vector (Boxed_vec128, mode), 1)
     | "%get_header" -> Primitive (Pget_header mode, 1)
-    | "%atomic_load" ->
-        Primitive ((Patomic_load {immediate_or_pointer=Pointer}), 1)
-    | "%atomic_set" ->
-        Primitive (Patomic_set {immediate_or_pointer=Pointer}, 2)
-    | "%atomic_exchange" ->
-        Primitive (Patomic_exchange {immediate_or_pointer=Pointer}, 2)
-    | "%atomic_compare_exchange" ->
-        Primitive (Patomic_compare_exchange {immediate_or_pointer=Pointer}, 3)
-    | "%atomic_cas" ->
-        Primitive (Patomic_compare_set {immediate_or_pointer=Pointer}, 3)
-    | "%atomic_fetch_add" -> Primitive (Patomic_fetch_add, 2)
-    | "%atomic_add" -> Primitive (Patomic_add, 2)
-    | "%atomic_sub" -> Primitive (Patomic_sub, 2)
-    | "%atomic_land" -> Primitive (Patomic_land, 2)
-    | "%atomic_lor" -> Primitive (Patomic_lor, 2)
-    | "%atomic_lxor" -> Primitive (Patomic_lxor, 2)
+    (* CR melse: field/loc variants of other primitives *)
+    | "%atomic_load" -> Atomic (Load, Ref)
+    | "%atomic_load_field" -> Atomic (Load, Field)
+    | "%atomic_load_loc" -> Atomic (Load, Loc)
+    | "%atomic_set"  -> Atomic (Set, Ref)
+    | "%atomic_exchange" -> Atomic (Exchange, Ref)
+    | "%atomic_exchange_field" -> Atomic (Exchange, Field)
+    | "%atomic_exchange_loc" -> Atomic (Exchange, Loc)
+    | "%atomic_compare_exchange" -> Atomic (Compare_exchange, Ref)
+    | "%atomic_cas" -> Atomic (Cas, Ref)
+    | "%atomic_cas_field" -> Atomic (Cas, Field)
+    | "%atomic_cas_loc" -> Atomic (Cas, Loc)
+    | "%atomic_fetch_add" -> Atomic (Faa, Ref)
+    | "%atomic_fetch_add_field" -> Atomic (Faa, Field)
+    | "%atomic_fetch_add_loc" -> Atomic (Faa, Loc)
+    | "%atomic_add" -> Atomic (Add, Ref)
+    | "%atomic_sub" -> Atomic (Sub, Ref)
+    | "%atomic_land" -> Atomic (Land, Ref)
+    | "%atomic_lor" -> Atomic (Lor, Ref)
+    | "%atomic_lxor" -> Atomic (Lxor, Ref)
     | "%runstack" ->
       if runtime5 then Primitive (Prunstack, 3) else Unsupported Prunstack
     | "%reperform" ->
@@ -1624,6 +1647,90 @@ let lambda_of_loc kind sloc =
                        ~include_zero_alloc:false sloc in
     Lconst (Const_immstring scope_name)
 
+let atomic_arity op (kind : atomic_kind) =
+  let arity_of_op =
+    match op with
+    | Load -> 1
+    | Set -> 2
+    | Exchange -> 2
+    | Compare_exchange -> 3
+    | Cas -> 3
+    | Faa -> 2
+    | Add | Sub | Lxor | Land | Lor -> 2
+  in
+  let extra_kind_arity =
+    match kind with
+    | Ref | Loc -> 0
+    | Field -> 1
+  in
+  arity_of_op + extra_kind_arity
+
+let lambda_of_atomic prim_name loc op (kind : atomic_kind) args =
+  if List.length args <> atomic_arity op kind then
+    raise (Error (to_location loc, Wrong_arity_builtin_primitive prim_name)) ;
+  let split = function
+    | [] ->
+        (* split is only called when [arity >= 1] *)
+        assert false
+    | first :: rest ->
+        first, rest
+  in
+  let prim =
+    match op with
+    | Load -> Patomic_load {immediate_or_pointer = Pointer}
+    | Set -> Patomic_set {immediate_or_pointer = Pointer}
+    | Compare_exchange -> Patomic_compare_exchange {immediate_or_pointer=Pointer}
+    | Exchange -> Patomic_exchange {immediate_or_pointer=Pointer}
+    | Cas -> Patomic_compare_set {immediate_or_pointer=Pointer}
+    | Faa -> Patomic_fetch_add
+    | Add -> Patomic_add
+    | Sub -> Patomic_sub
+    | Lor -> Patomic_lor
+    | Lxor -> Patomic_lxor
+    | Land -> Patomic_land
+  in
+  match kind with
+  | Ref ->
+      (* the primitive application
+           [Lprim(%atomic_exchange, [ref; v])]
+         becomes
+           [Lprim(caml_atomic_exchange_field, [ref; 0; v])]
+      *)
+      let ref_arg, rest = split args in
+      let args = ref_arg :: Lconst (Lambda.const_int 0) :: rest in
+      Lprim (prim, args, loc)
+  | Field ->
+      (* the primitive application
+           [Lprim(%atomic_exchange_field, [ptr; ofs; v])]
+         becomes
+           [Lprim(caml_atomic_exchange_field, [ptr; ofs; v])] *)
+      Lprim (prim, args, loc)
+  | Loc ->
+      (* CR melse: unboxed tuple? *)
+      (* the primitive application
+           [Lprim(%atomic_exchange_loc, [(ptr, ofs); v])]
+         becomes
+           [Lprim(caml_atomic_exchange_field, [ptr; ofs; v])]
+         and in the general case of a non-tuple expression <loc>
+           [Lprim(%atomic_exchange_loc, [loc; v])]
+         becomes
+           [Llet(p, loc,
+              Lprim(caml_atomic_exchange_field, [Field(p, 0); Field(p, 1); v]))]
+      *)
+      let loc_arg, rest = split args in
+      match loc_arg with
+      | Lprim (Pmakeblock _, [ptr; ofs], _argloc) ->
+          let args = ptr :: ofs :: rest in
+          Lprim (prim, args, loc)
+      | _ ->
+          let varg = Ident.create_local "atomic_arg" in
+          let ptr = Lprim (Pfield (0, Pointer, Reads_agree), [Lvar varg], loc) in
+          let ofs =
+            Lprim (Pfield (1, Immediate, Reads_agree), [Lvar varg], loc)
+          in
+          let args = ptr :: ofs :: rest in
+          Llet (Strict, Pvalue { raw_kind = Pgenval; nullable = Non_nullable }, varg, loc_arg, Lprim (prim, args, loc))
+
 let caml_restore_raw_backtrace =
   Lambda.simple_prim_on_values ~name:"caml_restore_raw_backtrace" ~arity:2
     ~alloc:false
@@ -1721,6 +1828,8 @@ let lambda_of_prim prim_name prim loc args arg_exps =
       Lprim (Ppeek layout, [ptr], loc)
   | Poke (Some layout), [ptr; new_value] ->
       Lprim (Ppoke layout, [ptr; new_value], loc)
+  | Atomic (op, kind), args ->
+      lambda_of_atomic prim_name loc op kind args
   | Unsupported prim, _ ->
       let exn =
         transl_extension_path loc (Lazy.force Env.initial)
@@ -1773,6 +1882,7 @@ let check_primitive_arity loc p =
     | Frame_pointers -> p.prim_arity = 0
     | Identity | Peek _ -> p.prim_arity = 1
     | Apply _ | Revapply _ | Poke _ -> p.prim_arity = 2
+    | Atomic (op, kind) -> p.prim_arity = atomic_arity op kind
     | Unsupported _ -> true
   in
   if not ok then raise(Error(loc, Wrong_arity_builtin_primitive p.prim_name))
@@ -1972,7 +2082,7 @@ let primitive_needs_event_after = function
   | Lazy_force _ | Send _ | Send_self _ | Send_cache _
   | Apply _ | Revapply _ -> true
   | Raise _ | Raise_with_backtrace | Loc _ | Frame_pointers | Identity
-  | Peek _ | Poke _ | Unsupported _ -> false
+  | Peek _ | Poke _ | Atomic _ | Unsupported _ -> false
 
 let transl_primitive_application loc p env ty ~poly_mode ~stack ~poly_sort
     path exp args arg_exps pos =

runtime/memory.c

@@ -372,7 +372,7 @@ CAMLprim value caml_atomic_compare_exchange_field (value obj, intnat field, valu
   }
 }
 
-CAMLexport value caml_atomic_cas_field (value obj, value vfield, value oldval, value newval)
+CAMLprim value caml_atomic_cas_field (value obj, value vfield, value oldval, value newval)
 {
   intnat field = Long_val(vfield);
   if (caml_atomic_compare_exchange_field(obj, field, oldval, newval) == oldval) {

testsuite/tests/atomic-locs/cmm.compilers.reference

@@ -17,7 +17,7 @@ cmm:
 (function camlCmm.standard_atomic_get_271 (r: val) (load_mut_atomic val r))
 
 (function camlCmm.standard_atomic_cas_294 (r: val oldv: val newv: val)
- (extcall "caml_atomic_cas" r oldv newv int,int,int->val))
+ (extcall "caml_atomic_cas_field" r 1 oldv newv int,int,int,int->val))
 
 (function camlCmm.entry ()
  (let standard_atomic_get "camlCmm.2"