diff --git a/middle_end/flambda2/validate/flambda2_core.ml b/middle_end/flambda2/validate/flambda2_core.ml index 46f69cb1722..22a55aaa53a 100644 --- a/middle_end/flambda2/validate/flambda2_core.ml +++ b/middle_end/flambda2/validate/flambda2_core.ml @@ -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) @@ -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 @@ -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 diff --git a/middle_end/flambda2/validate/flambda2_core.mli b/middle_end/flambda2/validate/flambda2_core.mli index 84103a6ef38..3361441a32c 100644 --- a/middle_end/flambda2/validate/flambda2_core.mli +++ b/middle_end/flambda2/validate/flambda2_core.mli @@ -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 @@ -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 diff --git a/middle_end/flambda2/validate/normalize.ml b/middle_end/flambda2/validate/normalize.ml index a990ce698ff..8a50405b901 100644 --- a/middle_end/flambda2/validate/normalize.ml +++ b/middle_end/flambda2/validate/normalize.ml @@ -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) @@ -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]. *) @@ -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 @@ -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)