Skip to content

Commit babb54a

Browse files
committed
Checkpoint on flambda2 implementation
1 parent 1236295 commit babb54a

File tree

3 files changed

+207
-33
lines changed

3 files changed

+207
-33
lines changed

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

Lines changed: 81 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,37 @@ let convert_array_kind (kind : L.array_kind) : converted_array_kind =
131131
| Pintarray -> Array_kind Immediates
132132
| Pfloatarray -> Array_kind Naked_floats
133133

134+
type converted_array_ref_kind =
135+
| Array_ref_kind of P.Array_ref_kind.t
136+
| Float_array_opt_dynamic_ref of Alloc_mode.For_allocations.t
137+
138+
let convert_array_ref_kind
139+
~current_region (kind : L.array_ref_kind) : converted_array_ref_kind =
140+
match kind with
141+
| Pgenarray_ref mode ->
142+
check_float_array_optimisation_enabled ();
143+
Float_array_opt_dynamic_ref
144+
(Alloc_mode.For_allocations.from_lambda ~current_region mode)
145+
| Paddrarray_ref -> Array_ref_kind Values
146+
| Pintarray_ref -> Array_ref_kind Immediates
147+
| Pfloatarray_ref mode ->
148+
Array_ref_kind
149+
(Naked_floats (Alloc_mode.For_allocations.from_lambda ~current_region mode))
150+
151+
type converted_array_set_kind =
152+
| Array_set_kind of P.Array_set_kind.t
153+
| Float_array_opt_dynamic_set of Alloc_mode.For_assignments.t
154+
155+
let convert_array_set_kind (kind : L.array_set_kind) : converted_array_set_kind =
156+
match kind with
157+
| Pgenarray_set mode ->
158+
check_float_array_optimisation_enabled ();
159+
Float_array_opt_dynamic_set (Alloc_mode.For_assignments.from_lambda mode)
160+
| Paddrarray_set mode ->
161+
Array_set_kind (Values (Alloc_mode.For_assignments.from_lambda mode))
162+
| Pintarray_set -> Array_set_kind Immediates
163+
| Pfloatarray_set -> Array_set_kind Naked_floats
164+
134165
type converted_duplicate_array_kind =
135166
| Duplicate_array_kind of P.Duplicate_array_kind.t
136167
| Float_array_opt_dynamic
@@ -491,20 +522,26 @@ let check_array_access ~dbg ~array ~index primitive : H.expr_primitive =
491522
~conditions:(array_access_validity_condition array index)
492523
~dbg
493524

494-
let array_load_unsafe ~array ~index (array_kind : P.Array_kind.t)
495-
~current_region : H.expr_primitive =
496-
match array_kind with
525+
let array_load_unsafe ~array ~index (array_ref_kind : P.Array_ref_kind.t)
526+
~current_region : H.expr_primitive =
527+
match array_ref_kind with
497528
| Immediates | Values ->
498-
Binary (Array_load (array_kind, Mutable), array, index)
499-
| Naked_floats ->
500-
box_float L.alloc_heap
529+
Binary (Array_load (array_ref_kind, Mutable), array, index)
530+
| Naked_floats mode ->
531+
box_float mode
501532
(Binary (Array_load (Naked_floats, Mutable), array, index))
502533
~current_region
503534

504-
let array_set_unsafe ~array ~index ~new_value (array_kind : P.Array_kind.t) :
535+
let array_set_unsafe ~array ~index ~new_value (array_set_kind : P.Array_set_kind.t) :
505536
H.expr_primitive =
506-
match array_kind with
507-
| Immediates | Values ->
537+
match array_set_kind with
538+
| Immediates
539+
Ternary
540+
( Array_set (array_kind, Assignment Alloc_mode.For_assignments.heap),
541+
array,
542+
index,
543+
new_value )
544+
| Values ->
508545
Ternary
509546
( Array_set (array_kind, Assignment Alloc_mode.For_assignments.heap),
510547
array,
@@ -518,17 +555,29 @@ let array_set_unsafe ~array ~index ~new_value (array_kind : P.Array_kind.t) :
518555
index,
519556
unbox_float new_value )
520557

