Only speculate during inlining if argument types have useful information

middle_end/flambda2/simplify/inlining/call_site_inlining_decision.ml

@@ -19,6 +19,8 @@
 open! Flambda.Import
 module DE = Downwards_env
 module DA = Downwards_acc
+module T = Flambda2_types
+module TE = T.Typing_env
 module UA = Upwards_acc
 module UE = Upwards_env
 
@@ -36,6 +38,7 @@ type t =
   | Missing_code
   | Definition_says_not_to_inline
   | Environment_says_never_inline
+  | Argument_types_not_useful
   | Unrolling_depth_exceeded
   | Max_inlining_depth_exceeded
   | Recursion_depth_exceeded
@@ -62,6 +65,8 @@ let [@ocamlformat "disable"] print ppf t =
     Format.fprintf ppf "Definition_says_not_to_inline"
   | Environment_says_never_inline ->
     Format.fprintf ppf "Environment_says_never_inline"
+  | Argument_types_not_useful ->
+    Format.fprintf ppf "Argument_types_not_useful"
   | Unrolling_depth_exceeded ->
     Format.fprintf ppf "Unrolling_depth_exceeded"
   | Max_inlining_depth_exceeded ->
@@ -112,7 +117,7 @@ let can_inline (t : t) : can_inline =
   match t with
   | Missing_code | Environment_says_never_inline | Max_inlining_depth_exceeded
   | Recursion_depth_exceeded | Speculatively_not_inline _
-  | Definition_says_not_to_inline ->
+  | Definition_says_not_to_inline | Argument_types_not_useful ->
     (* If there's an [@inlined] attribute on this, something's gone wrong *)
     Do_not_inline
       { warn_if_attribute_ignored = true; because_of_definition = true }
@@ -140,6 +145,9 @@ let report_reason fmt t =
        never@ be@ inlinable"
   | Environment_says_never_inline ->
     Format.fprintf fmt "the@ environment@ says@ never@ to@ inline"
+  | Argument_types_not_useful ->
+    Format.fprintf fmt
+      "there@ was@ no@ useful@ information@ about@ the@ arguments"
   | Unrolling_depth_exceeded ->
     Format.fprintf fmt "the@ maximum@ unrolling@ depth@ has@ been@ exceeded"
   | Max_inlining_depth_exceeded ->
@@ -251,18 +259,35 @@ let speculative_inlining dacc ~apply ~function_type ~simplify_expr ~return_arity
   in
   UA.cost_metrics uacc
 
+let argument_types_useful dacc apply =
+  if not
+       (Flambda_features.Inlining.speculative_inlining_only_if_arguments_useful
+          ())
+  then true
+  else
+    let typing_env = DE.typing_env (DA.denv dacc) in
+    List.exists
+      (fun simple ->
+        Simple.pattern_match simple
+          ~name:(fun name ~coercion:_ ->
+            let ty = TE.find typing_env name None in
+            not (T.is_unknown typing_env ty))
+          ~const:(fun _ -> true))
+      (Apply.args apply)
+
 let might_inline dacc ~apply ~function_type ~simplify_expr ~return_arity : t =
   let denv = DA.denv dacc in
   let env_prohibits_inlining = not (DE.can_inline denv) in
   match DE.find_code denv (FT.code_id function_type) with
   | None -> Missing_code
   | Some code ->
-    let function_decl_decision = Code.inlining_decision code in
-    if Function_decl_inlining_decision_type.must_be_inlined
-         function_decl_decision
+    let useful = argument_types_useful dacc apply in
+    let decision = Code.inlining_decision code in
+    if not useful
+    then Argument_types_not_useful
+    else if Function_decl_inlining_decision_type.must_be_inlined decision
     then Definition_says_inline
-    else if Function_decl_inlining_decision_type.cannot_be_inlined
-              function_decl_decision
+    else if Function_decl_inlining_decision_type.cannot_be_inlined decision
     then Definition_says_not_to_inline
     else if env_prohibits_inlining
     then Environment_says_never_inline

middle_end/flambda2/simplify/inlining/call_site_inlining_decision.mli

@@ -25,6 +25,7 @@ type t = private
   | Missing_code
   | Definition_says_not_to_inline
   | Environment_says_never_inline
+  | Argument_types_not_useful
   | Unrolling_depth_exceeded
   | Max_inlining_depth_exceeded
   | Recursion_depth_exceeded

middle_end/flambda2/types/expand_head.ml

@@ -54,6 +54,8 @@ module Expanded_type : sig
 
   val is_bottom : t -> bool
 
