Skip to content

Commit

Permalink
A couple optimizations around primitives.
Browse files Browse the repository at this point in the history 
1) If a primitive with only generative effects is unused, we can delete it.
2) Special cases for `begin_region` / `end_region` to delete unused regions
  • Loading branch information
@ccasin
ccasin committed May 27, 2024
1 parent 8e113b7 commit 99d56d3
Show file tree
Hide file tree
Showing 3 changed files with 179 additions and 65 deletions.
133 changes: 80 additions & 53 deletions middle_end/flambda2/validate/flambda2_core.ml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1092,11 +1092,17 @@ and ids_for_export_set_of_closures
value_slots function_decls_ids

and ids_for_export_prim (t : primitive) =
match t with
match[@warning "-4"] t with
| Nullary
(Invalid _ | Optimised_out _ | Probe_is_enabled _ | Begin_region
| Enter_inlined_apply _) ->
Ids_for_export.empty
| Unary (End_region, _region) ->
(* Following [simplify/flow/flow_acc.ml], uses of a region in [End_region] don't count
as uses. *)
(* CR: But is that actually OK for all uses of [ids_for_export], or should I have a
separate free_vars function for this? *)
Ids_for_export.empty
| Unary (prim, arg) ->
Ids_for_export.union
(P.ids_for_export_unary_primitive prim)
Expand Down Expand Up @@ -1306,6 +1312,14 @@ module Core_let = struct
Error Pattern_match_pair_error.Mismatched_let_bindings
)
)

let let_var_free_in : Bound_for_let.t -> core_exp -> bool = fun var exp ->
let free_vars = ids_for_export exp in
Bound_for_let.fold_all_bound_names var
~init:false
~var:(fun acc v -> acc || Variable.Set.mem (Bound_var.var v) free_vars.variables)
~symbol:(fun acc v -> acc || Symbol.Set.mem v free_vars.symbols)
~code_id:(fun acc v -> acc || Code_id.Set.mem v free_vars.code_ids)
end

module Core_continuation_handler = struct
Expand Down Expand Up @@ -1562,58 +1576,71 @@ let literal_contained (literal1 : literal) (literal2 : literal) : bool =
| (Simple _ | Cont _ | Slot (_, (Function_slot _ | Value_slot _))
| Res_cont (Never_returns | Return _) | Code_id _), _ -> false

let effects_and_coeffects (p : primitive) =
match p with
| Nullary prim ->
Flambda_primitive.effects_and_coeffects_of_nullary_primitive prim
| Unary (prim, _) ->
Flambda_primitive.effects_and_coeffects_of_unary_primitive prim
| Binary (prim, _, _) ->
Flambda_primitive.effects_and_coeffects_of_binary_primitive prim
| Ternary (prim, _, _, _) ->
Flambda_primitive.effects_and_coeffects_of_ternary_primitive prim
| Variadic (prim, _) ->
Flambda_primitive.effects_and_coeffects_of_variadic_primitive prim

let no_effects (p : primitive) =
match effects_and_coeffects p with
| No_effects, _, _ -> true
| ( (Only_generative_effects _ | Arbitrary_effects),
(No_coeffects | Has_coeffects),
_ ) ->
false

let no_effects (e : core_exp) : bool =
match must_be_prim e with
| None -> true
| Some p -> no_effects p

let can_inline (p : primitive) =
match effects_and_coeffects p with
| No_effects, No_coeffects, _ -> true
| Only_generative_effects _, No_coeffects, _ -> true
| ( (No_effects | Only_generative_effects _ | Arbitrary_effects),
(No_coeffects | Has_coeffects),
_ ) ->
false

let can_inline (e : core_exp) : bool =
match must_be_prim e with
| None -> true
| Some p -> can_inline p

let no_effects_or_coeffects (p : primitive) =
match effects_and_coeffects p with
| No_effects, No_coeffects, _ -> true
| ( (No_effects | Only_generative_effects _ | Arbitrary_effects),
(No_coeffects | Has_coeffects),
_ ) ->
false

