Apply CSE on updated primitive

Author: chambart

middle_end/flambda/simplify/simplify_primitive.ml

@@ -48,12 +48,12 @@ let try_cse dacc ~original_prim ~simplified_args_with_tys ~min_name_mode
   if not (Name_mode.equal min_name_mode Name_mode.normal) then Not_applied dacc
   else
     let result_var = VB.var result_var in
-    match apply_cse dacc ~original_prim with
+    let args = List.map fst simplified_args_with_tys in
+    match apply_cse dacc ~original_prim:(P.update_args original_prim args) with
     | Some replace_with ->
       let named = Named.create_simple replace_with in
       let ty = T.alias_type_of (P.result_kind' original_prim) replace_with in
       let env_extension = TEE.one_equation (Name.var result_var) ty in
-      let args = List.map fst simplified_args_with_tys in
       let simplified_named =
         let cost_metrics =
           Cost_metrics.notify_removed

middle_end/flambda/terms/flambda_primitive.ml

@@ -1638,6 +1638,19 @@ let args t =
   | Ternary (_, x0, x1, x2) -> [x0; x1; x2]
   | Variadic (_, xs) -> xs
 
+let update_args t args =
+  match t, args with
+  | Nullary _, [] -> t
+  | Unary (p, _), [x0] -> Unary (p, x0)
+  | Binary (p, _, _), [x0; x1] -> Binary (p, x0, x1)
+  | Ternary (p, _, _, _), [x0; x1; x2] -> Ternary (p, x0, x1, x2)
+  | Variadic (p, l), xs ->
+    assert(List.length l = List.length xs);
+    Variadic (p, xs)
+  | _, _ ->
+    Misc.fatal_errorf "Wrong arity for updating primitive %a with arguments %a"
+      print t Simple.List.print args
+
 let result_kind (t : t) =
   match t with
   | Nullary prim -> result_kind_of_nullary_primitive prim

middle_end/flambda/terms/flambda_primitive.mli

@@ -334,6 +334,8 @@ include Contains_ids.S with type t := t
 
 val args : t -> Simple.t list
 
+val update_args : t -> Simple.t list -> t
+
 (** Simpler version (e.g. for [Inlining_cost]), where only the actual
     primitive matters, not the arguments. *)
 module Without_args : sig