lambda: add support for more 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: gasche

asmcomp/cmm_helpers.ml

@@ -741,6 +741,9 @@ let array_indexing ?typ log2size ptr ofs dbg =
       Cop(add, [Cop(add, [ptr; lsl_const ofs (log2size - 1) dbg], dbg);
                     Cconst_int((-1) lsl (log2size - 1), dbg)], dbg)
 
+let field_address_computed ptr ofs dbg =
+  array_indexing log2_size_addr ptr ofs dbg
+
 let addr_array_ref arr ofs dbg =
   Cop(mk_load_mut Word_val,
     [array_indexing log2_size_addr arr ofs dbg], dbg)

asmcomp/cmm_helpers.mli

@@ -214,6 +214,10 @@ val set_field :
   expression -> int -> expression -> Lambda.initialization_or_assignment ->
   Debuginfo.t -> expression
 
+(** [field_address_computed ptr ofs dbg] returns an expression for the address
+    at offset [ofs] of the block pointed to by [ptr]. *)
+val field_address_computed : expression -> expression -> Debuginfo.t -> expression
+
 (** Load a block's header *)
 val get_header : expression -> Debuginfo.t -> expression
 

asmcomp/cmmgen.ml

@@ -831,8 +831,6 @@ and transl_prim_1 env p arg dbg =
        dbg)
   | Pdls_get ->
       Cop(Cdls_get, [transl env arg], dbg)
-  | Patomic_load ->
-      Cop(mk_load_atomic Word_val, [transl env arg], dbg)
   | Ppoll ->
     (Csequence (remove_unit (transl env arg),
                 return_unit dbg (Cop(Cpoll, [], dbg))))
@@ -853,7 +851,9 @@ and transl_prim_1 env p arg dbg =
     | Plslbint _ | Plsrbint _ | Pasrbint _ | Pbintcomp (_, _)
     | Pbigarrayref (_, _, _, _) | Pbigarrayset (_, _, _, _)
     | Pbigarraydim _ | Pstring_load _ | Pbytes_load _ | Pbytes_set _
-    | Pbigstring_load _ | Pbigstring_set _)
+    | Pbigstring_load _ | Pbigstring_set _
+    | Patomic_load
+    )
     ->
       fatal_errorf "Cmmgen.transl_prim_1: %a"
         Printclambda_primitives.primitive p
@@ -870,6 +870,12 @@ and transl_prim_2 env p arg1 arg2 dbg =
       let float_val = transl_unbox_float dbg env arg2 in
       setfloatfield n init ptr float_val dbg
 
+  | Patomic_load ->
+      let ptr = transl env arg1 in
+      let ofs = transl env arg2 in
+      Cop(mk_load_atomic Word_val,
+          [field_address_computed ptr ofs dbg], dbg)
+
   (* Boolean operations *)
   | Psequand ->
       let dbg' = Debuginfo.none in
@@ -1022,7 +1028,6 @@ and transl_prim_2 env p arg1 arg2 dbg =
                      [transl_unbox_int dbg env bi arg1;
                       transl_unbox_int dbg env bi arg2], dbg)) dbg
   | Prunstack | Pperform | Presume | Preperform | Pdls_get
-  | Patomic_load
   | Pnot | Pnegint | Pintoffloat | Pfloatofint | Pnegfloat
   | Pabsfloat | Pstringlength | Pbyteslength | Pbytessetu | Pbytessets
   | Pisint | Pbswap16 | Pint_as_pointer | Popaque | Pread_symbol _

bytecomp/bytegen.ml

@@ -488,7 +488,7 @@ let comp_primitive stack_info p sz args =
   | Pint_as_pointer -> Kccall("caml_int_as_pointer", 1)
   | Pbytes_to_string -> Kccall("caml_string_of_bytes", 1)
   | Pbytes_of_string -> Kccall("caml_bytes_of_string", 1)