521-
let[@inline always] match_on_array_kind ~array array_kind f : H.expr_primitive =
522-
match convert_array_kind array_kind with
523-
| Array_kind ((Immediates | Values) as array_kind) -> f array_kind
524-
| Array_kind Naked_floats -> f P.Array_kind.Naked_floats
525-
| Float_array_opt_dynamic ->
558+
let[@inline always] match_on_array_ref_kind ~current_region ~array array_ref_kind f
559+
: H.expr_primitive =
560+
match convert_array_ref_kind ~current_region array_ref_kind with
561+
| Array_ref_kind array_ref_kind -> f array_ref_kind
562+
| Float_array_opt_dynamic_ref mode ->
563+
(* CR keryan: we should push the ITE as low as possible to avoid duplicating
564+
too much *)
565+
If_then_else
566+
( Unary (Is_flat_float_array, array),
567+
f (P.Array_ref_kind.Naked_floats mode),
568+
f P.Array_ref_kind.Values )
569+
570+
let[@inline always] match_on_array_set_kind ~array array_ref_kind f
571+
: H.expr_primitive =
572+
match convert_array_set_kind array_ref_kind with
573+
| Array_set_kind array_set_kind -> f array_set_kind
574+
| Float_array_opt_dynamic_set mode ->
526575
(* CR keryan: we should push the ITE as low as possible to avoid duplicating
527576
too much *)
528577
If_then_else
529578
( Unary (Is_flat_float_array, array),
530-
f P.Array_kind.Naked_floats,
531-
f P.Array_kind.Values )
579+
f P.Array_set_kind.Naked_floats,
580+
f (P.Array_set_kind.Values mode) )
532581

533582
(* Safe arith (div/mod by zero) *)
534583
let checked_arith_op ~dbg (bi : Lambda.boxed_integer option) op mode arg1 arg2
@@ -969,22 +1018,22 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list)
9691018
| Pmodbint { size = Pnativeint; is_safe = Safe; mode }, [arg1; arg2] ->
9701019
checked_arith_op ~dbg (Some Pnativeint) Mod (Some mode) arg1 arg2
9711020
~current_region
972-
| Parrayrefu array_kind, [array; index] ->
1021+
| Parrayrefu array_ref_kind, [array; index] ->
9731022
(* For this and the following cases we will end up relying on the backend to
9741023
CSE the two accesses to the array's header word in the [Pgenarray]
9751024
case. *)
976-
match_on_array_kind ~array array_kind
1025+
match_on_array_ref_kind ~current_region ~array array_ref_kind
9771026
(array_load_unsafe ~array ~index ~current_region)
978-
| Parrayrefs array_kind, [array; index] ->
1027+
| Parrayrefs array_ref_kind, [array; index] ->
9791028
check_array_access ~dbg ~array ~index
980-
(match_on_array_kind ~array array_kind
1029+
(match_on_array_ref_kind ~current_region ~array array_ref_kind
9811030
(array_load_unsafe ~array ~index ~current_region))
982-
| Parraysetu array_kind, [array; index; new_value] ->
983-
match_on_array_kind ~array array_kind
1031+
| Parraysetu array_set_kind, [array; index; new_value] ->
1032+
match_on_array_set_kind ~array array_set_kind
9841033
(array_set_unsafe ~array ~index ~new_value)
985-
| Parraysets array_kind, [array; index; new_value] ->
1034+
| Parraysets array_set_kind, [array; index; new_value] ->
9861035
check_array_access ~dbg ~array ~index
987-
(match_on_array_kind ~array array_kind
1036+
(match_on_array_set_kind ~array array_set_kind
9881037
(array_set_unsafe ~array ~index ~new_value))
9891038
| Pbytessetu (* unsafe *), [bytes; index; new_value] ->
9901039
bytes_like_set_unsafe ~access_size:Eight Bytes bytes index new_value
@@ -1194,17 +1243,21 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list)
11941243
| Plsrbint _ | Pasrbint _ | Pfield_computed _ | Pdivbint _ | Pmodbint _
11951244
| Psetfloatfield _ | Pbintcomp _ | Pbigstring_load_16 _
11961245
| Pbigstring_load_32 _ | Pbigstring_load_64 _
1197-
| Parrayrefu (Pgenarray | Paddrarray | Pintarray | Pfloatarray)
1198-
| Parrayrefs (Pgenarray | Paddrarray | Pintarray | Pfloatarray)
1246+
| Parrayrefu
1247+
(Pgenarray_ref _ | Paddrarray_ref | Pintarray_ref | Pfloatarray_ref _)
1248+
| Parrayrefs
1249+
(Pgenarray_ref _ | Paddrarray_ref | Pintarray_ref | Pfloatarray_ref _)
11991250
| Pcompare_ints | Pcompare_floats | Pcompare_bints _ ),
12001251
([] | [_] | _ :: _ :: _ :: _) ) ->
12011252
Misc.fatal_errorf
12021253
"Closure_conversion.convert_primitive: Wrong arity for binary primitive \
12031254
%a (%a)"
12041255
Printlambda.primitive prim H.print_list_of_simple_or_prim args
12051256
| ( ( Psetfield_computed _ | Pbytessetu | Pbytessets
1206-
| Parraysetu (Pgenarray | Paddrarray | Pintarray | Pfloatarray)
1207-
| Parraysets (Pgenarray | Paddrarray | Pintarray | Pfloatarray)
1257+
| Parraysetu
1258+
(Pgenarray_set _ | Paddrarray_set _ | Pintarray_set | Pfloatarray_set)
1259+
| Parraysets
1260+
(Pgenarray_set _ | Paddrarray_set _ | Pintarray_set | Pfloatarray_set)
12081261
| Pbytes_set_16 _ | Pbytes_set_32 _ | Pbytes_set_64 _
12091262
| Pbigstring_set_16 _ | Pbigstring_set_32 _ | Pbigstring_set_64 _ ),
12101263
([] | [_] | [_; _] | _ :: _ :: _ :: _ :: _) ) ->

