Optimization debugging

src/analysis_and_optimization/Dataflow_utils.ml

@@ -24,7 +24,7 @@ let branching_traverse_statement stmt ~join ~init ~f =
     | IfElse (pred, then_s, else_s_opt) ->
         let s', c = f init then_s in
         Option.value_map else_s_opt
-          ~default:(s', IfElse (pred, c, None))
+          ~default:(join s' init, IfElse (pred, c, None))
           ~f:(fun else_s ->
             let s'', c' = f init else_s in
             (join s' s'', IfElse (pred, c, Some c')) )
@@ -109,7 +109,7 @@ let is_ctrl_flow pattern =
    set of a statement. It's advantageous to build them together because they both rely on
    some of the same Break, Continue and Return bookkeeping.
 *)
-let build_cf_graphs statement_map =
+let build_cf_graphs ?blocks_after_body:(blocks_after_body=true) statement_map =
   let rec build_cf_graph_rec (cf_parent : label option)
       ((in_state, in_map) : cf_state * (label, cf_edges) Map.Poly.t)
       (label : label) : cf_state * (label, cf_edges) Map.Poly.t =
@@ -119,16 +119,24 @@ let build_cf_graphs statement_map =
     let join (state1, map1) (state2, map2) =
       (join_cf_states state1 state2, union_maps_left map1 map2)
     in
+    (* This node is the parent of substatements, unless this is a Block, which
+       is visited after substatements *)
+    let substmt_preds =
+      match stmt with
+      | Block _ when blocks_after_body -> in_state.exits
+      | _ -> Set.Poly.singleton label
+    in
     (* The accumulated state after traversing substatements *)
     let substmt_state_unlooped, substmt_map =
       branching_fold_statement stmt ~join
-        ~init:({in_state with exits= Set.Poly.singleton label}, in_map)
+        ~init:({in_state with exits= substmt_preds}, in_map)
         ~f:(build_cf_graph_rec child_cf)
     in
     (* If the statement is a loop, we need to include the loop body exits as predecessors
          of the loop *)
     let substmt_state, predecessors =
-      let looped_state =
+      match stmt with
+      | For _ | While _ ->
         (* Loop statements are preceded by:
            1. The statements that come before the loop
            2. The natural exit points of the loop body
@@ -155,9 +163,13 @@ let build_cf_graphs statement_map =
             ]
         in
         ({substmt_state_unlooped with exits= loop_exits}, loop_predecessors)
