(**************************************************************************) (* *) (* OCaml *) (* *) (* Pierre Chambart, OCamlPro *) (* Mark Shinwell and Leo White, Jane Street Europe *) (* *) (* Copyright 2013--2016 OCamlPro SAS *) (* Copyright 2014--2016 Jane Street Group LLC *) (* *) (* All rights reserved. This file is distributed under the terms of *) (* the GNU Lesser General Public License version 2.1, with the *) (* special exception on linking described in the file LICENSE. *) (* *) (**************************************************************************) [@@@ocaml.warning "+a-4-9-30-40-41-42-66"] open! Int_replace_polymorphic_compare module E = Inline_and_simplify_aux.Env module B = Inlining_cost.Benefit module Definition = struct type t = | Existing_inner_free_var of Variable.t | Projection_from_existing_specialised_arg of Projection.t include Identifiable.Make (struct type nonrec t = t let compare t1 t2 = match t1, t2 with | Existing_inner_free_var var1, Existing_inner_free_var var2 -> Variable.compare var1 var2 | Projection_from_existing_specialised_arg proj1, Projection_from_existing_specialised_arg proj2 -> Projection.compare proj1 proj2 | Existing_inner_free_var _, _ -> -1 | _, Existing_inner_free_var _ -> 1 let equal t1 t2 = (compare t1 t2) = 0 let hash = Hashtbl.hash let print ppf t = match t with | Existing_inner_free_var var -> Format.fprintf ppf "Existing_inner_free_var %a" Variable.print var | Projection_from_existing_specialised_arg projection -> Format.fprintf ppf "Projection_from_existing_specialised_arg %a" Projection.print projection let output _ _ = failwith "Definition.output not yet implemented" end) end module What_to_specialise = struct type t = { (* [definitions] is indexed by (fun_var, group) *) definitions : Definition.t list Variable.Pair.Map.t; set_of_closures : Flambda.set_of_closures; make_direct_call_surrogates_for : Variable.Set.t; } let create ~set_of_closures = { definitions = Variable.Pair.Map.empty; set_of_closures; make_direct_call_surrogates_for = Variable.Set.empty; } let new_specialised_arg t ~fun_var ~group ~definition = let key = fun_var, group in let definitions = match Variable.Pair.Map.find key t.definitions with | exception Not_found -> [] | definitions -> definitions in let definitions = Variable.Pair.Map.add (fun_var, group) (definition :: definitions) t.definitions in { t with definitions; } let make_direct_call_surrogate_for t ~fun_var = match Variable.Map.find fun_var t.set_of_closures.function_decls.funs with | exception Not_found -> Misc.fatal_errorf "use_direct_call_surrogate_for: %a is not a fun_var \ from the given set of closures" Variable.print fun_var | _ -> { t with make_direct_call_surrogates_for = Variable.Set.add fun_var t.make_direct_call_surrogates_for; } end module W = What_to_specialise module type S = sig val pass_name : string val what_to_specialise : env:Inline_and_simplify_aux.Env.t -> set_of_closures:Flambda.set_of_closures -> What_to_specialise.t end module Processed_what_to_specialise = struct type for_one_function = { fun_var : Variable.t; function_decl : Flambda.function_declaration; make_direct_call_surrogates : bool; new_definitions_indexed_by_new_inner_vars : Definition.t Variable.Map.t; all_new_definitions : Definition.Set.t; new_inner_to_new_outer_vars : Variable.t Variable.Map.t; total_number_of_args : int; existing_specialised_args : Flambda.specialised_to Variable.Map.t; } type t = { set_of_closures : Flambda.set_of_closures; existing_definitions_via_spec_args_indexed_by_fun_var : Definition.Set.t Variable.Map.t; (* The following two maps' definitions have already been rewritten into their lifted form (i.e. they reference outer rather than inner variables). *) new_lifted_defns_indexed_by_new_outer_vars : Projection.t Variable.Map.t; new_outer_vars_indexed_by_new_lifted_defns : Variable.t Projection.Map.t; functions : for_one_function Variable.Map.t; make_direct_call_surrogates_for : Variable.Set.t; } let lift_projection t ~(projection : Projection.t) = (* The lifted definition must be in terms of outer variables, not inner variables. *) let find_outer_var inner_var = match Variable.Map.find inner_var t.set_of_closures.specialised_args with | (outer_var : Flambda.specialised_to) -> outer_var.var | exception Not_found -> Misc.fatal_errorf "find_outer_var: expected %a \ to be in [specialised_args], but it is \ not. The projection was: %a. Set of closures: %a" Variable.print inner_var Projection.print projection Flambda.print_set_of_closures t.set_of_closures in Projection.map_projecting_from projection ~f:find_outer_var let really_add_new_specialised_arg t ~group ~(definition : Definition.t) ~(for_one_function : for_one_function) = let fun_var = for_one_function.fun_var in (* We know here that a new specialised argument must be added. This needs a "new inner var" and a "new outer var". However if there is already a lifted projection being introduced around the set of closures (corresponding to another new specialised argument), we should re-use its "new outer var" to avoid duplication of projection definitions. Likewise if the definition is just [Existing_inner_free_var], in which case we can use the corresponding existing outer free variable. *) let new_outer_var, t = let existing_outer_var = match definition with | Existing_inner_free_var _ -> None | Projection_from_existing_specialised_arg projection -> let projection = lift_projection t ~projection in match Projection.Map.find projection t.new_outer_vars_indexed_by_new_lifted_defns with | new_outer_var -> Some new_outer_var | exception Not_found -> None in match existing_outer_var with | Some existing_outer_var -> existing_outer_var, t | None -> match definition with | Existing_inner_free_var existing_inner_var -> begin match Variable.Map.find existing_inner_var t.set_of_closures.free_vars with | exception Not_found -> Misc.fatal_errorf "really_add_new_specialised_arg: \ Existing_inner_free_var %a is not an inner free variable \ of %a in %a" Variable.print existing_inner_var Variable.print fun_var Flambda.print_set_of_closures t.set_of_closures | existing_outer_var -> existing_outer_var.var, t end | Projection_from_existing_specialised_arg projection -> let new_outer_var = Variable.rename group in let projection = lift_projection t ~projection in let new_outer_vars_indexed_by_new_lifted_defns = Projection.Map.add projection new_outer_var t.new_outer_vars_indexed_by_new_lifted_defns in let new_lifted_defns_indexed_by_new_outer_vars = Variable.Map.add new_outer_var projection t.new_lifted_defns_indexed_by_new_outer_vars in let t = { t with new_outer_vars_indexed_by_new_lifted_defns; new_lifted_defns_indexed_by_new_outer_vars; } in new_outer_var, t in let new_inner_var = Variable.rename group in let new_inner_to_new_outer_vars = Variable.Map.add new_inner_var new_outer_var for_one_function.new_inner_to_new_outer_vars in let for_one_function : for_one_function = { for_one_function with new_definitions_indexed_by_new_inner_vars = Variable.Map.add new_inner_var definition for_one_function.new_definitions_indexed_by_new_inner_vars; all_new_definitions = Definition.Set.add definition for_one_function.all_new_definitions; new_inner_to_new_outer_vars; total_number_of_args = for_one_function.total_number_of_args + 1; } in { t with functions = Variable.Map.add fun_var for_one_function t.functions; } let new_specialised_arg t ~fun_var ~group ~definition = let for_one_function : for_one_function = match Variable.Map.find fun_var t.functions with | exception Not_found -> begin match Variable.Map.find fun_var t.set_of_closures.function_decls.funs with | exception Not_found -> assert false | (function_decl : Flambda.function_declaration) -> let params = Parameter.Set.vars function_decl.params in let existing_specialised_args = Variable.Map.filter (fun inner_var _spec_to -> Variable.Set.mem inner_var params) t.set_of_closures.specialised_args in let make_direct_call_surrogates = Variable.Set.mem fun_var t.make_direct_call_surrogates_for in { fun_var; function_decl; make_direct_call_surrogates; new_definitions_indexed_by_new_inner_vars = Variable.Map.empty; all_new_definitions = Definition.Set.empty; new_inner_to_new_outer_vars = Variable.Map.empty; (* The "+ 1" is just in case there is a closure environment parameter added later. *) total_number_of_args = List.length function_decl.params + 1; existing_specialised_args; } end | for_one_function -> for_one_function in (* Determine whether there already exists an existing specialised argument that is known to be equal to the one proposed to this function. If so, use that instead. (Note that we also desire to dedup against any new specialised arguments added to the current function; but that happens automatically since [Extract_projections] returns a set.) *) let exists_already = match Variable.Map.find fun_var t.existing_definitions_via_spec_args_indexed_by_fun_var with | exception Not_found -> false | definitions -> Definition.Set.mem definition definitions in if exists_already then t else really_add_new_specialised_arg t ~group ~definition ~for_one_function let create ~env ~(what_to_specialise : W.t) = let existing_definitions_via_spec_args_indexed_by_fun_var = Variable.Map.map (fun (function_decl : Flambda.function_declaration) -> if function_decl.stub then Definition.Set.empty else let params = Parameter.Set.vars function_decl.params in Variable.Map.fold (fun inner_var (spec_to : Flambda.specialised_to) definitions -> if not (Variable.Set.mem inner_var params) then definitions else let definition : Definition.t = match spec_to.projection with | None -> Existing_inner_free_var inner_var | Some projection -> Projection_from_existing_specialised_arg projection in Definition.Set.add definition definitions) what_to_specialise.set_of_closures.specialised_args Definition.Set.empty) what_to_specialise.set_of_closures.function_decls.funs in let t : t = { set_of_closures = what_to_specialise.set_of_closures; existing_definitions_via_spec_args_indexed_by_fun_var; new_lifted_defns_indexed_by_new_outer_vars = Variable.Map.empty; new_outer_vars_indexed_by_new_lifted_defns = Projection.Map.empty; functions = Variable.Map.empty; make_direct_call_surrogates_for = what_to_specialise.make_direct_call_surrogates_for; } in (* It is important to limit the number of arguments added: if arguments end up being passed on the stack, tail call optimization will be disabled (see asmcomp/selectgen.ml). For each group of new specialised args provided by [T], either all or none of them will be added. (This is to avoid the situation where we add extra arguments but yet fail to eliminate an original one by stopping part-way through the specialised args addition.) *) let by_group = Variable.Pair.Map.fold (fun (fun_var, group) definitions by_group -> let fun_vars_and_definitions = match Variable.Map.find group by_group with | exception Not_found -> [] | fun_vars_and_definitions -> fun_vars_and_definitions in Variable.Map.add group ((fun_var, definitions)::fun_vars_and_definitions) by_group) what_to_specialise.definitions Variable.Map.empty in let module Backend = (val (E.backend env) : Backend_intf.S) in Variable.Map.fold (fun group fun_vars_and_definitions t -> let original_t = t in let t = (* Try adding all specialised args in the current group. *) List.fold_left (fun t (fun_var, definitions) -> List.fold_left (fun t definition -> new_specialised_arg t ~fun_var ~group ~definition) t definitions) t fun_vars_and_definitions in let some_function_has_too_many_args = Variable.Map.exists (fun _ (for_one_function : for_one_function) -> for_one_function.total_number_of_args > Backend.max_sensible_number_of_arguments) t.functions in if some_function_has_too_many_args then original_t (* drop this group *) else t) by_group t end module P = Processed_what_to_specialise let check_invariants ~pass_name ~(set_of_closures : Flambda.set_of_closures) ~original_set_of_closures = if !Clflags.flambda_invariant_checks then begin Variable.Map.iter (fun fun_var (function_decl : Flambda.function_declaration) -> let params = Parameter.Set.vars function_decl.params in Variable.Map.iter (fun inner_var (outer_var : Flambda.specialised_to) -> if Variable.Set.mem inner_var params then begin assert (not (Variable.Set.mem outer_var.var function_decl.free_variables)); match outer_var.projection with | None -> () | Some projection -> let from = Projection.projecting_from projection in if not (Variable.Set.mem from params) then begin Misc.fatal_errorf "Augment_specialised_args (%s): \ specialised argument (%a -> %a) references a \ projection variable that is not a specialised \ argument of the function %a. @ The set of closures \ before the transformation was:@ %a. @ The set of \ closures after the transformation was:@ %a." pass_name Variable.print inner_var Flambda.print_specialised_to outer_var Variable.print fun_var Flambda.print_set_of_closures original_set_of_closures Flambda.print_set_of_closures set_of_closures end end) set_of_closures.specialised_args) set_of_closures.function_decls.funs end module Make (T : S) = struct let () = Pass_wrapper.register ~pass_name:T.pass_name let rename_function_and_parameters ~fun_var ~(function_decl : Flambda.function_declaration) = let new_fun_var = Variable.rename fun_var in let params_renaming_list = List.map (fun param -> let new_param = Parameter.rename param in param, new_param) function_decl.params in let renamed_params = List.map snd params_renaming_list in let params_renaming = Variable.Map.of_list (List.map (fun (param, new_param) -> Parameter.var param, Parameter.var new_param) params_renaming_list) in new_fun_var, params_renaming, renamed_params let create_wrapper ~(for_one_function : P.for_one_function) ~benefit = let fun_var = for_one_function.fun_var in let function_decl = for_one_function.function_decl in (* To avoid increasing the free variables of the wrapper, for general cleanliness, we restate the definitions of the newly-specialised arguments in the wrapper itself in terms of the original specialised arguments. The variables bound to these definitions are called the "specialised args bound in the wrapper". Note that the domain of [params_renaming] is a (non-strict) superset of the "inner vars" of the original specialised args. *) let params = Parameter.Set.vars function_decl.params in let new_fun_var, params_renaming, wrapper_params = rename_function_and_parameters ~fun_var ~function_decl in let find_wrapper_param param = assert (Variable.Set.mem param params); match Variable.Map.find param params_renaming with | wrapper_param -> wrapper_param | exception Not_found -> Misc.fatal_errorf "find_wrapper_param: expected %a \ to be in [params_renaming], but it is not." Variable.print param in let new_inner_vars_to_spec_args_bound_in_the_wrapper_renaming = Variable.Map.mapi (fun new_inner_var _ -> Variable.rename new_inner_var) for_one_function.new_definitions_indexed_by_new_inner_vars in let spec_args_bound_in_the_wrapper = (* N.B.: in the order matching the new specialised argument parameters to the main function. *) Variable.Map.data new_inner_vars_to_spec_args_bound_in_the_wrapper_renaming in (* New definitions that project from existing specialised args need to be rewritten to use the corresponding specialised args of the wrapper. Definitions that are just equality to existing inner free variables do not need to be changed. Once this has been done the wrapper body can be constructed. We also need to rewrite definitions for any existing specialised args; these now have corresponding wrapper parameters that must also be specialised. *) let wrapper_body, benefit = let apply : Flambda.expr = Apply { func = new_fun_var; args = (Parameter.List.vars wrapper_params) @ spec_args_bound_in_the_wrapper; kind = Direct (Closure_id.wrap new_fun_var); dbg = Debuginfo.none; inline = Default_inline; specialise = Default_specialise; } in Variable.Map.fold (fun new_inner_var definition (wrapper_body, benefit) -> let definition : Definition.t = match (definition : Definition.t) with | Existing_inner_free_var _ -> definition | Projection_from_existing_specialised_arg projection -> Projection_from_existing_specialised_arg (Projection.map_projecting_from projection ~f:find_wrapper_param) in let benefit = match (definition : Definition.t) with | Existing_inner_free_var _ -> benefit | Projection_from_existing_specialised_arg projection -> B.add_projection projection benefit in match Variable.Map.find new_inner_var new_inner_vars_to_spec_args_bound_in_the_wrapper_renaming with | exception Not_found -> assert false | new_inner_var_of_wrapper -> let named : Flambda.named = match definition with | Existing_inner_free_var existing_inner_var -> Expr (Var existing_inner_var) | Projection_from_existing_specialised_arg projection -> Flambda_utils.projection_to_named projection in let wrapper_body = Flambda.create_let new_inner_var_of_wrapper named wrapper_body in (wrapper_body, benefit)) for_one_function.new_definitions_indexed_by_new_inner_vars (apply, benefit) in let rewritten_existing_specialised_args = Variable.Map.fold (fun inner_var (spec_to : Flambda.specialised_to) result -> let inner_var = find_wrapper_param inner_var in let projection = match spec_to.projection with | None -> None | Some projection -> Some (Projection.map_projecting_from projection ~f:find_wrapper_param) in let spec_to : Flambda.specialised_to = { var = spec_to.var; projection; } in Variable.Map.add inner_var spec_to result) for_one_function.existing_specialised_args Variable.Map.empty in let new_function_decl = Flambda.create_function_declaration ~params:wrapper_params ~body:wrapper_body ~stub:true ~dbg:Debuginfo.none ~inline:Default_inline ~specialise:Default_specialise ~is_a_functor:false ~closure_origin:function_decl.closure_origin in new_fun_var, new_function_decl, rewritten_existing_specialised_args, benefit let rewrite_function_decl (t : P.t) ~env ~duplicate_function ~(for_one_function : P.for_one_function) ~benefit = let set_of_closures = t.set_of_closures in let fun_var = for_one_function.fun_var in let function_decl = for_one_function.function_decl in let num_definitions = Variable.Map.cardinal for_one_function. new_definitions_indexed_by_new_inner_vars in if function_decl.stub || num_definitions < 1 || Variable.Map.mem fun_var set_of_closures.direct_call_surrogates then None else let new_fun_var, wrapper, rewritten_existing_specialised_args, benefit = create_wrapper ~for_one_function ~benefit in let new_specialised_args = Variable.Map.mapi (fun new_inner_var (definition : Definition.t) : Flambda.specialised_to -> assert (not (Variable.Map.mem new_inner_var set_of_closures.specialised_args)); match Variable.Map.find new_inner_var for_one_function.new_inner_to_new_outer_vars with | exception Not_found -> assert false | new_outer_var -> match definition with | Existing_inner_free_var _ -> { var = new_outer_var; projection = None; } | Projection_from_existing_specialised_arg projection -> let projecting_from = Projection.projecting_from projection in assert (Variable.Map.mem projecting_from set_of_closures.specialised_args); assert (Variable.Set.mem projecting_from (Parameter.Set.vars function_decl.params)); { var = new_outer_var; projection = Some projection; }) for_one_function.new_definitions_indexed_by_new_inner_vars in let specialised_args = Variable.Map.disjoint_union rewritten_existing_specialised_args new_specialised_args in let specialised_args, existing_function_decl = if not for_one_function.make_direct_call_surrogates then specialised_args, None else let function_decl, new_specialised_args = duplicate_function ~env ~set_of_closures ~fun_var ~new_fun_var in let specialised_args = Variable.Map.disjoint_union specialised_args new_specialised_args in specialised_args, Some function_decl in let all_params = let new_params = Variable.Set.elements (Variable.Map.keys for_one_function.new_inner_to_new_outer_vars) in let new_params = List.map Parameter.wrap new_params in function_decl.params @ new_params in let closure_origin = Closure_origin.create (Closure_id.wrap new_fun_var) in let rewritten_function_decl = Flambda.create_function_declaration ~params:all_params ~body:function_decl.body ~stub:function_decl.stub ~dbg:function_decl.dbg ~inline:function_decl.inline ~specialise:function_decl.specialise ~is_a_functor:function_decl.is_a_functor ~closure_origin in let funs, direct_call_surrogates = if for_one_function.make_direct_call_surrogates then let surrogate = Variable.rename fun_var in let funs = (* In this case, the original function declaration remains untouched up to alpha-equivalence. Direct calls to it (including inside the rewritten original function) will be replaced by calls to the surrogate (i.e. the wrapper) which will then be inlined. *) let existing_function_decl = match existing_function_decl with | Some decl -> decl | None -> assert false in Variable.Map.add new_fun_var rewritten_function_decl (Variable.Map.add surrogate wrapper (Variable.Map.add fun_var existing_function_decl Variable.Map.empty)) in let direct_call_surrogates = Variable.Map.add fun_var surrogate Variable.Map.empty in funs, direct_call_surrogates else let funs = Variable.Map.add new_fun_var rewritten_function_decl (Variable.Map.add fun_var wrapper Variable.Map.empty) in funs, Variable.Map.empty in let free_vars = Variable.Map.empty in Some (funs, free_vars, specialised_args, direct_call_surrogates, benefit) let add_lifted_projections_around_set_of_closures ~(set_of_closures : Flambda.set_of_closures) ~benefit ~new_lifted_defns_indexed_by_new_outer_vars = let body = Flambda_utils.name_expr ~name:Internal_variable_names.set_of_closures (Set_of_closures set_of_closures) in Variable.Map.fold (fun new_outer_var (projection : Projection.t) (expr, benefit) -> let named = Flambda_utils.projection_to_named projection in let benefit = B.add_projection projection benefit in let expr = Flambda.create_let new_outer_var named expr in expr, benefit) new_lifted_defns_indexed_by_new_outer_vars (body, benefit) let rewrite_set_of_closures_core ~env ~duplicate_function ~benefit ~(set_of_closures : Flambda.set_of_closures) = let what_to_specialise = P.create ~env ~what_to_specialise:(T.what_to_specialise ~env ~set_of_closures) in let original_set_of_closures = set_of_closures in let funs, free_vars, specialised_args, direct_call_surrogates, done_something, benefit = Variable.Map.fold (fun fun_var function_decl (funs, free_vars, specialised_args, direct_call_surrogates, done_something, benefit) -> match Variable.Map.find fun_var what_to_specialise.functions with | exception Not_found -> let funs = Variable.Map.add fun_var function_decl funs in funs, free_vars, specialised_args, direct_call_surrogates, done_something, benefit | (for_one_function : P.for_one_function) -> assert (Variable.equal fun_var for_one_function.fun_var); match rewrite_function_decl what_to_specialise ~env ~duplicate_function ~for_one_function ~benefit with | None -> let function_decl = for_one_function.function_decl in let funs = Variable.Map.add fun_var function_decl funs in funs, free_vars, specialised_args, direct_call_surrogates, done_something, benefit | Some (funs', free_vars', specialised_args', direct_call_surrogates', benefit) -> let funs = Variable.Map.disjoint_union funs funs' in let direct_call_surrogates = Variable.Map.disjoint_union direct_call_surrogates direct_call_surrogates' in let free_vars = Variable.Map.disjoint_union free_vars free_vars' in let specialised_args = Variable.Map.disjoint_union specialised_args specialised_args' in funs, free_vars, specialised_args, direct_call_surrogates, true, benefit) set_of_closures.function_decls.funs (Variable.Map.empty, set_of_closures.free_vars, set_of_closures.specialised_args, set_of_closures.direct_call_surrogates, false, benefit) in if not done_something then None else let function_decls = Flambda.update_function_declarations set_of_closures.function_decls ~funs in assert (Variable.Map.cardinal specialised_args >= Variable.Map.cardinal original_set_of_closures.specialised_args); let set_of_closures = Flambda.create_set_of_closures ~function_decls ~free_vars ~specialised_args ~direct_call_surrogates in if !Clflags.flambda_invariant_checks then begin check_invariants ~set_of_closures ~original_set_of_closures ~pass_name:T.pass_name end; let expr, benefit = add_lifted_projections_around_set_of_closures ~set_of_closures ~benefit ~new_lifted_defns_indexed_by_new_outer_vars: what_to_specialise.new_lifted_defns_indexed_by_new_outer_vars in Some (expr, benefit) let rewrite_set_of_closures ~env ~duplicate_function ~set_of_closures = Pass_wrapper.with_dump ~ppf_dump:(Inline_and_simplify_aux.Env.ppf_dump env) ~pass_name:T.pass_name ~input:set_of_closures ~print_input:Flambda.print_set_of_closures ~print_output:(fun ppf (expr, _) -> Flambda.print ppf expr) ~f:(fun () -> rewrite_set_of_closures_core ~env ~duplicate_function ~benefit:B.zero ~set_of_closures) end