middle_end/flambda2/terms/flambda_primitive.ml

Lines changed: 85 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,87 @@ module Array_kind = struct
9494
| Naked_floats -> Flambda_kind.With_subkind.naked_float
9595
end
9696

97+
module Array_ref_kind = struct
98+
type t =
99+
| Immediates
100+
| Values
101+
| Naked_floats of Alloc_mode.For_allocations.t
102+
103+
let [@ocamlformat "disable"] print ppf t =
104+
match t with
105+
| Immediates -> Format.pp_print_string ppf "Immediates"
106+
| Values -> Format.pp_print_string ppf "Values"
107+
| Naked_floats mode ->
108+
Format.fprintf ppf "@[<hov 1>(Naked_floats %a)@]"
109+
Alloc_mode.For_allocations.print mode
110+
111+
let compare = Stdlib.compare
112+
113+
let element_kind_for_set t =
114+
match t with
115+
| Immediates | Values -> K.value
116+
| Naked_floats _ -> K.naked_float
117+
118+
let element_kind_for_creation = element_kind_for_set
119+
120+
let element_kind_for_load = element_kind_for_set
121+
122+
let to_lambda t : Lambda.array_ref_kind =
123+
match t with
124+
| Immediates -> Pintarray_ref
125+
| Values -> Paddrarray_ref
126+
| Naked_floats mode ->
127+
Pfloatarray_ref (match mode with
128+
| Heap -> Lambda.alloc_heap
129+
| Local _ -> Lambda.alloc_local)
130+
131+
let element_kind t =
132+
match t with
133+
| Immediates -> Flambda_kind.With_subkind.tagged_immediate
134+
| Values -> Flambda_kind.With_subkind.any_value
135+
| Naked_floats _ -> Flambda_kind.With_subkind.naked_float
136+
end
137+
138+
module Array_set_kind = struct
139+
type t =
140+
| Immediates
141+
| Values of Alloc_mode.For_assignments.t
142+
| Naked_floats
143+
144+
let [@ocamlformat "disable"] print ppf t =
145+
match t with
146+
| Immediates -> Format.pp_print_string ppf "Immediates"
147+
| Values mode ->
148+
Format.fprintf ppf "@[<hov 1>(Values %a)@]"
149+
Alloc_mode.For_assignments.print mode
150+
| Naked_floats -> Format.fprintf ppf "Naked_floats"
151+
152+
let compare = Stdlib.compare
153+
154+
let element_kind_for_set t =
155+
match t with
156+
| Immediates | Values _ -> K.value
157+
| Naked_floats -> K.naked_float
158+
159+
let element_kind_for_creation = element_kind_for_set
160+
161+
let element_kind_for_load = element_kind_for_set
162+
163+
let to_lambda t : Lambda.array_set_kind =
164+
match t with
165+
| Immediates -> Pintarray_set
166+
| Values mode -> Paddrarray_set (match mode with
167+
| Heap -> Lambda.modify_heap
168+
| Local -> Lambda.modify_maybe_stack)
169+
| Naked_floats -> Pfloatarray_set
170+
171+
let element_kind t =
172+
match t with
173+
| Immediates -> Flambda_kind.With_subkind.tagged_immediate
174+
| Values _ -> Flambda_kind.With_subkind.any_value
175+
| Naked_floats -> Flambda_kind.With_subkind.naked_float
176+
end
177+
97178
module Duplicate_block_kind = struct
98179
type t =
99180
| Values of
@@ -1321,7 +1402,7 @@ let ids_for_export_binary_primitive p =
13211402