+  val is_unknown : t -> bool
+
   val to_type : t -> Type_grammar.t
 
   type descr = private
@@ -129,6 +131,9 @@ end = struct
   let is_bottom t =
     match t.descr with Bottom -> true | Unknown | Ok _ -> false
 
+  let is_unknown t =
+    match t.descr with Unknown -> true | Bottom | Ok _ -> false
+
   let of_non_alias_type ?coercion ty : t =
     match TG.descr ty with
     | Value Unknown -> create_unknown K.value
@@ -331,6 +336,8 @@ let[@inline always] get_canonical_simples_and_expand_heads ~left_env ~left_ty
 
 let is_bottom env t = ET.is_bottom (expand_head env t)
 
+let is_unknown env t = ET.is_unknown (expand_head env t)
+
 let missing_kind env free_names =
   Name_occurrences.fold_variables free_names ~init:false
     ~f:(fun missing_kind var ->

middle_end/flambda2/types/expand_head.mli

@@ -46,6 +46,8 @@ module Expanded_type : sig
 
   val is_bottom : t -> bool
 
+  val is_unknown : t -> bool
+
   val to_type : t -> Type_grammar.t
 
   type descr = private
@@ -87,6 +89,8 @@ val get_canonical_simples_and_expand_heads :
 
 val is_bottom : Typing_env.t -> Type_grammar.t -> bool
 
+val is_unknown : Typing_env.t -> Type_grammar.t -> bool
+
 val make_suitable_for_environment :
   Typing_env.t ->
   Type_grammar.t ->

middle_end/flambda2/types/flambda2_types.mli

@@ -453,6 +453,8 @@ val bottom_types_from_arity : Flambda_arity.t -> t list
     constructed (in other words, it is [Invalid]). *)
 val is_bottom : Typing_env.t -> t -> bool
 
+val is_unknown : Typing_env.t -> t -> bool
+
 val type_for_const : Reg_width_const.t -> t
 
 val kind_for_const : Reg_width_const.t -> Flambda_kind.t

middle_end/flambda2/ui/flambda_features.ml

@@ -125,6 +125,9 @@ module Inlining = struct
     match round_or_default with
     | Round round -> FH.get ~key:round !I.threshold
     | Default -> D.threshold
+
+  let speculative_inlining_only_if_arguments_useful () =
+    !Clflags.Flambda2.Inlining.speculative_inlining_only_if_arguments_useful
 end
 
 module Debug = struct

middle_end/flambda2/ui/flambda_features.mli

@@ -88,6 +88,8 @@ module Inlining : sig
   val large_function_size : round_or_default -> int
 
   val threshold : round_or_default -> float
+
+  val speculative_inlining_only_if_arguments_useful : unit -> bool
 end
 
 module Debug : sig

ocaml/driver/compenv.ml

@@ -439,6 +439,9 @@ let read_one_param ppf position name v =
     Float_arg_helper.parse v
       "Bad syntax in OCAMLPARAM for 'flambda2-inline-threshold'"
       Flambda2.Inlining.threshold
+  | "flambda2-speculative-inlining-only-if-arguments-useful" ->
+    set "flambda2-speculative-inlining-only-if-arguments-useful"
+      [Flambda2.Inlining.speculative_inlining_only_if_arguments_useful] v
   | "flambda2-treat-invalid-code-as-unreachable" ->
     set "flambda2-treat-invalid-code-as-unreachable"
       [Flambda2.treat_invalid_code_as_unreachable] v

ocaml/driver/main_args.ml

@@ -1160,6 +1160,24 @@ let mk_flambda2_inline_threshold f =
         \     more aggressive) (Flambda 2 only)"
       Clflags.Flambda2.Inlining.Default.threshold
 
+let mk_flambda2_speculative_inlining_only_if_arguments_useful f =
+  "-flambda2-speculative-inlining-only-if-arguments-useful", Arg.Unit f,
+    Printf.sprintf " Only\n\
+        \    perform speculative inlining if the Flambda type system has\n\
+        \    useful information about the argument(s) at the call site%s\n\
+        \    (Flambda 2 only)"
+      (format_default
+        Flambda2.Inlining.Default.speculative_inlining_only_if_arguments_useful)
+
+let mk_no_flambda2_speculative_inlining_only_if_arguments_useful f =
+  "-no-flambda2-speculative-inlining-only-if-arguments-useful", Arg.Unit f,
+    Printf.sprintf " Ignore\n\
+        \     whether the Flambda type system has useful information\n\
+        \     about the argument(s) at the call site when performing\n\
+        \     speculative inlining%s (Flambda 2 only)"
+      (format_not_default
+        Flambda2.Inlining.Default.speculative_inlining_only_if_arguments_useful)
+
 let mk_flambda2_treat_invalid_code_as_unreachable f =
   "-flambda2-treat-invalid-code-as-unreachable", Arg.Unit f,
   Printf.sprintf " Treat code deemed as\n\
@@ -1452,6 +1470,8 @@ module type Optcommon_options = sig
   val _flambda2_inline_small_function_size : string -> unit
   val _flambda2_inline_large_function_size : string -> unit
   val _flambda2_inline_threshold : string -> unit