let no_effects_or_coeffects (e : core_exp) : bool =
match must_be_prim e with
| None -> true
| Some p -> no_effects_or_coeffects p
module Effects = struct
let effects_and_coeffects (p : primitive) =
match p with
| Nullary prim ->
Flambda_primitive.effects_and_coeffects_of_nullary_primitive prim
| Unary (prim, _) ->
Flambda_primitive.effects_and_coeffects_of_unary_primitive prim
| Binary (prim, _, _) ->
Flambda_primitive.effects_and_coeffects_of_binary_primitive prim
| Ternary (prim, _, _, _) ->
Flambda_primitive.effects_and_coeffects_of_ternary_primitive prim
| Variadic (prim, _) ->
Flambda_primitive.effects_and_coeffects_of_variadic_primitive prim

let no_effects (p : primitive) =
match effects_and_coeffects p with
| No_effects, _, _ -> true
| ( (Only_generative_effects _ | Arbitrary_effects),
(No_coeffects | Has_coeffects),
_ ) ->
false

let no_effects (e : core_exp) : bool =
match must_be_prim e with
| None -> true
| Some p -> no_effects p

let no_effects_or_coeffects (p : primitive) =
match effects_and_coeffects p with
| No_effects, No_coeffects, _ -> true
| ( (No_effects | Only_generative_effects _ | Arbitrary_effects),
(No_coeffects | Has_coeffects),
_ ) ->
false

let no_effects_or_coeffects (e : core_exp) : bool =
match must_be_prim e with
| None -> true
| Some p -> no_effects_or_coeffects p