-      in
-      match stmt with
-      | For _ | While _ -> looped_state
+      | Block _ when blocks_after_body ->
+        (* Block statements are preceded by the natural exit points of the block
+           body *)
+        let block_predecessors = substmt_state_unlooped.exits in
+        (* Block exits are just the block node *)
+        let block_exits = Set.Poly.singleton label in
+        ({substmt_state_unlooped with exits= block_exits}, block_predecessors)
       | _ -> (substmt_state_unlooped, in_state.exits)
     in
     (* Some statements interact with the break/return/continue states
@@ -223,6 +235,6 @@ let build_cf_graph statement_map =
   cf_graph
 
 (** See interface file *)
-let build_predecessor_graph statement_map =
-  let exits, pred_graph, _ = build_cf_graphs statement_map in
+let build_predecessor_graph ?blocks_after_body:(blocks_after_body=true) statement_map =
+    let exits, pred_graph, _ = build_cf_graphs ~blocks_after_body statement_map in
   (exits, pred_graph)

src/analysis_and_optimization/Dataflow_utils.mli

@@ -6,8 +6,9 @@ val union_maps_left :
   ('a, 'b) Map.Poly.t -> ('a, 'b) Map.Poly.t -> ('a, 'b) Map.Poly.t
 (** Union maps, preserving the left element in a collision *)
 
-val build_cf_graphs :
-     (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t
+val build_cf_graphs
+  : ?blocks_after_body:bool
+  -> (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t
   -> label Set.Poly.t
      * (label, label Set.Poly.t) Map.Poly.t
      * (label, label Set.Poly.t) Map.Poly.t
@@ -33,8 +34,9 @@ val build_cf_graph :
    and return statements shouldn't affect other branches of execution.
 *)
 
-val build_predecessor_graph :
-     (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t
+val build_predecessor_graph
+  : ?blocks_after_body:bool
+  -> (label, (Expr.Typed.t, label) Stmt.Fixed.Pattern.t * 'm) Map.Poly.t
   -> label Set.Poly.t * (label, label Set.Poly.t) Map.Poly.t
 (**
    Building the predecessor graph requires a traversal with state that includes the

src/analysis_and_optimization/Dependence_analysis.ml

@@ -170,16 +170,6 @@ let prog_rhs_variables
   in
   union_map labels ~f:label_vars
 
-let rec var_declarations Stmt.Fixed.({pattern; _}) : string Set.Poly.t =
-  match pattern with
-  | Decl {decl_id; _} -> Set.Poly.singleton decl_id
-  | IfElse (_, s, None) | While (_, s) | For {body= s; _} -> var_declarations s
-  | IfElse (_, s1, Some s2) ->
-      Set.Poly.union (var_declarations s1) (var_declarations s2)
-  | Block slist | SList slist ->
-      Set.Poly.union_list (List.map ~f:var_declarations slist)
-  | _ -> Set.Poly.empty
-
 let stmt_uninitialized_variables (exceptions : string Set.Poly.t)
     (stmt : Stmt.Located.t) : (Location_span.t * string) Set.Poly.t =
   let flowgraph, flowgraph_to_mir =

src/analysis_and_optimization/Mir_utils.ml

@@ -2,6 +2,16 @@ open Core_kernel
 open Middle
 open Dataflow_types
 
+let rec var_declarations Stmt.Fixed.({pattern; _}) : string Set.Poly.t =
+  match pattern with
+  | Decl {decl_id; _} -> Set.Poly.singleton decl_id
+  | IfElse (_, s, None) | While (_, s) | For {body= s; _} -> var_declarations s
+  | IfElse (_, s1, Some s2) ->
+    Set.Poly.union (var_declarations s1) (var_declarations s2)
+  | Block slist | SList slist ->
+    Set.Poly.union_list (List.map ~f:var_declarations slist)
+  | _ -> Set.Poly.empty
+
 let rec map_rec_expr f e =
   let recurse = map_rec_expr f in
   Expr.Fixed.{e with pattern= f (Pattern.map recurse e.pattern)}
@@ -189,33 +199,51 @@ let rec summation_terms (Expr.Fixed.({pattern; _}) as rhs) =
 let stmt_of_block b =
   Stmt.Fixed.{pattern= SList b; meta= Stmt.Located.Meta.empty}
 
-let rec subst_expr m (Expr.Fixed.({pattern; _}) as e) =
-  match pattern with
-  | Var s -> ( match Map.find m s with Some e' -> e' | None -> e )
-  | _ -> Expr.Fixed.{e with pattern= Pattern.map (subst_expr m) pattern}
+let rec fn_subst_expr m e =
+  match m e with
+  | Some e' ->
+    (* let print_expr (e:Expr.Typed.t) = *)
+      (* [%sexp (e.pattern : Expr.Typed.Meta.t Expr.Fixed.t Expr.Fixed.Pattern.t)] |> Sexp.to_string *)
+    (* in *)
+    (* let _ = print_endline ("Replaced expr: " ^ print_expr e ^ " -> " ^ print_expr e') in *)
+    e'
+  | _ -> Expr.Fixed.{e with pattern= Pattern.map (fn_subst_expr m) e.pattern}
 
-let subst_idx m = Index.map (subst_expr m)
+let fn_subst_idx m = Index.map (fn_subst_expr m)
 
-let subst_stmt_base_helper g h b =
+let fn_subst_stmt_base_helper g h b =
   Stmt.Fixed.Pattern.(
     match b with
     | Assignment ((x, ut, l), e2) -> Assignment ((x, ut, List.map ~f:h l), g e2)
     | x -> map g (fun y -> y) x)
 
-let subst_stmt_base m = subst_stmt_base_helper (subst_expr m) (subst_idx m)
-let subst_stmt m = map_rec_stmt_loc (subst_stmt_base m)
+let fn_subst_stmt_base m = fn_subst_stmt_base_helper (fn_subst_expr m) (fn_subst_idx m)
+let fn_subst_stmt m = map_rec_stmt_loc (fn_subst_stmt_base m)
 
-let rec expr_subst_expr m e =
-  match Map.find m e with
-  | Some e' -> e'
-  | None ->
-      Expr.Fixed.{e with pattern= Pattern.map (expr_subst_expr m) e.pattern}
+let name_map m (e : Expr.Typed.t) =
+  match e.pattern with
+  | Var s ->
+    (match Map.Poly.find m s with
+     | Some s' -> Some {e with pattern = Var s'}
+     | None -> None)
+  | _ -> None
 
-let expr_subst_idx m = Index.map (expr_subst_expr m)
+let name_subst_stmt m = fn_subst_stmt (name_map m)
 
-let expr_subst_stmt_base m =
-  subst_stmt_base_helper (expr_subst_expr m) (expr_subst_idx m)
+let var_map m (e : Expr.Typed.t) =
+  match e.pattern with
+  | Var s -> Map.find m s
+  | _ -> None
+let subst_expr m e = fn_subst_expr (var_map m) e
+let subst_idx m = Index.map (subst_expr m)
+let subst_stmt_base m = fn_subst_stmt_base_helper (subst_expr m) (subst_idx m)
+let subst_stmt m = map_rec_stmt_loc (subst_stmt_base m)
 
+let expr_map m (e : Expr.Typed.t) = Map.find m e
+let expr_subst_expr m e = fn_subst_expr (expr_map m) e
+let expr_subst_idx m = Index.map (expr_subst_expr m)
+let expr_subst_stmt_base m =
+  fn_subst_stmt_base_helper (expr_subst_expr m) (expr_subst_idx m)
 let expr_subst_stmt m = map_rec_stmt_loc (expr_subst_stmt_base m)
 
 let rec expr_depth Expr.Fixed.({pattern; _}) =
@@ -278,12 +306,12 @@ let rec update_expr_ad_levels autodiffable_variables
       let e2 = update_expr_ad_levels autodiffable_variables e2 in
       { pattern= EOr (e1, e2)
       ; meta= {e.meta with adlevel= ad_level_sup [e1; e2]} }
-  | Indexed (e, i_list) ->
-      let e = update_expr_ad_levels autodiffable_variables e in
+  | Indexed (ixed, i_list) ->
+      let ixed = update_expr_ad_levels autodiffable_variables ixed in
       let i_list =
         List.map ~f:(update_idx_ad_levels autodiffable_variables) i_list
       in
-      { pattern= Indexed (e, i_list)
+      { pattern= Indexed (ixed, i_list)
       ; meta=
           { e.meta with
             adlevel= ad_level_sup (e :: List.concat_map ~f:Index.bounds i_list)

src/analysis_and_optimization/Mir_utils.mli

@@ -1,6 +1,8 @@
 open Core_kernel
 open Middle
 open Dataflow_types
+val var_declarations :
+  ('a, 'b) Stmt.Fixed.t -> string Set.Poly.t
 
 val map_rec_expr :
      (Expr.Typed.t Expr.Fixed.Pattern.t -> Expr.Typed.t Expr.Fixed.Pattern.t)
@@ -144,6 +146,11 @@ val subst_stmt :
   (string, Expr.Typed.t) Map.Poly.t -> Stmt.Located.t -> Stmt.Located.t
 (** Substitute variables occurring anywhere in a statement according to the provided Map. *)
 
+val name_subst_stmt :
+  (string, string) Map.Poly.t -> Stmt.Located.t -> Stmt.Located.t
+(** Substitute subexpressions occurring anywhere in a statement according to the provided Map. *)
+
+
 val expr_subst_expr :
   Expr.Typed.t Expr.Typed.Map.t -> Expr.Typed.t -> Expr.Typed.t
 (** Substitute subexpressions in an expression according to the provided Map, trying