104 lines
5.2 KiB
OCaml
104 lines
5.2 KiB
OCaml
(**************************************************************************)
|
|
(* *)
|
|
(* 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"]
|
|
|
|
module ASA = Augment_specialised_args
|
|
module W = ASA.What_to_specialise
|
|
|
|
module Transform = struct
|
|
let pass_name = "unbox-specialised-args"
|
|
let variable_suffix = ""
|
|
|
|
let precondition ~env:_ ~(set_of_closures : Flambda.set_of_closures) =
|
|
!Clflags.unbox_specialised_args
|
|
&& not (Variable.Map.is_empty set_of_closures.specialised_args)
|
|
|
|
let what_to_specialise ~env ~(set_of_closures : Flambda.set_of_closures) =
|
|
let what_to_specialise = W.create ~set_of_closures in
|
|
if not (precondition ~env ~set_of_closures) then
|
|
what_to_specialise
|
|
else
|
|
let projections_by_function =
|
|
Variable.Map.filter_map set_of_closures.function_decls.funs
|
|
~f:(fun _fun_var (function_decl : Flambda.function_declaration) ->
|
|
if function_decl.stub then None
|
|
else
|
|
Some (Extract_projections.from_function_decl ~env
|
|
~function_decl
|
|
~which_variables:set_of_closures.specialised_args))
|
|
in
|
|
(* CR-soon mshinwell: consider caching the Invariant_params *relation*
|
|
as well as the "_in_recursion" map *)
|
|
let invariant_params_flow =
|
|
Invariant_params.invariant_param_sources set_of_closures.function_decls
|
|
~backend:(Inline_and_simplify_aux.Env.backend env)
|
|
in
|
|
Variable.Map.fold (fun fun_var extractions what_to_specialise ->
|
|
Projection.Set.fold (fun (projection : Projection.t)
|
|
what_to_specialise ->
|
|
let group = Projection.projecting_from projection in
|
|
assert (Variable.Map.mem group set_of_closures.specialised_args);
|
|
let what_to_specialise =
|
|
W.new_specialised_arg what_to_specialise ~fun_var ~group
|
|
~definition:(Projection_from_existing_specialised_arg
|
|
projection)
|
|
in
|
|
match Variable.Map.find group invariant_params_flow with
|
|
| exception Not_found -> what_to_specialise
|
|
| flow ->
|
|
(* If for function [f] we would extract a projection expression
|
|
[e] from some specialised argument [x] of [f], and we know
|
|
from [Invariant_params] that a specialised argument [y] of
|
|
another function [g] flows to [x], we will add add [e] with
|
|
[y] substituted for [x] throughout as a newly-specialised
|
|
argument for [g]. This should help reduce the number of
|
|
simplification rounds required for mutually-recursive
|
|
functions. *)
|
|
Variable.Pair.Set.fold (fun (target_fun_var, target_spec_arg)
|
|
what_to_specialise ->
|
|
if Variable.equal fun_var target_fun_var
|
|
|| not (Variable.Map.mem target_spec_arg
|
|
set_of_closures.specialised_args)
|
|
then begin
|
|
what_to_specialise
|
|
end else begin
|
|
(* Rewrite the projection (that was in terms of an inner
|
|
specialised arg of [fun_var]) to be in terms of the
|
|
corresponding inner specialised arg of
|
|
[target_fun_var]. (The outer vars referenced in the
|
|
projection remain unchanged.) *)
|
|
let projection =
|
|
Projection.map_projecting_from projection
|
|
~f:(fun var ->
|
|
assert (Variable.equal var group);
|
|
target_spec_arg)
|
|
in
|
|
W.new_specialised_arg what_to_specialise
|
|
~fun_var:target_fun_var ~group
|
|
~definition:
|
|
(Projection_from_existing_specialised_arg projection)
|
|
end)
|
|
flow
|
|
what_to_specialise)
|
|
extractions
|
|
what_to_specialise)
|
|
projections_by_function
|
|
what_to_specialise
|
|
end
|
|
|
|
include ASA.Make (Transform)
|