type substitutability =
| Can_duplicate
(* Things with no co-effects and only generative effects can be substituted freely for
our analysis. Those for things with generative effects, it would be a bug for
flambda2 itself to do so, as this may increase allocation. *)
| Can_delete_if_unused
(* Things with co-effects and only generative effects. These can't be substituted,
because their co-effects mean that reordering them is a change in behavior. But if
their result is not used they can be deleted because they don't have observable
effects. *)
| No_substitutions
(* Things with real effects can't be moved or deleted at all. *)

let can_substitute (p : primitive) =
match effects_and_coeffects p with
| (No_effects | Only_generative_effects _), No_coeffects, _ -> Can_duplicate
| (No_effects | Only_generative_effects _), Has_coeffects, _ ->
Can_delete_if_unused
| Arbitrary_effects, (No_coeffects | Has_coeffects), _ -> No_substitutions

let can_substitute (e : core_exp) =
match must_be_prim e with
| None -> Can_duplicate (* CR ccasinghino: Is this right? *)
| Some p -> can_substitute p
end

let returns_unit (p : primitive) : bool =
match p with
Expand Down
27 changes: 23 additions & 4 deletions middle_end/flambda2/validate/flambda2_core.mli
View file
Original file line number Diff line number Diff line change
Expand Up @@ -216,6 +216,8 @@ module Core_let : sig
t -> t -> (Bound_for_let.t -> core_exp -> core_exp -> 'a) ->
(Bound_codelike.t -> Bound_codelike.t -> core_exp -> core_exp -> 'a) ->
('a, Pattern_match_pair_error.t) Result.t

val let_var_free_in : Bound_for_let.t -> core_exp -> bool
end

module Core_continuation_handler : sig
Expand Down Expand Up @@ -331,8 +333,25 @@ val static_const_group_fix :
val print_set_of_closures : Format.formatter -> set_of_closures -> unit
val literal_contained : literal -> literal -> bool

(* Effects *)
val no_effects_or_coeffects : core_exp -> bool
val no_effects : core_exp -> bool
val can_inline : core_exp -> bool
module Effects : sig
(* Effects *)
val no_effects_or_coeffects : core_exp -> bool
val no_effects : core_exp -> bool

type substitutability =
| Can_duplicate
(* Things with no co-effects and only generative effects can be substituted freely for
our analysis. Those for things with generative effects, it would be a bug for
flambda2 itself to do so, as this may increase allocation. *)
| Can_delete_if_unused
(* Things with co-effects and only generative effects. These can't be substituted,
because their co-effects mean that reordering them is a change in behavior. But if
their result is not used they can be deleted because they don't have observable
effects. *)
| No_substitutions
(* Things with real effects can't be moved or deleted at all. *)

val can_substitute : core_exp -> substitutability
end

val returns_unit : core_exp -> bool
84 changes: 76 additions & 8 deletions middle_end/flambda2/validate/normalize.ml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -339,6 +339,55 @@ let subst_params
Expr.create_named (Literal s))
() e

(* We want to remove [begin_region]/[end_region] pairs of the region is not used
for anything. This is accomplished in two steps:
1) When simplifying [let x = e1 in e2], if [x] is unused in [e2], we can drop
the whole let and just keep [e2]. And we don't count [end_region] as using
its region argument, so we will drop [begin_region] (which is [e1] here) this
way. That happens in [step_let]
2) But that leaves behind an [end_region] for a now undefined region in [e2].
This function checks for that and cleans it up. This is horribly inefficient
- flambda2 actually does a much smarter thing so it doesn't have to
re-traverse the term. *)
let remove_corresponding_end_region region_var e1 e2 =
let[@warning "-4"] is_let_end_region_of region_var e =
Core_let.pattern_match e ~f:(fun ~x:_ ~e1 ~e2 ->
match Expr.descr e1 with
| Named (Prim (Unary (End_region, region))) ->
begin match Expr.descr region with
| Named (Literal (Simple region)) ->
begin match Simple.must_be_var region with
| Some (region_var', _) when Variable.equal region_var region_var' ->
Some e2
| _ -> None
end
| _ -> None
end
| _ -> None)
in

let rec remove_end region_var e2 =
match Expr.descr e2 with
| Named _ | Invalid _-> e2
| Let e ->
begin match is_let_end_region_of region_var e with
| Some e -> e
| None -> let_fix (remove_end region_var) e
end
| Let_cont e -> let_cont_fix (remove_end region_var) e
| Apply e -> apply_fix (remove_end region_var) e
| Apply_cont e -> apply_cont_fix (remove_end region_var) e
| Lambda e -> lambda_fix (remove_end region_var) e
| Handler e -> handler_fix (remove_end region_var) e
| Switch e -> switch_fix (remove_end region_var) e
in
match[@warning "-4"] Expr.descr e1, region_var with
| Named (Prim (Nullary Begin_region)), Bound_for_let.Singleton v ->
remove_end (Bound_var.var v) e2
| _ -> e2

(* [LetCont-β]
e1 where k args = e2 ⟶ e1 [k \ λ args. e2] *)
let rec subst_cont (cont_e1: core_exp) (k: Bound_continuation.t)
Expand Down Expand Up @@ -482,15 +531,16 @@ and step_let let_abst body : core_exp =
step_named_for_let x e
| None -> x, step e1
in
if can_inline e1 then
(* [e1] can be inlined; i.e. it has no side effects. *)
match Effects.can_substitute e1 with
| Can_duplicate ->
(* [e1] can be substituted freely; i.e. it has only generative effects and
does not observe effects. *)
(* [Let-β]
let x = v in e1 ⟶ e2 [x\v] *)
subst_pattern ~bound:x ~let_body:e1 e2 |> step

(* Cannot be inlined; however, if it returns unit, then
value can be replaced with unit. *)
else
| No_substitutions ->
(* Cannot be substituted; however, if it returns unit, then value can be
replaced with unit. *)
let e2 =
if returns_unit e1 then
(* Substitute in unit value in place for [e2]. *)
Expand All @@ -500,7 +550,25 @@ and step_let let_abst body : core_exp =
else e2
in
Core_let.create ~x ~e1 ~e2:(step e2)
)
| Can_delete_if_unused ->
(* Can be deleted if its result is unused. We also apply the unit
optimization from above. *)
let e2 =
if returns_unit e1 then
(* Substitute in unit value in place for [e2]. *)
subst_pattern ~bound:x
~let_body:(Expr.create_named (Literal (Simple
(Simple.const Int_ids.Const.const_zero)))) e2
else e2
in
let e2 = step e2 in
let is_used = Core_let.let_var_free_in x e2 in
if is_used then
Core_let.create ~x ~e1 ~e2
else
(* A very special case for when we've removed a begin_region *)
remove_corresponding_end_region x e1 e2
)

and step_let_cont ({handler; body}:let_cont_expr) : core_exp =
Core_continuation_handler.pattern_match handler
Expand Down Expand Up @@ -822,7 +890,7 @@ and step_named_for_let (var: Bound_for_let.t) (body: named)
(Static (Bound_codelike.create var),
Expr.create_named (Static_consts consts)))
| _ -> (var, Expr.create_named (Static_consts consts))
)[@ocaml.warning "-4"]
)[@ocaml.warning "-4"]
| Prim v -> (var, Eval_prim.eval v)

and concretize_my_closure phi (slot : Function_slot.t)
Expand Down

0 comments on commit 99d56d3

Please sign in to comment.