+  val _flambda2_speculative_inlining_only_if_arguments_useful : unit -> unit
+  val _no_flambda2_speculative_inlining_only_if_arguments_useful : unit -> unit
 
   val _flambda2_inlining_report_bin : unit -> unit
 
@@ -1850,6 +1870,10 @@ struct
     mk_flambda2_inline_large_function_size
       F._flambda2_inline_large_function_size;
     mk_flambda2_inline_threshold F._flambda2_inline_threshold;
+    mk_flambda2_speculative_inlining_only_if_arguments_useful
+      F._flambda2_speculative_inlining_only_if_arguments_useful;
+    mk_no_flambda2_speculative_inlining_only_if_arguments_useful
+      F._no_flambda2_speculative_inlining_only_if_arguments_useful;
 
     mk_flambda2_inlining_report_bin F._flambda2_inlining_report_bin;
 
@@ -2035,6 +2059,10 @@ module Make_opttop_options (F : Opttop_options) = struct
     mk_flambda2_inline_large_function_size
       F._flambda2_inline_large_function_size;
     mk_flambda2_inline_threshold F._flambda2_inline_threshold;
+    mk_flambda2_speculative_inlining_only_if_arguments_useful
+      F._flambda2_speculative_inlining_only_if_arguments_useful;
+    mk_no_flambda2_speculative_inlining_only_if_arguments_useful
+      F._no_flambda2_speculative_inlining_only_if_arguments_useful;
 
     mk_flambda2_inlining_report_bin F._flambda2_inlining_report_bin;
 
@@ -2425,6 +2453,12 @@ module Default = struct
         "Syntax: -flambda2-inline-threshold <float> | <round>=<float>[,...]"
         Flambda2.Inlining.threshold
 
+    let _flambda2_speculative_inlining_only_if_arguments_useful =
+      set Flambda2.Inlining.speculative_inlining_only_if_arguments_useful
+
+    let _no_flambda2_speculative_inlining_only_if_arguments_useful =
+      clear Flambda2.Inlining.speculative_inlining_only_if_arguments_useful
+
     let _flambda2_inlining_report_bin = set Flambda2.Inlining.report_bin
 
     let _flambda2_unicode = set Flambda2.unicode

ocaml/driver/main_args.mli

@@ -257,6 +257,9 @@ module type Optcommon_options = sig
   val _flambda2_inline_large_function_size : string -> unit
   val _flambda2_inline_threshold : string -> unit
 
+  val _flambda2_speculative_inlining_only_if_arguments_useful : unit -> unit
+  val _no_flambda2_speculative_inlining_only_if_arguments_useful : unit -> unit
+
   val _flambda2_inlining_report_bin : unit -> unit
 
   val _flambda2_unicode : unit -> unit

ocaml/utils/clflags.ml

@@ -536,6 +536,8 @@ module Flambda2 = struct
       let large_function_size = 10
 
       let threshold = 10.
+
+      let speculative_inlining_only_if_arguments_useful = true
     end
 
     module F = Float_arg_helper
@@ -555,6 +557,9 @@ module Flambda2 = struct
 
     let threshold = ref (F.default Default.threshold)
 
+    let speculative_inlining_only_if_arguments_useful =
+      ref Default.speculative_inlining_only_if_arguments_useful
+
     let report_bin = ref false
 
     type inlining_arguments = {

ocaml/utils/clflags.mli

@@ -288,6 +288,8 @@ module Flambda2 : sig
       val large_function_size : int
 
       val threshold : float
+
+      val speculative_inlining_only_if_arguments_useful : bool
     end
 
     val max_depth : Int_arg_helper.parsed ref
@@ -304,6 +306,8 @@ module Flambda2 : sig
 
     val threshold : Float_arg_helper.parsed ref
 
+    val speculative_inlining_only_if_arguments_useful : bool ref
+
     val report_bin : bool ref
   end
 end