13221403
type ternary_primitive =
13231404
| Block_set of Block_access_kind.t * Init_or_assign.t
1324-
| Array_set of Array_kind.t * Init_or_assign.t
1405+
| Array_set of Array_set_kind.t * Init_or_assign.t
13251406
| Bytes_or_bigstring_set of bytes_like_value * string_accessor_width
13261407
| Bigarray_set of num_dimensions * Bigarray_kind.t * Bigarray_layout.t
13271408

@@ -1372,8 +1453,8 @@ let print_ternary_primitive ppf p =
13721453
fprintf ppf "(Block_set %a %a)" Block_access_kind.print kind
13731454
Init_or_assign.print init
13741455
| Array_set (kind, init) ->
1375-
fprintf ppf "(Array_set %a %a)" Array_kind.print kind Init_or_assign.print
1376-
init
1456+
fprintf ppf "(Array_set %a %a)"
1457+
Array_set_kind.print kind Init_or_assign.print init
13771458
| Bytes_or_bigstring_set (kind, string_accessor_width) ->
13781459
fprintf ppf "(Bytes_set %a %a)" print_bytes_like_value kind
13791460
print_string_accessor_width string_accessor_width
@@ -1389,7 +1470,7 @@ let args_kind_of_ternary_primitive p =
13891470
block_index_kind,
13901471
Block_access_kind.element_kind_for_set access_kind )
13911472
| Array_set (kind, _) ->
1392-
array_kind, array_index_kind, Array_kind.element_kind_for_set kind
1473+
array_kind, array_index_kind, Array_set_kind.element_kind_for_set kind
13931474
| Bytes_or_bigstring_set (Bytes, (Eight | Sixteen)) ->
13941475
string_or_bytes_kind, bytes_or_bigstring_index_kind, K.naked_immediate
13951476
| Bytes_or_bigstring_set (Bytes, Thirty_two) ->

middle_end/flambda2/terms/flambda_primitive.mli

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,44 @@ module Array_kind : sig
5353
val element_kind : t -> Flambda_kind.With_subkind.t
5454
end
5555

56+
module Array_ref_kind : sig
57+
type t =
58+
| Immediates (** An array consisting only of immediate values. *)
59+
| Values
60+
(** An array consisting of elements of kind [value]. With the float
61+
array optimisation enabled, such elements must never be [float]s. *)
62+
| Naked_floats of Alloc_mode.For_allocations.t
63+
(** An array consisting of naked floats, represented using
64+
[Double_array_tag]. *)
65+
66+
val print : Format.formatter -> t -> unit
67+
68+
val compare : t -> t -> int
69+
70+
val to_lambda : t -> Lambda.array_ref_kind
71+
72+
val element_kind : t -> Flambda_kind.With_subkind.t
73+
end
74+
75+
module Array_set_kind : sig
76+
type t =
77+
| Immediates (** An array consisting only of immediate values. *)
78+
| Values of Alloc_mode.For_assignments.t
79+
(** An array consisting of elements of kind [value]. With the float
80+
array optimisation enabled, such elements must never be [float]s. *)
81+
| Naked_floats
82+
(** An array consisting of naked floats, represented using
83+
[Double_array_tag]. *)
84+
85+
val print : Format.formatter -> t -> unit
86+
87+
val compare : t -> t -> int
88+
89+
val to_lambda : t -> Lambda.array_set_kind
90+
91+
val element_kind : t -> Flambda_kind.With_subkind.t
92+
end
93+
5694
module Duplicate_block_kind : sig
5795
type t =
5896
| Values of
@@ -338,6 +376,8 @@ type binary_float_arith_op =
338376
type binary_primitive =
339377
| Block_load of Block_access_kind.t * Mutability.t
340378
| Array_load of Array_kind.t * Mutability.t
379+
(** Doesn't need [Array_ref_kind.t] because it doesn't box any returned
380+
[float]s. *)
341381
| String_or_bigstring_load of string_like_value * string_accessor_width
342382
| Bigarray_load of num_dimensions * Bigarray_kind.t * Bigarray_layout.t
343383
| Phys_equal of equality_comparison
@@ -352,7 +392,7 @@ type binary_primitive =
352392
(** Primitives taking exactly three arguments. *)
353393
type ternary_primitive =
354394
| Block_set of Block_access_kind.t * Init_or_assign.t
355-
| Array_set of Array_kind.t * Init_or_assign.t
395+
| Array_set of Array_set_kind.t * Init_or_assign.t
356396
| Bytes_or_bigstring_set of bytes_like_value * string_accessor_width
357397
| Bigarray_set of num_dimensions * Bigarray_kind.t * Bigarray_layout.t
358398

0 commit comments

Comments
 (0)