Irreducible transform

lib/analysis/irreducible_loops.ml

@@ -38,23 +38,22 @@ module Make (G : GSig) = struct
   module VSet = Set.Make (G.V)
   module VMap = Map.Make (G.V)
 
+  let vset_pp fmt m =
+    Format.pp_print_string fmt
+      (VSet.to_string ~sep:", " ~start:"{" ~stop:"}" G.V.show m)
+
   (** Loop context information of a given block. *)
   type loop_info =
     | PrimaryHeader of {
         primary_header : G.V.t option;
             (** next-outermost header to this loop, if we are inside a loop *)
-        headers : VSet.t;
+        headers : VSet.t; [@printer vset_pp]
             (** the set of blocks which can be used to enter the loop. a header
                 is defined as dominating all non-header nodes of the loop. *)
-        nodes : VSet.t;
+        nodes : VSet.t; [@printer vset_pp]
             (** the set of blocks which are internal to the loop. this set forms
                 a strongly-connected component. note that a block may be
                 internal to multiple loops *)
-        entries : G.E.t Iter.t Lazy.t;
-            (** persistent iterator over edges which enter any header in this
-                loop *)
-        backedges : G.E.t Iter.t Lazy.t;
-            (** persistent iterator of back-edges to any header of this loop *)
       }  (** Designated primary header of a loop. *)
     | LoopParticipant of {
         primary_header : G.V.t;
@@ -64,8 +63,45 @@ module Make (G : GSig) = struct
         (** Member of a loop that is not a primary header (may be a non-primary
             header) *)
     | NonLoop
+  [@@deriving show { with_path = false }, eq]
 
   type block_info = { block : G.V.t; dfs_pos : int; loop : loop_info }
+  [@@deriving eq, show { with_path = false }]
+
+  let classify_block b =
+    match b with
+    | { loop = NonLoop } -> `NonLoop
+    | { loop = LoopParticipant _ } -> `LoopNode
+    | { loop = PrimaryHeader { headers } } when VSet.cardinal headers = 1 ->
+        `ReducibleHeader
+    | { loop = PrimaryHeader { headers } } -> `IrreducibleHeader
+
+  (** Accesses the Basil IR state to compute the set of entry edges originating
+      from outside the loop and going towards *any* header of the loop. *)
+  let compute_entries p = function
+    | { block; loop = PrimaryHeader { headers; nodes } } ->
+        Iter.(
+          VSet.to_iter headers
+          |> flat_map (fun h ->
+              G.pred p h |> List.to_iter
+              |> filter (fun a -> not @@ VSet.mem a nodes)
+              |> map (fun a -> G.find_edge p a h))
+          |> to_list)
+    | _ -> []
+
+  (** Accesses the Basil IR state to compute the set of back-edges. That is, the
+      set of edges originating from _inside_ the loop and going towards *any*
+      header of the loop. *)
+  let compute_backedges p = function
+    | { block; loop = PrimaryHeader { headers; nodes } } ->
+        Iter.(
+          VSet.to_iter headers
+          |> flat_map (fun h ->
+              G.pred p h |> List.to_iter
+              |> filter (fun a -> VSet.mem a nodes)
+              |> map (fun a -> G.find_edge p a h))
+          |> to_list)
+    | _ -> []
 
   module Implementation = struct
     (** Internal implementation of irreducible loop forest algorithm.
@@ -94,32 +130,6 @@ module Make (G : GSig) = struct
         - Loop: {:https://llvm.org/docs/LoopTerminology.html}
         - Cycle: {:https://llvm.org/docs/CycleTerminology.html} *)
 
-    open struct
-      (** Accesses the Basil IR state to compute the set of entry edges
-          originating from outside the loop and going towards *any* header of
-          the loop. *)
-      let compute_entries p block headers nodes =
-        lazy
-          (Iter.(
-             VSet.to_iter headers
-             |> flat_map (fun h ->
-                 diff (G.pred p h |> List.to_iter) (VSet.to_iter nodes))
-             |> map (fun src -> G.find_edge p src block))
-          |> Iter.persistent)
-
-      (** Accesses the Basil IR state to compute the set of back-edges. That is,
-          the set of edges originating from _inside_ the loop and going towards
-          *any* header of the loop. *)
-      let compute_backedges p block headers nodes =
-        lazy
-          (Iter.(
-             VSet.to_iter headers
-             |> flat_map (fun h ->
-                 inter (G.pred p h |> List.to_iter) (VSet.to_iter nodes))
-             |> map (fun src -> G.find_edge p src block))
-          |> Iter.persistent)
-    end
-
     type block_loop_state = {
       block : G.V.t;  (** the block which this loop information concerns. *)
       mutable iloop_header : G.V.t option;
@@ -133,11 +143,7 @@ module Make (G : GSig) = struct
               maintained even after visiting is finished. *)
       mutable is_traversed : bool;
           (** whether visiting this block has _started_. *)
-      mutable headers : VSet.t;
-          [@printer
-            fun fmt m ->
-              Format.pp_print_string fmt
-                (VSet.to_string ~sep:", " ~start:"{" ~stop:"}" G.V.show m)]
+      mutable headers : VSet.t; [@printer vset_pp]
           (** headers of the loop headed by this block. if this block heads a
               loop, this always contains `b` as the primary header. for
               irreducible loops, it also contains secondary headers. *)
@@ -155,13 +161,7 @@ module Make (G : GSig) = struct
       let loop =
         if is_loop_header e then
           PrimaryHeader
-            {
-              primary_header = e.iloop_header;
-              headers = e.headers;
-              nodes;
-              entries = compute_entries p e.block e.headers nodes;
-              backedges = compute_backedges p e.block e.headers nodes;
-            }
+            { primary_header = e.iloop_header; headers = e.headers; nodes }
         else if is_loop_participant e then
           LoopParticipant
             {
@@ -171,6 +171,7 @@ module Make (G : GSig) = struct
       in
       { block = e.block; dfs_pos = e.dfsp_pos_max; loop }
 
+    (** return loop forest in {b reverse} topological order.*)
     let compute_block_loop_info p (block_states : block_loop_state VMap.t) :
         block_info list =
       let open Iter in
@@ -284,10 +285,8 @@ module Make (G : GSig) = struct
     (** Sets [h] as the loop header for the block `b` and all containing loops.
     *)
     let tag_lhead s (b : block_loop_state) (h : block_loop_state option) =
-      print_endline "tag lhead";
       match h with
       | Some h when not @@ G.V.equal b.block h.block -> begin
-          print_endline (show_block_loop_state h);
           let rec loop ~(cur1 : block_loop_state) ~(cur2 : block_loop_state) =
             match cur1.iloop_header with
             | Some ih when G.V.equal ih cur2.block -> ()
@@ -363,18 +362,14 @@ module Make (G : GSig) = struct
             end
           | Return nh, [] ->
               (* this is the outermost call : we are done*)
-              raise (Recurse (Return nh, []))
+              raise_notrace (Recurse (Return nh, []))
         in
         (* iterate the remaining blocks *)
         iter_tails
           (fun it b ->
-            (*print_endline
-            (Printf.sprintf "%s %d"
-               (ID.to_string (b : block_loop_state).block)
-               (List.length it));*)
             match b with
             | { is_traversed = false } ->
-                raise
+                raise_notrace
                   (Recurse
                      ( Call { block = b; dfsp_pos = dfsp_pos + 1 },
                        ({ block = b0; dfsp_pos }, it) :: continuations ))
@@ -404,7 +399,7 @@ module Make (G : GSig) = struct
           it;
         b0.dfsp_pos <- 0;
         let result = Option.map st.l b0.iloop_header in
-        raise (Recurse (Return result, continuations))
+        raise_notrace (Recurse (Return result, continuations))
       in
       try run () with
       | Recurse (Return c, []) -> c
@@ -414,14 +409,23 @@ module Make (G : GSig) = struct
       print_endline
       @@ (VMap.to_iter st.loops
          |> Iter.to_string (fun (k, v) ->
-             Printf.sprintf "%s -> %s" (G.V.show k) (show_block_loop_state v)))
+             Printf.sprintf "%s -> %s\n" (G.V.show k) (show_block_loop_state v))
+         )
+
+    let dbg_show_r r =
+      print_endline
+      @@ (List.to_iter r
+         |> Iter.to_string (fun v -> Printf.sprintf "%s\n" (show_block_info v))
+         )
   end
 
+  (** Returns a list representing the loop forest in topological order. *)
   let solve g entry =
     let open Implementation in
     let st = create g in
     ignore @@ trav_loops st (Call { block = st.l entry; dfsp_pos = 1 }) [];
-    compute_block_loop_info g st.loops
+    let r = List.rev @@ compute_block_loop_info g st.loops in
+    r
 end
 
 module ProcIntra = struct
@@ -457,7 +461,7 @@ module ProcIntra = struct
       let open Option in
       Procedure.blocks_succ p src |> Iter.find_pred (fst %> ID.equal dest)
       |> function
-      | Some (e, b) -> (e, dest)
+      | Some _ -> (src, dest)
       | _ -> raise Not_found
 
     let succ p v = Procedure.blocks_succ p v |> Iter.map fst |> Iter.to_list
@@ -467,7 +471,8 @@ module ProcIntra = struct
   include Make (BlockGraph)
 
   (** Perform irreducible loop analysis and return a list containing the loop
-      information label for each block in the procedure . *)
+      information label for each block in the procedure. Returns all loop tags
+      for blocks in reverse-topological order. *)
   let solve_proc p =
     Procedure.get_entry_block p
     |> Option.flat_map_l (fun entry -> solve p entry)

lib/lang/block.ml

@@ -103,6 +103,26 @@ let map_fold_forwards ~(phi : 'acc -> 'v phi list -> 'acc * 'v phi list)
 let map ~phi f (b : ('v, 'e) t) : ('vv, 'ee) t =
   { stmts = Vector.map f b.stmts; phis = phi b.phis }
 
+(** Modify stmt list by creating a mutable copy of the underlying vector *)
+let fmap_stmts_copy (f : (('a, 'a, 'b) Stmt.t, 'c) Vector.t -> unit) b =
+  let v = Vector.copy b.stmts in
+  f v;
+  { b with stmts = Vector.freeze v }
+
+(** prepend statements to block statement list (copies underlying vector) *)
+let prepend_stmts (b : ('v, 'e) t) (nstmts : ('v, 'v, 'e) Stmt.t list) :
+    ('v, 'e) t =
+  let stmts = Vector.create () in
+  Vector.append_list stmts nstmts;
+  Vector.append stmts b.stmts;
+  let stmts = Vector.freeze stmts in
+  { b with stmts }
+
+(** append statements to block statement list (copies underlying vector) *)
+let append_stmts (b : ('v, 'e) t) (stmt : ('v, 'v, 'e) Stmt.t list) :
+    ('vv, 'ee) t =
+  fmap_stmts_copy (fun v -> Vector.append_list v stmt) b
+
 let flat_map ~phi f (b : ('v, 'e) t) : ('vv, 'ee) t =
   {
     stmts =

lib/lang/procedure.ml

@@ -338,15 +338,20 @@ let get_entry_block p =
         Some (List.hd id))
   with Not_found -> None
 
+(** Get the block for an id
+
+    @raise Not_found when the block does not exist. *)
+let find_block p id =
+  let open Edge in
+  let open G in
+  let g = graph p |> function Some e -> e | _ -> raise Not_found in
+  let _, e, _ = G.find_edge g (Begin id) (End id) in
+  match e with Block b -> b | Jump -> raise Not_found
+
 let get_block p id =
   let open Edge in
   let open G in
-  try
-    graph p
-    |> Option.flat_map (fun g ->
-        let _, e, _ = G.find_edge g (Begin id) (End id) in
-        match e with Block b -> Some b | Jump -> None)
-  with Not_found -> None
+  try Some (find_block p id) with Not_found -> None
 
 let decl_block_exn p name ?(phis = [])
     ~(stmts : ('var, 'var, 'expr) Stmt.t list) ?(successors = []) () =
@@ -365,6 +370,31 @@ let update_block p id (block : (Var.t, BasilExpr.t) Block.t) =
       let g = G.add_edge_e g (Begin id, Block block, End id) in
       g)
 
+let modify_succs p id ~remove ~add =
+  let open Edge in
+  let open G in
+  p
+  |> map_graph (fun g ->
+      let g =
+        G.succ_e g (End id)
+        |> List.filter
+             ( G.E.dst %> function
+               | Begin e -> List.exists (ID.equal e) remove
+               | _ -> false )
+        |> List.fold_left G.remove_edge_e g
+      in
+      let new_succs = List.map (fun s -> Vert.(End id, Jump, Begin s)) add in
+      List.fold_left G.add_edge_e g new_succs)
+
+let replace_block_succs p id succs =
+  let open Edge in
+  let open G in
+  p
+  |> map_graph (fun g ->
+      let g = G.succ_e g (End id) |> List.fold_left G.remove_edge_e g in
+      let new_succs = List.map (fun s -> Vert.(End id, Jump, Begin s)) succs in
+      List.fold_left G.add_edge_e g new_succs)
+
 let replace_edge p id (block : (Var.t, BasilExpr.t) Block.t) =
   update_block p id block
 

lib/passes.ml

@@ -146,6 +146,13 @@ module PassManager = struct
       doc = "Remove blocks unreachable from entry";
     }
 
+  let irreducible_loop =
+    {
+      name = "irreducible-loops";
+      apply = Proc Transforms.Irreducible_loop.transform;
+      doc = "Remove blocks unreachable from entry";
+    }
+
   let full_ssa =
     {
       name = "ssa";
@@ -210,6 +217,7 @@ module PassManager = struct
 
   let passes =
     [
+      irreducible_loop;
       cleanup_cfg;
       dfg_bool;
       dfg_ival_wint_product;

lib/transforms/dune

@@ -3,6 +3,7 @@
  (name transforms)
  (flags -w -27)
  (modules
+  irreducible_loop
   ssa
   cf_tx
   livevars