-  | Patomic_load -> Kccall("caml_atomic_load", 1)
+  | Patomic_load -> Kccall("caml_atomic_load_field", 2)
   | Pdls_get -> Kccall("caml_domain_dls_get", 1)
   | Ppoll -> Kccall("caml_process_pending_actions_with_root", 1)
   (* The cases below are handled in [comp_expr] before the [comp_primitive] call

lambda/translprim.ml

@@ -77,6 +77,20 @@ type loc_kind =
   | Loc_POS
   | Loc_FUNCTION
 
+type atomic_kind =
+  | Ref    (* operation on an atomic reference
+              (takes only a pointer) *)
+  | Record (* operation on an atomic record
+              (takes a pointer and an offset) *)
+  | Loc    (* operation on a first-class field
+              (takes a (pointer, offset) pair *)
+
+type atomic_op =
+  | Load
+  | Exchange
+  | Cas
+  | Fetch_add
+
 type prim =
   | Primitive of Lambda.primitive * int
   | External of Primitive.description
@@ -92,6 +106,7 @@ type prim =
   | Identity
   | Apply
   | Revapply
+  | Atomic of atomic_op * atomic_kind
 
 let used_primitives = Hashtbl.create 7
 let add_used_primitive loc env path =
@@ -114,12 +129,11 @@ let prim_sys_argv =
   Primitive.simple ~name:"caml_sys_argv" ~arity:1 ~alloc:true
 
 let prim_atomic_exchange =
-  Primitive.simple ~name:"caml_atomic_exchange" ~arity:2 ~alloc:false
+  Primitive.simple ~name:"caml_atomic_exchange_field" ~arity:3 ~alloc:false
 let prim_atomic_cas =
-  Primitive.simple ~name:"caml_atomic_cas" ~arity:3 ~alloc:false
+  Primitive.simple ~name:"caml_atomic_cas_field" ~arity:4 ~alloc:false
 let prim_atomic_fetch_add =
-  Primitive.simple ~name:"caml_atomic_fetch_add" ~arity:2 ~alloc:false
-
+  Primitive.simple ~name:"caml_atomic_fetch_add_field" ~arity:3 ~alloc:false
 
 let primitives_table =
   create_hashtable 57 [
@@ -372,10 +386,18 @@ let primitives_table =
     "%greaterequal", Comparison(Greater_equal, Compare_generic);
     "%greaterthan", Comparison(Greater_than, Compare_generic);
     "%compare", Comparison(Compare, Compare_generic);
-    "%atomic_load", Primitive (Patomic_load, 1);
-    "%atomic_exchange", External prim_atomic_exchange;
-    "%atomic_cas", External prim_atomic_cas;
-    "%atomic_fetch_add", External prim_atomic_fetch_add;
+    "%atomic_load", Atomic(Load, Ref);
+    "%atomic_exchange", Atomic(Exchange, Ref);
+    "%atomic_cas", Atomic(Cas, Ref);
+    "%atomic_fetch_add", Atomic(Fetch_add, Ref);
+    "%atomic_load_field", Atomic(Load, Record);
+    "%atomic_exchange_field", Atomic(Exchange, Record);
+    "%atomic_cas_field", Atomic(Cas, Record);
+    "%atomic_fetch_add_field", Atomic(Fetch_add, Record);
+    "%atomic_load_loc", Atomic(Load, Loc);
+    "%atomic_exchange_loc", Atomic(Exchange, Loc);
+    "%atomic_cas_loc", Atomic(Cas, Loc);
+    "%atomic_fetch_add_loc", Atomic(Fetch_add, Loc);
     "%runstack", Primitive (Prunstack, 3);
     "%reperform", Primitive (Preperform, 3);
     "%perform", Primitive (Pperform, 1);
@@ -658,6 +680,75 @@ let lambda_of_loc kind sloc =
     let scope_name = Debuginfo.Scoped_location.string_of_scoped_location sloc in
     Lconst (Const_immstring scope_name)
 
+let atomic_arity (op : atomic_op) (kind : atomic_kind) =
+  let arity_of_op = match op with
+    | Load -> 1
+    | Exchange -> 2
+    | Cas -> 3
+    | Fetch_add -> 2
+  in
+  let extra_kind_arity = match kind with
+    | Ref | Loc -> 0
+    | Record -> 1
+  in
+  arity_of_op + extra_kind_arity
+
+let lambda_of_atomic prim_name loc op 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 : atomic_op) with
+    | Load -> Patomic_load
+    | Exchange -> Pccall prim_atomic_exchange
+    | Cas -> Pccall prim_atomic_cas
+    | Fetch_add -> Pccall prim_atomic_fetch_add
+  in
+  match (kind : atomic_kind) with
+  | Record ->
+      (* the primitive application
+           [%atomic_exchange_field ptr ofs v]
+         becomes (in pseudo-code mixing C calls and OCaml expressions)
+           [caml_atomic_exchange_field(ptr, ofs, v)] *)
+      Lprim(prim, args, loc)
+  | Ref ->
+      (* the primitive application
+           [%atomic_exchange ref v]
+         becomes
+           [caml_atomic_exchange_field(ref, Val_long(0), v)] *)
+      let ref_arg, rest = split args in
+      let args = ref_arg :: Lconst (Lambda.const_int 0) :: rest in
+      Lprim(prim, args, loc)
+  | Loc ->
+      (* the primitive application
+           [%atomic_exchange_loc (ptr, ofs) v]
+         becomes
+           [caml_atomic_exchange_field(ptr, ofs, v)]
+         and in the general case of a non-tuple expression <loc>
+           [%atomic_exchange_loc <loc> v]
+         becomes
+           [let p = <loc> in
+            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, Immutable),
+                          [Lvar varg], loc) in
+          let ofs = Lprim(Pfield(1, Immediate, Immutable),
+                          [Lvar varg], loc) in
+          let args = ptr :: ofs :: rest in
+          Llet(Strict, Pgenval, varg, loc_arg,
+               Lprim(prim, args, loc))
+
 let caml_restore_raw_backtrace =
   Primitive.simple ~name:"caml_restore_raw_backtrace" ~arity:2 ~alloc:false
 
@@ -744,10 +835,13 @@ let lambda_of_prim prim_name prim loc args arg_exps =
         ap_inlined = Default_inline;
         ap_specialised = Default_specialise;
       }
+  | Atomic (op, kind), args ->
+      lambda_of_atomic prim_name loc op kind args
   | (Raise _ | Raise_with_backtrace
     | Lazy_force | Loc _ | Primitive _ | Comparison _
     | Send | Send_self | Send_cache | Frame_pointers | Identity
-    | Apply | Revapply), _ ->
+    | Apply | Revapply
+    ), _ ->
       raise(Error(to_location loc, Wrong_arity_builtin_primitive prim_name))
 
 let check_primitive_arity loc p =
@@ -766,6 +860,7 @@ let check_primitive_arity loc p =
     | Frame_pointers -> p.prim_arity = 0
     | Identity -> p.prim_arity = 1
     | Apply | Revapply -> p.prim_arity = 2
+    | Atomic (op, kind) -> p.prim_arity = atomic_arity op kind
   in
   if not ok then raise(Error(loc, Wrong_arity_builtin_primitive p.prim_name))
 
@@ -838,7 +933,11 @@ let primitive_needs_event_after = function
       lambda_primitive_needs_event_after (comparison_primitive comp knd)
   | Lazy_force | Send | Send_self | Send_cache
   | Apply | Revapply -> true
-  | Raise _ | Raise_with_backtrace | Loc _ | Frame_pointers | Identity -> false
+  | Raise _ | Raise_with_backtrace
+  | Loc _
+  | Frame_pointers | Identity
+  | Atomic (_, _)
+    -> false
 
 let transl_primitive_application loc p env ty path exp args arg_exps =
   let prim =

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

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